Skip to content

feat(mla): support native NoPE sparse MLA on SM120 #4791

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
lucamotz wants to merge 2 commits into from
Closed

feat(mla): support native NoPE sparse MLA on SM120 #4791

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
b69c34a
feat(mla): support native NoPE sparse MLA on SM120
lucamotz Aug 27, 2026
e2af235
docs(mla): clarify SM120 NoPE contracts
lucamotz Aug 28, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 12 additions & 8 deletions csrc/sparse_mla_sm120.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,10 +28,10 @@

// Sparse-MLA SM120 paged attention orchestrator (prefill-only).
//
// Decode for both DSV3_2 and DSV4 routes through the standalone
// Decode for DSV3_2, DSV4, GLM_NSA, and GLM53_NOPE routes through the standalone
// SparseMlaSm120DecodeDsv3_2 / SparseMlaSm120DecodeDsv4 entry points from
// Python (see flashinfer/sparse_mla_sm120.py). This entry point handles
// prefill dispatch for both model types (with optional dual cache for DSV4).
// Python (see flashinfer/mla/_sparse_mla_sm120.py). This entry point handles
// prefill dispatch for every model type (with optional dual cache for DSV4).
Comment thread
coderabbitai[bot] marked this conversation as resolved.

#include <cuda_runtime.h>
#include <flashinfer/attention/sparse_mla_sm120/model/model_type.h>
Expand Down Expand Up @@ -68,19 +68,20 @@ inline ModelType resolve_model_type(int d_qk, int64_t model_type) {
if (d_qk == 512) {
const auto mt = static_cast<ModelType>(
model_type == kAuto ? static_cast<int64_t>(ModelType::DSV4) : model_type);
TVM_FFI_ICHECK(mt == ModelType::DSV4)
<< "d_qk=512 supports only model_type auto or DSV4; got " << model_type;
TVM_FFI_ICHECK(mt == ModelType::DSV4 || mt == ModelType::GLM53_NOPE)
<< "d_qk=512 supports model_type auto, DSV4, or GLM53_NOPE; got " << model_type;
Comment thread
coderabbitai[bot] marked this conversation as resolved.
return mt;
}
TVM_FFI_ICHECK(false) << "Unsupported d_qk=" << d_qk
<< "; expected 576 (DSV3_2/GLM_NSA) or 512 (DSV4)";
<< "; expected 576 (DSV3_2/GLM_NSA) or 512 (DSV4/GLM53_NOPE)";
return ModelType::DSV4;
}

