flashinfer.attn_scores¶

Paged MQA logits (“attention scores”) kernels for Blackwell (SM100/SM103).

These compute, for every request and KV position, the per-head weighted sum of rectified query-key scores that a sparse-attention indexer uses to choose which KV tokens to keep:

\[\mathrm{logits}[t, p] = \sum_h w_{t,h} \cdot \mathrm{relu}\left( q_{t,h} \cdot k_p \right)\]

Two things about that expression are easy to get wrong:

The rectifier is applied per head, before the weighted sum – not to the summed score. Because the weights may be negative, the result may be too.
fp8_paged_mqa_logits() additionally multiplies by the per-token FP32 KV scale carried in the tail of each kv_fused row, so its result is \(s_p\) times the above. fp4_paged_mqa_logits() has no such per-position factor: MXFP4 block scales are folded into the dequantised values themselves.

Note

Within a request’s context the kernels write every position unconditionally and apply no causal or context mask, so callers must mask positions beyond each request’s context length themselves. A request with context_lens[b] == 0 is skipped entirely and its output row is never written – when passing out=, initialise it if you intend to read those rows.

Paged MQA Logits¶

`fp8_paged_mqa_logits`(q, kv_fused, weights, ...)	FP8 paged MQA logits for Blackwell (SM100).
`fp4_paged_mqa_logits`(q, sf_q, kv_fused, ...)	FP4 (MXFP4) paged MQA logits for Blackwell (SM100).

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.

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.

Scheduling and Setup¶

The persistent kernels need a per-call CTA work assignment. It is computed on the GPU by default, so the whole dispatch can be captured in a CUDA graph with no host round-trip. compute_paged_mqa_logits_schedule() lets a caller build it once and pass it back via schedule_meta=.

Warning

A reused schedule is only valid while ceil(context_lens / 256) is unchanged for every request. Recompute it whenever a context length crosses a 256-token boundary, including on every CUDA-graph replay.

`attn_scores.compute_paged_mqa_logits_schedule`(...)	Compute the CTA schedule tensor for paged MQA logits kernels.
`attn_scores.padded_context_len`(max_context_len)	Return the minimum allocated context dimension for paged MQA logits output.
`attn_scores.precompile_paged_mqa_logits`([...])	Pre-compile paged MQA logits kernels for common static configs.