diff --git a/csrc/sparse_mla_sm120.cu b/csrc/sparse_mla_sm120.cu index 2d24432b4a9..907d84b60c9 100644 --- a/csrc/sparse_mla_sm120.cu +++ b/csrc/sparse_mla_sm120.cu @@ -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). #include #include @@ -68,12 +68,12 @@ inline ModelType resolve_model_type(int d_qk, int64_t model_type) { if (d_qk == 512) { const auto mt = static_cast( model_type == kAuto ? static_cast(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; 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; } @@ -81,6 +81,7 @@ 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; @@ -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; diff --git a/csrc/sparse_mla_sm120_decode_dsv3_2.cu b/csrc/sparse_mla_sm120_decode_dsv3_2.cu index 742149245cf..8630039a0e2 100644 --- a/csrc/sparse_mla_sm120_decode_dsv3_2.cu +++ b/csrc/sparse_mla_sm120_decode_dsv3_2.cu @@ -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 @@ -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; - 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). @@ -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) @@ -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; diff --git a/csrc/sparse_mla_sm120_jit_binding.cu b/csrc/sparse_mla_sm120_jit_binding.cu index 5d59d0a4203..fef5b9ede10 100644 --- a/csrc/sparse_mla_sm120_jit_binding.cu +++ b/csrc/sparse_mla_sm120_jit_binding.cu @@ -178,10 +178,11 @@ void SparseMlaSm120DecodeDsv3_2(TensorView q, TensorView kv_cache, TensorView in const int num_heads = static_cast(q.size(1)); const int topk = static_cast(indices.size(-1)); const int d_qk = static_cast(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(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"); diff --git a/csrc/sparse_mla_sm120_prefill.cu b/csrc/sparse_mla_sm120_prefill.cu index 0c988caaae3..acbe28bdcea 100644 --- a/csrc/sparse_mla_sm120_prefill.cu +++ b/csrc/sparse_mla_sm120_prefill.cu @@ -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) @@ -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::D_QK == 576); - if (topk != 2048) return false; + static_assert(KVCacheTraits::D_QK == 576 || + (MT == ModelType::GLM53_NOPE && KVCacheTraits::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( + launch_prefill_sg( 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(Q, KV, indices, attn_sink, output, + launch_prefill_sg(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(Q, KV, indices, attn_sink, output, \ + launch_prefill_mg(Q, KV, indices, attn_sink, output, \ out_lse, sm_scale, num_tokens, \ stride_kv_block, topk_length_ptr, stream) @@ -432,6 +434,10 @@ bool sparse_mla_prefill_dispatch(ModelType mt, int num_heads, int topk, int page return dispatch_v32(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(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); diff --git a/flashinfer/mla/_core.py b/flashinfer/mla/_core.py index 493f453e424..18657b80894 100644 --- a/flashinfer/mla/_core.py +++ b/flashinfer/mla/_core.py @@ -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( @@ -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 @@ -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 @@ -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, diff --git a/flashinfer/mla/_sparse_mla_sm120.py b/flashinfer/mla/_sparse_mla_sm120.py index 67849b902d1..35240efe2bf 100644 --- a/flashinfer/mla/_sparse_mla_sm120.py +++ b/flashinfer/mla/_sparse_mla_sm120.py @@ -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 @@ -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 @@ -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 @@ -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 @@ -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 + ) + ) ) @@ -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 @@ -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 @@ -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`` @@ -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. diff --git a/include/flashinfer/attention/sparse_mla_sm120/common/fp8_quant.cuh b/include/flashinfer/attention/sparse_mla_sm120/common/fp8_quant.cuh index 8260f87af52..dfdc77b91ea 100644 --- a/include/flashinfer/attention/sparse_mla_sm120/common/fp8_quant.cuh +++ b/include/flashinfer/attention/sparse_mla_sm120/common/fp8_quant.cuh @@ -49,6 +49,7 @@ __device__ __forceinline__ void load_q_bf16_to_smem(bf16* q_nope_bf16, bf16* q_r const bf16* q_base, int valid_hpb = HPB) { using KV = KVCacheTraits; constexpr int D_NOPE = KV::D_NOPE; + constexpr int D_ROPE = KV::D_ROPE; constexpr int DIM = KV::D_QK; constexpr int BF16_STRIDE = KV::Q_NOPE_BF16_STRIDE; @@ -57,9 +58,12 @@ __device__ __forceinline__ void load_q_bf16_to_smem(bf16* q_nope_bf16, bf16* q_r q_nope_bf16[h * BF16_STRIDE + d] = (h < valid_hpb) ? q_base[h * DIM + d] : __float2bfloat16(0.f); } - for (int i = threadIdx.x; i < HPB * D_ROPE; i += _MATH_THREADS) { - int h = i / D_ROPE, d = i % D_ROPE; - q_rope[h * D_ROPE + d] = (h < valid_hpb) ? q_base[h * DIM + D_NOPE + d] : __float2bfloat16(0.f); + if constexpr (D_ROPE > 0) { + for (int i = threadIdx.x; i < HPB * D_ROPE; i += _MATH_THREADS) { + int h = i / D_ROPE, d = i % D_ROPE; + q_rope[h * D_ROPE + d] = + (h < valid_hpb) ? q_base[h * DIM + D_NOPE + d] : __float2bfloat16(0.f); + } } bar_sync_t<2, _MATH_THREADS>(); } @@ -70,6 +74,7 @@ __device__ __forceinline__ void quantize_q_to_smem(uint8_t* q_nope_fp8, float* q float* reduce_buf, int valid_hpb = HPB) { using KV = KVCacheTraits; constexpr int D_NOPE = KV::D_NOPE; + constexpr int D_ROPE = KV::D_ROPE; constexpr int Q_NOPE_STRIDE = KV::Q_NOPE_STRIDE; constexpr int QUANT_TILE = KV::QUANT_TILE; constexpr int NUM_SCALES = KV::NUM_SCALES; @@ -78,9 +83,12 @@ __device__ __forceinline__ void quantize_q_to_smem(uint8_t* q_nope_fp8, float* q float* amax = reduce_buf; // Step 1: copy Q rope to smem (only valid heads from gmem; zero-fill rest) - for (int i = threadIdx.x; i < HPB * D_ROPE; i += _MATH_THREADS) { - int h = i / D_ROPE, d = i % D_ROPE; - q_rope[h * D_ROPE + d] = (h < valid_hpb) ? q_base[h * DIM + D_NOPE + d] : __float2bfloat16(0.f); + if constexpr (D_ROPE > 0) { + for (int i = threadIdx.x; i < HPB * D_ROPE; i += _MATH_THREADS) { + int h = i / D_ROPE, d = i % D_ROPE; + q_rope[h * D_ROPE + d] = + (h < valid_hpb) ? q_base[h * DIM + D_NOPE + d] : __float2bfloat16(0.f); + } } // Step 2: init amax for (int i = threadIdx.x; i < HPB * NUM_SCALES; i += _MATH_THREADS) amax[i] = 0.f; diff --git a/include/flashinfer/attention/sparse_mla_sm120/common/kv_cache_io.cuh b/include/flashinfer/attention/sparse_mla_sm120/common/kv_cache_io.cuh index 86ae07cc6b0..780b7cb2166 100644 --- a/include/flashinfer/attention/sparse_mla_sm120/common/kv_cache_io.cuh +++ b/include/flashinfer/attention/sparse_mla_sm120/common/kv_cache_io.cuh @@ -34,13 +34,14 @@ // KV cache IO: gather BI entries from global KV pool to smem. // -// FlashMLA ABI: stride_kv_row = bytes_per_token (DSV3_2: 656, DSV4: 584). +// FlashMLA ABI: stride_kv_row = bytes_per_token (inline family: 656, DSV4: 584). // The IO stride used for address calculation is the DATA stride: -// DSV3_2: 656 (nope+scale+rope all contiguous, 656 % 16 = 0 ✓) -// DSV4: 576 (nope+rope only, footer scales excluded) -// 576 % 16 = 0 ✓ for cp.async.bulk +// inline family: 656 (data/scales/padding contiguous, 656 % 16 = 0 ✓) +// DSV4: 576 (nope+rope only, footer scales excluded) +// 576 % 16 = 0 ✓ for cp.async.bulk // -// DSV3_2 uses flat addressing: kv_ptr + global_idx * 656. +// DSV3_2, GLM_NSA, and GLM53_NOPE use flat addressing: +// kv_ptr + global_idx * 656. // DSV4 uses block-structured addressing (footer layout): // data: kv_ptr + block_idx * stride_kv_block + local_idx * 576 // scale: kv_ptr + block_idx * stride_kv_block + page_block_size * 576 + local_idx * 8 @@ -50,15 +51,15 @@ template struct KVIOTraits { using KV = KVCacheTraits; - // DSV3_2: IO_STRIDE = KV_GMEM_STRIDE = 656 (inline, bulk copy includes scale) + // Inline family: IO_STRIDE = KV_GMEM_STRIDE = 656 (bulk copy includes scale) // DSV4: IO_STRIDE = D_NOPE + D_ROPE*2 = 576 (footer, data portion only) static constexpr int IO_STRIDE = - KV::SCALE_IN_KV_SMEM ? KV::KV_GMEM_STRIDE : (KV::D_NOPE + D_ROPE * sizeof(bf16)); + KV::SCALE_IN_KV_SMEM ? KV::KV_GMEM_STRIDE : (KV::D_NOPE + KV::D_ROPE * sizeof(bf16)); static_assert(IO_STRIDE % 16 == 0, "IO stride must be 16B aligned for cp.async.bulk"); }; -// Bulk gather token nope data (and inline scales for DSV3_2) from global to smem. -// DSV3_2: flat addressing (idx * 656). DSV4: block-structured (footer layout). +// Bulk gather token nope data (and inline scales where present) from global to smem. +// Inline family: flat addressing (idx * 656). DSV4: block-structured (footer layout). template __device__ __forceinline__ void io_bulk_gather_tile(uint8_t* dst, const int32_t* indices, const uint8_t* __restrict__ kv_ptr, diff --git a/include/flashinfer/attention/sparse_mla_sm120/common/q_rope.cuh b/include/flashinfer/attention/sparse_mla_sm120/common/q_rope.cuh index 9ce8e2de3d9..2395aaedafe 100644 --- a/include/flashinfer/attention/sparse_mla_sm120/common/q_rope.cuh +++ b/include/flashinfer/attention/sparse_mla_sm120/common/q_rope.cuh @@ -41,28 +41,37 @@ // KV rope B operands are prefetched from global into registers BEFORE QK nope // MMA, so the ~300 cycle load latency overlaps with nope MMA compute. +template struct QRopeRegs { - uint32_t a[N_ROPE_CHUNKS][4]; + static constexpr int N_CHUNKS = KVCacheTraits::D_ROPE / 16; + uint32_t a[N_CHUNKS > 0 ? N_CHUNKS : 1][4]; }; +template struct KVRopePrefetch { - uint32_t b[N_ROPE_CHUNKS][2]; + static constexpr int N_CHUNKS = KVCacheTraits::D_ROPE / 16; + uint32_t b[N_CHUNKS > 0 ? N_CHUNKS : 1][2]; }; -__device__ __forceinline__ QRopeRegs preload_q_rope_regs(const bf16* q_rope_smem, int lane) { - QRopeRegs regs; +template +__device__ __forceinline__ QRopeRegs preload_q_rope_regs(const bf16* q_rope_smem, int lane) { + using KV = KVCacheTraits; + constexpr int N_CHUNKS = KV::D_ROPE / 16; + QRopeRegs regs{}; #pragma unroll - for (int ks = 0; ks < N_ROPE_CHUNKS; ks++) + for (int ks = 0; ks < N_CHUNKS; ks++) ldmatrix_load_A_bf16(regs.a[ks][0], regs.a[ks][1], regs.a[ks][2], regs.a[ks][3], - q_rope_smem + ks * 16, D_ROPE, lane); + q_rope_smem + ks * 16, KV::D_ROPE, lane); return regs; } -__device__ __forceinline__ KVRopePrefetch prefetch_kv_rope(const bf16* kv_rope_ptr, int lane) { +template +__device__ __forceinline__ KVRopePrefetch prefetch_kv_rope(const bf16* kv_rope_ptr, int lane) { + constexpr int N_CHUNKS = KVCacheTraits::D_ROPE / 16; const int tid = lane & 3; - KVRopePrefetch pf; + KVRopePrefetch pf{}; #pragma unroll - for (int ks = 0; ks < N_ROPE_CHUNKS; ks++) { + for (int ks = 0; ks < N_CHUNKS; ks++) { int ko = ks * 16; pf.b[ks][0] = *reinterpret_cast(kv_rope_ptr + ko + tid * 2); pf.b[ks][1] = *reinterpret_cast(kv_rope_ptr + ko + 8 + tid * 2); @@ -70,11 +79,13 @@ __device__ __forceinline__ KVRopePrefetch prefetch_kv_rope(const bf16* kv_rope_p return pf; } -__device__ __forceinline__ void compute_qk_rope(float qk[4], const QRopeRegs& qr, - const KVRopePrefetch& pf) { +template +__device__ __forceinline__ void compute_qk_rope(float qk[4], const QRopeRegs& qr, + const KVRopePrefetch& pf) { + constexpr int N_CHUNKS = KVCacheTraits::D_ROPE / 16; float ra[4] = {0.f, 0.f, 0.f, 0.f}; #pragma unroll - for (int ks = 0; ks < N_ROPE_CHUNKS; ks++) { + for (int ks = 0; ks < N_CHUNKS; ks++) { MmaBf16Result r = mma_bf16_m16n8k16(qr.a[ks][0], qr.a[ks][1], qr.a[ks][2], qr.a[ks][3], pf.b[ks][0], pf.b[ks][1], ra[0], ra[1], ra[2], ra[3]); ra[0] = r.d0; diff --git a/include/flashinfer/attention/sparse_mla_sm120/common/smem_layout.cuh b/include/flashinfer/attention/sparse_mla_sm120/common/smem_layout.cuh index 8d8a50d7bbb..da89b16a949 100644 --- a/include/flashinfer/attention/sparse_mla_sm120/common/smem_layout.cuh +++ b/include/flashinfer/attention/sparse_mla_sm120/common/smem_layout.cuh @@ -34,7 +34,7 @@ // Parameterized by ModelType and ComputeMode. // // Buffers (decode / prefill SG): -// q_nope_fp8, q_nope_sc, q_rope, kv_buf×2, [kv_scale_buf×2 for DSV4], +// q_nope_fp8, q_nope_sc, q_rope, kv_buf×2, [kv_scale_buf×2 for footer layouts], // reduce_buf, sum_reduce_buf (or union), m_smem, l_smem, // w_head_sc_all, w_fp8 (FP8 mode), v_trans, mbar_kv // @@ -50,13 +50,13 @@ struct SmemLayout { static constexpr size_t SMEM_Q_NOPE = BF16_Q ? HPB * KV::Q_NOPE_BF16_STRIDE * sizeof(bf16) : HPB * KV::Q_NOPE_STRIDE; static constexpr size_t SMEM_Q_SC = BF16_Q ? 0 : HPB * KV::NUM_SCALES * sizeof(float); - static constexpr size_t SMEM_Q_ROPE = HPB * D_ROPE * sizeof(bf16); + static constexpr size_t SMEM_Q_ROPE = HPB * KV::D_ROPE * sizeof(bf16); // KV double buffer static constexpr size_t SMEM_KV_BUF = BI * KV::KV_SMEM_STRIDE; // KV scale buffer: needed when bulk copy doesn't include scales. - // DSV3_2: copies 528B (nope+scale), scales in kv_smem → no extra buffer. + // Inline layouts copy nope+scale together, so scales already reside in kv_smem. // DSV4: copies 448B (nope only), scales at offset 576 → need separate buffer. static constexpr bool NEED_SCALE_BUF = (KV::KV_SMEM_COPY_BYTES < KV::KV_SCALE_GMEM_OFFSET + KV::SCALE_BYTES_PER_TOKEN); @@ -123,7 +123,7 @@ struct SmemLayoutMG { static constexpr int W_FP8_PARITIES = 2; static constexpr size_t SMEM_W_FP8_MG = W_FP8_PARITIES * N_HG * HPB * (BI + 16); // q_rope is only needed before the main loop; reuse the W_FP8 region. - static_assert(N_HG * HPB * D_ROPE * sizeof(bf16) <= SMEM_W_FP8_MG); + static_assert(N_HG * HPB * KV::D_ROPE * sizeof(bf16) <= SMEM_W_FP8_MG); static constexpr size_t SMEM_SCRATCH = 0; static constexpr size_t SMEM_MBAR_KV = 2 * sizeof(uint64_t); diff --git a/include/flashinfer/attention/sparse_mla_sm120/decode_dsv3_2_kernel.cuh b/include/flashinfer/attention/sparse_mla_sm120/decode_dsv3_2_kernel.cuh index f9ba123a55e..0ee0156e3ea 100644 --- a/include/flashinfer/attention/sparse_mla_sm120/decode_dsv3_2_kernel.cuh +++ b/include/flashinfer/attention/sparse_mla_sm120/decode_dsv3_2_kernel.cuh @@ -50,8 +50,9 @@ constexpr int DSV3_2_KV_BUF_COUNT = 2; constexpr int DSV3_2_ENTRIES_PER_WARP = DSV3_2_BI / DSV3_2_N_WARPS; // 8 constexpr int DSV3_2_QK_N_TILES = DSV3_2_ENTRIES_PER_WARP / 8; // 1 +template struct DecodeDsv3_2Smem { - using KV = KVCacheTraits; + using KV = KVCacheTraits; static constexpr int N_V_CHUNKS = KV::D_NOPE / KV::QUANT_TILE; static constexpr size_t SMEM_Q_ROPE = HPB * KV::D_ROPE * sizeof(bf16); @@ -126,7 +127,7 @@ __global__ void __launch_bounds__(DSV3_2_BLOCK_THREADS) sparse_mla_decode_dsv3_2 const int* __restrict__ topk_length_ptr, // [num_tokens] or null int num_tokens, int num_splits, int chunks_per_block, float sm_scale, size_t stride_kv_block) { using KV = KVCacheTraits; - static_assert(KV::D_QK == 576); + static_assert(KV::D_QK == 576 || (MT == ModelType::GLM53_NOPE && KV::D_QK == 512)); constexpr int D_NOPE = KV::D_NOPE; // 512 constexpr int D_ROPE_C = KV::D_ROPE; // 64 constexpr int D_QK = KV::D_QK; // 576 @@ -200,7 +201,7 @@ __global__ void __launch_bounds__(DSV3_2_BLOCK_THREADS) sparse_mla_decode_dsv3_2 // D_NOPE 512 + SCALE_BYTES_PER_TOKEN 16), so the QK / XV stages read // scales directly out of sm_kv_fp8. extern __shared__ __align__(16) char smem_raw[]; - auto sm = DecodeDsv3_2Smem::init(smem_raw); + auto sm = DecodeDsv3_2Smem::init(smem_raw); __shared__ bf16 sm_p_full[HPB][DSV3_2_BI]; // 2 KB static const int32_t* idx_base = indices + (size_t)t_idx * TOPK; @@ -249,8 +250,10 @@ __global__ void __launch_bounds__(DSV3_2_BLOCK_THREADS) sparse_mla_decode_dsv3_2 cp_async_bulk_g2s(kv_fp8_dst + (size_t)entry_idx * KV_SMEM_STRIDE, data_base, V2_BULK_NOPESC_BYTES, sm.mbar_full(buf)); // Bulk 2: RoPE (128 B) → sm_kv_rope slot. - cp_async_bulk_g2s(kv_rope_dst + (size_t)entry_idx * D_ROPE_C, data_base + KV_ROPE_OFFSET, - V2_BULK_ROPE_BYTES, sm.mbar_full(buf)); + if constexpr (D_ROPE_C > 0) { + cp_async_bulk_g2s(kv_rope_dst + (size_t)entry_idx * D_ROPE_C, data_base + KV_ROPE_OFFSET, + V2_BULK_ROPE_BYTES, sm.mbar_full(buf)); + } } }; diff --git a/include/flashinfer/attention/sparse_mla_sm120/model/kv_cache_traits.cuh b/include/flashinfer/attention/sparse_mla_sm120/model/kv_cache_traits.cuh index 3bdc6a69fad..c6e9e702b5e 100644 --- a/include/flashinfer/attention/sparse_mla_sm120/model/kv_cache_traits.cuh +++ b/include/flashinfer/attention/sparse_mla_sm120/model/kv_cache_traits.cuh @@ -92,6 +92,41 @@ struct KVCacheTraits : KVCacheTraits { static constexpr ScaleFormat SCALE_FORMAT = ScaleFormat::ARBITRARY_FP32; }; +template <> +struct KVCacheTraits { + // GLM-5.3-Flash is a native NoPE model. The absorbed query and latent KV + // dimensions are both 512; no positional-key lane exists. + static constexpr int D_NOPE = 512; + static constexpr int D_ROPE = 0; + static constexpr int D_QK = D_NOPE; + static constexpr int D_V = 512; + + static constexpr int QUANT_TILE = 128; + static constexpr int NUM_SCALES = D_NOPE / QUANT_TILE; + static constexpr ScaleFormat SCALE_FORMAT = ScaleFormat::ARBITRARY_FP32; + + // vLLM's fp8_ds_mla cache ABI remains 656 bytes/token. The first 528 + // bytes contain the 512 FP8 latent values plus four inline FP32 scales; + // the trailing 128 bytes are reserved padding and must never be treated as + // RoPE data by this specialization. + static constexpr bool SCALE_INLINE = true; + static constexpr int SCALE_BYTES_PER_TOKEN = NUM_SCALES * sizeof(float); + static constexpr int KV_GMEM_STRIDE = 656; + static constexpr int KV_SCALE_GMEM_OFFSET = D_NOPE; + static constexpr int KV_ROPE_GMEM_OFFSET = D_NOPE + SCALE_BYTES_PER_TOKEN; + static constexpr int KV_SMEM_STRIDE = D_NOPE + SCALE_BYTES_PER_TOKEN; + static constexpr int KV_SMEM_COPY_BYTES = KV_SMEM_STRIDE; + static constexpr bool SCALE_IN_KV_SMEM = true; + + static constexpr int Q_NOPE_STRIDE = D_NOPE + 16; + static constexpr int Q_NOPE_BF16_STRIDE = D_NOPE + 8; + static constexpr bool V_HAS_ROPE = false; + + __device__ static __forceinline__ uint8_t scale_to_ue8m0(float scale) { + return static_cast((__float_as_uint(scale) >> 23) & 0xFF); + } +}; + template <> struct KVCacheTraits { // Dimensions @@ -147,9 +182,8 @@ struct KVCacheTraits { static constexpr int HPB = 16; static constexpr int BI = 64; -// D_ROPE and D_V are shared across all supported models; the asserts below -// pin them to KVCacheTraits<...> so a new model with diverging values has -// to opt out explicitly. +// D_V is shared across all supported models. D_ROPE is the shared width for +// RoPE-bearing models; GLM53_NOPE explicitly specializes it to zero. static constexpr int D_ROPE = 64; static constexpr int D_V = 512; static_assert(KVCacheTraits::D_ROPE == D_ROPE); @@ -158,6 +192,8 @@ static_assert(KVCacheTraits::D_ROPE == D_ROPE); static_assert(KVCacheTraits::D_V == D_V); static_assert(KVCacheTraits::D_ROPE == D_ROPE); static_assert(KVCacheTraits::D_V == D_V); +static_assert(KVCacheTraits::D_ROPE == 0); +static_assert(KVCacheTraits::D_V == D_V); // Warp configuration static constexpr int N_MATH_WARPS = 8; diff --git a/include/flashinfer/attention/sparse_mla_sm120/model/model_type.h b/include/flashinfer/attention/sparse_mla_sm120/model/model_type.h index cdc6c24fea7..53af586c544 100644 --- a/include/flashinfer/attention/sparse_mla_sm120/model/model_type.h +++ b/include/flashinfer/attention/sparse_mla_sm120/model/model_type.h @@ -31,8 +31,10 @@ // ModelType determines KV cache layout, dimensions, and scale format. // DSV3_2: d_nope=512, power-of-2 FP32 scale inline, 656B/token // DSV4: d_nope=448, UE8M0 scale footer, 584B/token -// GLM_NSA: d_nope=512, arbitrary FP32 scale inline, 656B/token -enum class ModelType { DSV3_2, DSV4, GLM_NSA }; +// GLM_NSA: d_nope=512, d_rope=64, arbitrary FP32 scale inline, 656B/token +// GLM53_NOPE: d_nope=512, d_rope=0, arbitrary FP32 scale inline, +// 656B/token (the final 128 bytes are reserved cache padding) +enum class ModelType { DSV3_2, DSV4, GLM_NSA, GLM53_NOPE }; enum class ScaleFormat { POW2_FP32, UE8M0_BYTE, ARBITRARY_FP32 }; diff --git a/include/flashinfer/attention/sparse_mla_sm120/prefill_kernel.cuh b/include/flashinfer/attention/sparse_mla_sm120/prefill_kernel.cuh index 1d2b3440b88..81264da8ff6 100644 --- a/include/flashinfer/attention/sparse_mla_sm120/prefill_kernel.cuh +++ b/include/flashinfer/attention/sparse_mla_sm120/prefill_kernel.cuh @@ -53,11 +53,11 @@ // - No PDL (no dependent kernel) // // Template params (all constexpr): -// MT: ModelType (DSV3_2 / DSV4) +// MT: ModelType (DSV3_2 / DSV4 / GLM_NSA / GLM53_NOPE) // CM: ComputeMode (FP8 / BF16) for the QK MMA; XV is always FP8 // NUM_HEADS: 8, 16, 32, 64, 128 (NUM_HEADS < HPB=16 zero-pads + gates) -// TOPK: 128, 192, 256, 512, 1024, 2048 -// PAGE_BLOCK_SIZE: 64 (DSV3_2 and DSV4 both use the 64-token page layout) +// TOPK: 128, 192, 256, 512, 1024, 2048, 2176 +// PAGE_BLOCK_SIZE: 64 (all currently supported model types) // ============================================================================ struct PrefillColdParams { @@ -174,7 +174,7 @@ __global__ void __launch_bounds__(BLOCK_THREADS, 1) quantize_q_to_smem(sm.q_nope_fp8, sm.q_nope_sc, sm.q_rope, q_base, sm.reduce_buf, VALID_HPB); } - QRopeRegs q_rope_regs = preload_q_rope_regs(sm.q_rope, lane); + QRopeRegs q_rope_regs = preload_q_rope_regs(sm.q_rope, lane); for (int h = threadIdx.x; h < HPB; h += MATH_THREADS) sm.m_smem[h] = -1e30f; @@ -216,7 +216,7 @@ __global__ void __launch_bounds__(BLOCK_THREADS, 1) for (int i = threadIdx.x; i < CT::N_V_CHUNKS * HPB; i += MATH_THREADS) sm.w_head_sc_all[i] = 0.f; - KVRopePrefetch rope_pf = prefetch_kv_rope( + KVRopePrefetch rope_pf = prefetch_kv_rope( reinterpret_cast(entry_base[gid] + KV::KV_ROPE_GMEM_OFFSET), lane); // ── QK nope MMA ───────────────────── @@ -292,7 +292,7 @@ __global__ void __launch_bounds__(BLOCK_THREADS, 1) } // ── QK rope (BF16 MMA, uses prefetched B operands) ────── - compute_qk_rope(qk, q_rope_regs, rope_pf); + compute_qk_rope(qk, q_rope_regs, rope_pf); // ── Invalid index masking + topk_length overflow ───── { @@ -847,20 +847,20 @@ __device__ __forceinline__ void prefill_mg_impl( const bf16* q_base_g = Q + (size_t)s_i * NUM_HEADS * KV::D_QK + (size_t)(h_start + g * HPB) * KV::D_QK; if constexpr (CM == ComputeMode::BF16) { - load_q_bf16_to_smem(sm.q_nope_bf16(g), sm.q_rope() + g * HPB * D_ROPE, + load_q_bf16_to_smem(sm.q_nope_bf16(g), sm.q_rope() + g * HPB * KV::D_ROPE, q_base_g, VALID_HPB); } else { quantize_q_to_smem(sm.q_nope_fp8(g), sm.q_nope_sc(g), - sm.q_rope() + g * HPB * D_ROPE, q_base_g, + sm.q_rope() + g * HPB * KV::D_ROPE, q_base_g, sm.reduce_buf(), VALID_HPB); } } // Preload Q rope to registers for both groups - QRopeRegs q_rope_regs[MG_N_HG]; + QRopeRegs q_rope_regs[MG_N_HG]; #pragma unroll for (int g = 0; g < MG_N_HG; g++) - q_rope_regs[g] = preload_q_rope_regs(sm.q_rope() + g * HPB * D_ROPE, lane); + q_rope_regs[g] = preload_q_rope_regs(sm.q_rope() + g * HPB * KV::D_ROPE, lane); for (int i = threadIdx.x; i < MG_N_HG * HPB; i += MATH_THREADS) sm.m_smem()[i] = -1e30f; @@ -928,7 +928,7 @@ __device__ __forceinline__ void prefill_mg_impl( } } - KVRopePrefetch rope_pf = prefetch_kv_rope( + KVRopePrefetch rope_pf = prefetch_kv_rope( reinterpret_cast(entry_base_gid + KV::KV_ROPE_GMEM_OFFSET), lane); // Init per-group w_head_sc_all @@ -995,7 +995,7 @@ __device__ __forceinline__ void prefill_mg_impl( #pragma unroll for (int g = 0; g < 2; g++) { float* qk = qk_grp[g]; - compute_qk_rope(qk, q_rope_regs[g], rope_pf); + compute_qk_rope(qk, q_rope_regs[g], rope_pf); { int e0 = qk_nb + tid * 2, e1 = e0 + 1; @@ -1138,7 +1138,7 @@ __device__ __forceinline__ void prefill_mg_impl( } // QK rope (reuses prefetched B operands) - compute_qk_rope(qk, q_rope_regs[g], rope_pf); + compute_qk_rope(qk, q_rope_regs[g], rope_pf); // Invalid index masking + topk_length overflow. Dual splits per phase // (main: absolute ti*BI+e vs topk_len; extra: relative diff --git a/tests/attention/test_sparse_mla_sm120.py b/tests/attention/test_sparse_mla_sm120.py index cc55d9890fc..d760465d48b 100644 --- a/tests/attention/test_sparse_mla_sm120.py +++ b/tests/attention/test_sparse_mla_sm120.py @@ -179,6 +179,29 @@ def quantize_kv_glm_nsa(kv_bf16: torch.Tensor) -> torch.Tensor: return result.view(nb, bs, 1, bpt) +def quantize_kv_glm53_nope(kv_bf16: torch.Tensor) -> torch.Tensor: + """Pack native NoPE KV into the 656B ABI with arbitrary FP32 scales.""" + d_nope, tile_size, num_tiles = 512, 128, 4 + bpt = 656 + nb, bs, hk, d = kv_bf16.shape + assert d == d_nope and hk == 1 + nt = nb * bs + kv = kv_bf16.reshape(nt, d) + result = torch.zeros(nt, bpt, dtype=torch.uint8, device=kv.device) + + for ti in range(num_tiles): + tile = kv[:, ti * tile_size : (ti + 1) * tile_size].float() + scale = (tile.abs().amax(dim=-1).clamp(min=1e-4) / 448.0).to(torch.float32) + fp8 = (tile / scale.unsqueeze(-1)).clamp(-448, 448).to(torch.float8_e4m3fn) + result[:, ti * tile_size : (ti + 1) * tile_size] = fp8.view(torch.uint8) + result[:, d_nope + ti * 4 : d_nope + (ti + 1) * 4] = ( + scale.view(torch.float32).view(torch.uint8).view(nt, 4) + ) + + # Bytes 528:656 are reserved padding in the stable packed-cache ABI. + return result.view(nb, bs, 1, bpt) + + def _assert_has_non_pow2_inline_scales(packed: torch.Tensor) -> None: scales = packed.reshape(-1, 656)[:, 512:528].contiguous().view(torch.float32) log2_scales = scales.float().log2() @@ -962,6 +985,109 @@ def test_sparse_mla_sm120_prefill_glm_nsa_arbitrary_fp32(num_heads: int) -> None torch.testing.assert_close(out_lse, ref_lse, atol=5e-2, rtol=5e-2) +def test_sparse_mla_sm120_decode_glm53_nope() -> None: + torch.manual_seed(3) + device = torch.device("cuda") + d_qk = d_v = 512 + num_tokens, num_heads, topk = 4, 32, 2176 + page_block_size = 64 + num_blocks = 64 + s_kv = num_blocks * page_block_size + + kv_bf16 = ( + torch.randn( + num_blocks, page_block_size, 1, d_qk, device=device, dtype=torch.bfloat16 + ) + / 10.0 + ).clamp(-1, 1) + kv_packed = quantize_kv_glm53_nope(kv_bf16) + _assert_has_non_pow2_inline_scales(kv_packed) + kv_dequant = dequantize_kv_dsv3_2(kv_packed)[..., :d_qk] + + q = ( + torch.randn(num_tokens, num_heads, d_qk, device=device, dtype=torch.bfloat16) + / 10.0 + ).clamp(-1, 1) + indices = torch.randint( + 0, s_kv, (num_tokens, topk), device=device, dtype=torch.int32 + ) + indices[:, topk // 2 :] = -1 + sm_scale = d_qk**-0.5 + ref_out, ref_lse = _ref_sparse_attn(q, kv_dequant, indices, sm_scale, d_v) + + output = torch.zeros( + (num_tokens, num_heads, d_v), dtype=torch.bfloat16, device=device + ) + out_lse = torch.zeros((num_tokens, num_heads), dtype=torch.float32, device=device) + mid_out, mid_lse = _make_decode_scratch(num_tokens, num_heads, topk, d_v, device) + + sparse_mla_sm120_paged_attention( + q, + kv_packed, + indices, + output, + out_lse, + sm_scale, + d_v=d_v, + kv_scale_format="arbitrary_fp32", + mid_out=mid_out, + mid_lse=mid_lse, + ) + + torch.testing.assert_close(output, ref_out, atol=5e-2, rtol=5e-2) + torch.testing.assert_close(out_lse, ref_lse, atol=5e-2, rtol=5e-2) + + +def test_sparse_mla_sm120_prefill_glm53_nope() -> None: + torch.manual_seed(4) + device = torch.device("cuda") + d_qk = d_v = 512 + num_tokens, num_heads, topk = 65, 32, 2176 + page_block_size = 64 + num_blocks = 64 + s_kv = num_blocks * page_block_size + + kv_bf16 = ( + torch.randn( + num_blocks, page_block_size, 1, d_qk, device=device, dtype=torch.bfloat16 + ) + / 10.0 + ).clamp(-1, 1) + kv_packed = quantize_kv_glm53_nope(kv_bf16) + _assert_has_non_pow2_inline_scales(kv_packed) + kv_dequant = dequantize_kv_dsv3_2(kv_packed)[..., :d_qk] + + q = ( + torch.randn(num_tokens, num_heads, d_qk, device=device, dtype=torch.bfloat16) + / 10.0 + ).clamp(-1, 1) + indices = torch.randint( + 0, s_kv, (num_tokens, topk), device=device, dtype=torch.int32 + ) + indices[:, topk // 2 :] = -1 + sm_scale = d_qk**-0.5 + ref_out, ref_lse = _ref_sparse_attn(q, kv_dequant, indices, sm_scale, d_v) + + output = torch.zeros( + (num_tokens, num_heads, d_v), dtype=torch.bfloat16, device=device + ) + out_lse = torch.zeros((num_tokens, num_heads), dtype=torch.float32, device=device) + + sparse_mla_sm120_paged_attention( + q, + kv_packed, + indices, + output, + out_lse, + sm_scale, + d_v=d_v, + kv_scale_format="arbitrary_fp32", + ) + + torch.testing.assert_close(output, ref_out, atol=5e-2, rtol=5e-2) + torch.testing.assert_close(out_lse, ref_lse, atol=5e-2, rtol=5e-2) + + _DSV4_PREFILL_CONFIGS = [ (8, 128), (8, 192),