MLA¶

`mla_attention`	DSV4 MLA decode/verify megakernel up to the sparse attention, on SM100.
`mla_attention_plan`	Fill the plan slot of one layer kind in `workspace` and zero its barriers.
`mla_output_proj`	MLA output projection after an `mla_attention` launch, on SM100.
`mla_workspace`	Allocate the gmem workspace of `mla_attention_plan`, `mla_attention` and `mla_output_proj`.

mla_attention¶

mla_attention( q_or_x, hidden_states, kv_cache, page_size, positions, freqs_real, sm_scale, wq_a_weight, wq_a_sf, q_norm_weight, wq_b_weight, wq_b_sf, wkv_weight, wkv_sf, kv_norm_weight, swa_loc, attn_sink, workspace, q_eps=1e-06, kv_eps=1e-06, wkv_gate_weight=None, kv_score_ring=None, kv_gate_ring_row=None, c128_ape=None, c128_norm_weight=None, c128_freqs_real=None, c128_out_loc=None, c128_pool=None, c128_stride_kv_block=0, c128_page_block=0, c128_page_table=None, c128_attn_len=None, draft_tokens=0, reset_barriers=True, )¶

DSV4 MLA decode/verify megakernel up to the sparse attention, on SM100.

One persistent launch (one CTA per SM) runs x -> wq_a / wkv GEMMs -> q / kv RMSNorm + RoPE + KV cache write -> wq_b GEMM, the kv_gate GEMM and c128 compressor of HCA layers, and the sparse attention over the SWA window (plus the compressed c128 segment on HCA layers). mla_output_proj finishes the layer.

Only the DSV4 attention config is supported: 64 heads of head dim 512 (448 NoPE + 64 RoPE), q_lora rank 1024, kv_lora rank 512, value dim 512, c128 compress ratio 128, fp8 KV rows with one ue8m0 scale per 64 NoPE elements. The hidden size H is free but must be a multiple of 32.

Each token’s SWA window (its KV runs and length) comes from the plan slot of this layer kind in workspace: call mla_attention_plan for the kind first, once per forward, on the same stream with the same T.

Parameters:

