flashinfer.attn_scores¶

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

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 trailing scale region of each kv_fused block, 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. The asymmetry is algebraic – a scale uniform over a token’s whole K vector factors out of the dot product (and, being positive, commutes with the ReLU), while per-32-element-group scales on both operands do not.

The variants also differ beyond the scales. fp4_paged_mqa_logits() takes the extra q_sf scale tensor (plus sf_vec_size and is_kv_sf_interleaved) and packs two FP4 values per q byte. Only fp8_paged_mqa_logits() exposes acc_dtype (fp4’s MMA accumulator is fixed float32), and the default output_dtype is float32 (fp8) versus bfloat16 (fp4). fp8 accepts parametric shapes while fp4 requires exactly num_heads=64, head_dim=128. Finally, next_n=4 is single-pass on every supported arch for fp8, but for fp4 only on Rubin (SM107) – SM100/SM103 run it as two internal KV passes; build any caller schedule_meta with the schedule helper’s matching next_n/variant arguments. Paging arguments, the output contract, and out=/schedule_meta semantics are identical across the pair.

Note

Within each row’s scheduled extent the kernels write unconditionally and apply no causal or length mask; columns beyond that extent may never be written at all. Output row b*next_n + t is meaningful only for KV positions 0 .. seq_lens[b] - next_n + t inclusive (the newest slot t = next_n - 1 sees the whole sequence, each earlier slot one position fewer); callers must mask or slice past that per-slot limit themselves. A request with seq_lens[b] == 0 is skipped entirely and its output rows are 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/SM103) and Rubin (SM107).
`fp4_paged_mqa_logits`(q, q_sf, kv_fused, ...)	FP4 (MXFP4) paged MQA logits for Blackwell (SM100/SM103) and Rubin (SM107).

flashinfer.fp8_paged_mqa_logits(q: Tensor, kv_fused: Tensor, weights: Tensor, block_tables: Tensor, seq_lens: Tensor, max_seq_len: int, *, output_dtype: dtype = torch.float32, epi_dtype: dtype = torch.float32, acc_dtype: dtype = torch.float32, schedule_meta: Tensor | None = None, out: Tensor | None = None) → Tensor¶

FP8 paged MQA logits for Blackwell (SM100/SM103) and Rubin (SM107).

Computes weighted per-head ReLU attention scores against a paged FP8 KV cache – the DeepSeek-style sparse-attention indexer:

logits[b*next_n + t, pos] =          (stored as output_dtype)
    scale[pos] * sum_h weights[b*next_n + t, h] * relu(q[b,t,h,:] @ k[pos,:])
    └────────────── epilogue: epi_dtype ─────────────┘   └─ MMA: acc_dtype ───┘

The per-head dot product accumulates over head_dim inside the tensor core in acc_dtype; everything after it – ReLU, the weights multiply, the sum across heads, the scale[pos] multiply – runs on CUDA cores in epi_dtype; the result is converted to output_dtype on store. ReLU is applied per head, BEFORE the weighted sum, so negative weights can make the output negative. next_n = q.shape[1] is the number of query positions scored per request (speculative/MTP draft depth; use 1 for plain decode) – a dimension of q, not an argument; the kernel specialises on it automatically. Slot t = 0 is the oldest draft position and t = next_n - 1 the newest; seq_lens[b] is the KV length visible to the newest slot.

Example

Two requests, three draft positions each, 32-token KV blocks:

B, next_n, num_heads, head_dim = 2, 3, 64, 128
block_size, max_seq_len = 32, 8000
device = "cuda"

# Per-request KV lengths and a block table into the shared KV pool.
seq_lens = torch.tensor([3000, 8000], dtype=torch.int32, device=device)
# Table width: ceil(8000/32) = 250 blocks per request -- the
# natural width (the helper keeps the rule in one place).
blocks_per_seq = flashinfer.min_block_table_width(
    max_seq_len, block_size
)  # 250
num_blocks = B * blocks_per_seq
block_tables = torch.arange(
    num_blocks, dtype=torch.int32, device=device
).view(B, blocks_per_seq)

# Pack the fused KV pool: per block, ALL fp8 values first, then the
# per-token float32 scales (two contiguous regions, not per-token rows).
kv = torch.randn(num_blocks, block_size, head_dim, device=device).to(
    torch.float8_e4m3fn
)
kv_scales = torch.ones(num_blocks, block_size, device=device)
kv_fused = torch.cat(
    [kv.view(torch.uint8).flatten(1), kv_scales.view(torch.uint8).flatten(1)],
    dim=1,
).view(num_blocks, block_size, 1, head_dim + 4)

