flashinfer.fp8_paged_mqa_logits

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