flashinfer.gemm.fp4_qkv_qknorm_rope

flashinfer.gemm.fp4_qkv_qknorm_rope(a_packed, a_block_scale, a_global_scale, qkv_weight_packed, qkv_weight_block_scale, qkv_weight_global_scale, q_norm_weight, k_norm_weight, cos_sin, positions, *, num_q_heads, num_kv_heads, head_dim, is_neox=False, qkv_scale=None, out_dtype=torch.bfloat16, output_quant_scale=None, out=None)

Packed NVFP4 QKV projection with QK RMSNorm and RoPE.

positions[i] selects the row of cos_sin used for token i. Values outside [0, M) use row 0. Supported on SM100, SM103, and SM107. num_q_heads must equal num_kv_heads, head_dim must be 128, and is_neox must be false. Output is bfloat16.

Parameters:

  • a_packed – [M, K / 2] uint8 activation, two FP4 values per byte.

  • a_block_scale – 128x4 block scales for a_packed.

  • a_global_scale – Host float multiplied into the activation scale.

  • qkv_weight_packed – [N, K / 2] uint8 weight packed as three equal Q, K, and V groups. N is 3 * num_q_heads * head_dim.

  • qkv_weight_block_scale – 128x4 block scales for qkv_weight_packed.

  • qkv_weight_global_scale – Host float multiplied into the weight scale.

  • q_norm_weight – [head_dim] bfloat16 RMSNorm weight for Q.

  • k_norm_weight – [head_dim] bfloat16 RMSNorm weight for K.

  • cos_sin – [M, head_dim] float32 table. The first half of each row is cosine and the second half is sine.

  • positions – [M] int64 indices into cos_sin.

  • num_q_heads – Number of query heads. Must equal num_kv_heads.

  • num_kv_heads – Number of key and value heads.

  • head_dim – Head size. Only 128 is supported.

  • is_neox – RoPE layout. Only false, the interleaved-pair layout, is supported.

  • qkv_scale – Optional [3] float32 scale for the Q, K, and V groups.

  • out_dtype – Output dtype. Only torch.bfloat16 is supported.

  • output_quant_scale – Not used. Quantized output is not supported for this epilogue.

  • out – Optional preallocated [M, N] bfloat16 output.

Returns:

The [M, N] bfloat16 output tensor.