flashinfer.fp4_paged_mqa_logits

flashinfer.fp4_paged_mqa_logits(q: Tensor, sf_q: Tensor, kv_fused: Tensor, weights: Tensor, context_lens: Tensor, block_table: Tensor, max_context_len: int, *, sf_vec_size: int = 32, output_dtype: dtype = torch.bfloat16, epi_dtype: dtype = torch.float32, num_epi_subtiles: int = 1, is_kv_sf_interleaved: bool = False, schedule_meta: Tensor = None, out: Tensor = None) → Tensor

FP4 (MXFP4) paged MQA logits for Blackwell (SM100).

Parameters:

  • q – [batch_size, next_n, num_heads, head_dim/2] uint8 Two FP4 (E2M1) values packed per byte.

  • sf_q – [batch_size, next_n, num_heads] int32 Packed UE8M0 scale factors for each (token, head): head_dim/sf_vec_size scales of one byte each.

  • kv_fused –

    [num_blocks, block_size, 1, kv_row_bytes] uint8 kv_row_bytes = head_dim/2 + head_dim/sf_vec_size: the packed FP4 values of one token, then that token’s UE8M0 scale factors. Within a block the two regions are contiguous, not interleaved per token:

    [KV data: block_size * head_dim/2 bytes] [KV SF : block_size * head_dim/sf_vec_size bytes]

  • weights – [batch_size*next_n, num_heads] float32 per-head weights

  • context_lens – [batch_size] int32 (CUDA) No entry may exceed max_context_len.

  • block_table – [batch_size, max_blocks_per_seq] int32 (CUDA) Values are physical block indices into kv_fused’s dim 0. max_blocks_per_seq must be at least max_b round_up(context_lens[b], 128) // block_size – wider than ceil(context_lens[b] / block_size) when a length is not a multiple of 128 (context_len=257 with block_size=64 needs 6 columns, not 5). Extra entries may be any valid index (0). This is a hard precondition, not a checked argument: too few columns is a device-side out-of-bounds READ, i.e. undefined behaviour – it may return corrupt logits, or fault and poison the CUDA context.

  • max_context_len –

    int maximum KV sequence length; must be >= max(context_lens). The output row is sized from this while the schedule follows context_lens, so a smaller value is a device-side out-of-bounds WRITE, likewise undefined behaviour.

    Both preconditions above are the caller’s to satisfy. FLASHINFER_VALIDATE_INPUTS=1 raises on a violation instead, but it is a development aid only: it is off by default, and it is skipped during CUDA-graph capture because the device-to-host copy it needs is illegal there. Neither setting makes the kernel itself safe against a violated contract.

  • sf_vec_size – number of FP4 values sharing one UE8M0 scale factor. Determines both sf_q’s packing and the scale-factor bytes of each fused KV row (head_dim/sf_vec_size).

  • output_dtype – output tensor dtype (float32, float16, or bfloat16)

  • epi_dtype – epilogue dtype (float32, float16, or bfloat16)

  • num_epi_subtiles – epilogue subtile count (perf knob, default 1)

  • is_kv_sf_interleaved – declares how the scale-factor tail of each kv_fused block is ordered. False (the default) means token order. True means the block’s scale factors are split into _SF_PER_INT32 equal runs which are then round-robin interleaved, so slot k holds the scale factor of token (k % _SF_PER_INT32) * (block_size // _SF_PER_INT32) + k // _SF_PER_INT32. This describes the KV cache you pass in, so set it to match how that cache was written; supplying the interleaved order can be faster.

  • schedule_meta – optional pre-computed [num_sms+1, 2] int32 CTA schedule on CUDA. If None, computed from context_lens each call. Reusable only while the entire ceil(context_lens / 256) vector and the target device’s SM count are unchanged – a fixed batch size is not sufficient, since one sequence crossing a 256-token boundary changes it with every shape identical. Under CUDA-graph replay with changing lengths, recompute into the same buffer before launching. A stale schedule can hang the kernel. Use compute_paged_mqa_logits_schedule() to generate it.

  • out – optional pre-allocated output [batch_size*next_n, padded_ctx_len]. Use padded_context_len() to size it. Required for CUDA graph capture. May have more rows than batch_size*next_n, so one address-stable max-batch buffer can be shared across captures; only the first batch_size*next_n rows are written and returned.

Returns:

[batch_size*next_n, max_context_len] output_dtype

Return type:

logits

Restrictions of the current kernel:

The signature above is the general form. This kernel is specialised for a single problem shape; the constraints below are enforced here, and a future kernel may widen them without changing this signature.

num_heads must equal _FP4_REQUIRED_NUM_HEADS. head_dim must equal _FP4_REQUIRED_HEAD_DIM. next_n must be in 1.._FP4_MAX_NEXT_N. sf_vec_size must equal _FP4_SF_VEC_SIZE, and head_dim/sf_vec_size

must equal _SF_PER_INT32.

block_size must divide _COMPUTE_BLOCK_KV, with the quotient at

most _MAX_BLOCKS_PER_MMA.

num_epi_subtiles must divide num_heads, with the quotient a multiple of

_EPI_SUBTILE_UNROLL.

is_kv_sf_interleaved may be True only when

block_size == _FP4_SF_INTERLEAVE_BLOCK_SIZE.