q = torch.randn(B, next_n, num_heads, head_dim, device=device).to(
    torch.float8_e4m3fn
)
weights = torch.rand(B * next_n, num_heads, device=device)

logits = flashinfer.fp8_paged_mqa_logits(
    q, kv_fused, weights, block_tables, seq_lens, max_seq_len
)
# logits: [B*next_n, max_seq_len] float32.  Row b*next_n + t is
# valid for positions 0 .. seq_lens[b] - next_n + t; mask the
# rest before use.

# CUDA graphs / steady-state serving: pass an address-stable
# pre-allocated output sized with padded_seq_len() -- here
# padded_seq_len(8000) = 8192 columns; the returned logits
# is its [6, 8000] view.
out = torch.empty(
    B * next_n,
    flashinfer.padded_seq_len(max_seq_len),
    dtype=torch.float32,
    device=device,
)
logits = flashinfer.fp8_paged_mqa_logits(
    q, kv_fused, weights, block_tables, seq_lens, max_seq_len, out=out
)

Parameters:

q – [batch_size, next_n, num_heads, head_dim] float8_e4m3fn
kv_fused –
[num_blocks, block_size, 1, head_dim + 4] uint8 The paged KV pool; “fused” means each block stores its quantized values and their scale factors together in one uint8 buffer. The 4-D shape is a size contract only – the kernel reads each block’s flattened block_size*(head_dim+4) bytes as two contiguous regions, NOT as per-token rows:

bytes [0, block_size*head_dim): the FP8 values,
token-major (token i at offset i*head_dim);

bytes [block_size*head_dim, end): block_size
little-endian float32 per-token scales (token i’s scale at block_size*head_dim + 4*i).

Byte offsets within one block, drawn with the Example’s block_size=32, head_dim=128:
```
0                                        4096            4224
├──────────────── REGION 1 ───────────────┤─── REGION 2 ───┤
│           FP8 values, token-major       │  fp32 scales   │
│ [tok0: 128B][tok1: 128B]...[tok31: 128B]│[s0][s1]...[s31]│
└──────────────────────────────────────────────────────────┘
```
Token i’s 128 value bytes start at byte i*128; its 4-byte little-endian float32 scale is at 4096 + 4*i. A token’s data is split: its values sit in region 1, its scale in region 2 – NOT 132-byte per-token rows, though that wrong layout has the same total size. Built in one expression from value/scale tensors – see the Example above. (Per-token scaling; FP4 uses block scaling folded into dequantization instead.)
weights – [batch_size*next_n, num_heads] float32 Per-head mixing weights, not model parameters: output[row, pos] = sum_h weights[row, h] * relu(score_h). Row b*next_n + t pairs with q[b, t] – rows are ordered exactly like the output rows; pass a [batch_size, next_n, num_heads] tensor as w.view(batch_size * next_n, num_heads). Cast internally when epi_dtype is float16.
block_tables – [batch_size, max_blocks_per_seq] int32, on q’s device, each row’s entries contiguous (stride(1) == 1). The row stride is free, so a row-strided view such as block_tables[::next_n] (a per-draft-token table de-expanded without a copy) is accepted zero-copy. Values are physical block indices into kv_fused’s dim 0. max_blocks_per_seq must be at least min_block_table_width(max(seq_lens), block_size) = ceil(max(seq_lens) / block_size) – the natural width, one entry per KV block a request occupies, i.e. the table a paged-KV serving stack already keeps. The kernel never reads row b past its own ceil(seq_lens[b] / block_size) entries, so a shorter request’s trailing entries may hold anything. The width 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.
seq_lens – [batch_size] int32, on q’s device; any stride (a strided view such as seq_lens[::next_n] is accepted zero-copy). Per-request KV length; no entry may exceed max_seq_len. (The DeepGEMM paged-MQA API this module ports calls it context_lens.)
max_seq_len –
int maximum KV sequence length; must be >= max(seq_lens). The output row is sized from this while the schedule follows seq_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 – float32 (default) or float16. bfloat16 is deliberately unsupported: the FP8 epilogue has no bf16 branch and would silently miscompute. Note that fp4_paged_mqa_logits defaults to bfloat16 – pass output_dtype explicitly in code that uses both.
epi_dtype – epilogue accumulation dtype (float32 or float16). float16 halves the epilogue traffic but saturates at |logit| > 65504 – long-sequence accumulations can overflow to inf where float32 would not.
acc_dtype – MMA accumulator dtype (float32 or float16)
schedule_meta – optional pre-computed [num_sms+1, 2] int32 CTA schedule, contiguous, on q’s device. Omit it (None, recommended) unless profiling shows the per-call schedule computation (one small GPU kernel, a few microseconds) matters in your eager hot loop. A caller-managed schedule is reusable only while the contents of seq_lens and the target device’s SM count are unchanged – a fixed batch size is not sufficient, since the schedule depends on the length values themselves. Whenever seq_lens changes (including under CUDA-graph replay), 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 storage, PyTorch out= style: the function writes into it and the returned logits is its [batch_size*next_n, max_seq_len] view (see Returns); if None, the same padded storage is allocated internally each call. Shape at least [batch_size*next_n, padded_seq_len(max_seq_len)]: the kernel pads the column count to an internal store granularity and writes the padding unconditionally, so a narrower buffer is an out-of-bounds write – always size it with padded_seq_len() (see the Example). Must be on q’s device, dtype output_dtype, rows contiguous (stride(1) == 1) and non-overlapping (stride(0) >= padded_seq_len(max_seq_len): a row slice of a wider buffer is fine; an as_strided() pitch or an expand()ed row is rejected, since later rows would overwrite earlier ones). Required for CUDA graph capture. Bigger is fine – e.g. one address-stable max-batch buffer shared across captures; only the first batch_size*next_n rows are written.