q_or_x (torch.Tensor) – [T, 64 * 512] bfloat16, contiguous, T > 0. q staging buffer; on return it holds the normalized attention output O (before the inverse RoPE) of every token whose KV was not split. The other tokens’ partials stay in workspace; pass both to mla_output_proj.
hidden_states (torch.Tensor) – [T, H] bfloat16 layer input, unit column stride, 16-byte aligned rows (may be a column slice of a wider buffer).
kv_cache (torch.Tensor) – [num_pages, page_bytes] uint8 paged SWA KV cache. A page holds page_size rows of 576 bytes (448 float8_e4m3fn NoPE + 64 bfloat16 RoPE), then page_size 8-byte ue8m0 scale slots; page_bytes is a multiple of 576. Each token’s new row is written in place.
page_size (int) – Rows per kv_cache page, a power of two, as passed to mla_attention_plan.
positions (torch.Tensor) – [T] int32 position of each token.
freqs_real (torch.Tensor) – [max_pos, 64] float32 RoPE table, contiguous, 32 interleaved (cos, sin) pairs per position.
sm_scale (float) – Softmax scale of the QK scores.
wq_a_weight (torch.Tensor) – [1024, H] float8_e4m3fn wq_a weight.
wq_a_sf (torch.Tensor) – [1024, H / 32] uint8 SfAtom ue8m0 scale of wq_a_weight.
q_norm_weight (torch.Tensor) – [1024] bfloat16 RMSNorm gain of q_lora.
wq_b_weight (torch.Tensor) – [64 * 512, 1024] float8_e4m3fn wq_b weight.
wq_b_sf (torch.Tensor) – [64 * 512, 1024 / 32] uint8 SfAtom ue8m0 scale of wq_b_weight.
wkv_weight (torch.Tensor) – [512, H] float8_e4m3fn wkv weight.
wkv_sf (torch.Tensor) – [512, H / 32] uint8 SfAtom ue8m0 scale of wkv_weight.
kv_norm_weight (torch.Tensor) – [512] bfloat16 RMSNorm gain of the new KV row.
swa_loc (torch.Tensor) – [T] int32 kv_cache slot of each token’s new KV row.
attn_sink (torch.Tensor | None) – [64] float32 per-head sink logit added to the softmax denominator, or None for no sink.
workspace (torch.Tensor) – uint8 workspace from mla_workspace, planned for this layer kind by mla_attention_plan.
q_eps (float) – RMSNorm epsilon of q.
kv_eps (float) – RMSNorm epsilon of the KV rows (new and compressed).
wkv_gate_weight (torch.Tensor | None) – [2 * 512, H] bfloat16 kv_gate weight (compressor kv and score halves). This and the kv_* / c128_* arguments below are given on HCA layers and left None on SWA layers; wkv_gate_weight and c128_attn_len must be given together.
kv_score_ring (torch.Tensor | None) – [num_requests * 128, 2 * 512] float32 ring of each request’s last 128 kv_gate outputs ([kv | score]); updated in place.
kv_gate_ring_row (torch.Tensor | None) – [T] int32 kv_score_ring row of each token (request * 128 + position % 128).
c128_ape (torch.Tensor | None) – [128, 512] float32 absolute position bias added to the ring scores.
c128_norm_weight (torch.Tensor | None) – [512] bfloat16 RMSNorm gain of a compressed row.
c128_freqs_real (torch.Tensor | None) – [max_pos, 64] float32 RoPE table of compressed rows; a row fired at position p is roped at p - 127 (its window start).
c128_out_loc (torch.Tensor | None) – [T] int32 c128_pool slot written by a token whose position % 128 == 127 (the compressor fires there).
c128_pool (torch.Tensor | None) – uint8 paged compressed KV pool, same row and scale layout as kv_cache with c128_page_block rows per page.
c128_stride_kv_block (int) – Bytes per c128_pool page, a multiple of 576.
c128_page_block (int) – Rows per c128_pool page, a power of two.
c128_page_table (torch.Tensor | None) – [T, max_pages] int32 c128_pool pages of each token’s compressed segment.
c128_attn_len (torch.Tensor | None) – [T] int32 compressed rows each token attends over; None on SWA layers.
draft_tokens (int) – Tokens per request in verify, laid out consecutively at consecutive positions; 0 if unknown.
reset_barriers (bool) – True to zero the barrier banks at exit; pass False only if the next op on the stream (mla_output_proj) does it.

Returns:

None. Results are written in place.

mla_attention_plan¶

mla_attention_plan( kv_indices, topk_lengths, page_size, kv_run_step, c128_attn_len, workspace, )¶

Fill the plan slot of one layer kind in workspace and zero its barriers.

workspace keeps one plan per layer kind: SWA (c128_attn_len is None) and HCA. Run it once per forward per kind, before that kind’s mla_attention launches on the same stream with the same T; layers of one kind share the plan as long as every launch’s barriers are zeroed after it (by mla_attention itself, or by mla_output_proj).

The plan compresses each token’s SWA window into at most two runs of consecutive slots (a run ends where the next slot is not the previous one plus kv_run_step or lies on another page). A window of more than two runs (a padding or not yet written window) is cut to its first row, as sglang pads such rows to length 1. The block counts and the attention use these effective lengths.

Parameters:

kv_indices (torch.Tensor) – [T, topk] int32 SWA kv_cache slots of each token’s window, unit column stride, T > 0; columns past a token’s length are ignored (e.g. -1 padding).
topk_lengths (torch.Tensor | None) – [T] int32 contiguous SWA rows each token attends over, including its own new row, or None if every token attends over topk rows.
page_size (int) – Rows per kv_cache page, a power of two.
kv_run_step (int) – Slot step between consecutive window rows, 1 or -1 (a property of the slot table’s producer).
c128_attn_len (torch.Tensor | None) – [T] int32 compressed rows each token attends over, or None on SWA layers; selects the plan slot.
workspace (torch.Tensor) – uint8 workspace from mla_workspace.

Returns:

