flashinfer.fp8_paged_mqa_logits

flashinfer.fp8_paged_mqa_logits(q: Tensor, kv_fused: Tensor, weights: Tensor, context_lens: Tensor, block_table: Tensor, max_context_len: int, *, output_dtype: dtype = torch.float32, epi_dtype: dtype = torch.float32, acc_dtype: dtype = torch.float32, num_epi_subtiles: int = 1, schedule_meta: Tensor = None, out: Tensor = None) → Tensor

FP8 paged MQA logits for Blackwell (SM100).

Parameters:

  • q – [batch_size, next_n, num_heads, head_dim] float8_e4m3fn

  • kv_fused –

    [num_blocks, block_size, 1, kv_row_bytes] uint8 kv_row_bytes = head_dim + 4: the FP8 values of one token, then that token’s single float32 scale (4 bytes – per-token scaling, not block scaling as in FP4). Within a block the two regions are contiguous, not interleaved per token:

    [KV data: block_size * head_dim bytes] [KV SF : block_size * 4 bytes, one float32/token]

  • 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.

  • output_dtype – output tensor dtype (float32 or float16)

  • epi_dtype – epilogue accumulation dtype (float32 or float16)

  • acc_dtype – MMA accumulator dtype (float32 or float16)

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

  • 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], where padded_ctx_len >= max_context_len and is a multiple of SPLIT_KV=256. If None, allocated each call. 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

(a view of out when provided)

Return type:

logits

Restrictions of the current kernel:

The signature above is the general form. This kernel is parametric in num_heads and head_dim; the constraints below are enforced here, and a future kernel may widen them without changing this signature.

head_dim must be a multiple of _FP8_MMA_INST_K. next_n*num_heads must lie in [_MMA_N_MIN, _MMA_N_MAX] and be a multiple

of _MMA_N_MULTIPLE.

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.

head_dim, num_heads and next_n together must fit the shared memory

available per block on the device.