Returns:

[batch_size*next_n, max_seq_len] output_dtype: Always a view of the padded storage described under out: of the caller’s buffer when provided, else of the per-call internal allocation. Row b*next_n + t holds scores for KV positions 0 .. seq_lens[b] - next_n + t inclusive – the newest slot sees the whole sequence, each earlier slot one position fewer. Columns past that limit, and the padding columns of out beyond max_seq_len, are UNSPECIFIED: slice or mask by the limit before consuming (e.g. before a top-k). Rows of a request with seq_lens[b] == 0 are never written.

Return type:

logits

CUDA graph capture:

Pass a pre-allocated, address-stable out sized with padded_seq_len(); keep q / kv_fused / weights / seq_lens / block_tables at static addresses; leave schedule_meta=None – the schedule is then computed by a small GPU kernel inside the capture and re-derives itself from seq_lens’ current contents on every replay. If you manage the schedule yourself, recompute it into the same buffer before every replay in which the contents of seq_lens changed; a stale schedule can hang the kernel.

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 32. num_heads must be a multiple of 4. next_n*num_heads must be a multiple of 8 in [8, 256]. block_size (= kv_fused.shape[1]) must be 32, 64, or 128. head_dim, num_heads and next_n together must fit per-CTA shared

memory: on SM100, head_dim <= 192 fits at num_heads=64, next_n=1; 256 does not.

output size batch_size*next_n * padded_seq_len(max_seq_len) must: be < 2**31 elements – the kernels index the output with 32-bit offsets. With out=, batch_size*next_n * out.stride(0) must also be < 2**31 (a wider buffer widens the row stride the offsets multiply by).

flashinfer.fp4_paged_mqa_logits(q: Tensor, q_sf: Tensor, kv_fused: Tensor, weights: Tensor, block_tables: Tensor, seq_lens: Tensor, max_seq_len: int, *, sf_vec_size: int = 32, output_dtype: dtype = torch.bfloat16, epi_dtype: dtype = torch.float32, is_kv_sf_interleaved: bool = False, schedule_meta: Tensor | None = None, out: Tensor | None = None) → Tensor¶

FP4 (MXFP4) paged MQA logits for Blackwell (SM100/SM103) and Rubin (SM107).

Computes weighted per-head ReLU attention scores against a paged MXFP4 KV cache – the DeepSeek-style sparse-attention indexer:

logits[b*next_n + t, pos] =          (stored as output_dtype)
    sum_h weights[b*next_n + t, h] * relu(q[b,t,h,:] @ k[pos,:])
    └──────── epilogue: epi_dtype ───────┘   └─ MMA: float32 ─────┘