[T, 4] int32 view of the plan’s per-token KV runs in workspace – (slot0, len0, slot1, rows), row g < len0 of the window on slot slot0 + g, the others on slot1 + g - len0, rows the effective length. Valid until the next plan of this kind.

Return type:

torch.Tensor

mla_output_proj¶

mla_output_proj( q_or_x, positions, freqs_real, wo_a_weight, wo_a_sf, wo_b_weight, wo_b_sf, out, workspace, attn_sink=None, c128_attn_len=None, )¶

MLA output projection after an mla_attention launch, on SM100.

Computes out = wo_b(quant(wo_a(quant(rope^-1(combine(O)))))): the split-KV LSE combine and inverse RoPE of O, group-32 e4m3 / ue8m0 quantization, the grouped wo_a GEMM, group-32 quantization and the wo_b GEMM, as two PDL-chained kernels. It also zeroes the attention’s barrier banks, so the preceding mla_attention may pass reset_barriers=False.

The wo_a weight stacks G output groups: group g maps O columns [g * K, (g + 1) * K) of the flat [T, 64 * 512] O to its 2 * H / G rows of o_lora, with K = 64 * 512 / G. o_lora is twice the hidden size H. The fixed work split needs H a multiple of 512 and at most 4096, K a multiple of 512 and 2 * H / G a multiple of 256.

Parameters:

q_or_x (torch.Tensor) – [T, 64 * 512] bfloat16, contiguous, 1 <= T <= 640. The q_or_x of the attention launch, which holds O on return of mla_attention; read only.
positions (torch.Tensor) – [T] int32 token positions of the attention launch.
freqs_real (torch.Tensor) – [max_pos, 64] float32 RoPE table of the attention launch.
wo_a_weight (torch.Tensor) – [2 * H, 64 * 512 / G] float8_e4m3fn wo_a weight.
wo_a_sf (torch.Tensor) – [2 * H, 64 * 512 / G / 32] uint8 SfAtom ue8m0 scale of wo_a_weight.
wo_b_weight (torch.Tensor) – [H, 2 * H] float8_e4m3fn wo_b weight.
wo_b_sf (torch.Tensor) – [H, 2 * H / 32] uint8 SfAtom ue8m0 scale of wo_b_weight.
out (torch.Tensor) – [T, H] bfloat16, contiguous; written in place.
workspace (torch.Tensor) – uint8 workspace of the attention launch.
attn_sink (torch.Tensor | None) – [64] float32 attention sink of the attention launch, or None if it had none.
c128_attn_len (torch.Tensor | None) – c128_attn_len of the attention launch (None on SWA layers). Only whether it is None is used: it selects the SWA or HCA plan slot of workspace, so it must match the attention launch.

Returns:

None. The result is written in place to out.

mla_workspace¶

mla_workspace( max_tokens, wq_a_weight, wkv_weight, wo_a_weight, wo_b_weight, )¶

Allocate the gmem workspace of mla_attention_plan, mla_attention and mla_output_proj.

One workspace serves every layer of the model (SWA and HCA alike, same weight shapes), every CUDA graph and any T up to max_tokens. It holds the barriers, one plan per layer kind, the attention partials and activations, and the output projection’s buffers. Its barrier prefix is zeroed here, once; the kernels keep it zeroed.

Call order: mla_workspace once (before any CUDA-graph capture); per forward one mla_attention_plan per layer kind (required: the attention reads its KV runs and work list); per layer mla_attention then mla_output_proj.

Parameters:

max_tokens (int) – Largest T the workspace serves, in [1, 640].
wq_a_weight (torch.Tensor) – [1024, H] float8_e4m3fn wq_a weight; sets H.
wkv_weight (torch.Tensor) – [512, H] float8_e4m3fn wkv weight.
wo_a_weight (torch.Tensor) – [2 * H, 64 * 512 / G] float8_e4m3fn wo_a weight.
wo_b_weight (torch.Tensor) – [H, 2 * H] float8_e4m3fn wo_b weight.

Returns:

1-D uint8 workspace tensor on the device of wq_a_weight.

Return type:

torch.Tensor