inline int bytes_per_token(ModelType mt) {
switch (mt) {
case ModelType::DSV3_2:
case ModelType::GLM_NSA:
case ModelType::GLM53_NOPE:
return 656;
case ModelType::DSV4:
return 584;
Expand Down Expand Up @@ -270,8 +271,11 @@ void SparseMlaSm120PagedAttention(
stream);
TVM_FFI_ICHECK(ok) << "Unsupported sparse-MLA prefill configuration: "
<< "model="
<< (mt == ModelType::DSV3_2 ? "DSV3_2"
: (mt == ModelType::GLM_NSA ? "GLM_NSA" : "DSV4"))
<< (mt == ModelType::DSV3_2
? "DSV3_2"
: (mt == ModelType::GLM_NSA
? "GLM_NSA"
: (mt == ModelType::GLM53_NOPE ? "GLM53_NOPE" : "DSV4")))
<< " num_heads=" << num_heads << " topk=" << topk
<< " page_block_size=" << page_block_size << " topk_extra=" << extra_topk
<< " extra_page_block_size=" << extra_page_block_size;
Expand Down
25 changes: 16 additions & 9 deletions csrc/sparse_mla_sm120_decode_dsv3_2.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@
// num_heads ∈ {8, 16, 32, 64, 128}
// topk ∈ {128, 512, 1024, 2048}
// pbs = 64
// = 20 instantiations.
// plus the GLM53_NOPE (num_heads=32, topk=2176) specialization.

#include <cuda_runtime.h>

Expand Down Expand Up @@ -38,7 +38,7 @@ static bool launch_decode_dsv3_2_impl(const bf16* Q, const uint8_t* KV_cache,
int chunks_per_block_override, float sm_scale,
size_t stride_kv_block, cudaStream_t stream) {
using KV = KVCacheTraits<MT>;
static_assert(KV::D_QK == 576);
static_assert(KV::D_QK == 576 || (MT == ModelType::GLM53_NOPE && KV::D_QK == 512));
constexpr int H_BLOCKS = (NUM_HEADS + HPB - 1) / HPB;

// Dynamic smem layout (must match decode_dsv3_2_kernel.cuh exactly).
Expand Down Expand Up @@ -139,13 +139,15 @@ bool launch_sparse_mla_decode_dsv3_2(ModelType mt, int num_heads, int topk, int
Q, KV_cache, indices, mid_out, mid_lse, topk_length, output, out_lse, attn_sink, \
num_tokens, num_splits, chunks_per_block_override, sm_scale, stride_kv_block, stream); \
}
#define DSV3_2_DISPATCH(H, K) \
do { \
if (mt == ModelType::DSV3_2) { \
DSV3_2_DISPATCH_MT(ModelType::DSV3_2, H, K) \
} else if (mt == ModelType::GLM_NSA) { \
DSV3_2_DISPATCH_MT(ModelType::GLM_NSA, H, K) \
} \
#define DSV3_2_DISPATCH(H, K) \
do { \
if (mt == ModelType::DSV3_2) { \
DSV3_2_DISPATCH_MT(ModelType::DSV3_2, H, K) \
} else if (mt == ModelType::GLM_NSA) { \
DSV3_2_DISPATCH_MT(ModelType::GLM_NSA, H, K) \
} else if (mt == ModelType::GLM53_NOPE) { \
DSV3_2_DISPATCH_MT(ModelType::GLM53_NOPE, H, K) \
} \
} while (0);
DSV3_2_DISPATCH(8, 128)
DSV3_2_DISPATCH(8, 512)
Expand All @@ -167,6 +169,11 @@ bool launch_sparse_mla_decode_dsv3_2(ModelType mt, int num_heads, int topk, int
DSV3_2_DISPATCH(128, 512)
DSV3_2_DISPATCH(128, 1024)
DSV3_2_DISPATCH(128, 2048)
// GLM-5.3 combines its 2048 sparse selection with the 128-token
// indexer window. Keep this instantiation model-specific.
if (mt == ModelType::GLM53_NOPE) {
DSV3_2_DISPATCH_MT(ModelType::GLM53_NOPE, 32, 2176)
}
#undef DSV3_2_DISPATCH
#undef DSV3_2_DISPATCH_MT
return false;
Expand Down
7 changes: 4 additions & 3 deletions csrc/sparse_mla_sm120_jit_binding.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -178,10 +178,11 @@ void SparseMlaSm120DecodeDsv3_2(TensorView q, TensorView kv_cache, TensorView in
const int num_heads = static_cast<int>(q.size(1));
const int topk = static_cast<int>(indices.size(-1));
const int d_qk = static_cast<int>(q.size(2));
TVM_FFI_ICHECK_EQ(d_qk, 576) << "decode-dsv3_2 expects DSV3_2 layout (d_qk=576); got " << d_qk;
const auto mt = static_cast<ModelType>(model_type);
TVM_FFI_ICHECK(mt == ModelType::DSV3_2 || mt == ModelType::GLM_NSA)
<< "decode-dsv3_2 expects model_type DSV3_2 or GLM_NSA; got " << model_type;
TVM_FFI_ICHECK((d_qk == 576 && (mt == ModelType::DSV3_2 || mt == ModelType::GLM_NSA)) ||
(d_qk == 512 && mt == ModelType::GLM53_NOPE))
<< "decode-v32 expects DSV3_2/GLM_NSA d_qk=576 or GLM53_NOPE d_qk=512; got d_qk=" << d_qk
<< " model_type=" << model_type;

constexpr int BPT_DSV3_2 = 656;
const PagedKVLayout kv_layout = parse_paged_kv_layout(kv_cache, BPT_DSV3_2, "kv_cache");
Expand Down
18 changes: 12 additions & 6 deletions csrc/sparse_mla_sm120_prefill.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@
// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

// Sparse-MLA SM120 prefill. Single raw-pointer entry point that dispatches:
// - DSV3_2 / DSV4 model split
// - DSV3_2 / DSV4 / GLM_NSA / GLM53_NOPE model split
// - SG (single-group, 16 heads/CTA) for num_heads <= 16
// - MG (multi-group, 32 heads/CTA) for num_heads > 16
// - Dual-cache MG variants (DSV4 only)
Expand Down Expand Up @@ -216,28 +216,30 @@ inline bool dispatch_v32(int num_heads, int topk, const bf16* Q, const uint8_t*
const int32_t* indices, const float* attn_sink, bf16* output,
float* out_lse, float sm_scale, int num_tokens, size_t stride_kv_block,
const int* topk_length_ptr, cudaStream_t stream) {
static_assert(KVCacheTraits<MT>::D_QK == 576);
if (topk != 2048) return false;
static_assert(KVCacheTraits<MT>::D_QK == 576 ||
(MT == ModelType::GLM53_NOPE && KVCacheTraits<MT>::D_QK == 512));
constexpr int TOPK = MT == ModelType::GLM53_NOPE ? 2176 : 2048;
if (topk != TOPK) return false;

// PBS=64 matches the V32 decode (`decode_dsv3_2_kernel.cuh`). NH=8 covers
// small-TP shards; the SG kernel zero-pads invalid head slots up to HPB=16
// internally and gates write-back by VALID_HPB.
if (num_heads <= HPB) {
if (num_heads == 8) {
launch_prefill_sg<MT, ComputeMode::FP8, 8, 2048, 64>(
launch_prefill_sg<MT, ComputeMode::FP8, 8, TOPK, 64>(
Q, KV, indices, attn_sink, output, out_lse, sm_scale, num_tokens, stride_kv_block,
topk_length_ptr, stream);
return true;
}
if (num_heads != 16) return false;
launch_prefill_sg<MT, ComputeMode::FP8, 16, 2048, 64>(Q, KV, indices, attn_sink, output,
launch_prefill_sg<MT, ComputeMode::FP8, 16, TOPK, 64>(Q, KV, indices, attn_sink, output,
out_lse, sm_scale, num_tokens,
stride_kv_block, topk_length_ptr, stream);
return true;
}

#define DISPATCH_DSV3_2_MG(NH) \
launch_prefill_mg<MT, ComputeMode::FP8, NH, 2048, 64>(Q, KV, indices, attn_sink, output, \
launch_prefill_mg<MT, ComputeMode::FP8, NH, TOPK, 64>(Q, KV, indices, attn_sink, output, \
out_lse, sm_scale, num_tokens, \
stride_kv_block, topk_length_ptr, stream)

Expand Down Expand Up @@ -432,6 +434,10 @@ bool sparse_mla_prefill_dispatch(ModelType mt, int num_heads, int topk, int page
return dispatch_v32<ModelType::GLM_NSA>(num_heads, topk, Q, KV_cache, indices, attn_sink,
output, out_lse, sm_scale, num_tokens,
stride_kv_block, topk_length, stream);
case ModelType::GLM53_NOPE:
return dispatch_v32<ModelType::GLM53_NOPE>(num_heads, topk, Q, KV_cache, indices, attn_sink,
output, out_lse, sm_scale, num_tokens,
stride_kv_block, topk_length, stream);
case ModelType::DSV4:
return dispatch_dsv4_single(num_heads, topk, Q, KV_cache, indices, attn_sink, output, out_lse,
sm_scale, num_tokens, stride_kv_block, topk_length, stream);
Expand Down
34 changes: 24 additions & 10 deletions flashinfer/mla/_core.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -581,12 +581,20 @@ def _trtllm_batch_decode_sparse_mla_v32_sm120(
raise ValueError(
f"SM120 sparse MLA v32/GLM expects BF16 query, got {query.dtype}"
)
if kv_lora_rank != 512 or qk_rope_head_dim != 64 or query.size(-1) != 576:
rope_v32 = kv_lora_rank == 512 and qk_rope_head_dim == 64 and query.size(-1) == 576
glm53_nope = (
kv_lora_rank == 512
and qk_rope_head_dim == 0
and query.size(-1) == 512
and str(kv_scale_format).lower().replace("-", "_") == "arbitrary_fp32"
)
if not (rope_v32 or glm53_nope):
raise ValueError(
"SM120 sparse MLA v32/GLM expects kv_lora_rank=512, "
f"qk_rope_head_dim=64, and query head dim 576; got "
"SM120 sparse MLA expects either the v32/GLM_NSA 512+64 layout "
"or GLM-5.3 native NoPE 512+0 with arbitrary_fp32 scales; got "
f"kv_lora_rank={kv_lora_rank}, "
f"qk_rope_head_dim={qk_rope_head_dim}, query dim={query.size(-1)}"
f"qk_rope_head_dim={qk_rope_head_dim}, query dim={query.size(-1)}, "
f"kv_scale_format={kv_scale_format!r}"
)
if workspace_buffer.device != query.device:
raise ValueError(
Expand Down Expand Up @@ -3547,8 +3555,8 @@ def trtllm_batch_decode_with_kv_cache_mla(
``head_dim_qk = kv_lora_rank + qk_rope_head_dim``. When
``cum_seq_lens_q`` is provided, TRTLLM-GEN and monolithic CuTeDSL
instead accept compact ``[total_q, num_heads, head_dim_qk]`` input.
For the SM120/SM121 v32/GLM sparse backend, this must be BF16 with
``head_dim_qk == 576``.
For the SM120/SM121 sparse backend, this must be BF16 with
``head_dim_qk == 576`` for v32/GLM_NSA or ``512`` for GLM-5.3 NoPE.
kv_cache : torch.Tensor
For TRTLLM-GEN, CuteDSL, and XQA, the paged KV cache is
``[num_pages, page_size, kv_lora_rank + qk_rope_head_dim]`` or
Expand All @@ -3561,8 +3569,11 @@ def trtllm_batch_decode_with_kv_cache_mla(
by kernels that use semaphore state.
qk_nope_head_dim : int
Non-RoPE query dimension. Dense MLA paths commonly use ``128`` or
``64`` depending on model. The SM120/SM121 sparse v32/GLM backend
ignores this value and validates ``query.shape[-1] == 576`` instead.
``64`` depending on model. The SM120/SM121 packed sparse backend keeps
this legacy argument for API compatibility but ignores it; that path
validates ``kv_lora_rank``, ``qk_rope_head_dim``, and
``query.shape[-1]`` instead. GLM-5.3 uses the native
``qk_rope_head_dim=0`` / ``query.shape[-1]=512`` geometry.
kv_lora_rank : int
Latent KV rank. TRTLLM-GEN and SM120/SM121 sparse v32/GLM use ``512``.
qk_rope_head_dim : int
Expand Down Expand Up @@ -3684,8 +3695,11 @@ def trtllm_batch_decode_with_kv_cache_mla(
feature (e.g. ``sinks``).
kv_scale_format : str = "auto"
Scale semantics for the SM120/SM121 packed v32/GLM sparse backend.
``"auto"`` and ``"pow2_fp32"`` select DSv3.2 power-of-2 FP32 inline
scales; ``"arbitrary_fp32"`` selects GLM-style arbitrary FP32 inline scales.
For ``head_dim_qk=576``, ``"auto"`` and ``"pow2_fp32"`` select
DSv3.2 power-of-2 FP32 inline scales, while ``"arbitrary_fp32"``
selects GLM-NSA. For ``head_dim_qk=512``, ``"auto"`` selects DSv4's
footer scales and ``"arbitrary_fp32"`` selects GLM-5.3's inline
scales; ``"pow2_fp32"`` is unsupported.
Ignored by the ``trtllm-gen``, ``xqa``, and ``cute-dsl`` backends.
cum_seq_lens_q : Optional[torch.Tensor] = None
Cumulative query sequence lengths for variable-length query support,
Expand Down
61 changes: 41 additions & 20 deletions flashinfer/mla/_sparse_mla_sm120.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,8 +29,9 @@
"""Internal Sparse-MLA paged attention implementation for SM120.

Auto-dispatches between decode (num_tokens <= 64) and prefill (larger). Both
DSv3.2 (d_qk=576) and DSv4 (d_qk=512) decode go through dedicated warp-spec
standalone kernels; prefill is dispatched through the shared orchestrator.
the RoPE-bearing v32 family (d_qk=576), DSv4 (d_qk=512), and GLM-5.3 NoPE
(d_qk=512) decode through dedicated warp-spec standalone kernels; prefill is
dispatched through the shared orchestrator.

The user-facing sparse MLA entry points are
``flashinfer.mla.trtllm_batch_decode_sparse_mla_dsv4`` for DeepSeek V4 and
Expand Down Expand Up @@ -66,7 +67,7 @@

# Kernel-side constants. Mirrored from
# include/flashinfer/attention/sparse_mla_sm120/{arch,model}/*.cuh.
_D_V = 512 # value head dim (universal across DSV3_2 and DSV4)
_D_V = 512 # value head dim (shared by every supported model type)
_BI = 64 # KV partition tile size in candidates (BLOCK_SIZE_N)

# Decode/prefill cutoff: num_tokens > _DECODE_MAX_TOKENS routes to the
Expand Down Expand Up @@ -131,11 +132,13 @@
(128, 2048),
}
)
_DECODE_GLM53_NOPE_DISPATCH = frozenset({(32, 2176)})
_DECODE_DSV3_2_PAGE_BLOCK_SIZE = 64

_MODEL_TYPE_DSV3_2 = 0
_MODEL_TYPE_DSV4 = 1
_MODEL_TYPE_GLM_NSA = 2
_MODEL_TYPE_GLM53_NOPE = 3
_KV_SCALE_FORMATS = frozenset({"auto", "pow2_fp32", "arbitrary_fp32"})
_BPT_DSV3_2 = 656
_BPT_DSV4 = 584
Expand Down Expand Up @@ -170,17 +173,18 @@ def _resolve_model_type(d_qk: int, kv_scale_format: str) -> int:
return _MODEL_TYPE_GLM_NSA
return _MODEL_TYPE_DSV3_2
if d_qk == 512:
if fmt == "arbitrary_fp32":
return _MODEL_TYPE_GLM53_NOPE
if fmt != "auto":
raise ValueError(
"kv_scale_format is only configurable for d_qk=576; "
f"got d_qk=512 with kv_scale_format={kv_scale_format!r}"
f"unsupported d_qk=512 kv_scale_format={kv_scale_format!r}"
)
return _MODEL_TYPE_DSV4
raise ValueError(f"SM120 sparse-MLA supports d_qk=576 or d_qk=512, got d_qk={d_qk}")


def _bytes_per_token_for_model_type(model_type: int) -> int:
if model_type in (_MODEL_TYPE_DSV3_2, _MODEL_TYPE_GLM_NSA):
if model_type in (_MODEL_TYPE_DSV3_2, _MODEL_TYPE_GLM_NSA, _MODEL_TYPE_GLM53_NOPE):
return _BPT_DSV3_2
if model_type == _MODEL_TYPE_DSV4:
return _BPT_DSV4
Expand Down Expand Up @@ -227,14 +231,25 @@ def _packed_kv_page_block_size(


def _decode_dsv3_2_dispatchable(
num_tokens: int, num_heads: int, topk: int, d_qk: int, page_block_size: int
num_tokens: int,
num_heads: int,
topk: int,
d_qk: int,
page_block_size: int,
model_type: int,
) -> bool:
"""True iff decode-dsv3_2 supports this shape configuration."""
return (
num_tokens <= _DECODE_MAX_TOKENS
and d_qk == 576
and d_qk in (512, 576)
and page_block_size == _DECODE_DSV3_2_PAGE_BLOCK_SIZE
and (num_heads, topk) in _DECODE_DSV3_2_DISPATCH
and (
(num_heads, topk) in _DECODE_DSV3_2_DISPATCH
or (
model_type == _MODEL_TYPE_GLM53_NOPE
and (num_heads, topk) in _DECODE_GLM53_NOPE_DISPATCH
)
)
)


Expand Down Expand Up @@ -366,7 +381,10 @@ def _paged_attention(
if model_type in (
_MODEL_TYPE_DSV3_2,
_MODEL_TYPE_GLM_NSA,
) and _decode_dsv3_2_dispatchable(num_tokens, num_heads, topk, d_qk, kv_pbs):
_MODEL_TYPE_GLM53_NOPE,
) and _decode_dsv3_2_dispatchable(
num_tokens, num_heads, topk, d_qk, kv_pbs, model_type
):
num_splits = (topk + _BI - 1) // _BI
mid_out_view, mid_lse_view = _decode_scratch_views(
mid_out, mid_lse, num_tokens, num_heads, num_splits, d_v
Expand Down Expand Up @@ -449,8 +467,9 @@ def _sparse_mla_sm120_paged_attention(
----------
q : torch.Tensor
Query tensor, shape ``[num_tokens, num_heads, d_qk]``, dtype bf16.
``d_qk=576`` uses the V32-family inline-scale cache and
``d_qk=512`` uses the DSv4 footer-scale cache.
``d_qk=576`` uses a V32-family inline-scale cache. With ``d_qk=512``,
``kv_scale_format="auto"`` selects the DSv4 footer-scale cache and
``"arbitrary_fp32"`` selects the GLM-5.3 inline-scale cache.
kv_cache : torch.Tensor
Byte-packed paged main KV cache. Accepted forms are 3D
``[num_blocks, page_block_size, bytes]``, HND
Expand All @@ -469,11 +488,12 @@ def _sparse_mla_sm120_paged_attention(
sm_scale : float
Softmax scale (typically ``1 / sqrt(d_qk)``).
d_v : int
Value head dim. ``512`` for both DSV3_2 and DSV4 today.
Value head dim. ``512`` for every supported model type.
kv_scale_format : str
Scale semantics for ``d_qk=576``. ``"auto"`` and ``"pow2_fp32"``
select DSv3.2 power-of-2 FP32 inline scales; ``"arbitrary_fp32"``
selects GLM-style arbitrary FP32 inline scales.
Model/cache selector. For ``d_qk=576``, ``"auto"`` and
``"pow2_fp32"`` select DSv3.2 while ``"arbitrary_fp32"`` selects
GLM-NSA. For ``d_qk=512``, ``"auto"`` selects DSv4 and
``"arbitrary_fp32"`` selects GLM-5.3 NoPE.
topk_length : Optional[torch.Tensor]
Effective top-k length per query token, shape ``[num_tokens]``, dtype
int32. Required for sliding-window MLA near sequence start; ``None``
Expand Down Expand Up @@ -548,11 +568,12 @@ class _SparseMLAPagedAttentionRunner:
max_num_heads : Optional[int]
Optional worst-case ``num_heads``.
d_v : int
Value head dim. ``512`` for DSV3_2 / DSV4.
Value head dim. ``512`` for every supported model type.
kv_scale_format : str
Scale semantics for ``d_qk=576``. ``"auto"`` and ``"pow2_fp32"``
select DSv3.2 power-of-2 FP32 inline scales; ``"arbitrary_fp32"``
selects GLM-style arbitrary FP32 inline scales.
Model/cache selector. For ``d_qk=576``, ``"auto"`` and
``"pow2_fp32"`` select DSv3.2 while ``"arbitrary_fp32"`` selects
GLM-NSA. For ``d_qk=512``, ``"auto"`` selects DSv4 and
``"arbitrary_fp32"`` selects GLM-5.3 NoPE.
device : Optional[torch.device]
Allocation target. Defaults to the current CUDA device.

Expand Down
Loading
Loading