The per-head dot product accumulates in float32 inside the tensor core (not configurable here, unlike the FP8 variant’s acc_dtype); ReLU, the weights multiply, and the head sum run on CUDA cores in epi_dtype; the result is converted to output_dtype on store. The UE8M0 block scales are folded into dequantizing q and k, so unlike the FP8 variant there is no trailing per-token scale – a per-group scale does not factor out of the 128-long dot product the way fp8’s uniform per-token scale does. ReLU is applied per head, BEFORE the weighted sum, so negative weights can make the output negative. next_n = q.shape[1] is the number of query positions scored per request (speculative/MTP draft depth; use 1 for plain decode) – a dimension of q, not an argument. Slot t = 0 is the oldest draft position, t = next_n - 1 the newest; seq_lens[b] is the KV length visible to the newest slot. next_n may be 1..4: on Rubin (SM107+) every next_n runs as a single pass over the KV; on SM100/SM103, next_n=4 executes as two internal passes that each read the whole KV – numerics identical, roughly 2x the KV-bandwidth cost (~1.3x measured runtime at a 16K sequence length versus next_n=3), and a caller-supplied schedule_meta must then be built with this call’s next_n and variant=”fp4” (see schedule_meta below).

Example

Two requests, three draft positions each, 32-token KV blocks. Zero-filled quantized tensors are valid placeholders for a shape or latency check – real callers produce them with their MXFP4 quantizer (see tests/attn_scores/test_attn_scores.py::_per_token_cast_to_fp4 for a reference implementation):

B, next_n, num_heads, head_dim = 2, 3, 64, 128
block_size, max_seq_len = 32, 8000
device = "cuda"

# Per-request KV lengths and a block table into the shared KV pool.
seq_lens = torch.tensor([3000, 8000], dtype=torch.int32, device=device)
# Table width: ceil(8000/32) = 250 blocks per request -- the
# natural width (the helper keeps the rule in one place).
blocks_per_seq = flashinfer.min_block_table_width(
    max_seq_len, block_size
)  # 250
num_blocks = B * blocks_per_seq
block_tables = torch.arange(
    num_blocks, dtype=torch.int32, device=device
).view(B, blocks_per_seq)
weights = torch.rand(B * next_n, num_heads, device=device)

# Packed FP4 inputs (placeholder-quantized; see note above).
q = torch.zeros(
    B, next_n, num_heads, head_dim // 2, dtype=torch.uint8, device=device
)
q_sf = torch.zeros(B, next_n, num_heads, dtype=torch.int32, device=device)
kv_fused = torch.zeros(
    num_blocks, block_size, 1, head_dim // 2 + 4,
    dtype=torch.uint8, device=device,
)

logits = flashinfer.fp4_paged_mqa_logits(
    q, q_sf, kv_fused, weights, block_tables, seq_lens, max_seq_len
)
# logits: [B*next_n, max_seq_len] = [6, 8000] bfloat16 (fp4's
# default).  Row b*next_n + t is valid for positions
# 0 .. seq_lens[b] - next_n + t; mask the rest before use.

# CUDA graphs / steady-state serving: same out= pattern as fp8,
# but the dtype must match fp4's output_dtype (bfloat16 here).
out = torch.empty(
    B * next_n,
    flashinfer.padded_seq_len(max_seq_len),
    dtype=torch.bfloat16,
    device=device,
)
logits = flashinfer.fp4_paged_mqa_logits(
    q, q_sf, kv_fused, weights, block_tables, seq_lens,
    max_seq_len, out=out,
)

Parameters:

q – [batch_size, next_n, num_heads, head_dim/2] uint8 Two FP4 (E2M1) values packed per byte: byte j of a head’s row holds element 2j in the LOW nibble and element 2j+1 in the HIGH nibble.
q_sf – [batch_size, next_n, num_heads] int32 Packed UE8M0 scale factors for each (token, head): head_dim/sf_vec_size = 4 one-byte scales per int32. Byte k of each int32 (k-th least significant, little-endian) is the UE8M0 scale – the raw float32 exponent byte – of that (token, head)’s elements [k*sf_vec_size, (k+1)*sf_vec_size). Equivalent to a uint8 [batch_size, next_n, num_heads, 4] tensor viewed with .view(torch.int32).

kv_fused –

