flashinfer.diffusion_ops.minimax_h3_bf16_pre_attention

flashinfer.diffusion_ops.minimax_h3_bf16_pre_attention(x: Tensor, x_norm_weight: Tensor, adaln_scale: Tensor, adaln_shift: Tensor, adaln_index: Tensor, qkv_weight: Tensor, q_norm_weight: Tensor, k_norm_weight: Tensor, rope_cos_sin: Tensor, *, ulysses_degree: int, out: Tensor, eps: float = 1e-05, qk_eps: float | None = None, rope_positions: Tensor | None = None, workspace: Tensor | None = None) → Tensor

Run the BF16 pre-attention projection stage for MiniMax-H3.

The operation applies input RMSNorm, indexed AdaLN, a BF16 QKV projection, per-head Q/K RMSNorm, partial 3-D split-half NeoX RoPE, and a destination-major output pack, as two kernel launches on the current stream: the normalized and modulated activation is written to a BF16 [M, 5376] workspace, which the persistent QKV GEMM then streams through TMA while its epilogue applies the Q/K RMSNorm, RoPE and pack. The collective that consumes out is outside this operation.

The operands are the engine’s own tensors: AdaLN tables may be column chunks of a wider projection, indices are int64, and RoPE is a shared (cos_sin_cache, positions) pair. No copies are made on the host.

Parameters:

  • x (torch.Tensor) – Contiguous BF16 input with shape [M, 5376].

  • x_norm_weight (torch.Tensor) – Contiguous BF16 input RMSNorm weight with shape [5376].

  • adaln_scale (torch.Tensor) – BF16 AdaLN tables with shape [rows, 5376] for any rows >= 1. stride(1) must be 1; stride(0) may exceed 5376 (for example 6 * 5376 for a column chunk of the engine’s [rows, 6 * 5376] modulation projection) and must be a multiple of 8 elements (16 bytes); the data pointer must be 16-byte aligned. The tensors are passed to the kernel as they are, with their row count and row stride.

  • adaln_shift (torch.Tensor) – BF16 AdaLN tables with shape [rows, 5376] for any rows >= 1. stride(1) must be 1; stride(0) may exceed 5376 (for example 6 * 5376 for a column chunk of the engine’s [rows, 6 * 5376] modulation projection) and must be a multiple of 8 elements (16 bytes); the data pointer must be 16-byte aligned. The tensors are passed to the kernel as they are, with their row count and row stride.

  • adaln_index (torch.Tensor) – Contiguous int64 row indices with shape [M]. Values in [0, rows) select a table row; any other int64 value is guarded in the CUDA kernel and makes its output row all-zero instead of addressing outside the tables.

  • qkv_weight (torch.Tensor) – Contiguous BF16 fused QKV projection weight with physical shape [21504, 5376] in the engine-resident row order [qkv_kind, head, head_dim] ([q_all | k_all | v_all], the layout the model holds after loading the per-head-interleaved checkpoint).

  • q_norm_weight (torch.Tensor) – Contiguous BF16 per-head RMSNorm weights with shape [128].

  • k_norm_weight (torch.Tensor) – Contiguous BF16 per-head RMSNorm weights with shape [128].

  • rope_cos_sin (torch.Tensor) – Contiguous BF16 cache with shape [S, 96] and S >= 1. Columns [0, 48) hold frame/height/width cosine values and columns [48, 96) hold the corresponding sine values. Row m of the projection uses cache row rope_positions[m]. RoPE transforms Q/K dimensions [0, 96); dimensions [96, 128) pass through.

  • ulysses_degree (int) – Destination count, one of 1, 2, 4, or 8.

  • out (torch.Tensor) – Caller-owned contiguous BF16 destination with shape [P, M, 56 // P, 3, 128].

  • eps (float) – Epsilon of the input RMSNorm (runtime argument).

  • qk_eps (Optional[float]) – Epsilon of the per-head Q/K RMSNorms. None uses eps.

  • rope_positions (Optional[torch.Tensor]) – Contiguous int64 [M] cache row per token. Positions outside [0, S) are clamped on the device. None means the identity (row m uses cache row m, so S >= M is required); the identity is passed as a cached per-(M, device) arange tensor that is allocated on the first call for that M and reused afterwards, so the identity path is CUDA-graph safe once warmed up with one eager call of the same M.

  • workspace (Optional[torch.Tensor]) – Contiguous BF16 [M, 5376] scratch for the normalized and modulated activation (overwritten). None allocates one per call from the caching allocator; callers that pin memory or capture CUDA graphs may pass their own buffer.

Returns:

The same tensor passed as out.

Return type:

torch.Tensor

Notes

The CUDA kernel independently guards its AdaLN table loads, so malformed indices cannot form an out-of-bounds address. Valid indices preserve the MiniMax-H3 checkpoint semantics without a synchronizing host reduction.

This is a direct entry point for all supported destination counts (P in {1, 2, 4, 8}); the two-launch stage is faster than the segmented norm + cuBLAS + fused-postprocess chain at every production center on SM100a and SM103a, including P=1.