[num_blocks, block_size, 1, head_dim/2 + 4] uint8 The paged KV pool; “fused” means each block stores its packed FP4 values and their UE8M0 scale factors together in one uint8 buffer. The 4-D shape is a size contract only – the kernel reads each block’s flattened bytes as two contiguous regions, NOT as per-token rows:

bytes [0, block_size*head_dim/2): the packed FP4
values, token-major, same nibble rule as q;

bytes [block_size*head_dim/2, end): the UE8M0
scales, head_dim/sf_vec_size = 4 bytes per token, token-major by default (see is_kv_sf_interleaved).

Byte offsets within one block, drawn with the Example’s block_size=32, head_dim=128 (default token-major scales):

0                                     2048            2176
├────────────── REGION 1 ─────────────┤─── REGION 2 ───┤
│      packed FP4 values, token-major │  UE8M0 scales  │
│[tok0: 64B][tok1: 64B]...[tok31: 64B]│[t0][t1]...[t31]│
└──────────────────────────────────────────────────────┘

Token i’s 64 packed-value bytes start at byte i*64; its 4 UE8M0 scale bytes are at 2048 + 4*i. Zooming into one cell of each region:

[tok_i: 64B] = [e1|e0][e3|e2][e5|e4] ... [e127|e126]
                hi lo  hi lo  hi lo       hi    lo

  128 FP4 values squeezed two per byte: byte j
  holds element 2j in its LOW nibble (bits 0-3)
  and element 2j+1 in its HIGH nibble (bits 4-7)
  -- the same nibble rule as q.

[t_i: 4B] = [sf e0-31][sf e32-63][sf e64-95][sf e96-127]
             byte 0    byte 1     byte 2     byte 3

  4 UE8M0 scales squeezed into one little-endian
  int32, one per 32-element group: ascending byte
  offset = ascending group, so byte k (the k-th
  least significant of the int32) is the scale of
  elements [k*32, (k+1)*32).

Build it like the fp8 example’s torch.cat recipe, with the value/scale regions sized as above.

weights – [batch_size*next_n, num_heads] float32 Per-head mixing weights, not model parameters: output[row, pos] = sum_h weights[row, h] * relu(score_h). Row b*next_n + t pairs with q[b, t] – rows are ordered exactly like the output rows; pass a [batch_size, next_n, num_heads] tensor as w.view(batch_size * next_n, num_heads). Cast internally when epi_dtype is float16/bfloat16.
block_tables – [batch_size, max_blocks_per_seq] int32, on q’s device, each row’s entries contiguous (stride(1) == 1). The row stride is free, so a row-strided view such as block_tables[::next_n] (a per-draft-token table de-expanded without a copy) is accepted zero-copy. Values are physical block indices into kv_fused’s dim 0. max_blocks_per_seq must be at least min_block_table_width(max(seq_lens), block_size) = ceil(max(seq_lens) / block_size) – the natural width, one entry per KV block a request occupies, i.e. the table a paged-KV serving stack already keeps. The kernel never reads row b past its own ceil(seq_lens[b] / block_size) entries, so a shorter request’s trailing entries may hold anything. The width 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.
seq_lens – [batch_size] int32, on q’s device; any stride (a strided view such as seq_lens[::next_n] is accepted zero-copy). Per-request KV length; no entry may exceed max_seq_len. (The DeepGEMM paged-MQA API this module ports calls it context_lens.)
max_seq_len –
int maximum KV sequence length; must be >= max(seq_lens). The output row is sized from this while the schedule follows seq_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 (32). Determines both q_sf’s packing and the scale-factor bytes of each fused KV row.
output_dtype – bfloat16 (default), float32, or float16. The default differs from fp8_paged_mqa_logits (float32) – it mirrors what the upstream TensorRT-LLM integration binds for each variant; pass output_dtype explicitly in code that uses both.
epi_dtype – epilogue accumulation dtype (float32, float16, or bfloat16). float16 halves the epilogue traffic but saturates at |logit| > 65504 – long-sequence accumulations can overflow to inf where float32 would not. There is no acc_dtype knob: block-scaled FP4 MMA always accumulates in float32.
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 4 equal runs which are then round-robin interleaved, so slot k holds the scale factor of token (k % 4) * (block_size // 4) + k // 4. 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, contiguous, on q’s device. Omit it (None, recommended) unless profiling shows the per-call schedule computation (one small GPU kernel, a few microseconds) matters in your eager hot loop. Always generate it with compute_paged_mqa_logits_schedule(seq_lens, device=…, next_n=…, variant=”fp4”), mirroring this call’s arguments: for next_n=4 on SM100/SM103 the kernel internally restructures the problem (the two-pass split described above), and the helper applies the same internal policy, so the returned schedule is the right one on every architecture. A schedule built with mismatched arguments passes the always-on shape checks but can hang the kernel; the opt-in FLASHINFER_VALIDATE_INPUTS=1 freshness check catches it outside graph capture. A caller-managed schedule is reusable only while the contents of seq_lens and the target device’s SM count are unchanged – a fixed batch size is not sufficient, since the schedule depends on the length values themselves. Whenever seq_lens changes (including under CUDA-graph replay), 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 storage, PyTorch out= style: the function writes into it and the returned logits is its [batch_size*next_n, max_seq_len] view (see Returns); if None, the same padded storage is allocated internally each call. Shape at least [batch_size*next_n, padded_seq_len(max_seq_len)]: the kernel pads the column count to an internal store granularity and writes the padding unconditionally, so a narrower buffer is an out-of-bounds write – always size it with padded_seq_len() (see the Example). Must be on q’s device, dtype output_dtype, rows contiguous (stride(1) == 1) and non-overlapping (stride(0) >= padded_seq_len(max_seq_len): a row slice of a wider buffer is fine; an as_strided() pitch or an expand()ed row is rejected, since later rows would overwrite earlier ones). Required for CUDA graph capture. Bigger is fine – e.g. one address-stable max-batch buffer shared across captures; only the first batch_size*next_n rows are written.

Returns:

[batch_size*next_n, max_seq_len] output_dtype: Always a view of the padded storage described under out: of the caller’s buffer when provided, else of the per-call internal allocation. Row b*next_n + t holds scores for KV positions 0 .. seq_lens[b] - next_n + t inclusive – the newest slot sees the whole sequence, each earlier slot one position fewer. Columns past that limit, and the padding columns of out beyond max_seq_len, are UNSPECIFIED: slice or mask by the limit before consuming (e.g. before a top-k). Rows of a request with seq_lens[b] == 0 are never written.

Return type:

logits

CUDA graph capture:

Pass a pre-allocated, address-stable out sized with padded_seq_len(); keep q / q_sf / kv_fused / weights / seq_lens / block_tables at static addresses; leave schedule_meta=None – the schedule is then computed by a small GPU kernel inside the capture and re-derives itself from seq_lens’ current contents on every replay (including for the internal next_n=4 split on SM100/SM103). If you manage the schedule yourself, build it with this call’s next_n and variant=”fp4”, and recompute it into the same buffer before every replay in which the contents of seq_lens changed; a stale schedule can hang the kernel.

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 64. head_dim must equal 128. next_n must be in 1..4. schedule_meta must be generated with this call’s next_n and

variant=”fp4” (see schedule_meta above); on SM100/SM103 next_n=4 runs the internal two-pass split, which changes the schedule contents.

sf_vec_size must equal 32 (so head_dim/sf_vec_size = 4 scales per: token – exactly one int32 in q_sf).

block_size (= kv_fused.shape[1]) must be 32, 64, or 128. is_kv_sf_interleaved may be True only when block_size == 128. output size batch_size*next_n * padded_seq_len(max_seq_len) must

be < 2**31 elements – the kernels index the output with 32-bit offsets. With out=, batch_size*next_n * out.stride(0) must also be < 2**31 (a wider buffer widens the row stride the offsets multiply by).

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=. Pass the helper the same variant and next_n as the main call: it applies the kernel’s internal scheduling policy for that configuration, so the returned tensor (treat it as opaque) is correct on every architecture.

Warning

A reused schedule is only valid while the contents of seq_lens are unchanged for every request. Recompute it whenever seq_lens changes, including on every CUDA-graph replay.

`compute_paged_mqa_logits_schedule`(seq_lens)	Compute the CTA schedule tensor for paged MQA logits kernels.
`min_block_table_width`(seq_len, block_size)	Return the minimum `block_tables` column count for a sequence length.
`padded_seq_len`(max_seq_len)	Return the minimum column count of the paged MQA logits `out` buffer.
`precompile_paged_mqa_logits`([device, ...])	Pre-compile paged MQA logits kernels for common static configs.