minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[Refactor] Rename NSA → DSA: user-facing aliases, file/class/import rename (#25821)

Browse Source 
Co-authored-by: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Cheng Wan
2026-05-20 00:18:04 -07:00
committed by GitHub
co-authored by Claude Sonnet 4.6
parent da6d549ab2
commit 8131641bc6
162 changed files with 11298 additions and 10740 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.claude/skills/llm-torch-profiler-analysis/references/fuse-overlap-catalog.md
+11 -11
View File
@@ -33,7 +33,7 @@ rechecked recent merged and open optimization PRs through the GitHub CLI/API.
The vLLM torch.compile pass inventory is now split out in The vLLM torch.compile pass inventory is now split out in
[`vllm-torch-compile-fusions.md`](vllm-torch-compile-fusions.md). Stable [`vllm-torch-compile-fusions.md`](vllm-torch-compile-fusions.md). Stable
current-code families remain folded into the mainline rows below. New current-code families remain folded into the mainline rows below. New
status-sensitive rows were added for DeepSeek-V4, GLM5 NSA / PDL, NVFP4 MoE, status-sensitive rows were added for DeepSeek-V4, GLM5 DSA / PDL, NVFP4 MoE,
torch.compile decode, vLLM DSV4, vLLM ROCm WMMA, and vLLM GPU/CPU sync-removal torch.compile decode, vLLM DSV4, vLLM ROCm WMMA, and vLLM GPU/CPU sync-removal
work. Recheck PR state before treating an in-flight row as shipped. work. Recheck PR state before treating an in-flight row as shipped.
@@ -53,13 +53,13 @@ work. Recheck PR state before treating an in-flight row as shipped.
| Fused QK RoPE reshape + KV cache write | `fused_qk_rope_reshape_and_cache*`<br>RoPE followed by reshape / cache DtoD | `python/sglang/srt/layers/attention/utils.py::fused_qk_rope_reshape_and_cache` | One Triton kernel applies RoPE to Q / K, reshapes cache layout, and writes K / V directly to paged cache | Treat separate RoPE + reshape + cache-write ladders as an existing attention-prep fusion family. | | Fused QK RoPE reshape + KV cache write | `fused_qk_rope_reshape_and_cache*`<br>RoPE followed by reshape / cache DtoD | `python/sglang/srt/layers/attention/utils.py::fused_qk_rope_reshape_and_cache` | One Triton kernel applies RoPE to Q / K, reshapes cache layout, and writes K / V directly to paged cache | Treat separate RoPE + reshape + cache-write ladders as an existing attention-prep fusion family. |
| Fused RoPE + KV cache store | `fused_set_kv_buffer`<br>RoPE followed by KV-store, DtoD, or cache-write kernels | `python/sglang/jit_kernel/rope.py`<br>`python/sglang/srt/models/utils.py::enable_fused_set_kv_buffer` | Shared entrypoints can route to fused RoPE + KV-store or model-side `fused_set_kv_buffer` fast paths | Compare against the fused cache-store path before proposing a new KV rewrite. | | Fused RoPE + KV cache store | `fused_set_kv_buffer`<br>RoPE followed by KV-store, DtoD, or cache-write kernels | `python/sglang/jit_kernel/rope.py`<br>`python/sglang/srt/models/utils.py::enable_fused_set_kv_buffer` | Shared entrypoints can route to fused RoPE + KV-store or model-side `fused_set_kv_buffer` fast paths | Compare against the fused cache-store path before proposing a new KV rewrite. |
| Fused decode metadata setup | `normal_decode_set_metadata`<br>`cache_seqlens_int32`<br>`cu_seqlens_k`<br>`page_table`<br>`swa_page_table` | `python/sglang/srt/layers/attention/flashattention_backend.py::normal_decode_set_metadata` | Triton decode path fuses seq-len cast/add, prefix-sum, req-to-token gather, page-table divide, and optional SWA metadata build into 1-2 kernels | If decode exposes multiple tiny metadata kernels before attention, first compare against this existing fused metadata-prep path. | | Fused decode metadata setup | `normal_decode_set_metadata`<br>`cache_seqlens_int32`<br>`cu_seqlens_k`<br>`page_table`<br>`swa_page_table` | `python/sglang/srt/layers/attention/flashattention_backend.py::normal_decode_set_metadata` | Triton decode path fuses seq-len cast/add, prefix-sum, req-to-token gather, page-table divide, and optional SWA metadata build into 1-2 kernels | If decode exposes multiple tiny metadata kernels before attention, first compare against this existing fused metadata-prep path. |
| NSA fused metadata copy for graph replay | `fused_metadata_copy`<br>`fused_metadata_copy_multi`<br>`fused_nsa_cache_seqlens`<br>`fused_flashmla_metadata` | `python/sglang/jit_kernel/fused_metadata_copy.py` | CUDA graph replay path fuses multiple metadata copies into one kernel or one multi-destination kernel | Treat bursts of tiny metadata-copy kernels around NSA replay as a missed existing replay fusion. | | DSA fused metadata copy for graph replay | `fused_metadata_copy`<br>`fused_metadata_copy_multi`<br>`fused_dsa_cache_seqlens`<br>`fused_flashmla_metadata` | `python/sglang/jit_kernel/fused_metadata_copy.py` | CUDA graph replay path fuses multiple metadata copies into one kernel or one multi-destination kernel | Treat bursts of tiny metadata-copy kernels around DSA replay as a missed existing replay fusion. |
| DeepSeek MLA fused projection + norm + RoPE | `qkv_proj_with_rope_fused_weight`<br>`fused_qkv_a_proj_with_mqa`<br>`forward_absorb_fused_mla_rope*` | `python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_fused_rope_cpu.py`<br>`python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_fused_rope_rocm.py`<br>`python/sglang/srt/models/deepseek_v2.py` | CPU / ROCm paths fuse DeepSeek MLA projection packing with q / k norm, RoPE, and cache-oriented MLA prep | For DeepSeek MLA, split proj / norm / rope prep is usually an existing backend-specific fuse that did not fire. | | DeepSeek MLA fused projection + norm + RoPE | `qkv_proj_with_rope_fused_weight`<br>`fused_qkv_a_proj_with_mqa`<br>`forward_absorb_fused_mla_rope*` | `python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_fused_rope_cpu.py`<br>`python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_fused_rope_rocm.py`<br>`python/sglang/srt/models/deepseek_v2.py` | CPU / ROCm paths fuse DeepSeek MLA projection packing with q / k norm, RoPE, and cache-oriented MLA prep | For DeepSeek MLA, split proj / norm / rope prep is usually an existing backend-specific fuse that did not fire. |
| Fused QK RoPE concat + MLA cache write | `fused_qk_rope_cat_and_cache_mla`<br>`set_mla_kv_buffer` | `python/sglang/srt/layers/rocm_linear_utils.py`<br>`python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py` | ROCm MLA path can fuse Q / K RoPE packing, concat, and MLA cache write in one backend-specific op | On DeepSeek / MLA traces, separate RoPE-cat-cache steps are not automatically novel. | | Fused QK RoPE concat + MLA cache write | `fused_qk_rope_cat_and_cache_mla`<br>`set_mla_kv_buffer` | `python/sglang/srt/layers/rocm_linear_utils.py`<br>`python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py` | ROCm MLA path can fuse Q / K RoPE packing, concat, and MLA cache write in one backend-specific op | On DeepSeek / MLA traces, separate RoPE-cat-cache steps are not automatically novel. |
| Qwen3 decode fused QK norm + 3D mRoPE + KV cache write | `fused_qk_norm_mrope_3d_cache_pts_quant_shuffle`<br>`mrope`<br>decode cache write | `python/sglang/srt/models/qwen3.py` | ROCm / AITER decode path fuses QK norm, 3D mRoPE, and paged KV cache write | On Qwen3-style decode, separate norm + mRoPE + cache-store kernels are not a novel opportunity. | | Qwen3 decode fused QK norm + 3D mRoPE + KV cache write | `fused_qk_norm_mrope_3d_cache_pts_quant_shuffle`<br>`mrope`<br>decode cache write | `python/sglang/srt/models/qwen3.py` | ROCm / AITER decode path fuses QK norm, 3D mRoPE, and paged KV cache write | On Qwen3-style decode, separate norm + mRoPE + cache-store kernels are not a novel opportunity. |
| NPU fused split-QKV + RMSNorm + RoPE | `split_qkv_rmsnorm_rope` | `python/sglang/srt/models/llama.py`<br>`python/sglang/srt/models/qwen3.py`<br>`python/sglang/srt/models/qwen3_moe.py`<br>`python/sglang/srt/models/glm4_moe.py` | Ascend path fuses QKV split, Q / K RMSNorm, and RoPE in one op | On NPU traces, separate split / norm / rope kernels usually mean the fused path is unavailable or bypassed. | | NPU fused split-QKV + RMSNorm + RoPE | `split_qkv_rmsnorm_rope` | `python/sglang/srt/models/llama.py`<br>`python/sglang/srt/models/qwen3.py`<br>`python/sglang/srt/models/qwen3_moe.py`<br>`python/sglang/srt/models/glm4_moe.py` | Ascend path fuses QKV split, Q / K RMSNorm, and RoPE in one op | On NPU traces, separate split / norm / rope kernels usually mean the fused path is unavailable or bypassed. |
| Fused FP8 quantize + paged KV cache write | `trtllm_fp8_kv_kernel`<br>`fp8 kv cache write`<br>`paged KV cache write` | `python/sglang/srt/layers/attention/triton_ops/trtllm_fp8_kv_kernel.py` | TRTLLM MHA path fuses FP8 quantization, scale computation, and paged K / V cache write | If FP8 KV cache traces show standalone quant plus write kernels, first compare against this existing Triton fuse. | | Fused FP8 quantize + paged KV cache write | `trtllm_fp8_kv_kernel`<br>`fp8 kv cache write`<br>`paged KV cache write` | `python/sglang/srt/layers/attention/triton_ops/trtllm_fp8_kv_kernel.py` | TRTLLM MHA path fuses FP8 quantization, scale computation, and paged K / V cache write | If FP8 KV cache traces show standalone quant plus write kernels, first compare against this existing Triton fuse. |
| Fused MLA KV cache write + FP8 quant | `set_mla_kv_buffer_fp8_quant*`<br>`set_mla_kv_buffer_triton_fp8_quant` | `python/sglang/srt/mem_cache/utils.py`<br>`python/sglang/srt/mem_cache/memory_pool.py` | MLA / NSA KV pool path can quantize K and write directly into KV storage without a separate concat-and-quant chain | Treat standalone quant + KV-buffer write on MLA paths as missing existing fusion first. | | Fused MLA KV cache write + FP8 quant | `set_mla_kv_buffer_fp8_quant*`<br>`set_mla_kv_buffer_triton_fp8_quant` | `python/sglang/srt/mem_cache/utils.py`<br>`python/sglang/srt/mem_cache/memory_pool.py` | MLA / DSA KV pool path can quantize K and write directly into KV storage without a separate concat-and-quant chain | Treat standalone quant + KV-buffer write on MLA paths as missing existing fusion first. |
| Fused MoE router / top-k / softcapping | `FusedMoeRouter`<br>`fused_moe_router*`<br>router GEMM + `topk` + `tanh` | `python/sglang/srt/layers/moe/router.py` | Single fused router kernel covers router matmul, softcapping, and top-k selection | Treat exposed router matmul + softcap + top-k chains as an existing MoE fusion family. | | Fused MoE router / top-k / softcapping | `FusedMoeRouter`<br>`fused_moe_router*`<br>router GEMM + `topk` + `tanh` | `python/sglang/srt/layers/moe/router.py` | Single fused router kernel covers router matmul, softcapping, and top-k selection | Treat exposed router matmul + softcap + top-k chains as an existing MoE fusion family. |
| Fused MoE grouped-topk / gate kernels | `fused_topk_deepseek`<br>`moe_fused_gate`<br>`aiter_fused_topk`<br>`kimi_k2_moe_fused_gate` | `python/sglang/srt/layers/moe/topk.py` | CUDA / ROCm / FlashInfer kernels fuse bias, grouped-topk, renorm, and routed scaling into one gate op | Check backend / model eligibility before proposing a novel router-gate fusion. | | Fused MoE grouped-topk / gate kernels | `fused_topk_deepseek`<br>`moe_fused_gate`<br>`aiter_fused_topk`<br>`kimi_k2_moe_fused_gate` | `python/sglang/srt/layers/moe/topk.py` | CUDA / ROCm / FlashInfer kernels fuse bias, grouped-topk, renorm, and routed scaling into one gate op | Check backend / model eligibility before proposing a novel router-gate fusion. |
| Qwen-style shared-expert append into routed top-k output | `_append_shared_to_topk_output`<br>`fused_append_shared_experts_with_weights`<br>`num_fused_shared_experts` | `python/sglang/srt/models/qwen2_moe.py`<br>`python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe_triton_kernels.py` | Qwen-style MoE paths can append shared-expert ids and sigmoid gate weights to routed top-k output in one Triton kernel so the shared experts execute inside the fused MoE path | Treat routed top-k plus shared-expert pad / concat ladders as an existing MoE-prep fusion family first. | | Qwen-style shared-expert append into routed top-k output | `_append_shared_to_topk_output`<br>`fused_append_shared_experts_with_weights`<br>`num_fused_shared_experts` | `python/sglang/srt/models/qwen2_moe.py`<br>`python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe_triton_kernels.py` | Qwen-style MoE paths can append shared-expert ids and sigmoid gate weights to routed top-k output in one Triton kernel so the shared experts execute inside the fused MoE path | Treat routed top-k plus shared-expert pad / concat ladders as an existing MoE-prep fusion family first. |
@@ -67,8 +67,8 @@ work. Recheck PR state before treating an in-flight row as shipped.
| Fused MoE sum + all-reduce | routed MoE followed by explicit sum-reduce kernels | `python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py`<br>`python/sglang/srt/layers/moe/fused_moe_triton/fused_moe_triton_kernels.py` | `fuse_sum_all_reduce=True` path in the second MoE GEMM | Before inventing a new MoE reduction fuse, check whether `enable_fused_moe_sum_all_reduce` is simply off or the quant path is incompatible. | | Fused MoE sum + all-reduce | routed MoE followed by explicit sum-reduce kernels | `python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py`<br>`python/sglang/srt/layers/moe/fused_moe_triton/fused_moe_triton_kernels.py` | `fuse_sum_all_reduce=True` path in the second MoE GEMM | Before inventing a new MoE reduction fuse, check whether `enable_fused_moe_sum_all_reduce` is simply off or the quant path is incompatible. |
| Fused MoE activation + quant / re-quant | `silu_and_mul_*quant*`<br>`npu_dequant_swiglu_quant`<br>`swiglu_quant` | `python/sglang/srt/layers/moe/ep_moe/kernels.py`<br>`python/sglang/jit_kernel/nvfp4.py`<br>`python/sglang/srt/layers/moe/cutlass_w4a8_moe.py`<br>`python/sglang/srt/hardware_backend/npu/quantization/fused_moe_method_npu.py` | Quantized MoE backends fuse SwiGLU / SiLU-and-mul with FP8 / FP4 / NPU re-quant before the second expert GEMM | If MoE traces show standalone activation then quant kernels, first check whether the quantized fused path is missing. | | Fused MoE activation + quant / re-quant | `silu_and_mul_*quant*`<br>`npu_dequant_swiglu_quant`<br>`swiglu_quant` | `python/sglang/srt/layers/moe/ep_moe/kernels.py`<br>`python/sglang/jit_kernel/nvfp4.py`<br>`python/sglang/srt/layers/moe/cutlass_w4a8_moe.py`<br>`python/sglang/srt/hardware_backend/npu/quantization/fused_moe_method_npu.py` | Quantized MoE backends fuse SwiGLU / SiLU-and-mul with FP8 / FP4 / NPU re-quant before the second expert GEMM | If MoE traces show standalone activation then quant kernels, first check whether the quantized fused path is missing. |
| DeepSeek comm-prep fused RMSNorm + quant / flatten-quant | `fused_rms_fp8_group_quant`<br>`fused_rms_mxfp4_quant`<br>`fused_flatten_fp8_group_quant`<br>`fused_flatten_mxfp4_quant` | `python/sglang/srt/layers/communicator.py`<br>`python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py`<br>`python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py` | DeepSeek MLA / MHA ROCm paths fuse RMSNorm or flatten with FP8 / MXFP4 quantization for comm / attention prep | On DeepSeek quant traces, split norm + quant or flatten + quant is an existing family, not a new idea. | | DeepSeek comm-prep fused RMSNorm + quant / flatten-quant | `fused_rms_fp8_group_quant`<br>`fused_rms_mxfp4_quant`<br>`fused_flatten_fp8_group_quant`<br>`fused_flatten_mxfp4_quant` | `python/sglang/srt/layers/communicator.py`<br>`python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py`<br>`python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py` | DeepSeek MLA / MHA ROCm paths fuse RMSNorm or flatten with FP8 / MXFP4 quantization for comm / attention prep | On DeepSeek quant traces, split norm + quant or flatten + quant is an existing family, not a new idea. |
| NSA fused top-k transform / page-table build | `fast_topk_transform_fused`<br>`fast_topk_transform_ragged_fused` | `python/sglang/srt/layers/attention/nsa_backend.py` | NSA can fuse top-k selection with paged / ragged index transform instead of separate top-k plus metadata scatter | If NSA top-k metadata work is split, check `SGLANG_NSA_FUSE_TOPK` and backend support first. | | DSA fused top-k transform / page-table build | `fast_topk_transform_fused`<br>`fast_topk_transform_ragged_fused` | `python/sglang/srt/layers/attention/dsa_backend.py` | DSA can fuse top-k selection with paged / ragged index transform instead of separate top-k plus metadata scatter | If DSA top-k metadata work is split, check `SGLANG_DSA_FUSE_TOPK` and backend support first. |
| NSA fused quantize + indexed K-cache store | `fused_store_index_k_cache`<br>`act_quant`<br>`index_k_with_scale_buffer` | `python/sglang/jit_kernel/fused_store_index_cache.py`<br>`python/sglang/srt/layers/attention/nsa/nsa_indexer.py` | Single JIT kernel quantizes bf16 K to fp8 + scale and writes directly into NSA index cache | Treat split `act_quant` + buffer-store on CUDA as missing an existing fused store path. | | DSA fused quantize + indexed K-cache store | `fused_store_index_k_cache`<br>`act_quant`<br>`index_k_with_scale_buffer` | `python/sglang/jit_kernel/fused_store_index_cache.py`<br>`python/sglang/srt/layers/attention/dsa/dsa_indexer.py` | Single JIT kernel quantizes bf16 K to fp8 + scale and writes directly into DSA index cache | Treat split `act_quant` + buffer-store on CUDA as missing an existing fused store path. |
| Fused sampling temperature + softmax | `fused_temperature_softmax*` | `python/sglang/srt/layers/fused_sampling.py`<br>`python/sglang/srt/layers/sampler.py` | Triton single-pass / multi-pass kernels fuse temperature scaling and softmax during decode | Separate temp-divide + softmax at decode batch sizes is often a missed existing fusion. | | Fused sampling temperature + softmax | `fused_temperature_softmax*` | `python/sglang/srt/layers/fused_sampling.py`<br>`python/sglang/srt/layers/sampler.py` | Triton single-pass / multi-pass kernels fuse temperature scaling and softmax during decode | Separate temp-divide + softmax at decode batch sizes is often a missed existing fusion. |
| Fused logit softcap | `fused_softcap`<br>`final_logit_softcapping` | `python/sglang/srt/layers/elementwise.py`<br>`python/sglang/srt/layers/logits_processor.py` | Triton kernels fuse cast-to-float and softcap / tanh math for logits or generic elementwise softcapping | Treat exposed cast + softcap ladders as an existing Triton fuse family. | | Fused logit softcap | `fused_softcap`<br>`final_logit_softcapping` | `python/sglang/srt/layers/elementwise.py`<br>`python/sglang/srt/layers/logits_processor.py` | Triton kernels fuse cast-to-float and softcap / tanh math for logits or generic elementwise softcapping | Treat exposed cast + softcap ladders as an existing Triton fuse family. |
| Linear-attention packed projection reshuffle | `fused_qkvzba_split_reshape_cat*`<br>`qkvz_proj`<br>`ba_proj`<br>`qkvabz_proj`<br>`fused_qkvbfg_a_proj` | `python/sglang/jit_kernel/triton/gdn_fused_proj.py`<br>`python/sglang/srt/models/qwen3_next.py`<br>`python/sglang/srt/models/qwen3_5.py`<br>`python/sglang/srt/models/kimi_linear.py`<br>`python/sglang/srt/models/jet_nemotron.py` | GDN / Kimi / Jet-style linear-attn models pack multiple projections, then fuse split / reshape / cat into one kernel | Treat split reshape / transpose / cat ladders as an existing linear-attention fusion family. | | Linear-attention packed projection reshuffle | `fused_qkvzba_split_reshape_cat*`<br>`qkvz_proj`<br>`ba_proj`<br>`qkvabz_proj`<br>`fused_qkvbfg_a_proj` | `python/sglang/jit_kernel/triton/gdn_fused_proj.py`<br>`python/sglang/srt/models/qwen3_next.py`<br>`python/sglang/srt/models/qwen3_5.py`<br>`python/sglang/srt/models/kimi_linear.py`<br>`python/sglang/srt/models/jet_nemotron.py` | GDN / Kimi / Jet-style linear-attn models pack multiple projections, then fuse split / reshape / cat into one kernel | Treat split reshape / transpose / cat ladders as an existing linear-attention fusion family. |
@@ -90,7 +90,7 @@ work. Recheck PR state before treating an in-flight row as shipped.
| Llama4 shared branch vs routed branch overlap | shared expert branch plus routed MoE branch as adjacent windows | `python/sglang/srt/models/llama4.py` | shared expert on current stream, router + topk + routed experts on `alt_stream` | Use Llama4 as the first precedent for branch-level overlap in similar sparse models. | | Llama4 shared branch vs routed branch overlap | shared expert branch plus routed MoE branch as adjacent windows | `python/sglang/srt/models/llama4.py` | shared expert on current stream, router + topk + routed experts on `alt_stream` | Use Llama4 as the first precedent for branch-level overlap in similar sparse models. |
| ExaoneMoE shared experts vs router experts overlap | shared expert output and router-expert output form a two-branch window | `python/sglang/srt/models/exaone_moe.py::forward_normal_dual_stream` | shared experts on current stream, router + routed experts on `alt_stream`, explicit join before combine | This is an existing dual-stream MoE overlap family. | | ExaoneMoE shared experts vs router experts overlap | shared expert output and router-expert output form a two-branch window | `python/sglang/srt/models/exaone_moe.py::forward_normal_dual_stream` | shared experts on current stream, router + routed experts on `alt_stream`, explicit join before combine | This is an existing dual-stream MoE overlap family. |
| Grok residual-MoE branch overlap | dense MLP and block-sparse MoE branches in parallel | `python/sglang/srt/models/grok.py::moe_with_rmoe` | dense MLP on current stream, MoE on `alt_stream`, fused dual residual RMSNorm around boundaries | Treat exposed Grok branch overlap as an existing pattern. | | Grok residual-MoE branch overlap | dense MLP and block-sparse MoE branches in parallel | `python/sglang/srt/models/grok.py::moe_with_rmoe` | dense MLP on current stream, MoE on `alt_stream`, fused dual residual RMSNorm around boundaries | Treat exposed Grok branch overlap as an existing pattern. |
| NSA dual-stream overlap | Q-proj, K-proj, RoPE, cache-store, quantization in tight two-stream windows | `python/sglang/srt/layers/attention/nsa/nsa_indexer.py` | Q / K projection split, RoPE split, cache-store vs quantization overlap | NSA already contains several dual-stream overlap precedents. | | DSA dual-stream overlap | Q-proj, K-proj, RoPE, cache-store, quantization in tight two-stream windows | `python/sglang/srt/layers/attention/dsa/dsa_indexer.py` | Q / K projection split, RoPE split, cache-store vs quantization overlap | DSA already contains several dual-stream overlap precedents. |
| MoriEP async dispatch / combine comm stream | `MoriEP`<br>`_comm_stream`<br>`dispatch`<br>`combine`<br>`done_event` | `python/sglang/srt/layers/moe/token_dispatcher/moriep.py` | MoriEP can submit dispatch and combine onto a dedicated communication stream and synchronize only through events | Treat MoriEP comm / compute interleave as an existing MoE overlap family. | | MoriEP async dispatch / combine comm stream | `MoriEP`<br>`_comm_stream`<br>`dispatch`<br>`combine`<br>`done_event` | `python/sglang/srt/layers/moe/token_dispatcher/moriep.py` | MoriEP can submit dispatch and combine onto a dedicated communication stream and synchronize only through events | Treat MoriEP comm / compute interleave as an existing MoE overlap family. |
| Heterogeneous-TP staging scatter overlap | `scatter_stream`<br>`_scatter_stream`<br>`staging` | `python/sglang/srt/disaggregation/common/staging_handler.py`<br>`python/sglang/srt/disaggregation/common/staging_buffer.py` | decode-side staging scatter kernels can run on a dedicated stream while forward continues on the main stream | If decode traces show staging scatter kernels adjacent to forward kernels, classify them against this existing overlap family first. | | Heterogeneous-TP staging scatter overlap | `scatter_stream`<br>`_scatter_stream`<br>`staging` | `python/sglang/srt/disaggregation/common/staging_handler.py`<br>`python/sglang/srt/disaggregation/common/staging_buffer.py` | decode-side staging scatter kernels can run on a dedicated stream while forward continues on the main stream | If decode traces show staging scatter kernels adjacent to forward kernels, classify them against this existing overlap family first. |
| Generic `alt_stream` overlap families | `alt_stream` plus explicit `wait_stream` / `with torch.cuda.stream(...)` | `qwen2_moe.py`<br>`qwen3_moe.py`<br>`glm4_moe.py`<br>`bailing_moe.py`<br>`llada2.py`<br>`grok.py`<br>`olmo2.py`<br>`step3p5.py`<br>`longcat_flash.py`<br>`falcon_h1.py` | model-specific overlap on attention prep, MoE branches, or cache-store | Search these families before designing a new overlap scheme from scratch. | | Generic `alt_stream` overlap families | `alt_stream` plus explicit `wait_stream` / `with torch.cuda.stream(...)` | `qwen2_moe.py`<br>`qwen3_moe.py`<br>`glm4_moe.py`<br>`bailing_moe.py`<br>`llada2.py`<br>`grok.py`<br>`olmo2.py`<br>`step3p5.py`<br>`longcat_flash.py`<br>`falcon_h1.py` | model-specific overlap on attention prep, MoE branches, or cache-store | Search these families before designing a new overlap scheme from scratch. |
@@ -134,16 +134,16 @@ Stable entries should be folded into the mainline family rows above.
| Pattern | Trace keywords | Primary code | Existing path | Skill should conclude | | Pattern | Trace keywords | Primary code | Existing path | Skill should conclude |
| --- | --- | --- | --- | --- | | --- | --- | --- | --- | --- |
| PR `#21877` fused grouped down-GEMM + combine | `grouped_gemm_nt_masked`<br>`combine`<br>`fused grouped gemm combine` | `PR #21877`<br>`python/sglang/srt/layers/moe/ep_moe/flashinfer_cutedsl_moe.py`<br>`python/sglang/srt/layers/moe/token_dispatcher/deepep.py` | FlashInfer CuTeDSL kernel fuses the second expert GEMM with DeepEP low-latency combine | Treat this as a concrete upstream MoE fuse / overlap family, not a new thought experiment. | | PR `#21877` fused grouped down-GEMM + combine | `grouped_gemm_nt_masked`<br>`combine`<br>`fused grouped gemm combine` | `PR #21877`<br>`python/sglang/srt/layers/moe/ep_moe/flashinfer_cutedsl_moe.py`<br>`python/sglang/srt/layers/moe/token_dispatcher/deepep.py` | FlashInfer CuTeDSL kernel fuses the second expert GEMM with DeepEP low-latency combine | Treat this as a concrete upstream MoE fuse / overlap family, not a new thought experiment. |
| PR `#21889` fused BF16 to FP4 quant + paged KV write | `set_mla_kv_buffer_fp4_quant_kernel`<br>`fp4 kv cache` | `PR #21889`<br>`python/sglang/srt/mem_cache/utils.py` | Triton kernel writes FP4 NSA KV pages directly while quantizing BF16 input | If NSA FP4 KV paths are split into quant plus store, classify them as an in-flight upstream fuse family. | | PR `#21889` fused BF16 to FP4 quant + paged KV write | `set_mla_kv_buffer_fp4_quant_kernel`<br>`fp4 kv cache` | `PR #21889`<br>`python/sglang/srt/mem_cache/utils.py` | Triton kernel writes FP4 DSA KV pages directly while quantizing BF16 input | If DSA FP4 KV paths are split into quant plus store, classify them as an in-flight upstream fuse family. |
| PR `#21889` fused FP4 paged dequant to FP8 + page-table remap | `_dequant_fp4_to_fp8_paged_kernel`<br>`WRITE_PT`<br>`dequant_fp4_paged_decode` | `PR #21889`<br>`python/sglang/srt/layers/attention/nsa/dequant_fp4_to_fp8.py` | Triton kernel reads FP4 pages, writes FP8 directly, and can fuse decode-side page-table remap | Treat this as an upstream in-flight decode-prep fusion family. | | PR `#21889` fused FP4 paged dequant to FP8 + page-table remap | `_dequant_fp4_to_fp8_paged_kernel`<br>`WRITE_PT`<br>`dequant_fp4_paged_decode` | `PR #21889`<br>`python/sglang/srt/layers/attention/dsa/dequant_fp4_to_fp8.py` | Triton kernel reads FP4 pages, writes FP8 directly, and can fuse decode-side page-table remap | Treat this as an upstream in-flight decode-prep fusion family. |
| PR `#21491` FlashInfer TRTLLM FP8 MoE with fused shared experts | `num_fused_shared_experts`<br>`trtllm_fp8_block_scale_moe` | `PR #21491`<br>`python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py`<br>`python/sglang/srt/models/deepseek_v2.py` | FlashInfer TRTLLM FP8 MoE path can fuse shared experts inside the routed MoE kernel | On FP8 TRTLLM MoE discussions, treat fused shared experts as an upstream pattern that already has a concrete PR. | | PR `#21491` FlashInfer TRTLLM FP8 MoE with fused shared experts | `num_fused_shared_experts`<br>`trtllm_fp8_block_scale_moe` | `PR #21491`<br>`python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py`<br>`python/sglang/srt/models/deepseek_v2.py` | FlashInfer TRTLLM FP8 MoE path can fuse shared experts inside the routed MoE kernel | On FP8 TRTLLM MoE discussions, treat fused shared experts as an upstream pattern that already has a concrete PR. |
| PR `#22005` fused add + RMSNorm + per-token FP8 quant | `fused_add_rmsnorm_per_token_quant`<br>`per_token_quant_fp8` | `PR #22005`<br>`python/sglang/jit_kernel/csrc/elementwise/fused_add_rmsnorm_per_token_quant.cuh`<br>`python/sglang/jit_kernel/fused_add_rmsnorm_per_token_quant.py` | CUDA JIT kernel keeps normed values in registers and emits BF16 + FP8 outputs plus per-token scales | If FP8 online-quant traces show add+norm followed by per-token quant, treat this as an in-flight upstream CUDA fuse family. | | PR `#22005` fused add + RMSNorm + per-token FP8 quant | `fused_add_rmsnorm_per_token_quant`<br>`per_token_quant_fp8` | `PR #22005`<br>`python/sglang/jit_kernel/csrc/elementwise/fused_add_rmsnorm_per_token_quant.cuh`<br>`python/sglang/jit_kernel/fused_add_rmsnorm_per_token_quant.py` | CUDA JIT kernel keeps normed values in registers and emits BF16 + FP8 outputs plus per-token scales | If FP8 online-quant traces show add+norm followed by per-token quant, treat this as an in-flight upstream CUDA fuse family. |
| PR `#20667` Qwen3.5 fused QK norm + RoPE + KV cache write | `fused_qk_norm_rope_cache_pts_quant_shuffle`<br>`fused_qk_norm_mrope_3d_cache_pts_quant_shuffle`<br>`rotary_dim` | `PR #20667`<br>`python/sglang/srt/models/qwen3_5.py`<br>`python/sglang/srt/models/utils.py` | ROCm / AITER path fuses Q / K RMSNorm, partial or 3D RoPE, and direct KV cache write for Qwen3.5 attention | Treat split QK-norm + RoPE + cache-store on Qwen3.5 as a concrete in-flight upstream family, not a novel idea. | | PR `#20667` Qwen3.5 fused QK norm + RoPE + KV cache write | `fused_qk_norm_rope_cache_pts_quant_shuffle`<br>`fused_qk_norm_mrope_3d_cache_pts_quant_shuffle`<br>`rotary_dim` | `PR #20667`<br>`python/sglang/srt/models/qwen3_5.py`<br>`python/sglang/srt/models/utils.py` | ROCm / AITER path fuses Q / K RMSNorm, partial or 3D RoPE, and direct KV cache write for Qwen3.5 attention | Treat split QK-norm + RoPE + cache-store on Qwen3.5 as a concrete in-flight upstream family, not a novel idea. |
| PR `#22392` CUTLASS FP8 GEMM replacing nvjet | `cutlass_scaled_mm`<br>`fp8_scaled_mm`<br>`nvjet`<br>`cudaMemsetAsync` | `PR #22392`<br>`sgl-kernel/python/sgl_kernel/gemm.py`<br>`python/sglang/srt/layers/quantization/fp8_utils.py` | Runtime replacement swaps nvjet FP8 GEMMs for CUTLASS kernels, removing per-launch memset bubbles and extra output-copy kernels | Treat nvjet GEMM + memset bubble ladders as an in-flight SGLang linear-kernel family before calling them novel. | | PR `#22392` CUTLASS FP8 GEMM replacing nvjet | `cutlass_scaled_mm`<br>`fp8_scaled_mm`<br>`nvjet`<br>`cudaMemsetAsync` | `PR #22392`<br>`sgl-kernel/python/sgl_kernel/gemm.py`<br>`python/sglang/srt/layers/quantization/fp8_utils.py` | Runtime replacement swaps nvjet FP8 GEMMs for CUTLASS kernels, removing per-launch memset bubbles and extra output-copy kernels | Treat nvjet GEMM + memset bubble ladders as an in-flight SGLang linear-kernel family before calling them novel. |
| PR `#18612` NVFP4 CUTLASS MoE fused SiLU+Mul+quant | `silu_and_mul_scaled_nvfp4`<br>`nvfp4 expert quant`<br>`cutlass moe` | `PR #18612`<br>`python/sglang/srt/layers/moe/cutlass_w4a8_moe.py`<br>`python/sglang/jit_kernel/nvfp4.py` | Fuses MoE activation epilogue and NVFP4 expert quantization before the CUTLASS MoE second GEMM | Treat split SiLU+Mul then NVFP4 expert quant in CUTLASS MoE traces as an in-flight upstream SGLang family. | | PR `#18612` NVFP4 CUTLASS MoE fused SiLU+Mul+quant | `silu_and_mul_scaled_nvfp4`<br>`nvfp4 expert quant`<br>`cutlass moe` | `PR #18612`<br>`python/sglang/srt/layers/moe/cutlass_w4a8_moe.py`<br>`python/sglang/jit_kernel/nvfp4.py` | Fuses MoE activation epilogue and NVFP4 expert quantization before the CUTLASS MoE second GEMM | Treat split SiLU+Mul then NVFP4 expert quant in CUTLASS MoE traces as an in-flight upstream SGLang family. |
| PR `#22918` FlashInfer per-token NVFP4 MoE | `per_token_nvfp4`<br>`trtllm_fp4_block_scale_moe`<br>`FlashInfer MoE` | `PR #22918`<br>`python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py` | Adds FlashInfer-backed per-token NVFP4 MoE execution so expert quant/dequant work can move into the fused MoE backend | Treat standalone per-token NVFP4 MoE support kernels as a candidate missing backend-selection path, not an automatically novel kernel idea. | | PR `#22918` FlashInfer per-token NVFP4 MoE | `per_token_nvfp4`<br>`trtllm_fp4_block_scale_moe`<br>`FlashInfer MoE` | `PR #22918`<br>`python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py` | Adds FlashInfer-backed per-token NVFP4 MoE execution so expert quant/dequant work can move into the fused MoE backend | Treat standalone per-token NVFP4 MoE support kernels as a candidate missing backend-selection path, not an automatically novel kernel idea. |
| PR `#22851` NSA top-k backend and FlashInfer / PyTorch top-k split | `nsa topk`<br>`flashinfer_topk`<br>`pytorch_topk`<br>`fast_topk_transform` | `PR #22851`<br>`python/sglang/srt/layers/attention/nsa_backend.py` | Makes NSA top-k backend selection explicit and aligns fused top-k transform with FlashInfer / PyTorch fallbacks | When NSA top-k dominates decode, first classify it as backend selection or fused-transform eligibility work. | | PR `#22851` DSA top-k backend and FlashInfer / PyTorch top-k split | `dsa topk`<br>`flashinfer_topk`<br>`pytorch_topk`<br>`fast_topk_transform` | `PR #22851`<br>`python/sglang/srt/layers/attention/dsa_backend.py` | Makes DSA top-k backend selection explicit and aligns fused top-k transform with FlashInfer / PyTorch fallbacks | When DSA top-k dominates decode, first classify it as backend selection or fused-transform eligibility work. |
| PR `#24125` GLM5 NSA decode CatArrayBatchedCopy removal | `CatArrayBatchedCopy`<br>`GLM-5`<br>`NSA`<br>`TileLang decode` | `PR #24125`<br>`python/sglang/srt/layers/attention/nsa_backend.py` | Skips redundant cat/copy work in the GLM5 NSA TileLang decode path | Treat cat/copy bursts in GLM5 NSA decode as a concrete in-flight cleanup opportunity. | | PR `#24125` GLM5 DSA decode CatArrayBatchedCopy removal | `CatArrayBatchedCopy`<br>`GLM-5`<br>`DSA`<br>`TileLang decode` | `PR #24125`<br>`python/sglang/srt/layers/attention/dsa_backend.py` | Skips redundant cat/copy work in the GLM5 DSA TileLang decode path | Treat cat/copy bursts in GLM5 DSA decode as a concrete in-flight cleanup opportunity. |
| PR `#24007` MoE LoRA virtual experts for csgmv backend | `csgmv`<br>`virtual experts`<br>`MoE LoRA`<br>`fused_moe_lora` | `PR #24007`<br>`python/sglang/srt/layers/lora_backend.py`<br>`python/sglang/srt/layers/moe` | Routes MoE LoRA adapter work through virtual experts so csgmv-style kernels can batch it instead of launching fragmented adapter work | Treat MoE-LoRA tiny-kernel ladders as an in-flight batching/fusion family. | | PR `#24007` MoE LoRA virtual experts for csgmv backend | `csgmv`<br>`virtual experts`<br>`MoE LoRA`<br>`fused_moe_lora` | `PR #24007`<br>`python/sglang/srt/layers/lora_backend.py`<br>`python/sglang/srt/layers/moe` | Routes MoE LoRA adapter work through virtual experts so csgmv-style kernels can batch it instead of launching fragmented adapter work | Treat MoE-LoRA tiny-kernel ladders as an in-flight batching/fusion family. |
| PR `#24150` torch.compile local decode support | `enable_torch_compile`<br>`local compile`<br>`decode compile`<br>`torchinductor` | `PR #24150`<br>`python/sglang/srt` | Extends SGLang torch.compile coverage to local decode regions, so Inductor-generated fusion may replace hand-authored tiny kernels | When decode traces show compiler-generated kernels or missing named fused kernels, check this in-flight compile path before calling the shape unsupported. | | PR `#24150` torch.compile local decode support | `enable_torch_compile`<br>`local compile`<br>`decode compile`<br>`torchinductor` | `PR #24150`<br>`python/sglang/srt` | Extends SGLang torch.compile coverage to local decode regions, so Inductor-generated fusion may replace hand-authored tiny kernels | When decode traces show compiler-generated kernels or missing named fused kernels, check this in-flight compile path before calling the shape unsupported. |
@@ -294,7 +294,7 @@ contain the same implementation.
| `enable_single_batch_overlap` | `python/sglang/srt/server_args.py` | Enables the SBO family. | | `enable_single_batch_overlap` | `python/sglang/srt/server_args.py` | Enables the SBO family. |
| `enable_fused_moe_sum_all_reduce` | `python/sglang/srt/server_args.py` | Enables fused MoE sum-reduce in the down path. | | `enable_fused_moe_sum_all_reduce` | `python/sglang/srt/server_args.py` | Enables fused MoE sum-reduce in the down path. |
| `SGLANG_BLACKWELL_OVERLAP_SHARED_EXPERTS_OUTSIDE_SBO` | `python/sglang/srt/environ.py` | Alters how DeepSeek-style shared-expert overlap behaves on Blackwell. | | `SGLANG_BLACKWELL_OVERLAP_SHARED_EXPERTS_OUTSIDE_SBO` | `python/sglang/srt/environ.py` | Alters how DeepSeek-style shared-expert overlap behaves on Blackwell. |
| `SGLANG_NSA_FUSE_TOPK` | `python/sglang/srt/environ.py` | Gates NSA fused top-k transform / page-table build. | | `SGLANG_DSA_FUSE_TOPK` | `python/sglang/srt/environ.py` | Gates DSA fused top-k transform / page-table build. |
| `SGLANG_DISAGG_STAGING_BUFFER` | `python/sglang/srt/environ.py` | Enables the heterogeneous-TP staging-buffer family and its overlap windows. | | `SGLANG_DISAGG_STAGING_BUFFER` | `python/sglang/srt/environ.py` | Enables the heterogeneous-TP staging-buffer family and its overlap windows. |
| `SGLANG_STAGING_USE_TORCH` | `python/sglang/srt/disaggregation/common/staging_buffer.py` | Forces torch fallback for staging gather / scatter, so Triton staging kernels may disappear by design. | | `SGLANG_STAGING_USE_TORCH` | `python/sglang/srt/disaggregation/common/staging_buffer.py` | Forces torch fallback for staging gather / scatter, so Triton staging kernels may disappear by design. |
| `SGLANG_VIT_ENABLE_CUDA_GRAPH` | `python/sglang/srt/environ.py` | Can intentionally disable vision `aux_stream` overlap. | | `SGLANG_VIT_ENABLE_CUDA_GRAPH` | `python/sglang/srt/environ.py` | Can intentionally disable vision `aux_stream` overlap. |
.claude/skills/llm-torch-profiler-analysis/references/overlap-catalog.md
+1 -1
View File
@@ -48,7 +48,7 @@ upstream overlap references as of this refresh.
| Llama4 shared branch vs routed branch overlap | shared expert branch plus routed MoE branch as adjacent windows | `python/sglang/srt/models/llama4.py` | shared expert on current stream, router + topk + routed experts on `alt_stream` | Use Llama4 as the first precedent for branch-level overlap in similar sparse models. | | Llama4 shared branch vs routed branch overlap | shared expert branch plus routed MoE branch as adjacent windows | `python/sglang/srt/models/llama4.py` | shared expert on current stream, router + topk + routed experts on `alt_stream` | Use Llama4 as the first precedent for branch-level overlap in similar sparse models. |
| ExaoneMoE shared experts vs router experts overlap | shared expert output and router-expert output form a two-branch window | `python/sglang/srt/models/exaone_moe.py::forward_normal_dual_stream` | shared experts on current stream, router + routed experts on `alt_stream`, explicit join before combine | This is an existing dual-stream MoE overlap family. | | ExaoneMoE shared experts vs router experts overlap | shared expert output and router-expert output form a two-branch window | `python/sglang/srt/models/exaone_moe.py::forward_normal_dual_stream` | shared experts on current stream, router + routed experts on `alt_stream`, explicit join before combine | This is an existing dual-stream MoE overlap family. |
| Grok residual-MoE branch overlap | dense MLP and block-sparse MoE branches in parallel | `python/sglang/srt/models/grok.py::moe_with_rmoe` | dense MLP on current stream, MoE on `alt_stream`, fused dual residual RMSNorm around boundaries | Treat exposed Grok branch overlap as an existing pattern. | | Grok residual-MoE branch overlap | dense MLP and block-sparse MoE branches in parallel | `python/sglang/srt/models/grok.py::moe_with_rmoe` | dense MLP on current stream, MoE on `alt_stream`, fused dual residual RMSNorm around boundaries | Treat exposed Grok branch overlap as an existing pattern. |
| NSA dual-stream overlap | Q-proj, K-proj, RoPE, cache-store, quantization in tight two-stream windows | `python/sglang/srt/layers/attention/nsa/nsa_indexer.py` | Q / K projection split, RoPE split, cache-store vs quantization overlap | NSA already contains several dual-stream overlap precedents. | | DSA dual-stream overlap | Q-proj, K-proj, RoPE, cache-store, quantization in tight two-stream windows | `python/sglang/srt/layers/attention/dsa/dsa_indexer.py` | Q / K projection split, RoPE split, cache-store vs quantization overlap | DSA already contains several dual-stream overlap precedents. |
| MoriEP async dispatch / combine comm stream | `MoriEP`<br>`_comm_stream`<br>`dispatch`<br>`combine`<br>`done_event` | `python/sglang/srt/layers/moe/token_dispatcher/moriep.py` | MoriEP can submit dispatch and combine onto a dedicated communication stream and synchronize only through events | Treat MoriEP comm / compute interleave as an existing MoE overlap family. | | MoriEP async dispatch / combine comm stream | `MoriEP`<br>`_comm_stream`<br>`dispatch`<br>`combine`<br>`done_event` | `python/sglang/srt/layers/moe/token_dispatcher/moriep.py` | MoriEP can submit dispatch and combine onto a dedicated communication stream and synchronize only through events | Treat MoriEP comm / compute interleave as an existing MoE overlap family. |
| Generic `alt_stream` overlap families | `alt_stream` plus explicit `wait_stream` / `with torch.cuda.stream(...)` | `qwen2_moe.py`<br>`qwen3_moe.py`<br>`glm4_moe.py`<br>`bailing_moe.py`<br>`llada2.py`<br>`grok.py`<br>`olmo2.py`<br>`step3p5.py`<br>`longcat_flash.py`<br>`falcon_h1.py` | model-specific overlap on attention prep, MoE branches, or cache-store | Search these families before designing a new overlap scheme from scratch. | | Generic `alt_stream` overlap families | `alt_stream` plus explicit `wait_stream` / `with torch.cuda.stream(...)` | `qwen2_moe.py`<br>`qwen3_moe.py`<br>`glm4_moe.py`<br>`bailing_moe.py`<br>`llada2.py`<br>`grok.py`<br>`olmo2.py`<br>`step3p5.py`<br>`longcat_flash.py`<br>`falcon_h1.py` | model-specific overlap on attention prep, MoE branches, or cache-store | Search these families before designing a new overlap scheme from scratch. |
.claude/skills/llm-torch-profiler-analysis/scripts/triage_kernel_helpers.py
+9 -9
View File
@@ -530,16 +530,16 @@ FUSION_PATTERN_REGISTRY: Tuple[FusionPatternSpec, ...] = (
likely_share=0.5, likely_share=0.5,
), ),
FusionPatternSpec( FusionPatternSpec(
pattern="NSA fused metadata copy for graph replay", pattern="DSA fused metadata copy for graph replay",
candidate_path="python/sglang/jit_kernel/fused_metadata_copy.py", candidate_path="python/sglang/jit_kernel/fused_metadata_copy.py",
active_keywords=( active_keywords=(
"fused_metadata_copy", "fused_metadata_copy",
"fused_metadata_copy_multi", "fused_metadata_copy_multi",
"fused_nsa_cache_seqlens", "fused_dsa_cache_seqlens",
"fused_flashmla_metadata", "fused_flashmla_metadata",
), ),
rationale_hint=( rationale_hint=(
"NSA replay metadata copies are already fused into one-kernel" " families." "DSA replay metadata copies are already fused into one-kernel" " families."
), ),
min_share=0.02, min_share=0.02,
likely_share=0.2, likely_share=0.2,
@@ -744,23 +744,23 @@ FUSION_PATTERN_REGISTRY: Tuple[FusionPatternSpec, ...] = (
likely_share=1.5, likely_share=1.5,
), ),
FusionPatternSpec( FusionPatternSpec(
pattern="NSA fused top-k transform / page-table build", pattern="DSA fused top-k transform / page-table build",
candidate_path="python/sglang/srt/layers/attention/nsa_backend.py", candidate_path="python/sglang/srt/layers/attention/dsa_backend.py",
active_keywords=( active_keywords=(
"fast_topk_transform_fused", "fast_topk_transform_fused",
"fast_topk_transform_ragged_fused", "fast_topk_transform_ragged_fused",
), ),
rationale_hint=( rationale_hint=(
"NSA top-k metadata preparation already has fused transform kernels." "DSA top-k metadata preparation already has fused transform kernels."
), ),
min_share=0.05, min_share=0.05,
likely_share=0.3, likely_share=0.3,
), ),
FusionPatternSpec( FusionPatternSpec(
pattern="NSA fused quantize + indexed K-cache store", pattern="DSA fused quantize + indexed K-cache store",
candidate_path=( candidate_path=(
"python/sglang/jit_kernel/fused_store_index_cache.py" "python/sglang/jit_kernel/fused_store_index_cache.py"
"<br>python/sglang/srt/layers/attention/nsa/nsa_indexer.py" "<br>python/sglang/srt/layers/attention/dsa/dsa_indexer.py"
), ),
active_keywords=("fused_store_index_k_cache",), active_keywords=("fused_store_index_k_cache",),
split_groups=( split_groups=(
@@ -768,7 +768,7 @@ FUSION_PATTERN_REGISTRY: Tuple[FusionPatternSpec, ...] = (
("index_k", "cache", "store"), ("index_k", "cache", "store"),
), ),
rationale_hint=( rationale_hint=(
"NSA already has a fused quantize-and-indexed-store kernel family." "DSA already has a fused quantize-and-indexed-store kernel family."
), ),
min_share=0.2, min_share=0.2,
likely_share=1.0, likely_share=1.0,
.github/workflows/nightly-test-amd-rocm720.yml
+2 -2
View File
@@ -742,7 +742,7 @@ jobs:
bash scripts/ci/amd/amd_ci_install_dependency.sh --skip-test-time-deps bash scripts/ci/amd/amd_ci_install_dependency.sh --skip-test-time-deps
bash scripts/ci/amd/amd_ci_exec.sh pip install git+https://github.com/huggingface/transformers.git@96f807a33b75 bash scripts/ci/amd/amd_ci_exec.sh pip install git+https://github.com/huggingface/transformers.git@96f807a33b75
- name: Accuracy Test ROCm 7.2 (8-GPU GLM-5.1 NSA) - name: Accuracy Test ROCm 7.2 (8-GPU GLM-5.1 DSA)
timeout-minutes: 120 timeout-minutes: 120
run: | run: |
> github_summary.md # Clear summary file > github_summary.md # Clear summary file
@@ -1414,7 +1414,7 @@ jobs:
bash scripts/ci/amd/amd_ci_exec.sh pip install tabulate bash scripts/ci/amd/amd_ci_exec.sh pip install tabulate
bash scripts/ci/amd/amd_ci_exec.sh pip install git+https://github.com/huggingface/transformers.git@96f807a33b75 bash scripts/ci/amd/amd_ci_exec.sh pip install git+https://github.com/huggingface/transformers.git@96f807a33b75
- name: Accuracy Test MI35x ROCm 7.2 (8-GPU GLM-5.1 NSA) - name: Accuracy Test MI35x ROCm 7.2 (8-GPU GLM-5.1 DSA)
timeout-minutes: 180 timeout-minutes: 180
run: | run: |
> github_summary.md # Clear summary file > github_summary.md # Clear summary file
.github/workflows/nightly-test-amd.yml
+2 -2
View File
@@ -744,7 +744,7 @@ jobs:
bash scripts/ci/amd/amd_ci_install_dependency.sh bash scripts/ci/amd/amd_ci_install_dependency.sh
bash scripts/ci/amd/amd_ci_exec.sh pip install git+https://github.com/huggingface/transformers.git@96f807a33b75 bash scripts/ci/amd/amd_ci_exec.sh pip install git+https://github.com/huggingface/transformers.git@96f807a33b75
- name: Accuracy Test (8-GPU GLM-5.1 NSA) - name: Accuracy Test (8-GPU GLM-5.1 DSA)
timeout-minutes: 120 timeout-minutes: 120
run: | run: |
> github_summary.md # Clear summary file > github_summary.md # Clear summary file
@@ -1419,7 +1419,7 @@ jobs:
bash scripts/ci/amd/amd_ci_exec.sh pip install tabulate bash scripts/ci/amd/amd_ci_exec.sh pip install tabulate
bash scripts/ci/amd/amd_ci_exec.sh pip install git+https://github.com/huggingface/transformers.git@96f807a33b75 bash scripts/ci/amd/amd_ci_exec.sh pip install git+https://github.com/huggingface/transformers.git@96f807a33b75
- name: Accuracy Test MI35x (8-GPU GLM-5.1 NSA) - name: Accuracy Test MI35x (8-GPU GLM-5.1 DSA)
timeout-minutes: 180 timeout-minutes: 180
run: | run: |
> github_summary.md # Clear summary file > github_summary.md # Clear summary file
docs/advanced_features/attention_backend.md
+4 -4
View File
@@ -50,7 +50,7 @@ Multimodal attention is selected by `--mm-attention-backend`. The "MultiModal" c
```{note} ```{note}
- FlashAttention 4 supports both prefill and decode on SM90 (Hopper) and SM100 (Blackwell). FA4 MLA supports `page_size = 1`; FA4 MHA requires `page_size = 128`. On SM100, this is auto-enforced by the server; on SM90, users must set `--page-size 128` manually. - FlashAttention 4 supports both prefill and decode on SM90 (Hopper) and SM100 (Blackwell). FA4 MLA supports `page_size = 1`; FA4 MHA requires `page_size = 128`. On SM100, this is auto-enforced by the server; on SM90, users must set `--page-size 128` manually.
- NSA is specifically designed for [DeepSeek V3.2 DSA](https://lmsys.org/blog/2025-09-29-deepseek-V32/). See the [DSA Attention Backend (NSA)](#dsa-attention-backend-nsa) section and [DeepSeek V3.2 deployment guide](../basic_usage/deepseek_v32.md) for details. - DSA is specifically designed for [DeepSeek V3.2](https://lmsys.org/blog/2025-09-29-deepseek-V32/). See the [DSA Attention Backend](#dsa-attention-backend) section and [DeepSeek V3.2 deployment guide](../basic_usage/deepseek_v32.md) for details.
``` ```
```{warning} ```{warning}
@@ -107,11 +107,11 @@ GDN models are hybrid: the full-attention layers still require a standard `--att
- **Other CUDA (Hopper, Ampere, etc.)**: auto-selection works; no special constraints. - **Other CUDA (Hopper, Ampere, etc.)**: auto-selection works; no special constraints.
``` ```
### DSA Attention Backend (NSA) ### DSA Attention Backend
DSA (Deepseek Sparse Attention) is a native sparse attention mechanism used by [DeepSeek V3.2](https://lmsys.org/blog/2025-09-29-deepseek-V32/). It is activated automatically when the model architecture requires it and is selected via `--attention-backend nsa`. DSA (Deepseek Sparse Attention) is a native sparse attention mechanism used by [DeepSeek V3.2](https://lmsys.org/blog/2025-09-29-deepseek-V32/). It is activated automatically when the model architecture requires it and is selected via `--attention-backend dsa` (the legacy alias `--attention-backend nsa` is deprecated and kept for one release).
Internally, the NSA backend dispatches to different sub-backends for prefill and decode phases. You can override these with `--nsa-prefill-backend` and `--nsa-decode-backend`: Internally, the DSA backend dispatches to different sub-backends for prefill and decode phases. You can override these with `--dsa-prefill-backend` and `--dsa-decode-backend` (the `--nsa-prefill-backend` / `--nsa-decode-backend` aliases are deprecated):
| **Sub-backend** | **Prefill** | **Decode** | **Notes** | | **Sub-backend** | **Prefill** | **Decode** | **Notes** |
|-----------------------|-------------|------------|-----------------------------------------------| |-----------------------|-------------|------------|-----------------------------------------------|
docs/advanced_features/hisparse_guide.md
+2 -2
View File
@@ -90,7 +90,7 @@ python3 -m sglang.launch_server \
--tp-size 8 --dp-size 8 --enable-dp-attention \ --tp-size 8 --dp-size 8 --enable-dp-attention \
--mem-fraction-static 0.85 \ --mem-fraction-static 0.85 \
--kv-cache-dtype bfloat16 \ --kv-cache-dtype bfloat16 \
--nsa-decode-backend flashmla_sparse \ --dsa-decode-backend flashmla_sparse \
--disaggregation-mode decode \ --disaggregation-mode decode \
--disaggregation-ib-device mlx5_0,mlx5_1,mlx5_2,mlx5_3 \ --disaggregation-ib-device mlx5_0,mlx5_1,mlx5_2,mlx5_3 \
--dist-init-addr 127.0.0.1:5757 \ --dist-init-addr 127.0.0.1:5757 \
@@ -123,7 +123,7 @@ python3 -m sglang.bench_serving \
- The prefill instance does not need `--enable-hisparse`; it is unaware of HiSparse. - The prefill instance does not need `--enable-hisparse`; it is unaware of HiSparse.
- On the decode instance, the following flags are **required** for HiSparse: - On the decode instance, the following flags are **required** for HiSparse:
- `--kv-cache-dtype bfloat16` — currently only bfloat16 KV cache is supported (more dtypes planned). - `--kv-cache-dtype bfloat16` — currently only bfloat16 KV cache is supported (more dtypes planned).
- `--nsa-decode-backend flashmla_sparse` — currently only `flashmla_sparse` backend is supported. - `--dsa-decode-backend flashmla_sparse` — currently only `flashmla_sparse` backend is supported.
- `--enable-hisparse` — enables HiSparse. - `--enable-hisparse` — enables HiSparse.
- `--hisparse-config` — HiSparse configuration (top_k, device_buffer_size, host_to_device_ratio). - `--hisparse-config` — HiSparse configuration (top_k, device_buffer_size, host_to_device_ratio).
- `host_to_device_ratio` should be configured based on the host machine's available memory. For example: - `host_to_device_ratio` should be configured based on the host machine's available memory. For example:
docs/advanced_features/server_arguments.md
+9 -9
View File
@@ -265,14 +265,14 @@ Please consult the documentation below and [server_args.py](https://github.com/s
## Kernel Backends (Attention, Sampling, Grammar, GEMM) ## Kernel Backends (Attention, Sampling, Grammar, GEMM)
| Argument | Description | Defaults | Options | | Argument | Description | Defaults | Options |
| --- | --- | --- | --- | | --- | --- | --- | --- |
| `--attention-backend` | Choose the kernels for attention layers. | `None` | `triton`, `torch_native`, `flex_attention`, `nsa`, `cutlass_mla`, `fa3`, `fa4`, `flashinfer`, `flashmla`, `trtllm_mla`, `trtllm_mha`, `dual_chunk_flash_attn`, `aiter`, `wave`, `intel_amx`, `ascend` | | `--attention-backend` | Choose the kernels for attention layers. | `None` | `triton`, `torch_native`, `flex_attention`, `dsa` (canonical; `nsa` is a deprecated alias), `cutlass_mla`, `fa3`, `fa4`, `flashinfer`, `flashmla`, `trtllm_mla`, `trtllm_mha`, `dual_chunk_flash_attn`, `aiter`, `wave`, `intel_amx`, `ascend` |
| `--prefill-attention-backend` | Choose the kernels for prefill attention layers (have priority over --attention-backend). | `None` | `triton`, `torch_native`, `flex_attention`, `nsa`, `cutlass_mla`, `fa3`, `fa4`, `flashinfer`, `flashmla`, `trtllm_mla`, `trtllm_mha`, `dual_chunk_flash_attn`, `aiter`, `wave`, `intel_amx`, `ascend` | | `--prefill-attention-backend` | Choose the kernels for prefill attention layers (have priority over --attention-backend). | `None` | `triton`, `torch_native`, `flex_attention`, `dsa` (canonical; `nsa` is a deprecated alias), `cutlass_mla`, `fa3`, `fa4`, `flashinfer`, `flashmla`, `trtllm_mla`, `trtllm_mha`, `dual_chunk_flash_attn`, `aiter`, `wave`, `intel_amx`, `ascend` |
| `--decode-attention-backend` | Choose the kernels for decode attention layers (have priority over --attention-backend). | `None` | `triton`, `torch_native`, `flex_attention`, `nsa`, `cutlass_mla`, `fa3`, `fa4`, `flashinfer`, `flashmla`, `trtllm_mla`, `trtllm_mha`, `dual_chunk_flash_attn`, `aiter`, `wave`, `intel_amx`, `ascend` | | `--decode-attention-backend` | Choose the kernels for decode attention layers (have priority over --attention-backend). | `None` | `triton`, `torch_native`, `flex_attention`, `dsa` (canonical; `nsa` is a deprecated alias), `cutlass_mla`, `fa3`, `fa4`, `flashinfer`, `flashmla`, `trtllm_mla`, `trtllm_mha`, `dual_chunk_flash_attn`, `aiter`, `wave`, `intel_amx`, `ascend` |
| `--sampling-backend` | Choose the kernels for sampling layers. | `None` | `flashinfer`, `pytorch`, `ascend` | | `--sampling-backend` | Choose the kernels for sampling layers. | `None` | `flashinfer`, `pytorch`, `ascend` |
| `--grammar-backend` | Choose the backend for grammar-guided decoding. | `None` | `xgrammar`, `outlines`, `llguidance`, `none` | | `--grammar-backend` | Choose the backend for grammar-guided decoding. | `None` | `xgrammar`, `outlines`, `llguidance`, `none` |
| `--mm-attention-backend` | Set multimodal attention backend. | `None` | `sdpa`, `fa3`, `fa4`, `triton_attn`, `ascend_attn`, `aiter_attn` | | `--mm-attention-backend` | Set multimodal attention backend. | `None` | `sdpa`, `fa3`, `fa4`, `triton_attn`, `ascend_attn`, `aiter_attn` |
| `--nsa-prefill-backend` | Choose the NSA backend for the prefill stage (overrides `--attention-backend` when running DeepSeek NSA-style attention). | `flashmla_sparse` | `flashmla_sparse`, `flashmla_kv`, `flashmla_auto`, `fa3`, `tilelang`, `aiter`, `trtllm` | | `--dsa-prefill-backend` | Choose the DSA backend for the prefill stage (overrides `--attention-backend` when running DeepSeek DSA-style attention). `--nsa-prefill-backend` is a deprecated alias. | `flashmla_sparse` | `flashmla_sparse`, `flashmla_kv`, `flashmla_auto`, `fa3`, `tilelang`, `aiter`, `trtllm` |
| `--nsa-decode-backend` | Choose the NSA backend for the decode stage when running DeepSeek NSA-style attention. Overrides `--attention-backend` for decoding. | `fa3` | `flashmla_sparse`, `flashmla_kv`, `fa3`, `tilelang`, `aiter`, `trtllm` | | `--dsa-decode-backend` | Choose the DSA backend for the decode stage when running DeepSeek DSA-style attention. Overrides `--attention-backend` for decoding. `--nsa-decode-backend` is a deprecated alias. | `fa3` | `flashmla_sparse`, `flashmla_kv`, `fa3`, `tilelang`, `aiter`, `trtllm` |
| `--fp8-gemm-backend` | Choose the runner backend for Blockwise FP8 GEMM operations. Options: 'auto' (default, auto-selects based on hardware), 'deep_gemm' (JIT-compiled; enabled by default on NVIDIA Hopper (SM90) and Blackwell (SM100) when DeepGEMM is installed), 'flashinfer_trtllm' (FlashInfer TRTLLM backend; SM100/SM103 only), 'flashinfer_cutlass' (FlashInfer CUTLASS backend, SM120 only), 'flashinfer_deepgemm' (Hopper SM90 only, uses swapAB optimization for small M dimensions in decoding), 'cutlass' (optimal for Hopper/Blackwell GPUs and high-throughput), 'triton' (fallback, widely compatible), 'aiter' (ROCm only).| `auto` | `auto`, `deep_gemm`, `flashinfer_trtllm`, `flashinfer_cutlass`, `flashinfer_deepgemm`, `cutlass`, `triton`, `aiter` | | `--fp8-gemm-backend` | Choose the runner backend for Blockwise FP8 GEMM operations. Options: 'auto' (default, auto-selects based on hardware), 'deep_gemm' (JIT-compiled; enabled by default on NVIDIA Hopper (SM90) and Blackwell (SM100) when DeepGEMM is installed), 'flashinfer_trtllm' (FlashInfer TRTLLM backend; SM100/SM103 only), 'flashinfer_cutlass' (FlashInfer CUTLASS backend, SM120 only), 'flashinfer_deepgemm' (Hopper SM90 only, uses swapAB optimization for small M dimensions in decoding), 'cutlass' (optimal for Hopper/Blackwell GPUs and high-throughput), 'triton' (fallback, widely compatible), 'aiter' (ROCm only).| `auto` | `auto`, `deep_gemm`, `flashinfer_trtllm`, `flashinfer_cutlass`, `flashinfer_deepgemm`, `cutlass`, `triton`, `aiter` |
| `--fp4-gemm-backend` | Choose the runner backend for NVFP4 GEMM operations. Options: 'flashinfer_cutlass' (default), 'auto' (auto-selects between flashinfer_cudnn/flashinfer_cutlass based on CUDA/cuDNN version), 'flashinfer_cudnn' (FlashInfer cuDNN backend, optimal on CUDA 13+ with cuDNN 9.15+), 'flashinfer_trtllm' (FlashInfer TensorRT-LLM backend, requires different weight preparation with shuffling). All backends are from FlashInfer; when FlashInfer is unavailable, sgl-kernel CUTLASS is used as an automatic fallback.| `flashinfer_cutlass` | `auto`, `flashinfer_cudnn`, `flashinfer_cutlass`, `flashinfer_trtllm` | | `--fp4-gemm-backend` | Choose the runner backend for NVFP4 GEMM operations. Options: 'flashinfer_cutlass' (default), 'auto' (auto-selects between flashinfer_cudnn/flashinfer_cutlass based on CUDA/cuDNN version), 'flashinfer_cudnn' (FlashInfer cuDNN backend, optimal on CUDA 13+ with cuDNN 9.15+), 'flashinfer_trtllm' (FlashInfer TensorRT-LLM backend, requires different weight preparation with shuffling). All backends are from FlashInfer; when FlashInfer is unavailable, sgl-kernel CUTLASS is used as an automatic fallback.| `flashinfer_cutlass` | `auto`, `flashinfer_cudnn`, `flashinfer_cutlass`, `flashinfer_trtllm` |
| `--disable-flashinfer-autotune` | Flashinfer autotune is enabled by default. Set this flag to disable the autotune. | `False` | bool flag (set to enable) | | `--disable-flashinfer-autotune` | Flashinfer autotune is enabled by default. Set this flag to disable the autotune. | `False` | bool flag (set to enable) |
@@ -463,8 +463,8 @@ Please consult the documentation below and [server_args.py](https://github.com/s
| `--enable-deterministic-inference` | Enable deterministic inference mode with batch invariant ops. | `False` | bool flag (set to enable) | | `--enable-deterministic-inference` | Enable deterministic inference mode with batch invariant ops. | `False` | bool flag (set to enable) |
| `--rl-on-policy-target` | The training system that SGLang needs to match for true on-policy. | `None` | `fsdp` | | `--rl-on-policy-target` | The training system that SGLang needs to match for true on-policy. | `None` | `fsdp` |
| `--enable-attn-tp-input-scattered` | Allow input of attention to be scattered when only using tensor parallelism, to reduce the computational load of operations such as qkv latent. | `False` | bool flag (set to enable) | | `--enable-attn-tp-input-scattered` | Allow input of attention to be scattered when only using tensor parallelism, to reduce the computational load of operations such as qkv latent. | `False` | bool flag (set to enable) |
| `--enable-nsa-prefill-context-parallel` | Enable context parallelism used in the long sequence prefill phase of DeepSeek v3.2. | `False` | bool flag (set to enable) | | `--enable-dsa-prefill-context-parallel` | Enable context parallelism used in the long sequence prefill phase of DeepSeek v3.2. (`--enable-nsa-prefill-context-parallel` is a deprecated alias.) | `False` | bool flag (set to enable) |
| `--nsa-prefill-cp-mode` | Token splitting mode for the prefill phase of DeepSeek v3.2 under context parallelism. Optional values: `round-robin-split`(default),`in-seq-split`. `round-robin-split` distributes tokens across ranks based on `token_idx % cp_size`. It supports multi-batch prefill, fused MoE, and FP8 KV cache. | `in-seq-split` | `in-seq-split`, `round-robin-split` | | `--dsa-prefill-cp-mode` | Token splitting mode for the prefill phase of DeepSeek v3.2 under context parallelism. Optional values: `round-robin-split`(default),`in-seq-split`. `round-robin-split` distributes tokens across ranks based on `token_idx % cp_size`. It supports multi-batch prefill, fused MoE, and FP8 KV cache. (`--nsa-prefill-cp-mode` is a deprecated alias.) | `in-seq-split` | `in-seq-split`, `round-robin-split` |
| `--enable-fused-qk-norm-rope` | Enable fused qk normalization and rope rotary embedding. | `False` | bool flag (set to enable) | | `--enable-fused-qk-norm-rope` | Enable fused qk normalization and rope rotary embedding. | `False` | bool flag (set to enable) |
| `--enable-precise-embedding-interpolation` | Enable corner alignment for resize of embeddings grid to ensure more accurate(but slower) evaluation of interpolated embedding values. | `False` | bool flag (set to enable) | | `--enable-precise-embedding-interpolation` | Enable corner alignment for resize of embeddings grid to ensure more accurate(but slower) evaluation of interpolated embedding values. | `False` | bool flag (set to enable) |
@@ -569,5 +569,5 @@ Please consult the documentation below and [server_args.py](https://github.com/s
| `--crash-on-nan` | Crash the server on nan logprobs. | `False` | Type: str | | `--crash-on-nan` | Crash the server on nan logprobs. | `False` | Type: str |
| `--hybrid-kvcache-ratio` | Mix ratio in [0,1] between uniform and hybrid kv buffers (0.0 = pure uniform: swa_size / full_size = 1)(1.0 = pure hybrid: swa_size / full_size = local_attention_size / context_length) | `None` | Optional[float] | | `--hybrid-kvcache-ratio` | Mix ratio in [0,1] between uniform and hybrid kv buffers (0.0 = pure uniform: swa_size / full_size = 1)(1.0 = pure hybrid: swa_size / full_size = local_attention_size / context_length) | `None` | Optional[float] |
| `--load-watch-interval` | The interval of load watching in seconds. | `0.1` | Type: float | | `--load-watch-interval` | The interval of load watching in seconds. | `0.1` | Type: float |
| `--nsa-prefill` | Choose the NSA backend for the prefill stage (overrides `--attention-backend` when running DeepSeek NSA-style attention). | `flashmla_sparse` | `flashmla_sparse`, `flashmla_decode`, `fa3`, `tilelang`, `aiter` | | `--nsa-prefill` | Deprecated alias for `--dsa-prefill-backend`. Choose the DSA backend for the prefill stage (overrides `--attention-backend` when running DeepSeek DSA-style attention). | `flashmla_sparse` | `flashmla_sparse`, `flashmla_decode`, `fa3`, `tilelang`, `aiter` |
| `--nsa-decode` | Choose the NSA backend for the decode stage when running DeepSeek NSA-style attention. Overrides `--attention-backend` for decoding. | `flashmla_kv` | `flashmla_prefill`, `flashmla_kv`, `fa3`, `tilelang`, `aiter` | | `--nsa-decode` | Deprecated alias for `--dsa-decode-backend`. Choose the DSA backend for the decode stage when running DeepSeek DSA-style attention. Overrides `--attention-backend` for decoding. | `flashmla_kv` | `flashmla_prefill`, `flashmla_kv`, `fa3`, `tilelang`, `aiter` |
docs/basic_usage/deepseek_v32.md
+17 -17
View File
@@ -53,7 +53,7 @@ python -m sglang.launch_server --model deepseek-ai/DeepSeek-V3.2-Exp --tp 8 --ep
python -m sglang.launch_server --model deepseek-ai/DeepSeek-V3.2-Exp --tp 8 python -m sglang.launch_server --model deepseek-ai/DeepSeek-V3.2-Exp --tp 8
# Launch with TP on MI30x/MI35x # Launch with TP on MI30x/MI35x
python3 -m sglang.launch_server --model deepseek-ai/DeepSeek-V3.2-Exp --tp 8 --nsa-prefill-backend tilelang --nsa-decode-backend tilelang python3 -m sglang.launch_server --model deepseek-ai/DeepSeek-V3.2-Exp --tp 8 --dsa-prefill-backend tilelang --dsa-decode-backend tilelang
``` ```
To serve GLM-5, just replace the `--model` argument with `zai-org/GLM-5-FP8`. To serve GLM-5, just replace the `--model` argument with `zai-org/GLM-5-FP8`.
@@ -61,9 +61,9 @@ To serve GLM-5, just replace the `--model` argument with `zai-org/GLM-5-FP8`.
### Configuration Tips ### Configuration Tips
- **DP Attention**: To enable [DP Attention](../advanced_features/dp_dpa_smg_guide.md), please include `--enable-dp-attention --dp <dp-size>` in command. DP Attention is better for large concurrency scenarios. - **DP Attention**: To enable [DP Attention](../advanced_features/dp_dpa_smg_guide.md), please include `--enable-dp-attention --dp <dp-size>` in command. DP Attention is better for large concurrency scenarios.
- **TP Attention**: Launching with TP attention is also supported. TP attention is better for low latency scenarios. - **TP Attention**: Launching with TP attention is also supported. TP attention is better for low latency scenarios.
- **Short-sequence MHA prefill (adaptive)**: For short prefill sequences (default threshold: **2048 tokens**), the NSA backend uses standard MHA automatically (no extra flags). On H200 (SM90) this path uses the FlashAttention variable-length kernel; on B200 (SM100) it uses TRT-LLM ragged MHA. MHA uses `MHA_ONE_SHOT` for best performance, which computes multi-head attention over all tokens (both cached prefix and newly extended tokens) in a single kernel invocation, avoiding the overhead of chunked KV cache processing. This achieves optimal throughput for short sequences where total sequence length fits within the chunk capacity limit. - **Short-sequence MHA prefill (adaptive)**: For short prefill sequences (default threshold: **2048 tokens**), the DSA backend uses standard MHA automatically (no extra flags). On H200 (SM90) this path uses the FlashAttention variable-length kernel; on B200 (SM100) it uses TRT-LLM ragged MHA. MHA uses `MHA_ONE_SHOT` for best performance, which computes multi-head attention over all tokens (both cached prefix and newly extended tokens) in a single kernel invocation, avoiding the overhead of chunked KV cache processing. This achieves optimal throughput for short sequences where total sequence length fits within the chunk capacity limit.
- **MHA prefill threshold relaxation**: To apply MHA attention to requests longer than 2048 tokens, please set the flag `SGLANG_NSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD` to a value larger than 2048. As threshold grows larger, the prefill performance can be improved, but at the cost of potential accuracy drop. - **MHA prefill threshold relaxation**: To apply MHA attention to requests longer than 2048 tokens, please set the flag `SGLANG_DSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD` to a value larger than 2048 (`SGLANG_NSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD` is a deprecated alias). As threshold grows larger, the prefill performance can be improved, but at the cost of potential accuracy drop.
- **Choices of Attention Kernels**: The attention backend is automatically set to `nsa` attention backend for DeepSeek V3.2 model. In this backend, different kernels for sparse prefilling/decoding are implemented, which can be specified by `--nsa-prefill-backend` and `--nsa-decode-backend` server arguments. The choices of nsa prefill/decode attention kernels include: - **Choices of Attention Kernels**: The attention backend is automatically set to `dsa` attention backend for DeepSeek V3.2 model (the deprecated `nsa` alias also works). In this backend, different kernels for sparse prefilling/decoding are implemented, which can be specified by `--dsa-prefill-backend` and `--dsa-decode-backend` server arguments (the deprecated `--nsa-prefill-backend` / `--nsa-decode-backend` aliases also work). The choices of dsa prefill/decode attention kernels include:
- `flashmla_sparse`: `flash_mla_sparse_fwd` kernel from `flash_mla` library. Can run on both Hopper and Blackwell GPUs. It requires bf16 q, kv inputs. - `flashmla_sparse`: `flash_mla_sparse_fwd` kernel from `flash_mla` library. Can run on both Hopper and Blackwell GPUs. It requires bf16 q, kv inputs.
- `flashmla_kv`: `flash_mla_with_kvcache` kernel from `flash_mla` library. Can run on both Hopper and Blackwell GPUs. It requires bf16 q, fp8 k_cache inputs. - `flashmla_kv`: `flash_mla_with_kvcache` kernel from `flash_mla` library. Can run on both Hopper and Blackwell GPUs. It requires bf16 q, fp8 k_cache inputs.
- `flashmla_auto`: enables automatic selection of either `flashmla_sparse` or `flashmla_kv` kernel for prefill based on KV cache dtype, hardware, and heuristics. With BF16 KV cache, `flashmla_sparse` is always used on both Hopper and Blackwell. With FP8 KV cache: On Hopper (SM90), it unconditionally uses `flashmla_kv`; On Blackwell (SM100), it uses `flashmla_sparse` when `total_kv_tokens < total_q_tokens * 512`, otherwise falls back to `flashmla_kv`. The heuristics may need to be tuned if the performance of either kernel changes significantly. - `flashmla_auto`: enables automatic selection of either `flashmla_sparse` or `flashmla_kv` kernel for prefill based on KV cache dtype, hardware, and heuristics. With BF16 KV cache, `flashmla_sparse` is always used on both Hopper and Blackwell. With FP8 KV cache: On Hopper (SM90), it unconditionally uses `flashmla_kv`; On Blackwell (SM100), it uses `flashmla_sparse` when `total_kv_tokens < total_q_tokens * 512`, otherwise falls back to `flashmla_kv`. The heuristics may need to be tuned if the performance of either kernel changes significantly.
@@ -319,11 +319,11 @@ DeepSeek-V3.2-Speciale:
**Note: This feature is only verified on Hopper machines** **Note: This feature is only verified on Hopper machines**
For context parallel in DeepSeek V3.2 model, we provide two different modes of splitting tokens, which can be controlled with argument `--nsa-prefill-cp-mode`. For context parallel in DeepSeek V3.2 model, we provide two different modes of splitting tokens, which can be controlled with argument `--dsa-prefill-cp-mode` (the deprecated `--nsa-prefill-cp-mode` alias also works).
### In sequence splitting ### In sequence splitting
The first mode can be enabled by `--nsa-prefill-cp-mode in-seq-split`. This mode implements context parallel for DSA by splitting the sequence uniformly between context parallel ranks. At attention stage, each cp rank computes the indexer results of sharded sequence, and collects the whole kv cache through all gather operator. Add `attn_cp_size` for communication group for context parallel. The first mode can be enabled by `--dsa-prefill-cp-mode in-seq-split`. This mode implements context parallel for DSA by splitting the sequence uniformly between context parallel ranks. At attention stage, each cp rank computes the indexer results of sharded sequence, and collects the whole kv cache through all gather operator. Add `attn_cp_size` for communication group for context parallel.
Note that the in-sequence splitting mode has the following restrictions: Note that the in-sequence splitting mode has the following restrictions:
- The batch size is restricted to 1 for prefill batches - The batch size is restricted to 1 for prefill batches
@@ -335,12 +335,12 @@ For more details, please refer to PR https://github.com/sgl-project/sglang/pull/
Example: Example:
```bash ```bash
# In-seq splitting mode launched with EP + DP # In-seq splitting mode launched with EP + DP
python -m sglang.launch_server --model deepseek-ai/DeepSeek-V3.2-Exp --tp 8 --ep 8 --dp 2 --enable-dp-attention --enable-nsa-prefill-context-parallel --attn-cp-size 4 --nsa-prefill-cp-mode in-seq-split --max-running-requests 32 python -m sglang.launch_server --model deepseek-ai/DeepSeek-V3.2-Exp --tp 8 --ep 8 --dp 2 --enable-dp-attention --enable-dsa-prefill-context-parallel --attn-cp-size 4 --dsa-prefill-cp-mode in-seq-split --max-running-requests 32
``` ```
### Round robin splitting (default setting) ### Round robin splitting (default setting)
This mode can be enabled by specifying the parameter `--nsa-prefill-cp-mode round-robin-split`, which distributes tokens across ranks based on `token_idx % cp_size`. This mode can be enabled by specifying the parameter `--dsa-prefill-cp-mode round-robin-split`, which distributes tokens across ranks based on `token_idx % cp_size`.
In this scenario, compared to the in-sequence splitting method, it additionally supports the fused MoE backend (the fused MoE backend may deliver better performance than DeepEP in single-machine scenarios), FP8 KV-cache, and multi-batch prefill inference. However, it cannot be enabled with DP attention together. In this scenario, compared to the in-sequence splitting method, it additionally supports the fused MoE backend (the fused MoE backend may deliver better performance than DeepEP in single-machine scenarios), FP8 KV-cache, and multi-batch prefill inference. However, it cannot be enabled with DP attention together.
@@ -349,7 +349,7 @@ For more details, please refer to PR https://github.com/sgl-project/sglang/pull/
Example usage: Example usage:
```bash ```bash
# Launch with FusedMoe + CP8 # Launch with FusedMoe + CP8
python -m sglang.launch_server --model deepseek-ai/DeepSeek-V3.2-Exp --tp 8 --enable-nsa-prefill-context-parallel --attn-cp-size 8 --nsa-prefill-cp-mode round-robin-split --max-running-requests 32 python -m sglang.launch_server --model deepseek-ai/DeepSeek-V3.2-Exp --tp 8 --enable-dsa-prefill-context-parallel --attn-cp-size 8 --dsa-prefill-cp-mode round-robin-split --max-running-requests 32
``` ```
### Pipeline Parallel + Context Parallel (PP + CP) ### Pipeline Parallel + Context Parallel (PP + CP)
@@ -372,9 +372,9 @@ python3 -m sglang.launch_server \
--dist-init-addr <HEAD_NODE_IP>:62001 \ --dist-init-addr <HEAD_NODE_IP>:62001 \
--tp 8 --pp-size 2 \ --tp 8 --pp-size 2 \
--dp-size 1 --moe-dense-tp-size 1 \ --dp-size 1 --moe-dense-tp-size 1 \
--enable-nsa-prefill-context-parallel \ --enable-dsa-prefill-context-parallel \
--attn-cp-size 8 \ --attn-cp-size 8 \
--nsa-prefill-cp-mode round-robin-split \ --dsa-prefill-cp-mode round-robin-split \
--trust-remote-code \ --trust-remote-code \
--disable-radix-cache \ --disable-radix-cache \
--mem-fraction-static 0.8 \ --mem-fraction-static 0.8 \
@@ -396,9 +396,9 @@ python3 -m sglang.launch_server \
--dist-init-addr <HEAD_NODE_IP>:62001 \ --dist-init-addr <HEAD_NODE_IP>:62001 \
--tp 8 --pp-size 2 \ --tp 8 --pp-size 2 \
--dp-size 1 --moe-dense-tp-size 1 \ --dp-size 1 --moe-dense-tp-size 1 \
--enable-nsa-prefill-context-parallel \ --enable-dsa-prefill-context-parallel \
--attn-cp-size 8 \ --attn-cp-size 8 \
--nsa-prefill-cp-mode round-robin-split \ --dsa-prefill-cp-mode round-robin-split \
--trust-remote-code \ --trust-remote-code \
--disable-radix-cache \ --disable-radix-cache \
--mem-fraction-static 0.8 \ --mem-fraction-static 0.8 \
@@ -424,9 +424,9 @@ python -m sglang.launch_server \
--dist-init-addr <PREFILL_HEAD_IP>:20102 \ --dist-init-addr <PREFILL_HEAD_IP>:20102 \
--tp 8 --pp-size 2 \ --tp 8 --pp-size 2 \
--dp-size 1 --moe-dense-tp-size 1 \ --dp-size 1 --moe-dense-tp-size 1 \
--enable-nsa-prefill-context-parallel \ --enable-dsa-prefill-context-parallel \
--attn-cp-size 8 \ --attn-cp-size 8 \
--nsa-prefill-cp-mode round-robin-split \ --dsa-prefill-cp-mode round-robin-split \
--disaggregation-ib-device mlx5_bond_0,mlx5_bond_1,mlx5_bond_2,mlx5_bond_3 \ --disaggregation-ib-device mlx5_bond_0,mlx5_bond_1,mlx5_bond_2,mlx5_bond_3 \
--trust-remote-code \ --trust-remote-code \
--disable-radix-cache \ --disable-radix-cache \
@@ -450,9 +450,9 @@ python -m sglang.launch_server \
--dist-init-addr <PREFILL_HEAD_IP>:20102 \ --dist-init-addr <PREFILL_HEAD_IP>:20102 \
--tp 8 --pp-size 2 \ --tp 8 --pp-size 2 \
--dp-size 1 --moe-dense-tp-size 1 \ --dp-size 1 --moe-dense-tp-size 1 \
--enable-nsa-prefill-context-parallel \ --enable-dsa-prefill-context-parallel \
--attn-cp-size 8 \ --attn-cp-size 8 \
--nsa-prefill-cp-mode round-robin-split \ --dsa-prefill-cp-mode round-robin-split \
--disaggregation-ib-device mlx5_bond_0,mlx5_bond_1,mlx5_bond_2,mlx5_bond_3 \ --disaggregation-ib-device mlx5_bond_0,mlx5_bond_1,mlx5_bond_2,mlx5_bond_3 \
--trust-remote-code \ --trust-remote-code \
--disable-radix-cache \ --disable-radix-cache \
docs/platforms/ascend/ascend_npu_best_practice.md
+2 -2
View File
@@ -1113,8 +1113,8 @@ do
--nnodes 2 --node-rank $i \ --nnodes 2 --node-rank $i \
--disaggregation-bootstrap-port 8995 \ --disaggregation-bootstrap-port 8995 \
--moe-dense-tp-size 1 \ --moe-dense-tp-size 1 \
--enable-nsa-prefill-context-parallel \ --enable-dsa-prefill-context-parallel \
--nsa-prefill-cp-mode in-seq-split \ --dsa-prefill-cp-mode in-seq-split \
--attn-cp-size 32 \ --attn-cp-size 32 \
--speculative-algorithm NEXTN --speculative-num-steps 1 --speculative-eagle-topk 1 --speculative-num-draft-tokens 2 \ --speculative-algorithm NEXTN --speculative-num-steps 1 --speculative-eagle-topk 1 --speculative-num-draft-tokens 2 \
--dist-init-addr ${P_IP[0]}:10000 --dist-init-addr ${P_IP[0]}:10000
docs/platforms/ascend/ascend_npu_environment_variables.md
+1 -1
View File
@@ -9,7 +9,7 @@ This document provides a list of commonly used environment variables and aims to
|--------------------------------------------------|-------------------------------------------------------------------------------------------------------------------------------------------------------------|---------------| |--------------------------------------------------|-------------------------------------------------------------------------------------------------------------------------------------------------------------|---------------|
| `SGLANG_NPU_USE_MLAPO` | Adopts the `MLAPO` fusion operator in attention <br/> preprocessing stage of the MLA model. | `false` | | `SGLANG_NPU_USE_MLAPO` | Adopts the `MLAPO` fusion operator in attention <br/> preprocessing stage of the MLA model. | `false` |
| `SGLANG_USE_FIA_NZ` | Reshapes KV Cache for FIA NZ format.<br/> `SGLANG_USE_FIA_NZ` must be enabled with `SGLANG_NPU_USE_MLAPO` | `false` | | `SGLANG_USE_FIA_NZ` | Reshapes KV Cache for FIA NZ format.<br/> `SGLANG_USE_FIA_NZ` must be enabled with `SGLANG_NPU_USE_MLAPO` | `false` |
| `SGLANG_NPU_USE_MULTI_STREAM` | Enable dual-stream computation of shared experts <br/> and routing experts in DeepSeek models.<br/> Enable dual-stream computation in DeepSeek NSA Indexer. | `false` | | `SGLANG_NPU_USE_MULTI_STREAM` | Enable dual-stream computation of shared experts <br/> and routing experts in DeepSeek models.<br/> Enable dual-stream computation in DeepSeek DSA Indexer. | `false` |
| `SGLANG_NPU_DISABLE_ACL_FORMAT_WEIGHT` | Disable cast model weight tensor to a specific NPU <br/> ACL format. | `false` | | `SGLANG_NPU_DISABLE_ACL_FORMAT_WEIGHT` | Disable cast model weight tensor to a specific NPU <br/> ACL format. | `false` |
| `SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK` | The maximum number of dispatched tokens on each rank. | `128` | | `SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK` | The maximum number of dispatched tokens on each rank. | `128` |
docs/platforms/ascend/ascend_npu_support_features.md
+3 -3
View File
@@ -204,8 +204,8 @@ click [Server Arguments](https://docs.sglang.io/advanced_features/server_argumen
| `--sampling-backend` | `None` | `pytorch`,<br/>`ascend` | A2, A3 | | `--sampling-backend` | `None` | `pytorch`,<br/>`ascend` | A2, A3 |
| `--grammar-backend` | `None` | `xgrammar` | A2, A3 | | `--grammar-backend` | `None` | `xgrammar` | A2, A3 |
| `--mm-attention-backend` | `None` | `ascend_attn` | A2, A3 | | `--mm-attention-backend` | `None` | `ascend_attn` | A2, A3 |
| `--nsa-prefill-backend` | `flashmla_sparse` | `flashmla_sparse`,<br/> `flashmla_decode`,<br/>`fa3`,<br/> `tilelang`,<br/> `aiter` | Special for GPU | | `--dsa-prefill-backend` | `flashmla_sparse` | `flashmla_sparse`,<br/> `flashmla_decode`,<br/>`fa3`,<br/> `tilelang`,<br/> `aiter` | Special for GPU |
| `--nsa-decode-backend` | `fa3` | `flashmla_prefill`,<br/> `flashmla_kv`,<br/> `fa3`,<br/>`tilelang`,<br/> `aiter` | Special for GPU | | `--dsa-decode-backend` | `fa3` | `flashmla_prefill`,<br/> `flashmla_kv`,<br/> `fa3`,<br/>`tilelang`,<br/> `aiter` | Special for GPU |
| `--fp8-gemm-backend` | `auto` | `auto`,<br/> `deep_gemm`,<br/> `flashinfer_trtllm`,<br/>`flashinfer_cutlass`,<br/>`flashinfer_deepgemm`,<br/>`cutlass`,<br/> `triton`,<br/> `aiter` | Special for GPU | | `--fp8-gemm-backend` | `auto` | `auto`,<br/> `deep_gemm`,<br/> `flashinfer_trtllm`,<br/>`flashinfer_cutlass`,<br/>`flashinfer_deepgemm`,<br/>`cutlass`,<br/> `triton`,<br/> `aiter` | Special for GPU |
| `--disable-flashinfer-`<br/>`autotune` | `False` | bool flag<br/> (set to enable) | Special for GPU | | `--disable-flashinfer-`<br/>`autotune` | `False` | bool flag<br/> (set to enable) | Special for GPU |
@@ -371,7 +371,7 @@ click [Server Arguments](https://docs.sglang.io/advanced_features/server_argumen
| `--rl-on-policy-target` | `None` | `fsdp` | Planned | | `--rl-on-policy-target` | `None` | `fsdp` | Planned |
| `--enable-layerwise-`<br/>`nvtx-marker` | `False` | bool flag<br/> (set to enable) | Special for GPU | | `--enable-layerwise-`<br/>`nvtx-marker` | `False` | bool flag<br/> (set to enable) | Special for GPU |
| `--enable-attn-tp-`<br/>`input-scattered` | `False` | bool flag<br/> (set to enable) | Experimental | | `--enable-attn-tp-`<br/>`input-scattered` | `False` | bool flag<br/> (set to enable) | Experimental |
| `--enable-nsa-prefill-`<br/>`context-parallel` | `False` | bool flag<br/> (set to enable) | A2, A3 | | `--enable-dsa-prefill-`<br/>`context-parallel` | `False` | bool flag<br/> (set to enable) | A2, A3 |
| `--enable-fused-qk-`<br/>`norm-rope` | `False` | bool flag<br/> (set to enable) | Special for GPU | | `--enable-fused-qk-`<br/>`norm-rope` | `False` | bool flag<br/> (set to enable) | Special for GPU |
## Dynamic batch tokenizer ## Dynamic batch tokenizer
docs/platforms/plugin.md
+1 -1
View File
@@ -234,7 +234,7 @@ python -c "from sglang.srt.platforms import current_platform; print(current_plat
| `get_graph_runner_cls()` | `raise NotImplementedError` | Graph Runner class | | `get_graph_runner_cls()` | `raise NotImplementedError` | Graph Runner class |
| `get_mha_kv_pool_cls()` | `raise NotImplementedError` | MHA KV cache pool class | | `get_mha_kv_pool_cls()` | `raise NotImplementedError` | MHA KV cache pool class |
| `get_mla_kv_pool_cls()` | `raise NotImplementedError` | MLA KV cache pool class | | `get_mla_kv_pool_cls()` | `raise NotImplementedError` | MLA KV cache pool class |
| `get_nsa_kv_pool_cls()` | `raise NotImplementedError` | NSA KV cache pool class (DeepSeek V3.2) | | `get_dsa_kv_pool_cls()` | `raise NotImplementedError` | DSA KV cache pool class (DeepSeek V3.2) |
| `get_paged_allocator_cls()` | `raise NotImplementedError` | Paged allocator class | | `get_paged_allocator_cls()` | `raise NotImplementedError` | Paged allocator class |
| `get_piecewise_backend_cls()` | `raise NotImplementedError` | Piecewise compilation backend class | | `get_piecewise_backend_cls()` | `raise NotImplementedError` | Piecewise compilation backend class |
| `get_compile_backend(mode)` | `"inductor"` | Compilation backend string | | `get_compile_backend(mode)` | `"inductor"` | Compilation backend string |
docs/references/environment_variables.md
+4 -4
View File
@@ -88,16 +88,16 @@ SGLang supports various environment variables that can be used to configure its
| `SGLANG_MORI_POST_BATCH_SIZE` | Number of RDMA work requests posted in a single batch to each QP | `-1` | | `SGLANG_MORI_POST_BATCH_SIZE` | Number of RDMA work requests posted in a single batch to each QP | `-1` |
| `SGLANG_MORI_NUM_WORKERS` | Number of worker threads in the RDMA executor thread pool | `1` | | `SGLANG_MORI_NUM_WORKERS` | Number of worker threads in the RDMA executor thread pool | `1` |
## NSA Backend Configuration (For DeepSeek V3.2) ## DSA Backend Configuration (For DeepSeek V3.2)
<!-- # Environment variable to control mtp precomputing of metadata for multi-step speculative decoding --> <!-- # Environment variable to control mtp precomputing of metadata for multi-step speculative decoding -->
| Environment Variable | Description | Default Value | | Environment Variable | Description | Default Value |
| --- | --- | --- | | --- | --- | --- |
| `SGLANG_NSA_FUSE_TOPK` | Fuse the operation of picking topk logits and picking topk indices from page table | `true` | | `SGLANG_DSA_FUSE_TOPK` | Fuse the operation of picking topk logits and picking topk indices from page table (`SGLANG_NSA_FUSE_TOPK` is a deprecated alias) | `true` |
| `SGLANG_NSA_ENABLE_MTP_PRECOMPUTE_METADATA` | Precompute metadata that can be shared among different draft steps when MTP is enabled | `true` | | `SGLANG_DSA_ENABLE_MTP_PRECOMPUTE_METADATA` | Precompute metadata that can be shared among different draft steps when MTP is enabled (`SGLANG_NSA_ENABLE_MTP_PRECOMPUTE_METADATA` is a deprecated alias) | `true` |
| `SGLANG_USE_FUSED_METADATA_COPY` | Control whether to use fused metadata copy kernel for cuda graph replay | `true` | | `SGLANG_USE_FUSED_METADATA_COPY` | Control whether to use fused metadata copy kernel for cuda graph replay | `true` |
| `SGLANG_NSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD` | When the maximum kv len in current prefill batch exceeds this value, the sparse mla kernel will be applied, else it falls back to dense MHA implementation. Default to the index topk of model (2048 for DeepSeek V3.2) | `2048` | | `SGLANG_DSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD` | When the maximum kv len in current prefill batch exceeds this value, the sparse mla kernel will be applied, else it falls back to dense MHA implementation. Default to the index topk of model (2048 for DeepSeek V3.2) (`SGLANG_NSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD` is a deprecated alias) | `2048` |
## Memory Management ## Memory Management
docs_new/cookbook/autoregressive/DeepSeek/DeepSeek-V3_2.mdx
+9 -9
View File
@@ -349,11 +349,11 @@ sglang_args=$(echo serve \
--moe-a2a-backend deepep --ep-size 16 \ --moe-a2a-backend deepep --ep-size 16 \
--page-size 128 \ --page-size 128 \
--chunked-prefill-size 16384 \ --chunked-prefill-size 16384 \
--attention-backend nsa \ --attention-backend dsa \
--nsa-prefill-backend flashmla_sparse \ --dsa-prefill-backend flashmla_sparse \
--nsa-decode-backend flashmla_sparse \ --dsa-decode-backend flashmla_sparse \
--enable-nsa-prefill-context-parallel \ --enable-dsa-prefill-context-parallel \
--nsa-prefill-cp-mode round-robin-split \ --dsa-prefill-cp-mode round-robin-split \
--cuda-graph-max-bs 128 \ --cuda-graph-max-bs 128 \
--max-running-requests 128 \ --max-running-requests 128 \
--trust-remote-code --host "0.0.0.0" --port 30000 \ --trust-remote-code --host "0.0.0.0" --port 30000 \
@@ -361,7 +361,7 @@ sglang_args=$(echo serve \
--context-length 65536 \ --context-length 65536 \
--allow-auto-truncate --enable-metrics \ --allow-auto-truncate --enable-metrics \
--tool-call-parser deepseekv32 --reasoning-parser deepseek-v3 \ --tool-call-parser deepseekv32 --reasoning-parser deepseek-v3 \
--served-model-name DeepSeek-V3.2-nsa-pp-cp-ep-dp --served-model-name DeepSeek-V3.2-dsa-pp-cp-ep-dp
) )
sglang_args=($sglang_args) sglang_args=($sglang_args)
@@ -382,20 +382,20 @@ dp_config=" \
" "
cp_config=" \ cp_config=" \
--enable-nsa-prefill-context-parallel \ --enable-dsa-prefill-context-parallel \
" "
if [ "$dp" -eq 1 ]; then if [ "$dp" -eq 1 ]; then
cp_config=" \ cp_config=" \
$cp_config \ $cp_config \
--nsa-prefill-cp-mode round-robin-split \ --dsa-prefill-cp-mode round-robin-split \
" "
else else
cp_config=" \ cp_config=" \
$cp_config \ $cp_config \
--nsa-prefill-cp-mode in-seq-split \ --dsa-prefill-cp-mode in-seq-split \
" "
fi fi
docs_new/cookbook/autoregressive/GLM/GLM-5.1.mdx
+6 -6
View File
@@ -79,7 +79,7 @@ import { GLM51Deployment } from '/src/snippets/autoregressive/glm-51-deployment.
</tbody> </tbody>
</table> </table>
- **AMD GPUs**: Both BF16 and FP8 checkpoints are supported on MI300X/MI325X/MI355X at tp=8. Use `--nsa-prefill-backend tilelang --nsa-decode-backend tilelang` for the NSA attention backend. Add `--chunked-prefill-size 131072` and `--watchdog-timeout 1200` (20 minutes for weight loading). FP8 uses approximately half the memory of BF16 (~89 GB/GPU vs ~175 GB/GPU). EAGLE speculative decoding is not currently supported on AMD for GLM-5.1. - **AMD GPUs**: Both BF16 and FP8 checkpoints are supported on MI300X/MI325X/MI355X at tp=8. Use `--dsa-prefill-backend tilelang --dsa-decode-backend tilelang` for the DSA attention backend. Add `--chunked-prefill-size 131072` and `--watchdog-timeout 1200` (20 minutes for weight loading). FP8 uses approximately half the memory of BF16 (~89 GB/GPU vs ~175 GB/GPU). EAGLE speculative decoding is not currently supported on AMD for GLM-5.1.
- **GB300**: Only the FP8 checkpoint is recommended on GB300, with `tp=4`. For high-throughput DP attention on GB300, use `--dp 4`. - **GB300**: Only the FP8 checkpoint is recommended on GB300, with `tp=4`. For high-throughput DP attention on GB300, use `--dp 4`.
- For other configuration tips, please refer to [DeepSeek V3.2 documentation](../../../docs/basic_usage/deepseek_v32). GLM-5.1 and DeepSeek V3.2 share the same model structure, so the optimization techniques between these two models are also common (MTP, DSA kernel, Context Parallel...). - For other configuration tips, please refer to [DeepSeek V3.2 documentation](../../../docs/basic_usage/deepseek_v32). GLM-5.1 and DeepSeek V3.2 share the same model structure, so the optimization techniques between these two models are also common (MTP, DSA kernel, Context Parallel...).
- Use `--json-model-override-args '{"index_topk_pattern": "FFSFSSSFSSFFFSSSFFFSFSSSSSSFFSFFSFFSSFFFFFFSFFFFFSFFSSSSSSFSFFFSFSSSFSFFSFFSSS"}'` for GLM-5.1-FP8 if you want to enable the [IndexCache](https://github.com/THUDM/IndexCache) method. This feature is supported through [this PR](https://github.com/sgl-project/sglang/pull/21405) and introduces only a small accuracy loss. However, if you are running rigorous accuracy evaluations, it is not recommended to enable this feature. - Use `--json-model-override-args '{"index_topk_pattern": "FFSFSSSFSSFFFSSSFFFSFSSSSSSFFSFFSFFSSFFFFFFSFFFFFSFFSSSSSSFSFFFSFSSSFSFFSFFSSS"}'` for GLM-5.1-FP8 if you want to enable the [IndexCache](https://github.com/THUDM/IndexCache) method. This feature is supported through [this PR](https://github.com/sgl-project/sglang/pull/21405) and introduces only a small accuracy loss. However, if you are running rigorous accuracy evaluations, it is not recommended to enable this feature.
@@ -116,8 +116,8 @@ sglang serve \
--trust-remote-code \ --trust-remote-code \
--tool-call-parser glm47 \ --tool-call-parser glm47 \
--reasoning-parser glm45 \ --reasoning-parser glm45 \
--nsa-prefill-backend tilelang \ --dsa-prefill-backend tilelang \
--nsa-decode-backend tilelang \ --dsa-decode-backend tilelang \
--chunked-prefill-size 131072 \ --chunked-prefill-size 131072 \
--mem-fraction-static 0.80 \ --mem-fraction-static 0.80 \
--watchdog-timeout 1200 \ --watchdog-timeout 1200 \
@@ -132,8 +132,8 @@ sglang serve \
--model-path zai-org/GLM-5.1 \ --model-path zai-org/GLM-5.1 \
--tp 8 \ --tp 8 \
--trust-remote-code \ --trust-remote-code \
--nsa-prefill-backend tilelang \ --dsa-prefill-backend tilelang \
--nsa-decode-backend tilelang \ --dsa-decode-backend tilelang \
--chunked-prefill-size 131072 \ --chunked-prefill-size 131072 \
--mem-fraction-static 0.80 \ --mem-fraction-static 0.80 \
--watchdog-timeout 1200 \ --watchdog-timeout 1200 \
@@ -627,7 +627,7 @@ Average accuracy: 0.877
#### 5.3.1 GSM8K Benchmark (MI325/MI35x) #### 5.3.1 GSM8K Benchmark (MI325/MI35x)
- MI325/MI35x Test (GLM-5.1 BF16, `tp=8`, TileLang NSA backends) - MI325/MI35x Test (GLM-5.1 BF16, `tp=8`, TileLang DSA backends)
```bash Command ```bash Command
python3 benchmark/gsm8k/bench_sglang.py --num-questions 200 python3 benchmark/gsm8k/bench_sglang.py --num-questions 200
docs_new/cookbook/autoregressive/GLM/GLM-5.mdx
+5 -5
View File
@@ -85,9 +85,9 @@ import { GLM5Deployment } from '/src/snippets/autoregressive/glm-5-deployment.js
</tbody> </tbody>
</table> </table>
- **B200 (FP8)**: Use `--ep 1 --attention-backend nsa --nsa-decode-backend trtllm --nsa-prefill-backend trtllm --moe-runner-backend flashinfer_trtllm --enable-flashinfer-allreduce-fusion` for optimized NSA and MoE backends on Blackwell. Also add `--quantization fp8` for FP8 weight quantization. - **B200 (FP8)**: Use `--ep 1 --attention-backend dsa --dsa-decode-backend trtllm --dsa-prefill-backend trtllm --moe-runner-backend flashinfer_trtllm --enable-flashinfer-allreduce-fusion` for optimized DSA and MoE backends on Blackwell. Also add `--quantization fp8` for FP8 weight quantization.
- **AMD GPUs**: Use `--nsa-prefill-backend tilelang --nsa-decode-backend tilelang` for the NSA attention backend. Add `--chunked-prefill-size 131072` and `--watchdog-timeout 1200` (20 minutes for weight loading). EAGLE speculative decoding is not currently supported on AMD for GLM-5. - **AMD GPUs**: Use `--dsa-prefill-backend tilelang --dsa-decode-backend tilelang` for the DSA attention backend. Add `--chunked-prefill-size 131072` and `--watchdog-timeout 1200` (20 minutes for weight loading). EAGLE speculative decoding is not currently supported on AMD for GLM-5.
- For other configuration tips, please refer to [DeepSeek V3.2 documentation](../../../docs/basic_usage/deepseek_v32). GLM-5 and DeepSeek V3.2 share the same model structure, so the optimization techniques between these two models are also common (MTP, DSA kernel, Context Parallel...). - For other configuration tips, please refer to [DeepSeek V3.2 documentation](../../../docs/basic_usage/deepseek_v32). GLM-5 and DeepSeek V3.2 share the same model structure, so the optimization techniques between these two models are also common (MTP, DSA kernel, Context Parallel...).
- Use `--json-model-override-args '{"index_topk_pattern": "FFSFSSSFSSFFFSSSFFFSFSSSSSSFFSFFSFFSSFFFFFFSFFFFFSFFSSSSSSFSFFFSFSSSFSFFSFFSSS"}'` for GLM-5-FP8 if you want to enable the [IndexCache](https://github.com/THUDM/IndexCache) method. This feature is supported through [this PR](https://github.com/sgl-project/sglang/pull/21405) and introduces only a small accuracy loss. However, if you are running rigorous accuracy evaluations, it is not recommended to enable this feature. - Use `--json-model-override-args '{"index_topk_pattern": "FFSFSSSFSSFFFSSSFFFSFSSSSSSFFSFFSFFSSFFFFFFSFFFFFSFFSSSSSSFSFFFSFSSSFSFFSFFSSS"}'` for GLM-5-FP8 if you want to enable the [IndexCache](https://github.com/THUDM/IndexCache) method. This feature is supported through [this PR](https://github.com/sgl-project/sglang/pull/21405) and introduces only a small accuracy loss. However, if you are running rigorous accuracy evaluations, it is not recommended to enable this feature.
@@ -124,8 +124,8 @@ sglang serve \
--model-path zai-org/GLM-5 \ --model-path zai-org/GLM-5 \
--tp 8 \ --tp 8 \
--trust-remote-code \ --trust-remote-code \
--nsa-prefill-backend tilelang \ --dsa-prefill-backend tilelang \
--nsa-decode-backend tilelang \ --dsa-decode-backend tilelang \
--chunked-prefill-size 131072 \ --chunked-prefill-size 131072 \
--mem-fraction-static 0.80 \ --mem-fraction-static 0.80 \
--watchdog-timeout 1200 \ --watchdog-timeout 1200 \
@@ -653,7 +653,7 @@ Average accuracy: 0.877
#### 5.3.1 GSM8K Benchmark (MI325/MI35x) #### 5.3.1 GSM8K Benchmark (MI325/MI35x)
- MI325/MI35x Test (GLM-5 BF16, `tp=8`, TileLang NSA backends) - MI325/MI35x Test (GLM-5 BF16, `tp=8`, TileLang DSA backends)
```bash Command ```bash Command
python3 benchmark/gsm8k/bench_sglang.py --num-questions 200 python3 benchmark/gsm8k/bench_sglang.py --num-questions 200
docs_new/docs/advanced_features/attention_backend.mdx
+4 -4
View File
@@ -280,7 +280,7 @@ Multimodal attention is selected by `--mm-attention-backend`. The "MultiModal" c
<Note> <Note>
- FlashAttention 4 supports both prefill and decode on SM90 (Hopper) and SM100 (Blackwell). FA4 MLA supports `page_size = 1`; FA4 MHA requires `page_size = 128`. On SM100, this is auto-enforced by the server; on SM90, users must set `--page-size 128` manually. - FlashAttention 4 supports both prefill and decode on SM90 (Hopper) and SM100 (Blackwell). FA4 MLA supports `page_size = 1`; FA4 MHA requires `page_size = 128`. On SM100, this is auto-enforced by the server; on SM90, users must set `--page-size 128` manually.
- NSA is specifically designed for [DeepSeek V3.2 DSA](https://lmsys.org/blog/2025-09-29-deepseek-V32/). See the [DSA Attention Backend (NSA)](#dsa-attention-backend-nsa) section and [DeepSeek V3.2 deployment guide](../basic_usage/deepseek_v32) for details. - DSA is specifically designed for [DeepSeek V3.2](https://lmsys.org/blog/2025-09-29-deepseek-V32/). See the [DSA Attention Backend](#dsa-attention-backend) section and [DeepSeek V3.2 deployment guide](../basic_usage/deepseek_v32) for details.
</Note> </Note>
<Warning> <Warning>
@@ -378,11 +378,11 @@ GDN models are hybrid: the full-attention layers still require a standard `--att
- **Other CUDA (Hopper, Ampere, etc.)**: auto-selection works; no special constraints. - **Other CUDA (Hopper, Ampere, etc.)**: auto-selection works; no special constraints.
</Warning> </Warning>
### DSA Attention Backend (NSA) ### DSA Attention Backend
DSA (Deepseek Sparse Attention) is a native sparse attention mechanism used by [DeepSeek V3.2](https://lmsys.org/blog/2025-09-29-deepseek-V32/). It is activated automatically when the model architecture requires it and is selected via `--attention-backend nsa`. DSA (DeepSeek Sparse Attention) is a native sparse attention mechanism used by [DeepSeek V3.2](https://lmsys.org/blog/2025-09-29-deepseek-V32/). It is activated automatically when the model architecture requires it and is selected via `--attention-backend dsa` (deprecated alias: `nsa`).
Internally, the NSA backend dispatches to different sub-backends for prefill and decode phases. You can override these with `--nsa-prefill-backend` and `--nsa-decode-backend`: Internally, the DSA backend dispatches to different sub-backends for prefill and decode phases. You can override these with `--dsa-prefill-backend` and `--dsa-decode-backend`:
<table style={{width: "100%", borderCollapse: "collapse", tableLayout: "fixed"}}> <table style={{width: "100%", borderCollapse: "collapse", tableLayout: "fixed"}}>
<colgroup> <colgroup>
docs_new/docs/advanced_features/hisparse_guide.mdx
+2 -2
View File
@@ -142,7 +142,7 @@ python3 -m sglang.launch_server \
--tp-size 8 --dp-size 8 --enable-dp-attention \ --tp-size 8 --dp-size 8 --enable-dp-attention \
--mem-fraction-static 0.85 \ --mem-fraction-static 0.85 \
--kv-cache-dtype bfloat16 \ --kv-cache-dtype bfloat16 \
--nsa-decode-backend flashmla_sparse \ --dsa-decode-backend flashmla_sparse \
--disaggregation-mode decode \ --disaggregation-mode decode \
--disaggregation-ib-device mlx5_0,mlx5_1,mlx5_2,mlx5_3 \ --disaggregation-ib-device mlx5_0,mlx5_1,mlx5_2,mlx5_3 \
--dist-init-addr 127.0.0.1:5757 \ --dist-init-addr 127.0.0.1:5757 \
@@ -175,7 +175,7 @@ python3 -m sglang.bench_serving \
- The prefill instance does not need `--enable-hisparse`; it is unaware of HiSparse. - The prefill instance does not need `--enable-hisparse`; it is unaware of HiSparse.
- On the decode instance, the following flags are **required** for HiSparse: - On the decode instance, the following flags are **required** for HiSparse:
- `--kv-cache-dtype bfloat16` — currently only bfloat16 KV cache is supported (more dtypes planned). - `--kv-cache-dtype bfloat16` — currently only bfloat16 KV cache is supported (more dtypes planned).
- `--nsa-decode-backend flashmla_sparse` — currently only `flashmla_sparse` backend is supported. - `--dsa-decode-backend flashmla_sparse` — currently only `flashmla_sparse` backend is supported.
- `--enable-hisparse` — enables HiSparse. - `--enable-hisparse` — enables HiSparse.
- `--hisparse-config` — HiSparse configuration (top_k, device_buffer_size, host_to_device_ratio). - `--hisparse-config` — HiSparse configuration (top_k, device_buffer_size, host_to_device_ratio).
- `host_to_device_ratio` should be configured based on the host machine's available memory. For example: - `host_to_device_ratio` should be configured based on the host machine's available memory. For example:
docs_new/docs/advanced_features/server_arguments.mdx
+11 -11
View File
@@ -1162,19 +1162,19 @@ Please consult the documentation below and [server_args.py](https://github.com/s
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--attention-backend`</td> <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--attention-backend`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Choose the kernels for attention layers.</td> <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Choose the kernels for attention layers.</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`None`</td> <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`None`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>triton</code>, <code>torch_native</code>, <code>flex_attention</code>, <code>nsa</code>, <code>cutlass_mla</code>, <code>fa3</code>, <code>fa4</code>, <code>flashinfer</code>, <code>flashmla</code>, <code>trtllm_mla</code>, <code>trtllm_mha</code>, <code>dual_chunk_flash_attn</code>, <code>aiter</code>, <code>wave</code>, <code>intel_amx</code>, <code>ascend</code></td> <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>triton</code>, <code>torch_native</code>, <code>flex_attention</code>, <code>dsa</code>, <code>cutlass_mla</code>, <code>fa3</code>, <code>fa4</code>, <code>flashinfer</code>, <code>flashmla</code>, <code>trtllm_mla</code>, <code>trtllm_mha</code>, <code>dual_chunk_flash_attn</code>, <code>aiter</code>, <code>wave</code>, <code>intel_amx</code>, <code>ascend</code></td>
</tr> </tr>
<tr> <tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--prefill-attention-backend`</td> <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--prefill-attention-backend`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Choose the kernels for prefill attention layers (have priority over --attention-backend).</td> <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Choose the kernels for prefill attention layers (have priority over --attention-backend).</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`None`</td> <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`None`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>triton</code>, <code>torch_native</code>, <code>flex_attention</code>, <code>nsa</code>, <code>cutlass_mla</code>, <code>fa3</code>, <code>fa4</code>, <code>flashinfer</code>, <code>flashmla</code>, <code>trtllm_mla</code>, <code>trtllm_mha</code>, <code>dual_chunk_flash_attn</code>, <code>aiter</code>, <code>wave</code>, <code>intel_amx</code>, <code>ascend</code></td> <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>triton</code>, <code>torch_native</code>, <code>flex_attention</code>, <code>dsa</code>, <code>cutlass_mla</code>, <code>fa3</code>, <code>fa4</code>, <code>flashinfer</code>, <code>flashmla</code>, <code>trtllm_mla</code>, <code>trtllm_mha</code>, <code>dual_chunk_flash_attn</code>, <code>aiter</code>, <code>wave</code>, <code>intel_amx</code>, <code>ascend</code></td>
</tr> </tr>
<tr> <tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--decode-attention-backend`</td> <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--decode-attention-backend`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Choose the kernels for decode attention layers (have priority over --attention-backend).</td> <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Choose the kernels for decode attention layers (have priority over --attention-backend).</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`None`</td> <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`None`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>triton</code>, <code>torch_native</code>, <code>flex_attention</code>, <code>nsa</code>, <code>cutlass_mla</code>, <code>fa3</code>, <code>fa4</code>, <code>flashinfer</code>, <code>flashmla</code>, <code>trtllm_mla</code>, <code>trtllm_mha</code>, <code>dual_chunk_flash_attn</code>, <code>aiter</code>, <code>wave</code>, <code>intel_amx</code>, <code>ascend</code></td> <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>triton</code>, <code>torch_native</code>, <code>flex_attention</code>, <code>dsa</code>, <code>cutlass_mla</code>, <code>fa3</code>, <code>fa4</code>, <code>flashinfer</code>, <code>flashmla</code>, <code>trtllm_mla</code>, <code>trtllm_mha</code>, <code>dual_chunk_flash_attn</code>, <code>aiter</code>, <code>wave</code>, <code>intel_amx</code>, <code>ascend</code></td>
</tr> </tr>
<tr> <tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--sampling-backend`</td> <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--sampling-backend`</td>
@@ -1195,14 +1195,14 @@ Please consult the documentation below and [server_args.py](https://github.com/s
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>sdpa</code>, <code>fa3</code>, <code>fa4</code>, <code>triton_attn</code>, <code>ascend_attn</code>, <code>aiter_attn</code></td> <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>sdpa</code>, <code>fa3</code>, <code>fa4</code>, <code>triton_attn</code>, <code>ascend_attn</code>, <code>aiter_attn</code></td>
</tr> </tr>
<tr> <tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--nsa-prefill-backend`</td> <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--dsa-prefill-backend`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Choose the NSA backend for the prefill stage (overrides `--attention-backend` when running DeepSeek NSA-style attention).</td> <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>DSA backend for the prefill stage (overrides `--attention-backend` when running DeepSeek DSA-style attention).</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`flashmla_sparse`</td> <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`flashmla_sparse`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>flashmla_sparse</code>, <code>flashmla_kv</code>, <code>flashmla_auto</code>, <code>fa3</code>, <code>tilelang</code>, <code>aiter</code>, <code>trtllm</code></td> <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>flashmla_sparse</code>, <code>flashmla_kv</code>, <code>flashmla_auto</code>, <code>fa3</code>, <code>tilelang</code>, <code>aiter</code>, <code>trtllm</code></td>
</tr> </tr>
<tr> <tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--nsa-decode-backend`</td> <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--dsa-decode-backend`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Choose the NSA backend for the decode stage when running DeepSeek NSA-style attention. Overrides `--attention-backend` for decoding.</td> <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>DSA backend for the decode stage when running DeepSeek DSA-style attention. Overrides `--attention-backend` for decoding.</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`fa3`</td> <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`fa3`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>flashmla_sparse</code>, <code>flashmla_kv</code>, <code>fa3</code>, <code>tilelang</code>, <code>aiter</code>, <code>trtllm</code></td> <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>flashmla_sparse</code>, <code>flashmla_kv</code>, <code>fa3</code>, <code>tilelang</code>, <code>aiter</code>, <code>trtllm</code></td>
</tr> </tr>
@@ -2278,13 +2278,13 @@ Please consult the documentation below and [server_args.py](https://github.com/s
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>bool flag (set to enable)</td> <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>bool flag (set to enable)</td>
</tr> </tr>
<tr> <tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>--enable-nsa-prefill-context-parallel</code></td> <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>--enable-dsa-prefill-context-parallel</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Enable context parallelism used in the long sequence prefill phase of DeepSeek v3.2.</td> <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Enable context parallelism used in the long sequence prefill phase of DeepSeek v3.2.</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>False</code></td> <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>False</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>bool flag (set to enable)</td> <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>bool flag (set to enable)</td>
</tr> </tr>
<tr> <tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>--nsa-prefill-cp-mode</code></td> <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>--dsa-prefill-cp-mode</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Token splitting mode for the prefill phase of DeepSeek v3.2 under context parallelism. Optional values: <code>round-robin-split</code>(default),<code>in-seq-split</code>. <code>round-robin-split</code> distributes tokens across ranks based on <code>token_idx % cp_size</code>. It supports multi-batch prefill, fused MoE, and FP8 KV cache.</td> <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Token splitting mode for the prefill phase of DeepSeek v3.2 under context parallelism. Optional values: <code>round-robin-split</code>(default),<code>in-seq-split</code>. <code>round-robin-split</code> distributes tokens across ranks based on <code>token_idx % cp_size</code>. It supports multi-batch prefill, fused MoE, and FP8 KV cache.</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>in-seq-split</code></td> <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>in-seq-split</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>in-seq-split</code>, <code>round-robin-split</code></td> <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>in-seq-split</code>, <code>round-robin-split</code></td>
@@ -2841,13 +2841,13 @@ Please consult the documentation below and [server_args.py](https://github.com/s
</tr> </tr>
<tr> <tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--nsa-prefill`</td> <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--nsa-prefill`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Choose the NSA backend for the prefill stage (overrides `--attention-backend` when running DeepSeek NSA-style attention).</td> <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Deprecated alias for `--dsa-prefill-backend`.</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`flashmla_sparse`</td> <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`flashmla_sparse`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>`flashmla_sparse`, `flashmla_decode`, `fa3`, `tilelang`, `aiter`</td> <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>`flashmla_sparse`, `flashmla_decode`, `fa3`, `tilelang`, `aiter`</td>
</tr> </tr>
<tr> <tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--nsa-decode`</td> <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--nsa-decode`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Choose the NSA backend for the decode stage when running DeepSeek NSA-style attention. Overrides `--attention-backend` for decoding.</td> <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Deprecated alias for `--dsa-decode-backend`.</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`flashmla_kv`</td> <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`flashmla_kv`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>`flashmla_prefill`, `flashmla_kv`, `fa3`, `tilelang`, `aiter`</td> <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>`flashmla_prefill`, `flashmla_kv`, `fa3`, `tilelang`, `aiter`</td>
</tr> </tr>
docs_new/docs/basic_usage/deepseek_v32.mdx
+17 -17
View File
@@ -56,7 +56,7 @@ python -m sglang.launch_server --model deepseek-ai/DeepSeek-V3.2-Exp --tp 8 --ep
python -m sglang.launch_server --model deepseek-ai/DeepSeek-V3.2-Exp --tp 8 python -m sglang.launch_server --model deepseek-ai/DeepSeek-V3.2-Exp --tp 8
# Launch with TP on MI30x/MI35x # Launch with TP on MI30x/MI35x
python3 -m sglang.launch_server --model deepseek-ai/DeepSeek-V3.2-Exp --tp 8 --nsa-prefill-backend tilelang --nsa-decode-backend tilelang python3 -m sglang.launch_server --model deepseek-ai/DeepSeek-V3.2-Exp --tp 8 --dsa-prefill-backend tilelang --dsa-decode-backend tilelang
``` ```
To serve GLM-5, just replace the `--model` argument with `zai-org/GLM-5-FP8`. To serve GLM-5, just replace the `--model` argument with `zai-org/GLM-5-FP8`.
@@ -64,9 +64,9 @@ To serve GLM-5, just replace the `--model` argument with `zai-org/GLM-5-FP8`.
### Configuration Tips ### Configuration Tips
- **DP Attention**: To enable [DP Attention](../advanced_features/dp_dpa_smg_guide), please include `--enable-dp-attention --dp <dp-size>` in command. DP Attention is better for large concurrency scenarios. - **DP Attention**: To enable [DP Attention](../advanced_features/dp_dpa_smg_guide), please include `--enable-dp-attention --dp <dp-size>` in command. DP Attention is better for large concurrency scenarios.
- **TP Attention**: Launching with TP attention is also supported. TP attention is better for low latency scenarios. - **TP Attention**: Launching with TP attention is also supported. TP attention is better for low latency scenarios.
- **Short-sequence MHA prefill (adaptive)**: For short prefill sequences (default threshold: **2048 tokens**), the NSA backend uses standard MHA automatically (no extra flags). On H200 (SM90) this path uses the FlashAttention variable-length kernel; on B200 (SM100) it uses TRT-LLM ragged MHA. MHA uses `MHA_ONE_SHOT` for best performance, which computes multi-head attention over all tokens (both cached prefix and newly extended tokens) in a single kernel invocation, avoiding the overhead of chunked KV cache processing. This achieves optimal throughput for short sequences where total sequence length fits within the chunk capacity limit. - **Short-sequence MHA prefill (adaptive)**: For short prefill sequences (default threshold: **2048 tokens**), the DSA backend uses standard MHA automatically (no extra flags). On H200 (SM90) this path uses the FlashAttention variable-length kernel; on B200 (SM100) it uses TRT-LLM ragged MHA. MHA uses `MHA_ONE_SHOT` for best performance, which computes multi-head attention over all tokens (both cached prefix and newly extended tokens) in a single kernel invocation, avoiding the overhead of chunked KV cache processing. This achieves optimal throughput for short sequences where total sequence length fits within the chunk capacity limit.
- **MHA prefill threshold relaxation**: To apply MHA attention to requests longer than 2048 tokens, please set the flag `SGLANG_NSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD` to a value larger than 2048. As threshold grows larger, the prefill performance can be improved, but at the cost of potential accuracy drop. - **MHA prefill threshold relaxation**: To apply MHA attention to requests longer than 2048 tokens, please set the flag `SGLANG_DSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD` to a value larger than 2048 (`SGLANG_NSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD` is a deprecated alias). As threshold grows larger, the prefill performance can be improved, but at the cost of potential accuracy drop.
- **Choices of Attention Kernels**: The attention backend is automatically set to `nsa` attention backend for DeepSeek V3.2 model. In this backend, different kernels for sparse prefilling/decoding are implemented, which can be specified by `--nsa-prefill-backend` and `--nsa-decode-backend` server arguments. The choices of nsa prefill/decode attention kernels include: - **Choices of Attention Kernels**: The attention backend is automatically set to `dsa` attention backend for DeepSeek V3.2 model. In this backend, different kernels for sparse prefilling/decoding are implemented, which can be specified by `--dsa-prefill-backend` and `--dsa-decode-backend` server arguments. The choices of dsa prefill/decode attention kernels include:
- `flashmla_sparse`: `flash_mla_sparse_fwd` kernel from `flash_mla` library. Can run on both Hopper and Blackwell GPUs. It requires bf16 q, kv inputs. - `flashmla_sparse`: `flash_mla_sparse_fwd` kernel from `flash_mla` library. Can run on both Hopper and Blackwell GPUs. It requires bf16 q, kv inputs.
- `flashmla_kv`: `flash_mla_with_kvcache` kernel from `flash_mla` library. Can run on both Hopper and Blackwell GPUs. It requires bf16 q, fp8 k_cache inputs. - `flashmla_kv`: `flash_mla_with_kvcache` kernel from `flash_mla` library. Can run on both Hopper and Blackwell GPUs. It requires bf16 q, fp8 k_cache inputs.
- `flashmla_auto`: enables automatic selection of either `flashmla_sparse` or `flashmla_kv` kernel for prefill based on KV cache dtype, hardware, and heuristics. With BF16 KV cache, `flashmla_sparse` is always used on both Hopper and Blackwell. With FP8 KV cache: On Hopper (SM90), it unconditionally uses `flashmla_kv`; On Blackwell (SM100), it uses `flashmla_sparse` when `total_kv_tokens < total_q_tokens * 512`, otherwise falls back to `flashmla_kv`. The heuristics may need to be tuned if the performance of either kernel changes significantly. - `flashmla_auto`: enables automatic selection of either `flashmla_sparse` or `flashmla_kv` kernel for prefill based on KV cache dtype, hardware, and heuristics. With BF16 KV cache, `flashmla_sparse` is always used on both Hopper and Blackwell. With FP8 KV cache: On Hopper (SM90), it unconditionally uses `flashmla_kv`; On Blackwell (SM100), it uses `flashmla_sparse` when `total_kv_tokens < total_q_tokens * 512`, otherwise falls back to `flashmla_kv`. The heuristics may need to be tuned if the performance of either kernel changes significantly.
@@ -444,11 +444,11 @@ DeepSeek-V3.2-Speciale:
**Note: This feature is only verified on Hopper machines** **Note: This feature is only verified on Hopper machines**
For context parallel in DeepSeek V3.2 model, we provide two different modes of splitting tokens, which can be controlled with argument `--nsa-prefill-cp-mode`. For context parallel in DeepSeek V3.2 model, we provide two different modes of splitting tokens, which can be controlled with argument `--dsa-prefill-cp-mode`.
### In sequence splitting ### In sequence splitting
The first mode can be enabled by `--nsa-prefill-cp-mode in-seq-split`. This mode implements context parallel for DSA by splitting the sequence uniformly between context parallel ranks. At attention stage, each cp rank computes the indexer results of sharded sequence, and collects the whole kv cache through all gather operator. Add `attn_cp_size` for communication group for context parallel. The first mode can be enabled by `--dsa-prefill-cp-mode in-seq-split`. This mode implements context parallel for DSA by splitting the sequence uniformly between context parallel ranks. At attention stage, each cp rank computes the indexer results of sharded sequence, and collects the whole kv cache through all gather operator. Add `attn_cp_size` for communication group for context parallel.
Note that the in-sequence splitting mode has the following restrictions: Note that the in-sequence splitting mode has the following restrictions:
- The batch size is restricted to 1 for prefill batches - The batch size is restricted to 1 for prefill batches
@@ -460,12 +460,12 @@ For more details, please refer to PR https://github.com/sgl-project/sglang/pull/
Example: Example:
```bash Command ```bash Command
# In-seq splitting mode launched with EP + DP # In-seq splitting mode launched with EP + DP
python -m sglang.launch_server --model deepseek-ai/DeepSeek-V3.2-Exp --tp 8 --ep 8 --dp 2 --enable-dp-attention --enable-nsa-prefill-context-parallel --attn-cp-size 4 --nsa-prefill-cp-mode in-seq-split --max-running-requests 32 python -m sglang.launch_server --model deepseek-ai/DeepSeek-V3.2-Exp --tp 8 --ep 8 --dp 2 --enable-dp-attention --enable-dsa-prefill-context-parallel --attn-cp-size 4 --dsa-prefill-cp-mode in-seq-split --max-running-requests 32
``` ```
### Round robin splitting (default setting) ### Round robin splitting (default setting)
This mode can be enabled by specifying the parameter `--nsa-prefill-cp-mode round-robin-split`, which distributes tokens across ranks based on `token_idx % cp_size`. This mode can be enabled by specifying the parameter `--dsa-prefill-cp-mode round-robin-split`, which distributes tokens across ranks based on `token_idx % cp_size`.
In this scenario, compared to the in-sequence splitting method, it additionally supports the fused MoE backend (the fused MoE backend may deliver better performance than DeepEP in single-machine scenarios), FP8 KV-cache, and multi-batch prefill inference. However, it cannot be enabled with DP attention together. In this scenario, compared to the in-sequence splitting method, it additionally supports the fused MoE backend (the fused MoE backend may deliver better performance than DeepEP in single-machine scenarios), FP8 KV-cache, and multi-batch prefill inference. However, it cannot be enabled with DP attention together.
@@ -474,7 +474,7 @@ For more details, please refer to PR https://github.com/sgl-project/sglang/pull/
Example usage: Example usage:
```bash Command ```bash Command
# Launch with FusedMoe + CP8 # Launch with FusedMoe + CP8
python -m sglang.launch_server --model deepseek-ai/DeepSeek-V3.2-Exp --tp 8 --enable-nsa-prefill-context-parallel --attn-cp-size 8 --nsa-prefill-cp-mode round-robin-split --max-running-requests 32 python -m sglang.launch_server --model deepseek-ai/DeepSeek-V3.2-Exp --tp 8 --enable-dsa-prefill-context-parallel --attn-cp-size 8 --dsa-prefill-cp-mode round-robin-split --max-running-requests 32
``` ```
### Pipeline Parallel + Context Parallel (PP + CP) ### Pipeline Parallel + Context Parallel (PP + CP)
@@ -497,9 +497,9 @@ python3 -m sglang.launch_server \
--dist-init-addr <HEAD_NODE_IP>:62001 \ --dist-init-addr <HEAD_NODE_IP>:62001 \
--tp 8 --pp-size 2 \ --tp 8 --pp-size 2 \
--dp-size 1 --moe-dense-tp-size 1 \ --dp-size 1 --moe-dense-tp-size 1 \
--enable-nsa-prefill-context-parallel \ --enable-dsa-prefill-context-parallel \
--attn-cp-size 8 \ --attn-cp-size 8 \
--nsa-prefill-cp-mode round-robin-split \ --dsa-prefill-cp-mode round-robin-split \
--trust-remote-code \ --trust-remote-code \
--disable-radix-cache \ --disable-radix-cache \
--mem-fraction-static 0.8 \ --mem-fraction-static 0.8 \
@@ -521,9 +521,9 @@ python3 -m sglang.launch_server \
--dist-init-addr <HEAD_NODE_IP>:62001 \ --dist-init-addr <HEAD_NODE_IP>:62001 \
--tp 8 --pp-size 2 \ --tp 8 --pp-size 2 \
--dp-size 1 --moe-dense-tp-size 1 \ --dp-size 1 --moe-dense-tp-size 1 \
--enable-nsa-prefill-context-parallel \ --enable-dsa-prefill-context-parallel \
--attn-cp-size 8 \ --attn-cp-size 8 \
--nsa-prefill-cp-mode round-robin-split \ --dsa-prefill-cp-mode round-robin-split \
--trust-remote-code \ --trust-remote-code \
--disable-radix-cache \ --disable-radix-cache \
--mem-fraction-static 0.8 \ --mem-fraction-static 0.8 \
@@ -549,9 +549,9 @@ python -m sglang.launch_server \
--dist-init-addr <PREFILL_HEAD_IP>:20102 \ --dist-init-addr <PREFILL_HEAD_IP>:20102 \
--tp 8 --pp-size 2 \ --tp 8 --pp-size 2 \
--dp-size 1 --moe-dense-tp-size 1 \ --dp-size 1 --moe-dense-tp-size 1 \
--enable-nsa-prefill-context-parallel \ --enable-dsa-prefill-context-parallel \
--attn-cp-size 8 \ --attn-cp-size 8 \
--nsa-prefill-cp-mode round-robin-split \ --dsa-prefill-cp-mode round-robin-split \
--disaggregation-ib-device mlx5_bond_0,mlx5_bond_1,mlx5_bond_2,mlx5_bond_3 \ --disaggregation-ib-device mlx5_bond_0,mlx5_bond_1,mlx5_bond_2,mlx5_bond_3 \
--trust-remote-code \ --trust-remote-code \
--disable-radix-cache \ --disable-radix-cache \
@@ -575,9 +575,9 @@ python -m sglang.launch_server \
--dist-init-addr <PREFILL_HEAD_IP>:20102 \ --dist-init-addr <PREFILL_HEAD_IP>:20102 \
--tp 8 --pp-size 2 \ --tp 8 --pp-size 2 \
--dp-size 1 --moe-dense-tp-size 1 \ --dp-size 1 --moe-dense-tp-size 1 \
--enable-nsa-prefill-context-parallel \ --enable-dsa-prefill-context-parallel \
--attn-cp-size 8 \ --attn-cp-size 8 \
--nsa-prefill-cp-mode round-robin-split \ --dsa-prefill-cp-mode round-robin-split \
--disaggregation-ib-device mlx5_bond_0,mlx5_bond_1,mlx5_bond_2,mlx5_bond_3 \ --disaggregation-ib-device mlx5_bond_0,mlx5_bond_1,mlx5_bond_2,mlx5_bond_3 \
--trust-remote-code \ --trust-remote-code \
--disable-radix-cache \ --disable-radix-cache \
docs_new/docs/hardware-platforms/ascend-npus/ascend_npu_best_practice.mdx
+2 -2
View File
@@ -1598,8 +1598,8 @@ do
--nnodes 2 --node-rank $i \ --nnodes 2 --node-rank $i \
--disaggregation-bootstrap-port 8995 \ --disaggregation-bootstrap-port 8995 \
--moe-dense-tp-size 1 \ --moe-dense-tp-size 1 \
--enable-nsa-prefill-context-parallel \ --enable-dsa-prefill-context-parallel \
--nsa-prefill-cp-mode in-seq-split \ --dsa-prefill-cp-mode in-seq-split \
--attn-cp-size 32 \ --attn-cp-size 32 \
--speculative-algorithm NEXTN --speculative-num-steps 1 --speculative-eagle-topk 1 --speculative-num-draft-tokens 2 \ --speculative-algorithm NEXTN --speculative-num-steps 1 --speculative-eagle-topk 1 --speculative-num-draft-tokens 2 \
--dist-init-addr ${P_IP[0]}:10000 --dist-init-addr ${P_IP[0]}:10000
docs_new/docs/hardware-platforms/ascend-npus/ascend_npu_environment_variables.mdx
+1 -1
View File
@@ -29,7 +29,7 @@ This document provides a list of commonly used environment variables and aims to
</tr> </tr>
<tr> <tr>
<td><code>SGLANG_NPU_USE_MULTI_STREAM</code></td> <td><code>SGLANG_NPU_USE_MULTI_STREAM</code></td>
<td>Enable dual-stream computation of shared experts <br/> and routing experts in DeepSeek models.<br/> Enable dual-stream computation in DeepSeek NSA Indexer.</td> <td>Enable dual-stream computation of shared experts <br/> and routing experts in DeepSeek models.<br/> Enable dual-stream computation in DeepSeek DSA Indexer.</td>
<td><code>false</code></td> <td><code>false</code></td>
</tr> </tr>
<tr> <tr>
docs_new/docs/hardware-platforms/ascend-npus/ascend_npu_optimization.mdx
+6 -6
View File
@@ -187,18 +187,18 @@ These arguments and environment variables are critical for tuning prefill perfor
<td>true</td> <td>true</td>
</tr> </tr>
<tr> <tr>
<td>`--enable-nsa-prefill-context-parallel`</td> <td>`--enable-dsa-prefill-context-parallel`</td>
<td><strong>(DeepSeek V3.2 NSA-specific)</strong> Enables context parallelism for the long-sequence prefill phase of DeepSeek V3.2 with NSA (Native Sparse Attention). Distributes the sequence across CP ranks to parallelize the computationally expensive NSA prefill for ultra-long contexts.</td> <td><strong>(DeepSeek V3.2 DSA-specific)</strong> Enables context parallelism for the long-sequence prefill phase of DeepSeek V3.2 with DSA (DeepSeek Sparse Attention). Distributes the sequence across CP ranks to parallelize the computationally expensive DSA prefill for ultra-long contexts.</td>
<td>Enabled</td> <td>Enabled</td>
</tr> </tr>
<tr> <tr>
<td>`--nsa-prefill-cp-mode`</td> <td>`--dsa-prefill-cp-mode`</td>
<td><strong>(DeepSeek V3.2 NSA-specific)</strong> Controls how the long sequence is split across context parallel ranks: `in-seq-split` divides each sequence uniformly across CP ranks, optimal for single-request prefill. `round-robin-split` (code default) distributes tokens by index mod CP size, supporting multi-batch prefill. Only effective when `--enable-nsa-prefill-context-parallel` is enabled.</td> <td><strong>(DeepSeek V3.2 DSA-specific)</strong> Controls how the long sequence is split across context parallel ranks: `in-seq-split` divides each sequence uniformly across CP ranks, optimal for single-request prefill. `round-robin-split` (code default) distributes tokens by index mod CP size, supporting multi-batch prefill. Only effective when `--enable-dsa-prefill-context-parallel` is enabled.</td>
<td>`in-seq-split`</td> <td>`in-seq-split`</td>
</tr> </tr>
<tr> <tr>
<td>`--attn-cp-size`</td> <td>`--attn-cp-size`</td>
<td>Specifies the context parallelism group size for attention computation. Larger values distribute the sequence across more ranks, reducing per-rank memory and compute at the cost of increased communication. For models with NSA, this controls the CP size for sparse attention prefill. Set to the number of available devices for maximum parallelization.</td> <td>Specifies the context parallelism group size for attention computation. Larger values distribute the sequence across more ranks, reducing per-rank memory and compute at the cost of increased communication. For models with DSA, this controls the CP size for sparse attention prefill. Set to the number of available devices for maximum parallelization.</td>
<td>`32`</td> <td>`32`</td>
</tr> </tr>
</tbody> </tbody>
@@ -360,7 +360,7 @@ The following environment variables are used in other best practice configuratio
<tr> <tr>
<td>`HCCL_OP_EXPANSION_MODE=AIV`</td> <td>`HCCL_OP_EXPANSION_MODE=AIV`</td>
<td>Configures the HCCL communication algorithm scheduling to use AIV (Ascend Intelligent Vision) expansion mode, which can improve communication efficiency for certain collective operations.</td> <td>Configures the HCCL communication algorithm scheduling to use AIV (Ascend Intelligent Vision) expansion mode, which can improve communication efficiency for certain collective operations.</td>
<td>Used in Qwen MoE and R1 non-NSA configurations</td> <td>Used in Qwen MoE and R1 non-DSA configurations</td>
</tr> </tr>
<tr> <tr>
<td>`SGLANG_NPU_FUSED_MOE_MODE`</td> <td>`SGLANG_NPU_FUSED_MOE_MODE`</td>
docs_new/docs/hardware-platforms/ascend-npus/ascend_npu_support_features.mdx
+3 -3
View File
@@ -1225,13 +1225,13 @@ click [Server Arguments](../../advanced_features/server_arguments).
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>A2, A3</td> <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>A2, A3</td>
</tr> </tr>
<tr> <tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--nsa-prefill-backend`</td> <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--dsa-prefill-backend`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>`flashmla_sparse`</td> <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>`flashmla_sparse`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>flashmla_sparse</code>,<br/> <code>flashmla_decode</code>,<br/><code>fa3</code>,<br/> <code>tilelang</code>,<br/> <code>aiter</code></td> <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>flashmla_sparse</code>,<br/> <code>flashmla_decode</code>,<br/><code>fa3</code>,<br/> <code>tilelang</code>,<br/> <code>aiter</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Special for GPU</td> <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Special for GPU</td>
</tr> </tr>
<tr> <tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--nsa-decode-backend`</td> <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--dsa-decode-backend`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>`fa3`</td> <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>`fa3`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>flashmla_prefill</code>,<br/> <code>flashmla_kv</code>,<br/> <code>fa3</code>,<br/><code>tilelang</code>,<br/> <code>aiter</code></td> <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>flashmla_prefill</code>,<br/> <code>flashmla_kv</code>,<br/> <code>fa3</code>,<br/><code>tilelang</code>,<br/> <code>aiter</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Special for GPU</td> <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Special for GPU</td>
@@ -2226,7 +2226,7 @@ click [Server Arguments](../../advanced_features/server_arguments).
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Experimental</td> <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Experimental</td>
</tr> </tr>
<tr> <tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--enable-nsa-prefill-`<br/>`context-parallel`</td> <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--enable-dsa-prefill-`<br/>`context-parallel`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>`False`</td> <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>`False`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>bool flag<br/> (set to enable)</td> <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>bool flag<br/> (set to enable)</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>A2, A3</td> <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>A2, A3</td>
docs_new/docs/hardware-platforms/plugin.mdx
+2 -2
View File
@@ -512,9 +512,9 @@ python -c "from sglang.srt.platforms import current_platform; print(current_plat
<td>MLA KV cache pool class</td> <td>MLA KV cache pool class</td>
</tr> </tr>
<tr> <tr>
<td><code>get_nsa_kv_pool_cls()</code></td> <td><code>get_dsa_kv_pool_cls()</code></td>
<td><code>raise NotImplementedError</code></td> <td><code>raise NotImplementedError</code></td>
<td>NSA KV cache pool class (DeepSeek V3.2)</td> <td>DSA KV cache pool class (DeepSeek V3.2)</td>
</tr> </tr>
<tr> <tr>
<td><code>get_paged_allocator_cls()</code></td> <td><code>get_paged_allocator_cls()</code></td>
docs_new/docs/references/environment_variables.mdx
+7 -7
View File
@@ -393,7 +393,7 @@ SGLang supports various environment variables that can be used to configure its
</tbody> </tbody>
</table> </table>
## NSA Backend Configuration (For DeepSeek V3.2) ## DSA Backend Configuration (For DeepSeek V3.2)
{/* # Environment variable to control mtp precomputing of metadata for multi-step speculative decoding */} {/* # Environment variable to control mtp precomputing of metadata for multi-step speculative decoding */}
@@ -412,13 +412,13 @@ SGLang supports various environment variables that can be used to configure its
</thead> </thead>
<tbody> <tbody>
<tr> <tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_NSA_FUSE_TOPK</code></td> <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_DSA_FUSE_TOPK</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Fuse the operation of picking topk logits and picking topk indices from page table</td> <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Fuse the operation of picking topk logits and picking topk indices from page table. <code>SGLANG_NSA_FUSE_TOPK</code> is a deprecated alias.</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>true</code></td> <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>true</code></td>
</tr> </tr>
<tr> <tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_NSA_ENABLE_MTP_PRECOMPUTE_METADATA</code></td> <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_DSA_ENABLE_MTP_PRECOMPUTE_METADATA</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Precompute metadata that can be shared among different draft steps when MTP is enabled</td> <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Precompute metadata that can be shared among different draft steps when MTP is enabled. <code>SGLANG_NSA_ENABLE_MTP_PRECOMPUTE_METADATA</code> is a deprecated alias.</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>true</code></td> <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>true</code></td>
</tr> </tr>
<tr> <tr>
@@ -427,8 +427,8 @@ SGLang supports various environment variables that can be used to configure its
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>true</code></td> <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>true</code></td>
</tr> </tr>
<tr> <tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_NSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD</code></td> <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_DSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>When the maximum kv len in current prefill batch exceeds this value, the sparse mla kernel will be applied, else it falls back to dense MHA implementation. Default to the index topk of model (2048 for DeepSeek V3.2)</td> <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>When the maximum kv len in current prefill batch exceeds this value, the sparse mla kernel will be applied, else it falls back to dense MHA implementation. Default to the index topk of model (2048 for DeepSeek V3.2). <code>SGLANG_NSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD</code> is a deprecated alias.</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>2048</code></td> <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>2048</code></td>
</tr> </tr>
</tbody> </tbody>
docs_new/src/snippets/autoregressive/deepseek-v32-deployment.jsx
+2 -2
View File
@@ -196,8 +196,8 @@ export const DeepSeekV32Deployment = () => {
// Hardware platform specific parameters // Hardware platform specific parameters
if (isAMD) { if (isAMD) {
cmd += ' \\\n --trust-remote-code'; cmd += ' \\\n --trust-remote-code';
cmd += ' \\\n --nsa-prefill-backend tilelang'; cmd += ' \\\n --dsa-prefill-backend tilelang';
cmd += ' \\\n --nsa-decode-backend tilelang'; cmd += ' \\\n --dsa-decode-backend tilelang';
cmd += ' \\\n --cuda-graph-max-bs 64'; cmd += ' \\\n --cuda-graph-max-bs 64';
} }
docs_new/src/snippets/autoregressive/deepseek-v4-deployment.jsx
+2 -2
View File
@@ -650,8 +650,8 @@ export const DeepSeekV4Deployment = () => {
flags.push(` --tp ${tp}`); flags.push(` --tp ${tp}`);
if (multinode) flags.push(...multiNodeFlags(nnodes)); if (multinode) flags.push(...multiNodeFlags(nnodes));
flags.push(" --moe-a2a-backend deepep"); flags.push(" --moe-a2a-backend deepep");
flags.push(" --enable-nsa-prefill-context-parallel"); flags.push(" --enable-dsa-prefill-context-parallel");
flags.push(" --nsa-prefill-cp-mode round-robin-split"); flags.push(" --dsa-prefill-cp-mode round-robin-split");
flags.push(" --chunked-prefill-size 16384"); flags.push(" --chunked-prefill-size 16384");
// GB300 big CP needs higher mem-fraction-static: Pro 1.6T weights at // GB300 big CP needs higher mem-fraction-static: Pro 1.6T weights at
// tp=4 are ~224 GB/card on a 273 GB GB300, so 0.78 leaves a negative // tp=4 are ~224 GB/card on a 273 GB GB300, so 0.78 leaves a negative
docs_new/src/snippets/autoregressive/glm-5-deployment.jsx
+9 -9
View File
@@ -155,13 +155,13 @@ export const GLM5Deployment = () => {
cmd += ` --model-path ${modelName}`; cmd += ` --model-path ${modelName}`;
cmd += ` \\\n --tp ${tpValue}`; cmd += ` \\\n --tp ${tpValue}`;
// NVFP4 B200: trtllm NSA backends, flashinfer fusion, FP8 KV cache. // NVFP4 B200: trtllm DSA backends, flashinfer fusion, FP8 KV cache.
if (isNVFP4) { if (isNVFP4) {
cmd += ' \\\n --trust-remote-code'; cmd += ' \\\n --trust-remote-code';
cmd += ' \\\n --quantization modelopt_fp4'; cmd += ' \\\n --quantization modelopt_fp4';
cmd += ' \\\n --kv-cache-dtype fp8_e4m3'; cmd += ' \\\n --kv-cache-dtype fp8_e4m3';
cmd += ' \\\n --nsa-decode-backend trtllm'; cmd += ' \\\n --dsa-decode-backend trtllm';
cmd += ' \\\n --nsa-prefill-backend trtllm'; cmd += ' \\\n --dsa-prefill-backend trtllm';
cmd += ' \\\n --moe-runner-backend flashinfer_trtllm'; cmd += ' \\\n --moe-runner-backend flashinfer_trtllm';
cmd += ' \\\n --enable-flashinfer-allreduce-fusion'; cmd += ' \\\n --enable-flashinfer-allreduce-fusion';
cmd += ' \\\n --enable-dp-lm-head'; cmd += ' \\\n --enable-dp-lm-head';
@@ -174,11 +174,11 @@ export const GLM5Deployment = () => {
return cmd; return cmd;
} }
// AMD-specific: NSA tilelang backend. // AMD-specific: DSA tilelang backend.
if (isAMD) { if (isAMD) {
cmd += ' \\\n --trust-remote-code'; cmd += ' \\\n --trust-remote-code';
cmd += ' \\\n --nsa-prefill-backend tilelang'; cmd += ' \\\n --dsa-prefill-backend tilelang';
cmd += ' \\\n --nsa-decode-backend tilelang'; cmd += ' \\\n --dsa-decode-backend tilelang';
cmd += ' \\\n --chunked-prefill-size 131072'; cmd += ' \\\n --chunked-prefill-size 131072';
cmd += ' \\\n --watchdog-timeout 1200'; cmd += ' \\\n --watchdog-timeout 1200';
} }
@@ -199,9 +199,9 @@ export const GLM5Deployment = () => {
if (hardware === 'b200' && effectiveQuant === 'fp8') { if (hardware === 'b200' && effectiveQuant === 'fp8') {
cmd += ' \\\n --ep 1'; cmd += ' \\\n --ep 1';
cmd += ' \\\n --quantization fp8'; cmd += ' \\\n --quantization fp8';
cmd += ' \\\n --attention-backend nsa'; cmd += ' \\\n --attention-backend dsa';
cmd += ' \\\n --nsa-decode-backend trtllm'; cmd += ' \\\n --dsa-decode-backend trtllm';
cmd += ' \\\n --nsa-prefill-backend trtllm'; cmd += ' \\\n --dsa-prefill-backend trtllm';
cmd += ' \\\n --moe-runner-backend flashinfer_trtllm'; cmd += ' \\\n --moe-runner-backend flashinfer_trtllm';
cmd += ' \\\n --enable-flashinfer-allreduce-fusion'; cmd += ' \\\n --enable-flashinfer-allreduce-fusion';
} }
docs_new/src/snippets/autoregressive/glm-51-deployment.jsx
+2 -2
View File
@@ -152,8 +152,8 @@ export const GLM51Deployment = () => {
if (isAMD) { if (isAMD) {
cmd += ' \\\n --trust-remote-code'; cmd += ' \\\n --trust-remote-code';
cmd += ' \\\n --nsa-prefill-backend tilelang'; cmd += ' \\\n --dsa-prefill-backend tilelang';
cmd += ' \\\n --nsa-decode-backend tilelang'; cmd += ' \\\n --dsa-decode-backend tilelang';
cmd += ' \\\n --chunked-prefill-size 131072'; cmd += ' \\\n --chunked-prefill-size 131072';
cmd += ' \\\n --watchdog-timeout 1200'; cmd += ' \\\n --watchdog-timeout 1200';
} }
python/sglang/jit_kernel/csrc/nsa/fused_store_index_cache.cuh → python/sglang/jit_kernel/csrc/dsa/fused_store_index_cache.cuh
View File
python/sglang/jit_kernel/csrc/elementwise/fused_metadata_copy.cuh
+46 -46
View File
@@ -1,10 +1,10 @@
/* /*
* Fused metadata copy kernel for NSA backend CUDA graph replay. * Fused metadata copy kernel for DSA backend CUDA graph replay.
* JIT-compiled version for python/sglang/jit_kernel. * JIT-compiled version for python/sglang/jit_kernel.
* *
* OVERVIEW: * OVERVIEW:
* This kernel fuses multiple tensor copy operations (cache_seqlens, cu_seqlens_k, * This kernel fuses multiple tensor copy operations (cache_seqlens, cu_seqlens_k,
* page_table, nsa metadata, and optional FlashMLA metadata) into single kernel * page_table, dsa metadata, and optional FlashMLA metadata) into single kernel
* launches, significantly reducing kernel launch overhead and improving CUDA * launches, significantly reducing kernel launch overhead and improving CUDA
* graph replay performance during inference. * graph replay performance during inference.
* *
@@ -37,7 +37,7 @@
#include <algorithm> // for std::min #include <algorithm> // for std::min
#include <cuda_runtime.h> #include <cuda_runtime.h>
// Forward mode enum (must match Python ForwardMode in sglang/srt/layers/attention/nsa_backend.py) // Forward mode enum (must match Python ForwardMode in sglang/srt/layers/attention/dsa_backend.py)
enum ForwardModeEnum { DECODE = 0, TARGET_VERIFY = 1, DRAFT_EXTEND = 2 }; enum ForwardModeEnum { DECODE = 0, TARGET_VERIFY = 1, DRAFT_EXTEND = 2 };
/** /**
@@ -49,9 +49,9 @@ struct SourcePointers {
const int32_t* __restrict__ cache_seqlens; // [bs] sequence lengths in cache const int32_t* __restrict__ cache_seqlens; // [bs] sequence lengths in cache
const int32_t* __restrict__ cu_seqlens_k; // [bs+1] cumulative sequence lengths const int32_t* __restrict__ cu_seqlens_k; // [bs+1] cumulative sequence lengths
const int32_t* __restrict__ page_indices; // page table indices const int32_t* __restrict__ page_indices; // page table indices
const int32_t* __restrict__ nsa_cache_seqlens; // NSA-specific cache lengths const int32_t* __restrict__ dsa_cache_seqlens; // DSA-specific cache lengths
const int32_t* __restrict__ seqlens_expanded; // expanded sequence lengths (TARGET_VERIFY/DRAFT_EXTEND only) const int32_t* __restrict__ seqlens_expanded; // expanded sequence lengths (TARGET_VERIFY/DRAFT_EXTEND only)
const int32_t* __restrict__ nsa_cu_seqlens_k; // NSA cumulative sequence lengths const int32_t* __restrict__ dsa_cu_seqlens_k; // DSA cumulative sequence lengths
const int32_t* __restrict__ real_page_table; // optional real page table const int32_t* __restrict__ real_page_table; // optional real page table
const int32_t* __restrict__ flashmla_num_splits; // optional FlashMLA split counts const int32_t* __restrict__ flashmla_num_splits; // optional FlashMLA split counts
const int32_t* __restrict__ flashmla_metadata; // optional FlashMLA metadata const int32_t* __restrict__ flashmla_metadata; // optional FlashMLA metadata
@@ -66,9 +66,9 @@ struct DestinationPointers {
int32_t* __restrict__ cache_seqlens; // [bs] sequence lengths in cache int32_t* __restrict__ cache_seqlens; // [bs] sequence lengths in cache
int32_t* __restrict__ cu_seqlens_k; // [bs+1] cumulative sequence lengths int32_t* __restrict__ cu_seqlens_k; // [bs+1] cumulative sequence lengths
int32_t* __restrict__ page_table_1; // page table (note: different name from source) int32_t* __restrict__ page_table_1; // page table (note: different name from source)
int32_t* __restrict__ nsa_cache_seqlens; // NSA-specific cache lengths int32_t* __restrict__ dsa_cache_seqlens; // DSA-specific cache lengths
int32_t* __restrict__ seqlens_expanded; // expanded sequence lengths (TARGET_VERIFY/DRAFT_EXTEND only) int32_t* __restrict__ seqlens_expanded; // expanded sequence lengths (TARGET_VERIFY/DRAFT_EXTEND only)
int32_t* __restrict__ nsa_cu_seqlens_k; // NSA cumulative sequence lengths int32_t* __restrict__ dsa_cu_seqlens_k; // DSA cumulative sequence lengths
int32_t* __restrict__ real_page_table; // optional real page table int32_t* __restrict__ real_page_table; // optional real page table
int32_t* __restrict__ flashmla_num_splits; // optional FlashMLA split counts int32_t* __restrict__ flashmla_num_splits; // optional FlashMLA split counts
int32_t* __restrict__ flashmla_metadata; // optional FlashMLA metadata int32_t* __restrict__ flashmla_metadata; // optional FlashMLA metadata
@@ -189,26 +189,26 @@ __global__ void fused_metadata_copy_kernel(const FusedMetadataCopyParams __grid_
} }
} }
// Branch 3: NSA metadata copy (different loop sizes per mode) // Branch 3: DSA metadata copy (different loop sizes per mode)
if (forward_mode == 0) { // DECODE if (forward_mode == 0) { // DECODE
#pragma unroll 8 #pragma unroll 8
for (int i = tid; i < bs; i += total_threads) { for (int i = tid; i < bs; i += total_threads) {
dst.nsa_cache_seqlens[i] = src.nsa_cache_seqlens[i]; dst.dsa_cache_seqlens[i] = src.dsa_cache_seqlens[i];
} }
#pragma unroll 8 #pragma unroll 8
for (int i = tid; i < bs; i += total_threads) { for (int i = tid; i < bs; i += total_threads) {
dst.nsa_cu_seqlens_k[i + 1] = src.nsa_cu_seqlens_k[i + 1]; dst.dsa_cu_seqlens_k[i + 1] = src.dsa_cu_seqlens_k[i + 1];
} }
} else { // TARGET_VERIFY or DRAFT_EXTEND } else { // TARGET_VERIFY or DRAFT_EXTEND
#pragma unroll 4 #pragma unroll 4
for (int i = tid; i < seqlens_expanded_size; i += total_threads) { for (int i = tid; i < seqlens_expanded_size; i += total_threads) {
dst.nsa_cache_seqlens[i] = src.nsa_cache_seqlens[i]; dst.dsa_cache_seqlens[i] = src.dsa_cache_seqlens[i];
} }
#pragma unroll 4 #pragma unroll 4
for (int i = tid; i < seqlens_expanded_size; i += total_threads) { for (int i = tid; i < seqlens_expanded_size; i += total_threads) {
dst.nsa_cu_seqlens_k[i + 1] = src.nsa_cu_seqlens_k[i + 1]; dst.dsa_cu_seqlens_k[i + 1] = src.dsa_cu_seqlens_k[i + 1];
} }
} }
@@ -309,22 +309,22 @@ __global__ void fused_metadata_copy_multi_kernel(const FusedMetadataCopyMultiPar
dst2.page_table_1[row * page_table_1_stride + col] = val; dst2.page_table_1[row * page_table_1_stride + col] = val;
} }
// Copy nsa_cache_seqlens to all 3 backends // Copy dsa_cache_seqlens to all 3 backends
#pragma unroll 8 #pragma unroll 8
for (int i = tid; i < bs; i += total_threads) { for (int i = tid; i < bs; i += total_threads) {
int32_t val = src.nsa_cache_seqlens[i]; int32_t val = src.dsa_cache_seqlens[i];
dst0.nsa_cache_seqlens[i] = val; dst0.dsa_cache_seqlens[i] = val;
dst1.nsa_cache_seqlens[i] = val; dst1.dsa_cache_seqlens[i] = val;
dst2.nsa_cache_seqlens[i] = val; dst2.dsa_cache_seqlens[i] = val;
} }
// Copy NSA cu_seqlens to all 3 backends // Copy DSA cu_seqlens to all 3 backends
#pragma unroll 8 #pragma unroll 8
for (int i = tid; i < bs; i += total_threads) { for (int i = tid; i < bs; i += total_threads) {
int32_t val = src.nsa_cu_seqlens_k[i + 1]; int32_t val = src.dsa_cu_seqlens_k[i + 1];
dst0.nsa_cu_seqlens_k[i + 1] = val; dst0.dsa_cu_seqlens_k[i + 1] = val;
dst1.nsa_cu_seqlens_k[i + 1] = val; dst1.dsa_cu_seqlens_k[i + 1] = val;
dst2.nsa_cu_seqlens_k[i + 1] = val; dst2.dsa_cu_seqlens_k[i + 1] = val;
} }
// Copy real page table to all 3 backends // Copy real page table to all 3 backends
@@ -493,18 +493,18 @@ struct FusedMetadataCopyKernel {
run(const tvm::ffi::TensorView cache_seqlens_src, run(const tvm::ffi::TensorView cache_seqlens_src,
const tvm::ffi::TensorView cu_seqlens_k_src, const tvm::ffi::TensorView cu_seqlens_k_src,
const tvm::ffi::TensorView page_indices_src, const tvm::ffi::TensorView page_indices_src,
const tvm::ffi::TensorView nsa_cache_seqlens_src, const tvm::ffi::TensorView dsa_cache_seqlens_src,
const tvm::ffi::Optional<tvm::ffi::TensorView> seqlens_expanded_src, const tvm::ffi::Optional<tvm::ffi::TensorView> seqlens_expanded_src,
const tvm::ffi::TensorView nsa_cu_seqlens_k_src, const tvm::ffi::TensorView dsa_cu_seqlens_k_src,
const tvm::ffi::Optional<tvm::ffi::TensorView> real_page_table_src, const tvm::ffi::Optional<tvm::ffi::TensorView> real_page_table_src,
const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_num_splits_src, const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_num_splits_src,
const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_metadata_src, const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_metadata_src,
const tvm::ffi::TensorView cache_seqlens_dst, const tvm::ffi::TensorView cache_seqlens_dst,
const tvm::ffi::TensorView cu_seqlens_k_dst, const tvm::ffi::TensorView cu_seqlens_k_dst,
const tvm::ffi::TensorView page_table_1_dst, const tvm::ffi::TensorView page_table_1_dst,
const tvm::ffi::TensorView nsa_cache_seqlens_dst, const tvm::ffi::TensorView dsa_cache_seqlens_dst,
const tvm::ffi::Optional<tvm::ffi::TensorView> seqlens_expanded_dst, const tvm::ffi::Optional<tvm::ffi::TensorView> seqlens_expanded_dst,
const tvm::ffi::TensorView nsa_cu_seqlens_k_dst, const tvm::ffi::TensorView dsa_cu_seqlens_k_dst,
const tvm::ffi::Optional<tvm::ffi::TensorView> real_page_table_dst, const tvm::ffi::Optional<tvm::ffi::TensorView> real_page_table_dst,
const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_num_splits_dst, const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_num_splits_dst,
const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_metadata_dst, const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_metadata_dst,
@@ -522,9 +522,9 @@ struct FusedMetadataCopyKernel {
.cache_seqlens = unwrap_data_ptr<int32_t>(cache_seqlens_src, "cache_seqlens_src"), .cache_seqlens = unwrap_data_ptr<int32_t>(cache_seqlens_src, "cache_seqlens_src"),
.cu_seqlens_k = unwrap_data_ptr<int32_t>(cu_seqlens_k_src, "cu_seqlens_k_src"), .cu_seqlens_k = unwrap_data_ptr<int32_t>(cu_seqlens_k_src, "cu_seqlens_k_src"),
.page_indices = unwrap_data_ptr<int32_t>(page_indices_src, "page_indices_src"), .page_indices = unwrap_data_ptr<int32_t>(page_indices_src, "page_indices_src"),
.nsa_cache_seqlens = unwrap_data_ptr<int32_t>(nsa_cache_seqlens_src, "nsa_cache_seqlens_src"), .dsa_cache_seqlens = unwrap_data_ptr<int32_t>(dsa_cache_seqlens_src, "dsa_cache_seqlens_src"),
.seqlens_expanded = unwrap_optional_data_ptr<int32_t>(seqlens_expanded_src, "seqlens_expanded_src"), .seqlens_expanded = unwrap_optional_data_ptr<int32_t>(seqlens_expanded_src, "seqlens_expanded_src"),
.nsa_cu_seqlens_k = unwrap_data_ptr<int32_t>(nsa_cu_seqlens_k_src, "nsa_cu_seqlens_k_src"), .dsa_cu_seqlens_k = unwrap_data_ptr<int32_t>(dsa_cu_seqlens_k_src, "dsa_cu_seqlens_k_src"),
.real_page_table = unwrap_optional_data_ptr<int32_t>(real_page_table_src, "real_page_table_src"), .real_page_table = unwrap_optional_data_ptr<int32_t>(real_page_table_src, "real_page_table_src"),
.flashmla_num_splits = .flashmla_num_splits =
unwrap_optional_data_ptr<int32_t>(flashmla_num_splits_src, "flashmla_num_splits_src"), unwrap_optional_data_ptr<int32_t>(flashmla_num_splits_src, "flashmla_num_splits_src"),
@@ -535,9 +535,9 @@ struct FusedMetadataCopyKernel {
.cache_seqlens = unwrap_data_ptr_mut<int32_t>(cache_seqlens_dst, "cache_seqlens_dst"), .cache_seqlens = unwrap_data_ptr_mut<int32_t>(cache_seqlens_dst, "cache_seqlens_dst"),
.cu_seqlens_k = unwrap_data_ptr_mut<int32_t>(cu_seqlens_k_dst, "cu_seqlens_k_dst"), .cu_seqlens_k = unwrap_data_ptr_mut<int32_t>(cu_seqlens_k_dst, "cu_seqlens_k_dst"),
.page_table_1 = unwrap_data_ptr_mut<int32_t>(page_table_1_dst, "page_table_1_dst"), .page_table_1 = unwrap_data_ptr_mut<int32_t>(page_table_1_dst, "page_table_1_dst"),
.nsa_cache_seqlens = unwrap_data_ptr_mut<int32_t>(nsa_cache_seqlens_dst, "nsa_cache_seqlens_dst"), .dsa_cache_seqlens = unwrap_data_ptr_mut<int32_t>(dsa_cache_seqlens_dst, "dsa_cache_seqlens_dst"),
.seqlens_expanded = unwrap_optional_data_ptr_mut<int32_t>(seqlens_expanded_dst, "seqlens_expanded_dst"), .seqlens_expanded = unwrap_optional_data_ptr_mut<int32_t>(seqlens_expanded_dst, "seqlens_expanded_dst"),
.nsa_cu_seqlens_k = unwrap_data_ptr_mut<int32_t>(nsa_cu_seqlens_k_dst, "nsa_cu_seqlens_k_dst"), .dsa_cu_seqlens_k = unwrap_data_ptr_mut<int32_t>(dsa_cu_seqlens_k_dst, "dsa_cu_seqlens_k_dst"),
.real_page_table = unwrap_optional_data_ptr_mut<int32_t>(real_page_table_dst, "real_page_table_dst"), .real_page_table = unwrap_optional_data_ptr_mut<int32_t>(real_page_table_dst, "real_page_table_dst"),
.flashmla_num_splits = .flashmla_num_splits =
unwrap_optional_data_ptr_mut<int32_t>(flashmla_num_splits_dst, "flashmla_num_splits_dst"), unwrap_optional_data_ptr_mut<int32_t>(flashmla_num_splits_dst, "flashmla_num_splits_dst"),
@@ -605,32 +605,32 @@ struct FusedMetadataCopyMultiKernel {
run(const tvm::ffi::TensorView cache_seqlens_src, run(const tvm::ffi::TensorView cache_seqlens_src,
const tvm::ffi::TensorView cu_seqlens_k_src, const tvm::ffi::TensorView cu_seqlens_k_src,
const tvm::ffi::TensorView page_indices_src, const tvm::ffi::TensorView page_indices_src,
const tvm::ffi::TensorView nsa_cache_seqlens_src, const tvm::ffi::TensorView dsa_cache_seqlens_src,
const tvm::ffi::TensorView nsa_cu_seqlens_k_src, const tvm::ffi::TensorView dsa_cu_seqlens_k_src,
const tvm::ffi::Optional<tvm::ffi::TensorView> real_page_table_src, const tvm::ffi::Optional<tvm::ffi::TensorView> real_page_table_src,
const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_num_splits_src, const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_num_splits_src,
const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_metadata_src, const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_metadata_src,
const tvm::ffi::TensorView cache_seqlens_dst0, const tvm::ffi::TensorView cache_seqlens_dst0,
const tvm::ffi::TensorView cu_seqlens_k_dst0, const tvm::ffi::TensorView cu_seqlens_k_dst0,
const tvm::ffi::TensorView page_table_1_dst0, const tvm::ffi::TensorView page_table_1_dst0,
const tvm::ffi::TensorView nsa_cache_seqlens_dst0, const tvm::ffi::TensorView dsa_cache_seqlens_dst0,
const tvm::ffi::TensorView nsa_cu_seqlens_k_dst0, const tvm::ffi::TensorView dsa_cu_seqlens_k_dst0,
const tvm::ffi::Optional<tvm::ffi::TensorView> real_page_table_dst0, const tvm::ffi::Optional<tvm::ffi::TensorView> real_page_table_dst0,
const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_num_splits_dst0, const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_num_splits_dst0,
const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_metadata_dst0, const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_metadata_dst0,
const tvm::ffi::TensorView cache_seqlens_dst1, const tvm::ffi::TensorView cache_seqlens_dst1,
const tvm::ffi::TensorView cu_seqlens_k_dst1, const tvm::ffi::TensorView cu_seqlens_k_dst1,
const tvm::ffi::TensorView page_table_1_dst1, const tvm::ffi::TensorView page_table_1_dst1,
const tvm::ffi::TensorView nsa_cache_seqlens_dst1, const tvm::ffi::TensorView dsa_cache_seqlens_dst1,
const tvm::ffi::TensorView nsa_cu_seqlens_k_dst1, const tvm::ffi::TensorView dsa_cu_seqlens_k_dst1,
const tvm::ffi::Optional<tvm::ffi::TensorView> real_page_table_dst1, const tvm::ffi::Optional<tvm::ffi::TensorView> real_page_table_dst1,
const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_num_splits_dst1, const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_num_splits_dst1,
const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_metadata_dst1, const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_metadata_dst1,
const tvm::ffi::TensorView cache_seqlens_dst2, const tvm::ffi::TensorView cache_seqlens_dst2,
const tvm::ffi::TensorView cu_seqlens_k_dst2, const tvm::ffi::TensorView cu_seqlens_k_dst2,
const tvm::ffi::TensorView page_table_1_dst2, const tvm::ffi::TensorView page_table_1_dst2,
const tvm::ffi::TensorView nsa_cache_seqlens_dst2, const tvm::ffi::TensorView dsa_cache_seqlens_dst2,
const tvm::ffi::TensorView nsa_cu_seqlens_k_dst2, const tvm::ffi::TensorView dsa_cu_seqlens_k_dst2,
const tvm::ffi::Optional<tvm::ffi::TensorView> real_page_table_dst2, const tvm::ffi::Optional<tvm::ffi::TensorView> real_page_table_dst2,
const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_num_splits_dst2, const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_num_splits_dst2,
const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_metadata_dst2, const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_metadata_dst2,
@@ -647,9 +647,9 @@ struct FusedMetadataCopyMultiKernel {
.cache_seqlens = unwrap_data_ptr<int32_t>(cache_seqlens_src, "cache_seqlens_src"), .cache_seqlens = unwrap_data_ptr<int32_t>(cache_seqlens_src, "cache_seqlens_src"),
.cu_seqlens_k = unwrap_data_ptr<int32_t>(cu_seqlens_k_src, "cu_seqlens_k_src"), .cu_seqlens_k = unwrap_data_ptr<int32_t>(cu_seqlens_k_src, "cu_seqlens_k_src"),
.page_indices = unwrap_data_ptr<int32_t>(page_indices_src, "page_indices_src"), .page_indices = unwrap_data_ptr<int32_t>(page_indices_src, "page_indices_src"),
.nsa_cache_seqlens = unwrap_data_ptr<int32_t>(nsa_cache_seqlens_src, "nsa_cache_seqlens_src"), .dsa_cache_seqlens = unwrap_data_ptr<int32_t>(dsa_cache_seqlens_src, "dsa_cache_seqlens_src"),
.seqlens_expanded = nullptr, // Not used in multi-backend DECODE mode .seqlens_expanded = nullptr, // Not used in multi-backend DECODE mode
.nsa_cu_seqlens_k = unwrap_data_ptr<int32_t>(nsa_cu_seqlens_k_src, "nsa_cu_seqlens_k_src"), .dsa_cu_seqlens_k = unwrap_data_ptr<int32_t>(dsa_cu_seqlens_k_src, "dsa_cu_seqlens_k_src"),
.real_page_table = unwrap_optional_data_ptr<int32_t>(real_page_table_src, "real_page_table_src"), .real_page_table = unwrap_optional_data_ptr<int32_t>(real_page_table_src, "real_page_table_src"),
.flashmla_num_splits = .flashmla_num_splits =
unwrap_optional_data_ptr<int32_t>(flashmla_num_splits_src, "flashmla_num_splits_src"), unwrap_optional_data_ptr<int32_t>(flashmla_num_splits_src, "flashmla_num_splits_src"),
@@ -660,9 +660,9 @@ struct FusedMetadataCopyMultiKernel {
.cache_seqlens = unwrap_data_ptr_mut<int32_t>(cache_seqlens_dst0, "cache_seqlens_dst0"), .cache_seqlens = unwrap_data_ptr_mut<int32_t>(cache_seqlens_dst0, "cache_seqlens_dst0"),
.cu_seqlens_k = unwrap_data_ptr_mut<int32_t>(cu_seqlens_k_dst0, "cu_seqlens_k_dst0"), .cu_seqlens_k = unwrap_data_ptr_mut<int32_t>(cu_seqlens_k_dst0, "cu_seqlens_k_dst0"),
.page_table_1 = unwrap_data_ptr_mut<int32_t>(page_table_1_dst0, "page_table_1_dst0"), .page_table_1 = unwrap_data_ptr_mut<int32_t>(page_table_1_dst0, "page_table_1_dst0"),
.nsa_cache_seqlens = unwrap_data_ptr_mut<int32_t>(nsa_cache_seqlens_dst0, "nsa_cache_seqlens_dst0"), .dsa_cache_seqlens = unwrap_data_ptr_mut<int32_t>(dsa_cache_seqlens_dst0, "dsa_cache_seqlens_dst0"),
.seqlens_expanded = nullptr, .seqlens_expanded = nullptr,
.nsa_cu_seqlens_k = unwrap_data_ptr_mut<int32_t>(nsa_cu_seqlens_k_dst0, "nsa_cu_seqlens_k_dst0"), .dsa_cu_seqlens_k = unwrap_data_ptr_mut<int32_t>(dsa_cu_seqlens_k_dst0, "dsa_cu_seqlens_k_dst0"),
.real_page_table = unwrap_optional_data_ptr_mut<int32_t>(real_page_table_dst0, "real_page_table_dst0"), .real_page_table = unwrap_optional_data_ptr_mut<int32_t>(real_page_table_dst0, "real_page_table_dst0"),
.flashmla_num_splits = .flashmla_num_splits =
unwrap_optional_data_ptr_mut<int32_t>(flashmla_num_splits_dst0, "flashmla_num_splits_dst0"), unwrap_optional_data_ptr_mut<int32_t>(flashmla_num_splits_dst0, "flashmla_num_splits_dst0"),
@@ -674,9 +674,9 @@ struct FusedMetadataCopyMultiKernel {
.cache_seqlens = unwrap_data_ptr_mut<int32_t>(cache_seqlens_dst1, "cache_seqlens_dst1"), .cache_seqlens = unwrap_data_ptr_mut<int32_t>(cache_seqlens_dst1, "cache_seqlens_dst1"),
.cu_seqlens_k = unwrap_data_ptr_mut<int32_t>(cu_seqlens_k_dst1, "cu_seqlens_k_dst1"), .cu_seqlens_k = unwrap_data_ptr_mut<int32_t>(cu_seqlens_k_dst1, "cu_seqlens_k_dst1"),
.page_table_1 = unwrap_data_ptr_mut<int32_t>(page_table_1_dst1, "page_table_1_dst1"), .page_table_1 = unwrap_data_ptr_mut<int32_t>(page_table_1_dst1, "page_table_1_dst1"),
.nsa_cache_seqlens = unwrap_data_ptr_mut<int32_t>(nsa_cache_seqlens_dst1, "nsa_cache_seqlens_dst1"), .dsa_cache_seqlens = unwrap_data_ptr_mut<int32_t>(dsa_cache_seqlens_dst1, "dsa_cache_seqlens_dst1"),
.seqlens_expanded = nullptr, .seqlens_expanded = nullptr,
.nsa_cu_seqlens_k = unwrap_data_ptr_mut<int32_t>(nsa_cu_seqlens_k_dst1, "nsa_cu_seqlens_k_dst1"), .dsa_cu_seqlens_k = unwrap_data_ptr_mut<int32_t>(dsa_cu_seqlens_k_dst1, "dsa_cu_seqlens_k_dst1"),
.real_page_table = unwrap_optional_data_ptr_mut<int32_t>(real_page_table_dst1, "real_page_table_dst1"), .real_page_table = unwrap_optional_data_ptr_mut<int32_t>(real_page_table_dst1, "real_page_table_dst1"),
.flashmla_num_splits = .flashmla_num_splits =
unwrap_optional_data_ptr_mut<int32_t>(flashmla_num_splits_dst1, "flashmla_num_splits_dst1"), unwrap_optional_data_ptr_mut<int32_t>(flashmla_num_splits_dst1, "flashmla_num_splits_dst1"),
@@ -688,9 +688,9 @@ struct FusedMetadataCopyMultiKernel {
.cache_seqlens = unwrap_data_ptr_mut<int32_t>(cache_seqlens_dst2, "cache_seqlens_dst2"), .cache_seqlens = unwrap_data_ptr_mut<int32_t>(cache_seqlens_dst2, "cache_seqlens_dst2"),
.cu_seqlens_k = unwrap_data_ptr_mut<int32_t>(cu_seqlens_k_dst2, "cu_seqlens_k_dst2"), .cu_seqlens_k = unwrap_data_ptr_mut<int32_t>(cu_seqlens_k_dst2, "cu_seqlens_k_dst2"),
.page_table_1 = unwrap_data_ptr_mut<int32_t>(page_table_1_dst2, "page_table_1_dst2"), .page_table_1 = unwrap_data_ptr_mut<int32_t>(page_table_1_dst2, "page_table_1_dst2"),
.nsa_cache_seqlens = unwrap_data_ptr_mut<int32_t>(nsa_cache_seqlens_dst2, "nsa_cache_seqlens_dst2"), .dsa_cache_seqlens = unwrap_data_ptr_mut<int32_t>(dsa_cache_seqlens_dst2, "dsa_cache_seqlens_dst2"),
.seqlens_expanded = nullptr, .seqlens_expanded = nullptr,
.nsa_cu_seqlens_k = unwrap_data_ptr_mut<int32_t>(nsa_cu_seqlens_k_dst2, "nsa_cu_seqlens_k_dst2"), .dsa_cu_seqlens_k = unwrap_data_ptr_mut<int32_t>(dsa_cu_seqlens_k_dst2, "dsa_cu_seqlens_k_dst2"),
.real_page_table = unwrap_optional_data_ptr_mut<int32_t>(real_page_table_dst2, "real_page_table_dst2"), .real_page_table = unwrap_optional_data_ptr_mut<int32_t>(real_page_table_dst2, "real_page_table_dst2"),
.flashmla_num_splits = .flashmla_num_splits =
unwrap_optional_data_ptr_mut<int32_t>(flashmla_num_splits_dst2, "flashmla_num_splits_dst2"), unwrap_optional_data_ptr_mut<int32_t>(flashmla_num_splits_dst2, "flashmla_num_splits_dst2"),
python/sglang/jit_kernel/fused_metadata_copy.py
+39 -39
View File
@@ -1,5 +1,5 @@
""" """
Fused metadata copy kernel for NSA backend CUDA graph replay. Fused metadata copy kernel for DSA backend CUDA graph replay.
This module provides JIT-compiled CUDA kernels for fusing multiple tensor This module provides JIT-compiled CUDA kernels for fusing multiple tensor
copy operations into single kernel launches, reducing kernel launch overhead copy operations into single kernel launches, reducing kernel launch overhead
@@ -98,18 +98,18 @@ def fused_metadata_copy_cuda(
cache_seqlens_src: torch.Tensor, cache_seqlens_src: torch.Tensor,
cu_seqlens_k_src: torch.Tensor, cu_seqlens_k_src: torch.Tensor,
page_indices_src: torch.Tensor, page_indices_src: torch.Tensor,
nsa_cache_seqlens_src: torch.Tensor, dsa_cache_seqlens_src: torch.Tensor,
seqlens_expanded_src: Optional[torch.Tensor], seqlens_expanded_src: Optional[torch.Tensor],
nsa_cu_seqlens_k_src: torch.Tensor, dsa_cu_seqlens_k_src: torch.Tensor,
real_page_table_src: Optional[torch.Tensor], real_page_table_src: Optional[torch.Tensor],
flashmla_num_splits_src: Optional[torch.Tensor], flashmla_num_splits_src: Optional[torch.Tensor],
flashmla_metadata_src: Optional[torch.Tensor], flashmla_metadata_src: Optional[torch.Tensor],
cache_seqlens_dst: torch.Tensor, cache_seqlens_dst: torch.Tensor,
cu_seqlens_k_dst: torch.Tensor, cu_seqlens_k_dst: torch.Tensor,
page_table_1_dst: torch.Tensor, page_table_1_dst: torch.Tensor,
nsa_cache_seqlens_dst: torch.Tensor, dsa_cache_seqlens_dst: torch.Tensor,
seqlens_expanded_dst: Optional[torch.Tensor], seqlens_expanded_dst: Optional[torch.Tensor],
nsa_cu_seqlens_k_dst: torch.Tensor, dsa_cu_seqlens_k_dst: torch.Tensor,
real_page_table_dst: Optional[torch.Tensor], real_page_table_dst: Optional[torch.Tensor],
flashmla_num_splits_dst: Optional[torch.Tensor], flashmla_num_splits_dst: Optional[torch.Tensor],
flashmla_metadata_dst: Optional[torch.Tensor], flashmla_metadata_dst: Optional[torch.Tensor],
@@ -120,7 +120,7 @@ def fused_metadata_copy_cuda(
seqlens_expanded_size: int, seqlens_expanded_size: int,
) -> None: ) -> None:
""" """
Fused metadata copy kernel for NSA backend CUDA graph replay. Fused metadata copy kernel for DSA backend CUDA graph replay.
This function fuses multiple tensor copy operations into a single kernel launch, This function fuses multiple tensor copy operations into a single kernel launch,
reducing kernel launch overhead and improving performance. reducing kernel launch overhead and improving performance.
@@ -129,18 +129,18 @@ def fused_metadata_copy_cuda(
cache_seqlens_src: Source cache sequence lengths [bs] cache_seqlens_src: Source cache sequence lengths [bs]
cu_seqlens_k_src: Source cumulative sequence lengths [bs+1] cu_seqlens_k_src: Source cumulative sequence lengths [bs+1]
page_indices_src: Source page indices [rows, max_len] page_indices_src: Source page indices [rows, max_len]
nsa_cache_seqlens_src: Source NSA cache sequence lengths [size] dsa_cache_seqlens_src: Source DSA cache sequence lengths [size]
seqlens_expanded_src: Optional source expanded sequence lengths [size] (required for TARGET_VERIFY/DRAFT_EXTEND) seqlens_expanded_src: Optional source expanded sequence lengths [size] (required for TARGET_VERIFY/DRAFT_EXTEND)
nsa_cu_seqlens_k_src: Source NSA cumulative sequence lengths [size+1] dsa_cu_seqlens_k_src: Source DSA cumulative sequence lengths [size+1]
real_page_table_src: Optional source real page table [rows, cols] real_page_table_src: Optional source real page table [rows, cols]
flashmla_num_splits_src: Optional source FlashMLA num_splits [size+1] flashmla_num_splits_src: Optional source FlashMLA num_splits [size+1]
flashmla_metadata_src: Optional source FlashMLA metadata tensor flashmla_metadata_src: Optional source FlashMLA metadata tensor
cache_seqlens_dst: Destination cache sequence lengths [bs] cache_seqlens_dst: Destination cache sequence lengths [bs]
cu_seqlens_k_dst: Destination cumulative sequence lengths [bs+1] cu_seqlens_k_dst: Destination cumulative sequence lengths [bs+1]
page_table_1_dst: Destination page table [rows, stride] page_table_1_dst: Destination page table [rows, stride]
nsa_cache_seqlens_dst: Destination NSA cache sequence lengths [size] dsa_cache_seqlens_dst: Destination DSA cache sequence lengths [size]
seqlens_expanded_dst: Optional destination expanded sequence lengths [size] (required for TARGET_VERIFY/DRAFT_EXTEND) seqlens_expanded_dst: Optional destination expanded sequence lengths [size] (required for TARGET_VERIFY/DRAFT_EXTEND)
nsa_cu_seqlens_k_dst: Destination NSA cumulative sequence lengths [size+1] dsa_cu_seqlens_k_dst: Destination DSA cumulative sequence lengths [size+1]
real_page_table_dst: Optional destination real page table [rows, cols] real_page_table_dst: Optional destination real page table [rows, cols]
flashmla_num_splits_dst: Optional destination FlashMLA num_splits [size+1] flashmla_num_splits_dst: Optional destination FlashMLA num_splits [size+1]
flashmla_metadata_dst: Optional destination FlashMLA metadata tensor flashmla_metadata_dst: Optional destination FlashMLA metadata tensor
@@ -164,28 +164,28 @@ def fused_metadata_copy_cuda(
cache_seqlens_src = cache_seqlens_src.contiguous() cache_seqlens_src = cache_seqlens_src.contiguous()
cu_seqlens_k_src = cu_seqlens_k_src.contiguous() cu_seqlens_k_src = cu_seqlens_k_src.contiguous()
page_indices_src = page_indices_src.contiguous() page_indices_src = page_indices_src.contiguous()
nsa_cache_seqlens_src = nsa_cache_seqlens_src.contiguous() dsa_cache_seqlens_src = dsa_cache_seqlens_src.contiguous()
if seqlens_expanded_src is not None: if seqlens_expanded_src is not None:
seqlens_expanded_src = seqlens_expanded_src.contiguous() seqlens_expanded_src = seqlens_expanded_src.contiguous()
nsa_cu_seqlens_k_src = nsa_cu_seqlens_k_src.contiguous() dsa_cu_seqlens_k_src = dsa_cu_seqlens_k_src.contiguous()
# Call JIT-compiled kernel (None values are passed as Optional with no value) # Call JIT-compiled kernel (None values are passed as Optional with no value)
module.fused_metadata_copy( module.fused_metadata_copy(
cache_seqlens_src, cache_seqlens_src,
cu_seqlens_k_src, cu_seqlens_k_src,
page_indices_src, page_indices_src,
nsa_cache_seqlens_src, dsa_cache_seqlens_src,
seqlens_expanded_src, seqlens_expanded_src,
nsa_cu_seqlens_k_src, dsa_cu_seqlens_k_src,
real_page_table_src, real_page_table_src,
flashmla_num_splits_src, flashmla_num_splits_src,
flashmla_metadata_src, flashmla_metadata_src,
cache_seqlens_dst, cache_seqlens_dst,
cu_seqlens_k_dst, cu_seqlens_k_dst,
page_table_1_dst, page_table_1_dst,
nsa_cache_seqlens_dst, dsa_cache_seqlens_dst,
seqlens_expanded_dst, seqlens_expanded_dst,
nsa_cu_seqlens_k_dst, dsa_cu_seqlens_k_dst,
real_page_table_dst, real_page_table_dst,
flashmla_num_splits_dst, flashmla_num_splits_dst,
flashmla_metadata_dst, flashmla_metadata_dst,
@@ -200,32 +200,32 @@ def fused_metadata_copy_multi_cuda(
cache_seqlens_src: torch.Tensor, cache_seqlens_src: torch.Tensor,
cu_seqlens_k_src: torch.Tensor, cu_seqlens_k_src: torch.Tensor,
page_indices_src: torch.Tensor, page_indices_src: torch.Tensor,
nsa_cache_seqlens_src: torch.Tensor, dsa_cache_seqlens_src: torch.Tensor,
nsa_cu_seqlens_k_src: torch.Tensor, dsa_cu_seqlens_k_src: torch.Tensor,
real_page_table_src: Optional[torch.Tensor], real_page_table_src: Optional[torch.Tensor],
flashmla_num_splits_src: Optional[torch.Tensor], flashmla_num_splits_src: Optional[torch.Tensor],
flashmla_metadata_src: Optional[torch.Tensor], flashmla_metadata_src: Optional[torch.Tensor],
cache_seqlens_dst0: torch.Tensor, cache_seqlens_dst0: torch.Tensor,
cu_seqlens_k_dst0: torch.Tensor, cu_seqlens_k_dst0: torch.Tensor,
page_table_1_dst0: torch.Tensor, page_table_1_dst0: torch.Tensor,
nsa_cache_seqlens_dst0: torch.Tensor, dsa_cache_seqlens_dst0: torch.Tensor,
nsa_cu_seqlens_k_dst0: torch.Tensor, dsa_cu_seqlens_k_dst0: torch.Tensor,
real_page_table_dst0: Optional[torch.Tensor], real_page_table_dst0: Optional[torch.Tensor],
flashmla_num_splits_dst0: Optional[torch.Tensor], flashmla_num_splits_dst0: Optional[torch.Tensor],
flashmla_metadata_dst0: Optional[torch.Tensor], flashmla_metadata_dst0: Optional[torch.Tensor],
cache_seqlens_dst1: torch.Tensor, cache_seqlens_dst1: torch.Tensor,
cu_seqlens_k_dst1: torch.Tensor, cu_seqlens_k_dst1: torch.Tensor,
page_table_1_dst1: torch.Tensor, page_table_1_dst1: torch.Tensor,
nsa_cache_seqlens_dst1: torch.Tensor, dsa_cache_seqlens_dst1: torch.Tensor,
nsa_cu_seqlens_k_dst1: torch.Tensor, dsa_cu_seqlens_k_dst1: torch.Tensor,
real_page_table_dst1: Optional[torch.Tensor], real_page_table_dst1: Optional[torch.Tensor],
flashmla_num_splits_dst1: Optional[torch.Tensor], flashmla_num_splits_dst1: Optional[torch.Tensor],
flashmla_metadata_dst1: Optional[torch.Tensor], flashmla_metadata_dst1: Optional[torch.Tensor],
cache_seqlens_dst2: torch.Tensor, cache_seqlens_dst2: torch.Tensor,
cu_seqlens_k_dst2: torch.Tensor, cu_seqlens_k_dst2: torch.Tensor,
page_table_1_dst2: torch.Tensor, page_table_1_dst2: torch.Tensor,
nsa_cache_seqlens_dst2: torch.Tensor, dsa_cache_seqlens_dst2: torch.Tensor,
nsa_cu_seqlens_k_dst2: torch.Tensor, dsa_cu_seqlens_k_dst2: torch.Tensor,
real_page_table_dst2: Optional[torch.Tensor], real_page_table_dst2: Optional[torch.Tensor],
flashmla_num_splits_dst2: Optional[torch.Tensor], flashmla_num_splits_dst2: Optional[torch.Tensor],
flashmla_metadata_dst2: Optional[torch.Tensor], flashmla_metadata_dst2: Optional[torch.Tensor],
@@ -234,7 +234,7 @@ def fused_metadata_copy_multi_cuda(
seqlens_expanded_size: int, seqlens_expanded_size: int,
) -> None: ) -> None:
""" """
Multi-backend fused metadata copy kernel for NSA backend CUDA graph replay. Multi-backend fused metadata copy kernel for DSA backend CUDA graph replay.
This function copies metadata from one source to THREE destinations in a single This function copies metadata from one source to THREE destinations in a single
kernel launch, eliminating the overhead of 3 separate kernel calls. Currently kernel launch, eliminating the overhead of 3 separate kernel calls. Currently
@@ -244,16 +244,16 @@ def fused_metadata_copy_multi_cuda(
cache_seqlens_src: Source cache sequence lengths [bs] cache_seqlens_src: Source cache sequence lengths [bs]
cu_seqlens_k_src: Source cumulative sequence lengths [bs+1] cu_seqlens_k_src: Source cumulative sequence lengths [bs+1]
page_indices_src: Source page indices [bs, max_len] page_indices_src: Source page indices [bs, max_len]
nsa_cache_seqlens_src: Source NSA cache sequence lengths [bs] dsa_cache_seqlens_src: Source DSA cache sequence lengths [bs]
nsa_cu_seqlens_k_src: Source NSA cumulative sequence lengths [bs+1] dsa_cu_seqlens_k_src: Source DSA cumulative sequence lengths [bs+1]
real_page_table_src: Optional source real page table [bs, cols] real_page_table_src: Optional source real page table [bs, cols]
flashmla_num_splits_src: Optional source FlashMLA num_splits [bs+1] flashmla_num_splits_src: Optional source FlashMLA num_splits [bs+1]
flashmla_metadata_src: Optional source FlashMLA metadata tensor flashmla_metadata_src: Optional source FlashMLA metadata tensor
cache_seqlens_dst0-2: Destination cache sequence lengths for backends 0-2 cache_seqlens_dst0-2: Destination cache sequence lengths for backends 0-2
cu_seqlens_k_dst0-2: Destination cumulative sequence lengths for backends 0-2 cu_seqlens_k_dst0-2: Destination cumulative sequence lengths for backends 0-2
page_table_1_dst0-2: Destination page tables for backends 0-2 page_table_1_dst0-2: Destination page tables for backends 0-2
nsa_cache_seqlens_dst0-2: Destination NSA cache sequence lengths for backends 0-2 dsa_cache_seqlens_dst0-2: Destination DSA cache sequence lengths for backends 0-2
nsa_cu_seqlens_k_dst0-2: Destination NSA cumulative sequence lengths for backends 0-2 dsa_cu_seqlens_k_dst0-2: Destination DSA cumulative sequence lengths for backends 0-2
real_page_table_dst0-2: Optional destination real page tables for backends 0-2 real_page_table_dst0-2: Optional destination real page tables for backends 0-2
flashmla_num_splits_dst0-2: Optional destination FlashMLA num_splits for backends 0-2 flashmla_num_splits_dst0-2: Optional destination FlashMLA num_splits for backends 0-2
flashmla_metadata_dst0-2: Optional destination FlashMLA metadata tensors for backends 0-2 flashmla_metadata_dst0-2: Optional destination FlashMLA metadata tensors for backends 0-2
@@ -273,40 +273,40 @@ def fused_metadata_copy_multi_cuda(
cache_seqlens_src = cache_seqlens_src.contiguous() cache_seqlens_src = cache_seqlens_src.contiguous()
cu_seqlens_k_src = cu_seqlens_k_src.contiguous() cu_seqlens_k_src = cu_seqlens_k_src.contiguous()
page_indices_src = page_indices_src.contiguous() page_indices_src = page_indices_src.contiguous()
nsa_cache_seqlens_src = nsa_cache_seqlens_src.contiguous() dsa_cache_seqlens_src = dsa_cache_seqlens_src.contiguous()
nsa_cu_seqlens_k_src = nsa_cu_seqlens_k_src.contiguous() dsa_cu_seqlens_k_src = dsa_cu_seqlens_k_src.contiguous()
# Call JIT-compiled kernel (None values are passed as Optional with no value) # Call JIT-compiled kernel (None values are passed as Optional with no value)
module.fused_metadata_copy_multi( module.fused_metadata_copy_multi(
cache_seqlens_src, cache_seqlens_src,
cu_seqlens_k_src, cu_seqlens_k_src,
page_indices_src, page_indices_src,
nsa_cache_seqlens_src, dsa_cache_seqlens_src,
nsa_cu_seqlens_k_src, dsa_cu_seqlens_k_src,
real_page_table_src, real_page_table_src,
flashmla_num_splits_src, flashmla_num_splits_src,
flashmla_metadata_src, flashmla_metadata_src,
cache_seqlens_dst0, cache_seqlens_dst0,
cu_seqlens_k_dst0, cu_seqlens_k_dst0,
page_table_1_dst0, page_table_1_dst0,
nsa_cache_seqlens_dst0, dsa_cache_seqlens_dst0,
nsa_cu_seqlens_k_dst0, dsa_cu_seqlens_k_dst0,
real_page_table_dst0, real_page_table_dst0,
flashmla_num_splits_dst0, flashmla_num_splits_dst0,
flashmla_metadata_dst0, flashmla_metadata_dst0,
cache_seqlens_dst1, cache_seqlens_dst1,
cu_seqlens_k_dst1, cu_seqlens_k_dst1,
page_table_1_dst1, page_table_1_dst1,
nsa_cache_seqlens_dst1, dsa_cache_seqlens_dst1,
nsa_cu_seqlens_k_dst1, dsa_cu_seqlens_k_dst1,
real_page_table_dst1, real_page_table_dst1,
flashmla_num_splits_dst1, flashmla_num_splits_dst1,
flashmla_metadata_dst1, flashmla_metadata_dst1,
cache_seqlens_dst2, cache_seqlens_dst2,
cu_seqlens_k_dst2, cu_seqlens_k_dst2,
page_table_1_dst2, page_table_1_dst2,
nsa_cache_seqlens_dst2, dsa_cache_seqlens_dst2,
nsa_cu_seqlens_k_dst2, dsa_cu_seqlens_k_dst2,
real_page_table_dst2, real_page_table_dst2,
flashmla_num_splits_dst2, flashmla_num_splits_dst2,
flashmla_metadata_dst2, flashmla_metadata_dst2,
python/sglang/jit_kernel/fused_store_index_cache.py
+8 -8
View File
@@ -28,7 +28,7 @@ logger = logging.getLogger(__name__)
@cache_once @cache_once
def _jit_nsa_fused_store_module( def _jit_dsa_fused_store_module(
key_dtype: torch.dtype, indices_dtype: torch.dtype, page_size: int key_dtype: torch.dtype, indices_dtype: torch.dtype, page_size: int
) -> Module: ) -> Module:
""" """
@@ -39,13 +39,13 @@ def _jit_nsa_fused_store_module(
return load_jit( return load_jit(
"fused_store_index_k_cache", "fused_store_index_k_cache",
*args, *args,
cuda_files=["nsa/fused_store_index_cache.cuh"], cuda_files=["dsa/fused_store_index_cache.cuh"],
cuda_wrappers=[ cuda_wrappers=[
( (
"fused_store_index_k_cache", "fused_store_index_k_cache",
# - Float = bf16_t (sgl_kernel/type.cuh) # - Float = bf16_t (sgl_kernel/type.cuh)
# - IndicesT = int64_t (out_cache_loc is int64 in SGLang SetKAndS) # - IndicesT = int64_t (out_cache_loc is int64 in SGLang SetKAndS)
# - kPageSize = 64 (CUDA NSA) # - kPageSize = 64 (CUDA DSA)
f"FusedStoreCacheIndexerKernel<{args}>::run", f"FusedStoreCacheIndexerKernel<{args}>::run",
) )
], ],
@@ -53,15 +53,15 @@ def _jit_nsa_fused_store_module(
@cache_once @cache_once
def can_use_nsa_fused_store( def can_use_dsa_fused_store(
key_dtype: torch.dtype, indices_dtype: torch.dtype, page_size: int key_dtype: torch.dtype, indices_dtype: torch.dtype, page_size: int
) -> bool: ) -> bool:
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
try: try:
_jit_nsa_fused_store_module(key_dtype, indices_dtype, page_size) _jit_dsa_fused_store_module(key_dtype, indices_dtype, page_size)
return True return True
except Exception as e: except Exception as e:
logger.warning(f"Failed to load nsa fused store JIT kernel: {e}") logger.warning(f"Failed to load dsa fused store JIT kernel: {e}")
return False return False
@@ -73,7 +73,7 @@ def fused_store_index_k_cache(
page_size: int = 64, page_size: int = 64,
) -> None: ) -> None:
""" """
Fused: quantize bf16 key (N,128) -> fp8 + fp32 scale and write into NSATokenToKVPool.index_k_with_scale_buffer. Fused: quantize bf16 key (N,128) -> fp8 + fp32 scale and write into DSATokenToKVPool.index_k_with_scale_buffer.
key: (num_tokens, 128) bf16 (or reshapeable to it) key: (num_tokens, 128) bf16 (or reshapeable to it)
index_k_with_scale: (num_pages, 64*(128+4)) uint8 index_k_with_scale: (num_pages, 64*(128+4)) uint8
@@ -101,5 +101,5 @@ def fused_store_index_k_cache(
if not index_k_with_scale.is_contiguous(): if not index_k_with_scale.is_contiguous():
index_k_with_scale = index_k_with_scale.contiguous() index_k_with_scale = index_k_with_scale.contiguous()
module = _jit_nsa_fused_store_module(key.dtype, out_cache_loc.dtype, page_size) module = _jit_dsa_fused_store_module(key.dtype, out_cache_loc.dtype, page_size)
module.fused_store_index_k_cache(key, index_k_with_scale, out_cache_loc) module.fused_store_index_k_cache(key, index_k_with_scale, out_cache_loc)
python/sglang/jit_kernel/tests/test_fused_metadata_copy.py
+102 -102
View File
@@ -33,7 +33,7 @@ def create_test_metadata(
has_flashmla: bool = False, has_flashmla: bool = False,
device: str = "cuda", device: str = "cuda",
): ):
"""Create test metadata tensors matching NSA backend structure.""" """Create test metadata tensors matching DSA backend structure."""
# Basic tensors (always present) # Basic tensors (always present)
cache_seqlens_src = torch.randint( cache_seqlens_src = torch.randint(
1, max_len, (bs,), dtype=torch.int32, device=device 1, max_len, (bs,), dtype=torch.int32, device=device
@@ -44,28 +44,28 @@ def create_test_metadata(
page_indices_src = torch.randint( page_indices_src = torch.randint(
0, 1000, (bs, max_len), dtype=torch.int32, device=device 0, 1000, (bs, max_len), dtype=torch.int32, device=device
) )
nsa_cache_seqlens_src = torch.randint( dsa_cache_seqlens_src = torch.randint(
1, max_len, (seqlens_expanded_size,), dtype=torch.int32, device=device 1, max_len, (seqlens_expanded_size,), dtype=torch.int32, device=device
) )
seqlens_expanded_src = torch.randint( seqlens_expanded_src = torch.randint(
1, max_seqlen_k, (seqlens_expanded_size,), dtype=torch.int32, device=device 1, max_seqlen_k, (seqlens_expanded_size,), dtype=torch.int32, device=device
) )
nsa_cu_seqlens_k_src = torch.zeros( dsa_cu_seqlens_k_src = torch.zeros(
seqlens_expanded_size + 1, dtype=torch.int32, device=device seqlens_expanded_size + 1, dtype=torch.int32, device=device
) )
nsa_cu_seqlens_k_src[1:] = torch.cumsum(nsa_cache_seqlens_src, dim=0) dsa_cu_seqlens_k_src[1:] = torch.cumsum(dsa_cache_seqlens_src, dim=0)
# Destination tensors # Destination tensors
cache_seqlens_dst = torch.zeros(bs, dtype=torch.int32, device=device) cache_seqlens_dst = torch.zeros(bs, dtype=torch.int32, device=device)
cu_seqlens_k_dst = torch.zeros(bs + 1, dtype=torch.int32, device=device) cu_seqlens_k_dst = torch.zeros(bs + 1, dtype=torch.int32, device=device)
page_table_1_dst = torch.zeros((bs, max_len + 16), dtype=torch.int32, device=device) page_table_1_dst = torch.zeros((bs, max_len + 16), dtype=torch.int32, device=device)
nsa_cache_seqlens_dst = torch.zeros( dsa_cache_seqlens_dst = torch.zeros(
seqlens_expanded_size, dtype=torch.int32, device=device seqlens_expanded_size, dtype=torch.int32, device=device
) )
nsa_seqlens_expanded_dst = torch.zeros( dsa_seqlens_expanded_dst = torch.zeros(
seqlens_expanded_size, dtype=torch.int32, device=device seqlens_expanded_size, dtype=torch.int32, device=device
) )
nsa_cu_seqlens_k_dst = torch.zeros( dsa_cu_seqlens_k_dst = torch.zeros(
seqlens_expanded_size + 1, dtype=torch.int32, device=device seqlens_expanded_size + 1, dtype=torch.int32, device=device
) )
@@ -107,9 +107,9 @@ def create_test_metadata(
"cache_seqlens": cache_seqlens_src, "cache_seqlens": cache_seqlens_src,
"cu_seqlens_k": cu_seqlens_k_src, "cu_seqlens_k": cu_seqlens_k_src,
"page_indices": page_indices_src, "page_indices": page_indices_src,
"nsa_cache_seqlens": nsa_cache_seqlens_src, "dsa_cache_seqlens": dsa_cache_seqlens_src,
"seqlens_expanded": seqlens_expanded_src, "seqlens_expanded": seqlens_expanded_src,
"nsa_cu_seqlens_k": nsa_cu_seqlens_k_src, "dsa_cu_seqlens_k": dsa_cu_seqlens_k_src,
"real_page_table": real_page_table_src, "real_page_table": real_page_table_src,
"flashmla_num_splits": flashmla_num_splits_src, "flashmla_num_splits": flashmla_num_splits_src,
"flashmla_metadata": flashmla_metadata_src, "flashmla_metadata": flashmla_metadata_src,
@@ -118,9 +118,9 @@ def create_test_metadata(
"cache_seqlens": cache_seqlens_dst, "cache_seqlens": cache_seqlens_dst,
"cu_seqlens_k": cu_seqlens_k_dst, "cu_seqlens_k": cu_seqlens_k_dst,
"page_table_1": page_table_1_dst, "page_table_1": page_table_1_dst,
"nsa_cache_seqlens": nsa_cache_seqlens_dst, "dsa_cache_seqlens": dsa_cache_seqlens_dst,
"nsa_seqlens_expanded": nsa_seqlens_expanded_dst, "dsa_seqlens_expanded": dsa_seqlens_expanded_dst,
"nsa_cu_seqlens_k": nsa_cu_seqlens_k_dst, "dsa_cu_seqlens_k": dsa_cu_seqlens_k_dst,
"real_page_table": real_page_table_dst, "real_page_table": real_page_table_dst,
"flashmla_num_splits": flashmla_num_splits_dst, "flashmla_num_splits": flashmla_num_splits_dst,
"flashmla_metadata": flashmla_metadata_dst, "flashmla_metadata": flashmla_metadata_dst,
@@ -134,8 +134,8 @@ def reference_copy_decode(src, dst, max_len):
dst["cache_seqlens"].copy_(src["cache_seqlens"]) dst["cache_seqlens"].copy_(src["cache_seqlens"])
dst["cu_seqlens_k"][1:].copy_(src["cu_seqlens_k"][1:]) dst["cu_seqlens_k"][1:].copy_(src["cu_seqlens_k"][1:])
dst["page_table_1"][:, :max_len].copy_(src["page_indices"]) dst["page_table_1"][:, :max_len].copy_(src["page_indices"])
dst["nsa_cache_seqlens"].copy_(src["nsa_cache_seqlens"]) dst["dsa_cache_seqlens"].copy_(src["dsa_cache_seqlens"])
dst["nsa_cu_seqlens_k"][1 : bs + 1].copy_(src["nsa_cu_seqlens_k"][1 : bs + 1]) dst["dsa_cu_seqlens_k"][1 : bs + 1].copy_(src["dsa_cu_seqlens_k"][1 : bs + 1])
if src["real_page_table"] is not None: if src["real_page_table"] is not None:
rows, cols = src["real_page_table"].shape rows, cols = src["real_page_table"].shape
@@ -159,10 +159,10 @@ def reference_copy_target_verify(src, dst, max_seqlen_k, seqlens_expanded_size):
rows, cols = src["page_indices"].shape rows, cols = src["page_indices"].shape
dst["page_table_1"][:rows, :cols].copy_(src["page_indices"]) dst["page_table_1"][:rows, :cols].copy_(src["page_indices"])
dst["nsa_seqlens_expanded"][:seqlens_expanded_size].copy_(src["seqlens_expanded"]) dst["dsa_seqlens_expanded"][:seqlens_expanded_size].copy_(src["seqlens_expanded"])
dst["nsa_cache_seqlens"][:seqlens_expanded_size].copy_(src["nsa_cache_seqlens"]) dst["dsa_cache_seqlens"][:seqlens_expanded_size].copy_(src["dsa_cache_seqlens"])
dst["nsa_cu_seqlens_k"][1 : seqlens_expanded_size + 1].copy_( dst["dsa_cu_seqlens_k"][1 : seqlens_expanded_size + 1].copy_(
src["nsa_cu_seqlens_k"][1 : seqlens_expanded_size + 1] src["dsa_cu_seqlens_k"][1 : seqlens_expanded_size + 1]
) )
if src["real_page_table"] is not None: if src["real_page_table"] is not None:
@@ -187,10 +187,10 @@ def reference_copy_draft_extend(src, dst, max_seqlen_k, seqlens_expanded_size):
rows, cols = src["page_indices"].shape rows, cols = src["page_indices"].shape
dst["page_table_1"][:rows, :cols].copy_(src["page_indices"]) dst["page_table_1"][:rows, :cols].copy_(src["page_indices"])
dst["nsa_seqlens_expanded"][:seqlens_expanded_size].copy_(src["seqlens_expanded"]) dst["dsa_seqlens_expanded"][:seqlens_expanded_size].copy_(src["seqlens_expanded"])
dst["nsa_cache_seqlens"][:seqlens_expanded_size].copy_(src["nsa_cache_seqlens"]) dst["dsa_cache_seqlens"][:seqlens_expanded_size].copy_(src["dsa_cache_seqlens"])
dst["nsa_cu_seqlens_k"][1 : seqlens_expanded_size + 1].copy_( dst["dsa_cu_seqlens_k"][1 : seqlens_expanded_size + 1].copy_(
src["nsa_cu_seqlens_k"][1 : seqlens_expanded_size + 1] src["dsa_cu_seqlens_k"][1 : seqlens_expanded_size + 1]
) )
if src["real_page_table"] is not None: if src["real_page_table"] is not None:
@@ -233,13 +233,13 @@ def test_fused_metadata_copy_dtype_validation():
page_indices_src = torch.randint( page_indices_src = torch.randint(
0, 1000, (bs, max_len), dtype=torch.int32, device=device 0, 1000, (bs, max_len), dtype=torch.int32, device=device
) )
nsa_cache_seqlens_src = torch.randint( dsa_cache_seqlens_src = torch.randint(
1, max_len, (seqlens_expanded_size,), dtype=torch.int32, device=device 1, max_len, (seqlens_expanded_size,), dtype=torch.int32, device=device
) )
seqlens_expanded_src = torch.randint( seqlens_expanded_src = torch.randint(
1, max_seqlen_k, (seqlens_expanded_size,), dtype=torch.int32, device=device 1, max_seqlen_k, (seqlens_expanded_size,), dtype=torch.int32, device=device
) )
nsa_cu_seqlens_k_src = torch.zeros( dsa_cu_seqlens_k_src = torch.zeros(
seqlens_expanded_size + 1, dtype=torch.int32, device=device seqlens_expanded_size + 1, dtype=torch.int32, device=device
) )
@@ -247,13 +247,13 @@ def test_fused_metadata_copy_dtype_validation():
cache_seqlens_dst = torch.zeros(bs, dtype=torch.int32, device=device) cache_seqlens_dst = torch.zeros(bs, dtype=torch.int32, device=device)
cu_seqlens_k_dst = torch.zeros(bs + 1, dtype=torch.int32, device=device) cu_seqlens_k_dst = torch.zeros(bs + 1, dtype=torch.int32, device=device)
page_table_1_dst = torch.zeros((bs, max_len + 16), dtype=torch.int32, device=device) page_table_1_dst = torch.zeros((bs, max_len + 16), dtype=torch.int32, device=device)
nsa_cache_seqlens_dst = torch.zeros( dsa_cache_seqlens_dst = torch.zeros(
seqlens_expanded_size, dtype=torch.int32, device=device seqlens_expanded_size, dtype=torch.int32, device=device
) )
nsa_seqlens_expanded_dst = torch.zeros( dsa_seqlens_expanded_dst = torch.zeros(
seqlens_expanded_size, dtype=torch.int32, device=device seqlens_expanded_size, dtype=torch.int32, device=device
) )
nsa_cu_seqlens_k_dst = torch.zeros( dsa_cu_seqlens_k_dst = torch.zeros(
seqlens_expanded_size + 1, dtype=torch.int32, device=device seqlens_expanded_size + 1, dtype=torch.int32, device=device
) )
@@ -263,18 +263,18 @@ def test_fused_metadata_copy_dtype_validation():
cache_seqlens_src_wrong, # Wrong dtype: int64 cache_seqlens_src_wrong, # Wrong dtype: int64
cu_seqlens_k_src, cu_seqlens_k_src,
page_indices_src, page_indices_src,
nsa_cache_seqlens_src, dsa_cache_seqlens_src,
seqlens_expanded_src, seqlens_expanded_src,
nsa_cu_seqlens_k_src, dsa_cu_seqlens_k_src,
None, # real_page_table_src None, # real_page_table_src
None, # flashmla_num_splits_src None, # flashmla_num_splits_src
None, # flashmla_metadata_src None, # flashmla_metadata_src
cache_seqlens_dst, cache_seqlens_dst,
cu_seqlens_k_dst, cu_seqlens_k_dst,
page_table_1_dst, page_table_1_dst,
nsa_cache_seqlens_dst, dsa_cache_seqlens_dst,
nsa_seqlens_expanded_dst, dsa_seqlens_expanded_dst,
nsa_cu_seqlens_k_dst, dsa_cu_seqlens_k_dst,
None, # real_page_table_dst None, # real_page_table_dst
None, # flashmla_num_splits_dst None, # flashmla_num_splits_dst
None, # flashmla_metadata_dst None, # flashmla_metadata_dst
@@ -296,18 +296,18 @@ def test_fused_metadata_copy_dtype_validation():
cache_seqlens_src, cache_seqlens_src,
cu_seqlens_k_src, cu_seqlens_k_src,
page_indices_src, page_indices_src,
nsa_cache_seqlens_src, dsa_cache_seqlens_src,
seqlens_expanded_src, seqlens_expanded_src,
nsa_cu_seqlens_k_src, dsa_cu_seqlens_k_src,
None, None,
None, None,
None, None,
cache_seqlens_dst_wrong, # Wrong dtype: int64 cache_seqlens_dst_wrong, # Wrong dtype: int64
cu_seqlens_k_dst, cu_seqlens_k_dst,
page_table_1_dst, page_table_1_dst,
nsa_cache_seqlens_dst, dsa_cache_seqlens_dst,
nsa_seqlens_expanded_dst, dsa_seqlens_expanded_dst,
nsa_cu_seqlens_k_dst, dsa_cu_seqlens_k_dst,
None, None,
None, None,
None, None,
@@ -369,18 +369,18 @@ def test_fused_metadata_copy(bs, forward_mode, has_real_page_table, has_flashmla
data["src"]["cache_seqlens"], data["src"]["cache_seqlens"],
data["src"]["cu_seqlens_k"], data["src"]["cu_seqlens_k"],
data["src"]["page_indices"], data["src"]["page_indices"],
data["src"]["nsa_cache_seqlens"], data["src"]["dsa_cache_seqlens"],
data["src"]["seqlens_expanded"], data["src"]["seqlens_expanded"],
data["src"]["nsa_cu_seqlens_k"], data["src"]["dsa_cu_seqlens_k"],
data["src"]["real_page_table"], data["src"]["real_page_table"],
data["src"]["flashmla_num_splits"], data["src"]["flashmla_num_splits"],
data["src"]["flashmla_metadata"], data["src"]["flashmla_metadata"],
dst_fused["cache_seqlens"], dst_fused["cache_seqlens"],
dst_fused["cu_seqlens_k"], dst_fused["cu_seqlens_k"],
dst_fused["page_table_1"], dst_fused["page_table_1"],
dst_fused["nsa_cache_seqlens"], dst_fused["dsa_cache_seqlens"],
dst_fused["nsa_seqlens_expanded"], dst_fused["dsa_seqlens_expanded"],
dst_fused["nsa_cu_seqlens_k"], dst_fused["dsa_cu_seqlens_k"],
dst_fused["real_page_table"], dst_fused["real_page_table"],
dst_fused["flashmla_num_splits"], dst_fused["flashmla_num_splits"],
dst_fused["flashmla_metadata"], dst_fused["flashmla_metadata"],
@@ -402,14 +402,14 @@ def test_fused_metadata_copy(bs, forward_mode, has_real_page_table, has_flashmla
dst_ref["page_table_1"], dst_fused["page_table_1"] dst_ref["page_table_1"], dst_fused["page_table_1"]
), "page_table_1 mismatch" ), "page_table_1 mismatch"
assert torch.equal( assert torch.equal(
dst_ref["nsa_cache_seqlens"], dst_fused["nsa_cache_seqlens"] dst_ref["dsa_cache_seqlens"], dst_fused["dsa_cache_seqlens"]
), "nsa_cache_seqlens mismatch" ), "dsa_cache_seqlens mismatch"
assert torch.equal( assert torch.equal(
dst_ref["nsa_seqlens_expanded"], dst_fused["nsa_seqlens_expanded"] dst_ref["dsa_seqlens_expanded"], dst_fused["dsa_seqlens_expanded"]
), "nsa_seqlens_expanded mismatch" ), "dsa_seqlens_expanded mismatch"
assert torch.equal( assert torch.equal(
dst_ref["nsa_cu_seqlens_k"], dst_fused["nsa_cu_seqlens_k"] dst_ref["dsa_cu_seqlens_k"], dst_fused["dsa_cu_seqlens_k"]
), "nsa_cu_seqlens_k mismatch" ), "dsa_cu_seqlens_k mismatch"
if has_real_page_table: if has_real_page_table:
assert torch.equal( assert torch.equal(
@@ -458,18 +458,18 @@ def test_fused_metadata_copy_large_batch(bs):
data["src"]["cache_seqlens"], data["src"]["cache_seqlens"],
data["src"]["cu_seqlens_k"], data["src"]["cu_seqlens_k"],
data["src"]["page_indices"], data["src"]["page_indices"],
data["src"]["nsa_cache_seqlens"], data["src"]["dsa_cache_seqlens"],
data["src"]["seqlens_expanded"], data["src"]["seqlens_expanded"],
data["src"]["nsa_cu_seqlens_k"], data["src"]["dsa_cu_seqlens_k"],
data["src"]["real_page_table"], data["src"]["real_page_table"],
data["src"]["flashmla_num_splits"], data["src"]["flashmla_num_splits"],
data["src"]["flashmla_metadata"], data["src"]["flashmla_metadata"],
dst_fused["cache_seqlens"], dst_fused["cache_seqlens"],
dst_fused["cu_seqlens_k"], dst_fused["cu_seqlens_k"],
dst_fused["page_table_1"], dst_fused["page_table_1"],
dst_fused["nsa_cache_seqlens"], dst_fused["dsa_cache_seqlens"],
dst_fused["nsa_seqlens_expanded"], dst_fused["dsa_seqlens_expanded"],
dst_fused["nsa_cu_seqlens_k"], dst_fused["dsa_cu_seqlens_k"],
dst_fused["real_page_table"], dst_fused["real_page_table"],
dst_fused["flashmla_num_splits"], dst_fused["flashmla_num_splits"],
dst_fused["flashmla_metadata"], dst_fused["flashmla_metadata"],
@@ -510,13 +510,13 @@ def create_test_metadata_multi(
page_indices_src = torch.randint( page_indices_src = torch.randint(
0, 1000, (bs, max_len), dtype=torch.int32, device=device 0, 1000, (bs, max_len), dtype=torch.int32, device=device
) )
nsa_cache_seqlens_src = torch.randint( dsa_cache_seqlens_src = torch.randint(
1, max_len, (seqlens_expanded_size,), dtype=torch.int32, device=device 1, max_len, (seqlens_expanded_size,), dtype=torch.int32, device=device
) )
nsa_cu_seqlens_k_src = torch.zeros( dsa_cu_seqlens_k_src = torch.zeros(
seqlens_expanded_size + 1, dtype=torch.int32, device=device seqlens_expanded_size + 1, dtype=torch.int32, device=device
) )
nsa_cu_seqlens_k_src[1:] = torch.cumsum(nsa_cache_seqlens_src, dim=0) dsa_cu_seqlens_k_src[1:] = torch.cumsum(dsa_cache_seqlens_src, dim=0)
# Optional tensors # Optional tensors
real_page_table_src = None real_page_table_src = None
@@ -544,10 +544,10 @@ def create_test_metadata_multi(
page_table_1_dst = torch.zeros( page_table_1_dst = torch.zeros(
(bs, max_len + 16), dtype=torch.int32, device=device (bs, max_len + 16), dtype=torch.int32, device=device
) )
nsa_cache_seqlens_dst = torch.zeros( dsa_cache_seqlens_dst = torch.zeros(
seqlens_expanded_size, dtype=torch.int32, device=device seqlens_expanded_size, dtype=torch.int32, device=device
) )
nsa_cu_seqlens_k_dst = torch.zeros( dsa_cu_seqlens_k_dst = torch.zeros(
seqlens_expanded_size + 1, dtype=torch.int32, device=device seqlens_expanded_size + 1, dtype=torch.int32, device=device
) )
@@ -573,8 +573,8 @@ def create_test_metadata_multi(
"cache_seqlens_int32": cache_seqlens_dst, "cache_seqlens_int32": cache_seqlens_dst,
"cu_seqlens_k": cu_seqlens_k_dst, "cu_seqlens_k": cu_seqlens_k_dst,
"page_table_1": page_table_1_dst, "page_table_1": page_table_1_dst,
"nsa_cache_seqlens_int32": nsa_cache_seqlens_dst, "dsa_cache_seqlens_int32": dsa_cache_seqlens_dst,
"nsa_cu_seqlens_k": nsa_cu_seqlens_k_dst, "dsa_cu_seqlens_k": dsa_cu_seqlens_k_dst,
"real_page_table": real_page_table_dst, "real_page_table": real_page_table_dst,
"flashmla_num_splits": flashmla_num_splits_dst, "flashmla_num_splits": flashmla_num_splits_dst,
"flashmla_metadata": flashmla_metadata_dst, "flashmla_metadata": flashmla_metadata_dst,
@@ -585,8 +585,8 @@ def create_test_metadata_multi(
"cache_seqlens": cache_seqlens_src, "cache_seqlens": cache_seqlens_src,
"cu_seqlens_k": cu_seqlens_k_src, "cu_seqlens_k": cu_seqlens_k_src,
"page_indices": page_indices_src, "page_indices": page_indices_src,
"nsa_cache_seqlens": nsa_cache_seqlens_src, "dsa_cache_seqlens": dsa_cache_seqlens_src,
"nsa_cu_seqlens_k": nsa_cu_seqlens_k_src, "dsa_cu_seqlens_k": dsa_cu_seqlens_k_src,
"real_page_table": real_page_table_src, "real_page_table": real_page_table_src,
"flashmla_num_splits": flashmla_num_splits_src, "flashmla_num_splits": flashmla_num_splits_src,
"flashmla_metadata": flashmla_metadata_src, "flashmla_metadata": flashmla_metadata_src,
@@ -604,8 +604,8 @@ def reference_copy_for_loop(src, dst_list, bs, max_len):
dst["cache_seqlens_int32"].copy_(src["cache_seqlens"]) dst["cache_seqlens_int32"].copy_(src["cache_seqlens"])
dst["cu_seqlens_k"][1:].copy_(src["cu_seqlens_k"][1:]) dst["cu_seqlens_k"][1:].copy_(src["cu_seqlens_k"][1:])
dst["page_table_1"][:, :max_len].copy_(src["page_indices"]) dst["page_table_1"][:, :max_len].copy_(src["page_indices"])
dst["nsa_cache_seqlens_int32"].copy_(src["nsa_cache_seqlens"]) dst["dsa_cache_seqlens_int32"].copy_(src["dsa_cache_seqlens"])
dst["nsa_cu_seqlens_k"][1 : bs + 1].copy_(src["nsa_cu_seqlens_k"][1 : bs + 1]) dst["dsa_cu_seqlens_k"][1 : bs + 1].copy_(src["dsa_cu_seqlens_k"][1 : bs + 1])
if src["real_page_table"] is not None: if src["real_page_table"] is not None:
rows, cols = src["real_page_table"].shape rows, cols = src["real_page_table"].shape
@@ -641,10 +641,10 @@ def test_fused_metadata_copy_multi_dtype_validation():
page_indices_src = torch.randint( page_indices_src = torch.randint(
0, 1000, (bs, max_len), dtype=torch.int32, device=device 0, 1000, (bs, max_len), dtype=torch.int32, device=device
) )
nsa_cache_seqlens_src = torch.randint( dsa_cache_seqlens_src = torch.randint(
1, max_len, (seqlens_expanded_size,), dtype=torch.int32, device=device 1, max_len, (seqlens_expanded_size,), dtype=torch.int32, device=device
) )
nsa_cu_seqlens_k_src = torch.zeros( dsa_cu_seqlens_k_src = torch.zeros(
seqlens_expanded_size + 1, dtype=torch.int32, device=device seqlens_expanded_size + 1, dtype=torch.int32, device=device
) )
@@ -656,10 +656,10 @@ def test_fused_metadata_copy_multi_dtype_validation():
"page_table_1": torch.zeros( "page_table_1": torch.zeros(
(bs, max_len + 16), dtype=torch.int32, device=device (bs, max_len + 16), dtype=torch.int32, device=device
), ),
"nsa_cache_seqlens": torch.zeros( "dsa_cache_seqlens": torch.zeros(
seqlens_expanded_size, dtype=torch.int32, device=device seqlens_expanded_size, dtype=torch.int32, device=device
), ),
"nsa_cu_seqlens_k": torch.zeros( "dsa_cu_seqlens_k": torch.zeros(
seqlens_expanded_size + 1, dtype=torch.int32, device=device seqlens_expanded_size + 1, dtype=torch.int32, device=device
), ),
} }
@@ -674,8 +674,8 @@ def test_fused_metadata_copy_multi_dtype_validation():
cache_seqlens_src_wrong, # Wrong dtype: int64 cache_seqlens_src_wrong, # Wrong dtype: int64
cu_seqlens_k_src, cu_seqlens_k_src,
page_indices_src, page_indices_src,
nsa_cache_seqlens_src, dsa_cache_seqlens_src,
nsa_cu_seqlens_k_src, dsa_cu_seqlens_k_src,
None, # real_page_table_src None, # real_page_table_src
None, # flashmla_num_splits_src None, # flashmla_num_splits_src
None, # flashmla_metadata_src None, # flashmla_metadata_src
@@ -683,8 +683,8 @@ def test_fused_metadata_copy_multi_dtype_validation():
dst0["cache_seqlens"], dst0["cache_seqlens"],
dst0["cu_seqlens_k"], dst0["cu_seqlens_k"],
dst0["page_table_1"], dst0["page_table_1"],
dst0["nsa_cache_seqlens"], dst0["dsa_cache_seqlens"],
dst0["nsa_cu_seqlens_k"], dst0["dsa_cu_seqlens_k"],
None, None,
None, None,
None, None,
@@ -692,8 +692,8 @@ def test_fused_metadata_copy_multi_dtype_validation():
dst1["cache_seqlens"], dst1["cache_seqlens"],
dst1["cu_seqlens_k"], dst1["cu_seqlens_k"],
dst1["page_table_1"], dst1["page_table_1"],
dst1["nsa_cache_seqlens"], dst1["dsa_cache_seqlens"],
dst1["nsa_cu_seqlens_k"], dst1["dsa_cu_seqlens_k"],
None, None,
None, None,
None, None,
@@ -701,8 +701,8 @@ def test_fused_metadata_copy_multi_dtype_validation():
dst2["cache_seqlens"], dst2["cache_seqlens"],
dst2["cu_seqlens_k"], dst2["cu_seqlens_k"],
dst2["page_table_1"], dst2["page_table_1"],
dst2["nsa_cache_seqlens"], dst2["dsa_cache_seqlens"],
dst2["nsa_cu_seqlens_k"], dst2["dsa_cu_seqlens_k"],
None, None,
None, None,
None, None,
@@ -772,8 +772,8 @@ def test_fused_metadata_copy_multi(bs, has_real_page_table, has_flashmla):
data["src"]["cache_seqlens"], data["src"]["cache_seqlens"],
data["src"]["cu_seqlens_k"], data["src"]["cu_seqlens_k"],
data["src"]["page_indices"], data["src"]["page_indices"],
data["src"]["nsa_cache_seqlens"], data["src"]["dsa_cache_seqlens"],
data["src"]["nsa_cu_seqlens_k"], data["src"]["dsa_cu_seqlens_k"],
data["src"]["real_page_table"], data["src"]["real_page_table"],
data["src"]["flashmla_num_splits"], data["src"]["flashmla_num_splits"],
data["src"]["flashmla_metadata"], data["src"]["flashmla_metadata"],
@@ -781,8 +781,8 @@ def test_fused_metadata_copy_multi(bs, has_real_page_table, has_flashmla):
dst_fused_0["cache_seqlens_int32"], dst_fused_0["cache_seqlens_int32"],
dst_fused_0["cu_seqlens_k"], dst_fused_0["cu_seqlens_k"],
dst_fused_0["page_table_1"], dst_fused_0["page_table_1"],
dst_fused_0["nsa_cache_seqlens_int32"], dst_fused_0["dsa_cache_seqlens_int32"],
dst_fused_0["nsa_cu_seqlens_k"], dst_fused_0["dsa_cu_seqlens_k"],
dst_fused_0["real_page_table"], dst_fused_0["real_page_table"],
dst_fused_0["flashmla_num_splits"], dst_fused_0["flashmla_num_splits"],
dst_fused_0["flashmla_metadata"], dst_fused_0["flashmla_metadata"],
@@ -790,8 +790,8 @@ def test_fused_metadata_copy_multi(bs, has_real_page_table, has_flashmla):
dst_fused_1["cache_seqlens_int32"], dst_fused_1["cache_seqlens_int32"],
dst_fused_1["cu_seqlens_k"], dst_fused_1["cu_seqlens_k"],
dst_fused_1["page_table_1"], dst_fused_1["page_table_1"],
dst_fused_1["nsa_cache_seqlens_int32"], dst_fused_1["dsa_cache_seqlens_int32"],
dst_fused_1["nsa_cu_seqlens_k"], dst_fused_1["dsa_cu_seqlens_k"],
dst_fused_1["real_page_table"], dst_fused_1["real_page_table"],
dst_fused_1["flashmla_num_splits"], dst_fused_1["flashmla_num_splits"],
dst_fused_1["flashmla_metadata"], dst_fused_1["flashmla_metadata"],
@@ -799,8 +799,8 @@ def test_fused_metadata_copy_multi(bs, has_real_page_table, has_flashmla):
dst_fused_2["cache_seqlens_int32"], dst_fused_2["cache_seqlens_int32"],
dst_fused_2["cu_seqlens_k"], dst_fused_2["cu_seqlens_k"],
dst_fused_2["page_table_1"], dst_fused_2["page_table_1"],
dst_fused_2["nsa_cache_seqlens_int32"], dst_fused_2["dsa_cache_seqlens_int32"],
dst_fused_2["nsa_cu_seqlens_k"], dst_fused_2["dsa_cu_seqlens_k"],
dst_fused_2["real_page_table"], dst_fused_2["real_page_table"],
dst_fused_2["flashmla_num_splits"], dst_fused_2["flashmla_num_splits"],
dst_fused_2["flashmla_metadata"], dst_fused_2["flashmla_metadata"],
@@ -836,8 +836,8 @@ def test_fused_metadata_copy_multi(bs, has_real_page_table, has_flashmla):
"cache_seqlens_int32", "cache_seqlens_int32",
"cu_seqlens_k", "cu_seqlens_k",
"page_table_1", "page_table_1",
"nsa_cache_seqlens_int32", "dsa_cache_seqlens_int32",
"nsa_cu_seqlens_k", "dsa_cu_seqlens_k",
]: ]:
if not torch.equal(dst_ref[key], dst_fused[key]): if not torch.equal(dst_ref[key], dst_fused[key]):
diff = ( diff = (
@@ -965,32 +965,32 @@ def test_fused_metadata_copy_multi_large_batch(bs):
data["src"]["cache_seqlens"], data["src"]["cache_seqlens"],
data["src"]["cu_seqlens_k"], data["src"]["cu_seqlens_k"],
data["src"]["page_indices"], data["src"]["page_indices"],
data["src"]["nsa_cache_seqlens"], data["src"]["dsa_cache_seqlens"],
data["src"]["nsa_cu_seqlens_k"], data["src"]["dsa_cu_seqlens_k"],
data["src"]["real_page_table"], data["src"]["real_page_table"],
data["src"]["flashmla_num_splits"], data["src"]["flashmla_num_splits"],
data["src"]["flashmla_metadata"], data["src"]["flashmla_metadata"],
dst_fused_0["cache_seqlens_int32"], dst_fused_0["cache_seqlens_int32"],
dst_fused_0["cu_seqlens_k"], dst_fused_0["cu_seqlens_k"],
dst_fused_0["page_table_1"], dst_fused_0["page_table_1"],
dst_fused_0["nsa_cache_seqlens_int32"], dst_fused_0["dsa_cache_seqlens_int32"],
dst_fused_0["nsa_cu_seqlens_k"], dst_fused_0["dsa_cu_seqlens_k"],
dst_fused_0["real_page_table"], dst_fused_0["real_page_table"],
dst_fused_0["flashmla_num_splits"], dst_fused_0["flashmla_num_splits"],
dst_fused_0["flashmla_metadata"], dst_fused_0["flashmla_metadata"],
dst_fused_1["cache_seqlens_int32"], dst_fused_1["cache_seqlens_int32"],
dst_fused_1["cu_seqlens_k"], dst_fused_1["cu_seqlens_k"],
dst_fused_1["page_table_1"], dst_fused_1["page_table_1"],
dst_fused_1["nsa_cache_seqlens_int32"], dst_fused_1["dsa_cache_seqlens_int32"],
dst_fused_1["nsa_cu_seqlens_k"], dst_fused_1["dsa_cu_seqlens_k"],
dst_fused_1["real_page_table"], dst_fused_1["real_page_table"],
dst_fused_1["flashmla_num_splits"], dst_fused_1["flashmla_num_splits"],
dst_fused_1["flashmla_metadata"], dst_fused_1["flashmla_metadata"],
dst_fused_2["cache_seqlens_int32"], dst_fused_2["cache_seqlens_int32"],
dst_fused_2["cu_seqlens_k"], dst_fused_2["cu_seqlens_k"],
dst_fused_2["page_table_1"], dst_fused_2["page_table_1"],
dst_fused_2["nsa_cache_seqlens_int32"], dst_fused_2["dsa_cache_seqlens_int32"],
dst_fused_2["nsa_cu_seqlens_k"], dst_fused_2["dsa_cu_seqlens_k"],
dst_fused_2["real_page_table"], dst_fused_2["real_page_table"],
dst_fused_2["flashmla_num_splits"], dst_fused_2["flashmla_num_splits"],
dst_fused_2["flashmla_metadata"], dst_fused_2["flashmla_metadata"],
@@ -1013,32 +1013,32 @@ def test_fused_metadata_copy_multi_large_batch(bs):
data["src"]["cache_seqlens"], data["src"]["cache_seqlens"],
data["src"]["cu_seqlens_k"], data["src"]["cu_seqlens_k"],
data["src"]["page_indices"], data["src"]["page_indices"],
data["src"]["nsa_cache_seqlens"], data["src"]["dsa_cache_seqlens"],
data["src"]["nsa_cu_seqlens_k"], data["src"]["dsa_cu_seqlens_k"],
data["src"]["real_page_table"], data["src"]["real_page_table"],
data["src"]["flashmla_num_splits"], data["src"]["flashmla_num_splits"],
data["src"]["flashmla_metadata"], data["src"]["flashmla_metadata"],
dst_fused_0["cache_seqlens_int32"], dst_fused_0["cache_seqlens_int32"],
dst_fused_0["cu_seqlens_k"], dst_fused_0["cu_seqlens_k"],
dst_fused_0["page_table_1"], dst_fused_0["page_table_1"],
dst_fused_0["nsa_cache_seqlens_int32"], dst_fused_0["dsa_cache_seqlens_int32"],
dst_fused_0["nsa_cu_seqlens_k"], dst_fused_0["dsa_cu_seqlens_k"],
dst_fused_0["real_page_table"], dst_fused_0["real_page_table"],
dst_fused_0["flashmla_num_splits"], dst_fused_0["flashmla_num_splits"],
dst_fused_0["flashmla_metadata"], dst_fused_0["flashmla_metadata"],
dst_fused_1["cache_seqlens_int32"], dst_fused_1["cache_seqlens_int32"],
dst_fused_1["cu_seqlens_k"], dst_fused_1["cu_seqlens_k"],
dst_fused_1["page_table_1"], dst_fused_1["page_table_1"],
dst_fused_1["nsa_cache_seqlens_int32"], dst_fused_1["dsa_cache_seqlens_int32"],
dst_fused_1["nsa_cu_seqlens_k"], dst_fused_1["dsa_cu_seqlens_k"],
dst_fused_1["real_page_table"], dst_fused_1["real_page_table"],
dst_fused_1["flashmla_num_splits"], dst_fused_1["flashmla_num_splits"],
dst_fused_1["flashmla_metadata"], dst_fused_1["flashmla_metadata"],
dst_fused_2["cache_seqlens_int32"], dst_fused_2["cache_seqlens_int32"],
dst_fused_2["cu_seqlens_k"], dst_fused_2["cu_seqlens_k"],
dst_fused_2["page_table_1"], dst_fused_2["page_table_1"],
dst_fused_2["nsa_cache_seqlens_int32"], dst_fused_2["dsa_cache_seqlens_int32"],
dst_fused_2["nsa_cu_seqlens_k"], dst_fused_2["dsa_cu_seqlens_k"],
dst_fused_2["real_page_table"], dst_fused_2["real_page_table"],
dst_fused_2["flashmla_num_splits"], dst_fused_2["flashmla_num_splits"],
dst_fused_2["flashmla_metadata"], dst_fused_2["flashmla_metadata"],
python/sglang/jit_kernel/tests/test_fused_store_index_cache.py
+3 -3
View File
@@ -26,7 +26,7 @@ from sglang.test.ci.ci_register import register_cuda_ci
try: try:
from sglang.jit_kernel.fused_store_index_cache import ( from sglang.jit_kernel.fused_store_index_cache import (
can_use_nsa_fused_store, can_use_dsa_fused_store,
fused_store_index_k_cache, fused_store_index_k_cache,
) )
@@ -70,7 +70,7 @@ def _skip_if_unavailable(page_size: int = PAGE_SIZE):
pytest.skip("torch.float8_e4m3fn not available") pytest.skip("torch.float8_e4m3fn not available")
if not HAS_FUSED: if not HAS_FUSED:
pytest.skip("fused_store_index_cache not importable") pytest.skip("fused_store_index_cache not importable")
if not can_use_nsa_fused_store(torch.bfloat16, torch.int64, page_size): if not can_use_dsa_fused_store(torch.bfloat16, torch.int64, page_size):
pytest.skip("JIT kernel unavailable / failed to compile") pytest.skip("JIT kernel unavailable / failed to compile")
@@ -187,7 +187,7 @@ def _reference_quantize_and_store(
def _import_act_quant(): def _import_act_quant():
try: try:
from sglang.srt.layers.attention.nsa.triton_kernel import act_quant from sglang.srt.layers.attention.dsa.triton_kernel import act_quant
return act_quant return act_quant
except Exception: except Exception:
python/sglang/jit_kernel/tests/test_set_mla_kv_buffer.py
+1 -1
View File
@@ -75,7 +75,7 @@ def test_set_mla_kv_buffer_loc_dtypes(loc_dtype):
def test_set_mla_kv_buffer_uint8_byte_layout(): def test_set_mla_kv_buffer_uint8_byte_layout():
"""FP8 NSA byte-layout: cache_k_nope is uint8 with [fp8(512) | scales(16)] = 528, """FP8 DSA byte-layout: cache_k_nope is uint8 with [fp8(512) | scales(16)] = 528,
cache_k_rope is uint8 [128]; total payload = 656 bytes.""" cache_k_rope is uint8 [128]; total payload = 656 bytes."""
nope_bytes, rope_bytes = 528, 128 nope_bytes, rope_bytes = 528, 128
batch_size = 64 batch_size = 64
python/sglang/srt/arg_groups/deepseek_v4_hook.py
+3 -3
View File
@@ -54,13 +54,13 @@ def apply_deepseek_v4_defaults(server_args: "ServerArgs", model_arch: str) -> No
def validate_deepseek_v4_cp(server_args: "ServerArgs") -> None: def validate_deepseek_v4_cp(server_args: "ServerArgs") -> None:
"""Validate DeepSeek V4 context-parallel configuration.""" """Validate DeepSeek V4 context-parallel configuration."""
if not server_args.enable_nsa_prefill_context_parallel: if not server_args.enable_dsa_prefill_context_parallel:
return return
if server_args.nsa_prefill_cp_mode != "round-robin-split": if server_args.dsa_prefill_cp_mode != "round-robin-split":
raise ValueError( raise ValueError(
f"DeepSeekV4 only supports round-robin-split CP mode, " f"DeepSeekV4 only supports round-robin-split CP mode, "
f"got {server_args.nsa_prefill_cp_mode}" f"got {server_args.dsa_prefill_cp_mode}"
) )
server_args.enable_dp_attention = True server_args.enable_dp_attention = True
python/sglang/srt/arg_groups/hisparse_hook.py
+12 -12
View File
@@ -20,13 +20,13 @@ def _hisparse_default_backend(kv_cache_dtype: str) -> str:
return "flashmla_kv" if kv_cache_dtype == "fp8_e4m3" else "flashmla_sparse" return "flashmla_kv" if kv_cache_dtype == "fp8_e4m3" else "flashmla_sparse"
def apply_hisparse_nsa_backend_defaults( def apply_hisparse_dsa_backend_defaults(
server_args: "ServerArgs", server_args: "ServerArgs",
user_set_prefill: bool, user_set_prefill: bool,
user_set_decode: bool, user_set_decode: bool,
kv_cache_dtype: str, kv_cache_dtype: str,
) -> bool: ) -> bool:
"""Pick NSA backends for --enable-hisparse based on KV dtype. """Pick DSA backends for --enable-hisparse based on KV dtype.
BF16 KV -> flashmla_sparse, FP8 KV -> flashmla_kv. Returns True if hisparse BF16 KV -> flashmla_sparse, FP8 KV -> flashmla_kv. Returns True if hisparse
handled backend selection (caller should skip its own default logic). handled backend selection (caller should skip its own default logic).
@@ -36,29 +36,29 @@ def apply_hisparse_nsa_backend_defaults(
backend = _hisparse_default_backend(kv_cache_dtype) backend = _hisparse_default_backend(kv_cache_dtype)
if not user_set_prefill: if not user_set_prefill:
server_args.nsa_prefill_backend = backend server_args.dsa_prefill_backend = backend
if not user_set_decode: if not user_set_decode:
server_args.nsa_decode_backend = backend server_args.dsa_decode_backend = backend
logger.warning( logger.warning(
f"HiSparse enabled ({kv_cache_dtype}): using NSA backends " f"HiSparse enabled ({kv_cache_dtype}): using DSA backends "
f"prefill={server_args.nsa_prefill_backend}, decode={server_args.nsa_decode_backend}." f"prefill={server_args.dsa_prefill_backend}, decode={server_args.dsa_decode_backend}."
) )
return True return True
def validate_hisparse(server_args: "ServerArgs") -> None: def validate_hisparse(server_args: "ServerArgs") -> None:
"""Validate --enable-hisparse constraints (model class, radix cache, NSA backend).""" """Validate --enable-hisparse constraints (model class, radix cache, DSA backend)."""
if not server_args.enable_hisparse: if not server_args.enable_hisparse:
return return
from sglang.srt.configs.model_config import ( from sglang.srt.configs.model_config import (
is_deepseek_nsa, is_deepseek_dsa,
is_deepseek_v4, is_deepseek_v4,
) )
hf_config = server_args.get_model_config().hf_config hf_config = server_args.get_model_config().hf_config
is_v4_hisparse = is_deepseek_v4(hf_config) is_v4_hisparse = is_deepseek_v4(hf_config)
assert is_deepseek_nsa(hf_config) or is_v4_hisparse, ( assert is_deepseek_dsa(hf_config) or is_v4_hisparse, (
"--enable-hisparse is only supported for DSA (DeepSeek Sparse Attention) " "--enable-hisparse is only supported for DSA (DeepSeek Sparse Attention) "
"models (e.g., DeepSeek V3.2, GLM-5) and DeepSeek V4 now. " "models (e.g., DeepSeek V3.2, GLM-5) and DeepSeek V4 now. "
) )
@@ -83,13 +83,13 @@ def validate_hisparse(server_args: "ServerArgs") -> None:
server_args.kv_cache_dtype, {"flashmla_sparse", "flashmla_kv"} server_args.kv_cache_dtype, {"flashmla_sparse", "flashmla_kv"}
) )
for attr, label in [ for attr, label in [
("nsa_prefill_backend", "prefill"), ("dsa_prefill_backend", "prefill"),
("nsa_decode_backend", "decode"), ("dsa_decode_backend", "decode"),
]: ]:
backend = getattr(server_args, attr) backend = getattr(server_args, attr)
if backend is not None and backend not in allowed_backends: if backend is not None and backend not in allowed_backends:
raise ValueError( raise ValueError(
f"HiSparse with --kv-cache-dtype={server_args.kv_cache_dtype} requires " f"HiSparse with --kv-cache-dtype={server_args.kv_cache_dtype} requires "
f"--nsa-{label}-backend in {sorted(allowed_backends)}, " f"--dsa-{label}-backend in {sorted(allowed_backends)}, "
f"but got {backend}." f"but got {backend}."
) )
python/sglang/srt/configs/model_config.py
+12 -12
View File
@@ -99,7 +99,7 @@ def _hf_attr(config, name):
return getattr(config, name, None) return getattr(config, name, None)
def is_deepseek_nsa(config) -> bool: def is_deepseek_dsa(config) -> bool:
return ( return (
_hf_arch(config) _hf_arch(config)
in ( in (
@@ -121,31 +121,31 @@ def is_deepseek_v4(config) -> bool:
) )
def get_nsa_index_head_dim(config: PretrainedConfig) -> int: def get_dsa_index_head_dim(config: PretrainedConfig) -> int:
assert is_deepseek_nsa(config) or is_deepseek_v4(config) assert is_deepseek_dsa(config) or is_deepseek_v4(config)
return config.index_head_dim return config.index_head_dim
def get_nsa_index_topk(config: PretrainedConfig) -> int: def get_dsa_index_topk(config: PretrainedConfig) -> int:
assert is_deepseek_nsa(config) assert is_deepseek_dsa(config)
return config.index_topk return config.index_topk
def get_nsa_index_n_heads(config: PretrainedConfig) -> int: def get_dsa_index_n_heads(config: PretrainedConfig) -> int:
assert is_deepseek_nsa(config) assert is_deepseek_dsa(config)
return config.index_n_heads return config.index_n_heads
def get_num_indexer_layers(config) -> int: def get_num_indexer_layers(config) -> int:
"""Layer count for the global indexer-topk capturer's host buffer. """Layer count for the global indexer-topk capturer's host buffer.
NSA models (V3.2) instantiate an Indexer on every transformer layer. DSA models (V3.2) instantiate an Indexer on every transformer layer.
With index_topk_freq > 1 some layers reuse prev layer's topk; those still With index_topk_freq > 1 some layers reuse prev layer's topk; those still
get a slot (mirrored at the MLA call site). DSv4 has C4 indexers only on get a slot (mirrored at the MLA call site). DSv4 has C4 indexers only on
layers whose compress_ratio == 4. Other architectures: set layers whose compress_ratio == 4. Other architectures: set
num_indexer_layers on hf_text_config; 0 disables the capturer. num_indexer_layers on hf_text_config; 0 disables the capturer.
""" """
if is_deepseek_nsa(config): if is_deepseek_dsa(config):
return config.num_hidden_layers return config.num_hidden_layers
if is_deepseek_v4(config): if is_deepseek_v4(config):
compress_ratios = getattr(config, "compress_ratios", None) or [] compress_ratios = getattr(config, "compress_ratios", None) or []
@@ -329,7 +329,7 @@ class ModelConfig:
self.use_ngram_embedding = getattr(self.hf_config, "use_ngram_embedding", False) self.use_ngram_embedding = getattr(self.hf_config, "use_ngram_embedding", False)
self.is_piecewise_cuda_graph_disabled_model = ( self.is_piecewise_cuda_graph_disabled_model = (
is_piecewise_cuda_graph_disabled_model(self.hf_config.architectures) is_piecewise_cuda_graph_disabled_model(self.hf_config.architectures)
or is_deepseek_nsa(self.hf_text_config) or is_deepseek_dsa(self.hf_text_config)
) )
self.dtype = _get_and_verify_dtype(self.hf_text_config, dtype) self.dtype = _get_and_verify_dtype(self.hf_text_config, dtype)
@@ -622,8 +622,8 @@ class ModelConfig:
self.qk_rope_head_dim = self.hf_text_config.qk_rope_head_dim self.qk_rope_head_dim = self.hf_text_config.qk_rope_head_dim
self.v_head_dim = self.hf_text_config.v_head_dim self.v_head_dim = self.hf_text_config.v_head_dim
self.index_head_dim = ( self.index_head_dim = (
get_nsa_index_head_dim(self.hf_text_config) get_dsa_index_head_dim(self.hf_text_config)
if is_deepseek_nsa(self.hf_text_config) if is_deepseek_dsa(self.hf_text_config)
else None else None
) )
# Handle rope scaling # Handle rope scaling
python/sglang/srt/disaggregation/base/conn.py
+1 -1
View File
@@ -17,7 +17,7 @@ if TYPE_CHECKING:
class StateType(str, enum.Enum): class StateType(str, enum.Enum):
MAMBA = "mamba" MAMBA = "mamba"
SWA = "swa" SWA = "swa"
NSA = "nsa" DSA = "dsa"
@dataclasses.dataclass @dataclasses.dataclass
python/sglang/srt/disaggregation/decode.py
+3 -3
View File
@@ -954,7 +954,7 @@ class DecodePreallocQueue:
window_kv_indices_swa.cpu().numpy(), page_size window_kv_indices_swa.cpu().numpy(), page_size
) )
def _nsa_payload(): def _dsa_payload():
kv_indices_full = self.req_to_token_pool.req_to_token[ kv_indices_full = self.req_to_token_pool.req_to_token[
decode_req.req.req_pool_idx, :seq_len decode_req.req.req_pool_idx, :seq_len
] ]
@@ -971,8 +971,8 @@ class DecodePreallocQueue:
state_indices.append(_mamba_payload()) state_indices.append(_mamba_payload())
elif st == StateType.SWA: elif st == StateType.SWA:
state_indices.append(_swa_payload()) state_indices.append(_swa_payload())
elif st == StateType.NSA: elif st == StateType.DSA:
state_indices.append(_nsa_payload()) state_indices.append(_dsa_payload())
else: else:
state_indices.append(None) state_indices.append(None)
python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py
+1 -1
View File
@@ -238,7 +238,7 @@ class DecodeKVCacheOffloadManager:
kv_committed_len = req.pop_committed_kv_cache() kv_committed_len = req.pop_committed_kv_cache()
start = start_offset start = start_offset
end = kv_committed_len end = kv_committed_len
# Free the incremental part of the request (NSA-aware) # Free the incremental part of the request (DSA-aware)
kv_indices = self.req_to_token_pool.req_to_token[req.req_pool_idx, start:end] kv_indices = self.req_to_token_pool.req_to_token[req.req_pool_idx, start:end]
self.token_to_kv_pool_allocator.free(kv_indices) self.token_to_kv_pool_allocator.free(kv_indices)
python/sglang/srt/disaggregation/mooncake/conn.py
+1 -1
View File
@@ -1012,7 +1012,7 @@ class MooncakeKVManager(CommonKVManager):
) )
or rc or rc
) )
elif st in (StateType.SWA, StateType.NSA): elif st in (StateType.SWA, StateType.DSA):
if ( if (
target_rank_registration_info is not None target_rank_registration_info is not None
and not self.is_mla_backend and not self.is_mla_backend
python/sglang/srt/disaggregation/mori/conn.py
+3 -3
View File
@@ -960,8 +960,8 @@ class MoriKVManager(CommonKVManager):
return self._send_mamba_state( return self._send_mamba_state(
peer_info, src_state_indices, dst_state_indices peer_info, src_state_indices, dst_state_indices
) )
elif state_type in ("swa", "nsa"): elif state_type in ("swa", "dsa"):
return self._send_swa_nsa_state( return self._send_swa_dsa_state(
peer_info, src_state_indices, dst_state_indices, state_type peer_info, src_state_indices, dst_state_indices, state_type
) )
else: else:
@@ -1056,7 +1056,7 @@ class MoriKVManager(CommonKVManager):
return statuses return statuses
def _send_swa_nsa_state( def _send_swa_dsa_state(
self, self,
peer_info: KVArgsRegisterInfo, peer_info: KVArgsRegisterInfo,
src_state_indices: npt.NDArray[np.int32], src_state_indices: npt.NDArray[np.int32],
python/sglang/srt/disaggregation/nixl/conn.py
+1 -1
View File
@@ -1541,7 +1541,7 @@ class NixlKVManager(CommonKVManager):
dst_gpu_id, dst_gpu_id,
comp_notif, comp_notif,
) )
elif st in (StateType.SWA, StateType.NSA): elif st in (StateType.SWA, StateType.DSA):
if not self.is_mla_backend and self.attn_tp_size != decode_tp_size: if not self.is_mla_backend and self.attn_tp_size != decode_tp_size:
raise RuntimeError( raise RuntimeError(
f"PD Disaggregation does NOT support PD different TP sizes for non-MLA {st.upper()} hybrid models yet." f"PD Disaggregation does NOT support PD different TP sizes for non-MLA {st.upper()} hybrid models yet."
python/sglang/srt/disaggregation/prefill.py
+3 -3
View File
@@ -818,7 +818,7 @@ class SchedulerDisaggregationPrefillMixin:
window_kv_indices_swa.cpu().numpy(), page_size window_kv_indices_swa.cpu().numpy(), page_size
) )
def _nsa_payload(): def _dsa_payload():
kv_indices_full = self.req_to_token_pool.req_to_token[ kv_indices_full = self.req_to_token_pool.req_to_token[
req.req_pool_idx, :seq_len req.req_pool_idx, :seq_len
] ]
@@ -833,8 +833,8 @@ class SchedulerDisaggregationPrefillMixin:
state_indices.append(_mamba_payload()) state_indices.append(_mamba_payload())
elif st == StateType.SWA: elif st == StateType.SWA:
state_indices.append(_swa_payload()) state_indices.append(_swa_payload())
elif st == StateType.NSA: elif st == StateType.DSA:
state_indices.append(_nsa_payload()) state_indices.append(_dsa_payload())
else: else:
state_indices.append(None) state_indices.append(None)
python/sglang/srt/disaggregation/utils.py
+4 -4
View File
@@ -567,7 +567,7 @@ def setup_state_kv_args(
from sglang.srt.disaggregation.base.conn import StateType from sglang.srt.disaggregation.base.conn import StateType
from sglang.srt.hardware_backend.npu.memory_pool_npu import NPUMLATokenToKVPool from sglang.srt.hardware_backend.npu.memory_pool_npu import NPUMLATokenToKVPool
from sglang.srt.mem_cache.base_swa_memory_pool import BaseSWAKVPool from sglang.srt.mem_cache.base_swa_memory_pool import BaseSWAKVPool
from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool, NSATokenToKVPool from sglang.srt.mem_cache.memory_pool import DSATokenToKVPool, HybridLinearKVPool
kv_args.state_types = [] kv_args.state_types = []
kv_args.state_data_ptrs = [] kv_args.state_data_ptrs = []
@@ -593,9 +593,9 @@ def setup_state_kv_args(
append_state_component( append_state_component(
kv_args, StateType.MAMBA, data_ptrs, data_lens, item_lens, dim kv_args, StateType.MAMBA, data_ptrs, data_lens, item_lens, dim
) )
elif isinstance(token_to_kv_pool, (NSATokenToKVPool, NPUMLATokenToKVPool)): elif isinstance(token_to_kv_pool, (DSATokenToKVPool, NPUMLATokenToKVPool)):
if draft_token_to_kv_pool is not None and isinstance( if draft_token_to_kv_pool is not None and isinstance(
draft_token_to_kv_pool, NSATokenToKVPool draft_token_to_kv_pool, DSATokenToKVPool
): ):
( (
draft_data_ptrs, draft_data_ptrs,
@@ -612,7 +612,7 @@ def setup_state_kv_args(
kv_args.total_kv_layers = total_kv_layers kv_args.total_kv_layers = total_kv_layers
else: else:
append_state_component( append_state_component(
kv_args, StateType.NSA, data_ptrs, data_lens, item_lens kv_args, StateType.DSA, data_ptrs, data_lens, item_lens
) )
if ( if (
python/sglang/srt/environ.py
+46 -4
View File
@@ -134,6 +134,41 @@ class EnvInt(EnvField):
raise ValueError(f'"{value}" is not a valid integer value') raise ValueError(f'"{value}" is not a valid integer value')
class _DeprecatedEnvFallback:
"""Mixin for EnvField subclasses: if the canonical env var is not set,
check *deprecated_name* and emit DeprecationWarning before reading it.
Usage:
SGLANG_DSA_FUSE_TOPK = EnvBoolWithAlias(True, deprecated_name="SGLANG_NSA_FUSE_TOPK")
"""
def __init__(self, default: Any, deprecated_name: str):
super().__init__(default)
self.deprecated_name = deprecated_name
def get(self) -> Any:
if os.getenv(self.name) is None:
fallback = os.getenv(self.deprecated_name)
if fallback is not None:
warnings.warn(
f"Environment variable '{self.deprecated_name}' is deprecated; "
f"use '{self.name}' instead. "
"The alias will be removed in a future release.",
DeprecationWarning,
stacklevel=2,
)
os.environ[self.name] = fallback
return super().get()
class EnvBoolWithAlias(_DeprecatedEnvFallback, EnvBool):
pass
class EnvIntWithAlias(_DeprecatedEnvFallback, EnvInt):
pass
class EnvFloat(EnvField): class EnvFloat(EnvField):
def parse(self, value: str) -> float: def parse(self, value: str) -> float:
try: try:
@@ -428,11 +463,18 @@ class Envs:
SGLANG_NIXL_EP_BF16_DISPATCH = EnvBool(False) SGLANG_NIXL_EP_BF16_DISPATCH = EnvBool(False)
SGLANG_NIXL_EP_NUM_MAX_DISPATCH_TOKENS_PER_RANK = EnvInt(128) SGLANG_NIXL_EP_NUM_MAX_DISPATCH_TOKENS_PER_RANK = EnvInt(128)
# NSA Backend # DSA Backend (canonical names; fall back to SGLANG_NSA_* with deprecation warning)
SGLANG_NSA_FUSE_TOPK = EnvBool(True) SGLANG_DSA_FUSE_TOPK = EnvBoolWithAlias(True, deprecated_name="SGLANG_NSA_FUSE_TOPK")
SGLANG_NSA_ENABLE_MTP_PRECOMPUTE_METADATA = EnvBool(True) SGLANG_DSA_ENABLE_MTP_PRECOMPUTE_METADATA = EnvBoolWithAlias(
True, deprecated_name="SGLANG_NSA_ENABLE_MTP_PRECOMPUTE_METADATA"
)
SGLANG_DSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD = EnvIntWithAlias(
2048, deprecated_name="SGLANG_NSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD"
)
SGLANG_DSA_HIP_DISABLE_PRESHUFFLE = EnvBoolWithAlias(
False, deprecated_name="SGLANG_NSA_HIP_DISABLE_PRESHUFFLE"
)
SGLANG_USE_FUSED_METADATA_COPY = EnvBool(True) SGLANG_USE_FUSED_METADATA_COPY = EnvBool(True)
SGLANG_NSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD = EnvInt(2048)
# sgl-kernel # sgl-kernel
SGLANG_SKIP_SGL_KERNEL_VERSION_CHECK = EnvBool(False) SGLANG_SKIP_SGL_KERNEL_VERSION_CHECK = EnvBool(False)
python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py
+2 -2
View File
@@ -20,7 +20,7 @@ from sglang.srt.hardware_backend.npu.attention.mla_preprocess import (
is_mla_preprocess_enabled, is_mla_preprocess_enabled,
) )
from sglang.srt.layers.attention.base_attn_backend import AttentionBackend from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
from sglang.srt.layers.attention.nsa.utils import is_nsa_enable_prefill_cp from sglang.srt.layers.attention.dsa.utils import is_dsa_enable_prefill_cp
from sglang.srt.layers.dp_attention import get_attention_tp_size from sglang.srt.layers.dp_attention import get_attention_tp_size
from sglang.srt.layers.radix_attention import AttentionType from sglang.srt.layers.radix_attention import AttentionType
from sglang.srt.layers.utils.cp_utils import cp_all_gather_rerange_kv_cache from sglang.srt.layers.utils.cp_utils import cp_all_gather_rerange_kv_cache
@@ -932,7 +932,7 @@ class AscendAttnBackend(AttentionBackend):
if ( if (
is_prefill is_prefill
and is_nsa_enable_prefill_cp() and is_dsa_enable_prefill_cp()
and forward_batch.attn_cp_metadata is not None and forward_batch.attn_cp_metadata is not None
): ):
attn_out = self.do_cp_balance_attn( attn_out = self.do_cp_balance_attn(
python/sglang/srt/hardware_backend/npu/graph_runner/eagle_draft_extend_npu_graph_runner.py
+2 -2
View File
@@ -20,7 +20,7 @@ from typing import TYPE_CHECKING
import torch import torch
from sglang.srt.configs.model_config import is_deepseek_nsa from sglang.srt.configs.model_config import is_deepseek_dsa
from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.speculative.eagle_draft_extend_cuda_graph_runner import ( from sglang.srt.speculative.eagle_draft_extend_cuda_graph_runner import (
EAGLEDraftExtendCudaGraphRunner, EAGLEDraftExtendCudaGraphRunner,
@@ -59,7 +59,7 @@ class EAGLEDraftExtendNpuGraphRunner(EAGLEDraftExtendCudaGraphRunner):
) )
def _replay(self, forward_batch: ForwardBatch): def _replay(self, forward_batch: ForwardBatch):
if not is_deepseek_nsa(self.model_runner.model_config.hf_config): if not is_deepseek_dsa(self.model_runner.model_config.hf_config):
seq_lens = forward_batch.seq_lens_cpu.tolist() + [0] * ( seq_lens = forward_batch.seq_lens_cpu.tolist() + [0] * (
self.bs - self.raw_bs self.bs - self.raw_bs
) )
python/sglang/srt/hardware_backend/npu/graph_runner/eagle_draft_npu_graph_runner.py
+2 -2
View File
@@ -22,7 +22,7 @@ from typing import TYPE_CHECKING, Dict, Union
import numpy as np import numpy as np
import torch import torch
from sglang.srt.configs.model_config import AttentionArch, is_deepseek_nsa from sglang.srt.configs.model_config import AttentionArch, is_deepseek_dsa
from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.speculative.eagle_draft_cuda_graph_runner import ( from sglang.srt.speculative.eagle_draft_cuda_graph_runner import (
EAGLEDraftCudaGraphRunner, EAGLEDraftCudaGraphRunner,
@@ -96,7 +96,7 @@ class EAGLEDraftNpuGraphRunner(EAGLEDraftCudaGraphRunner):
def _replay(self, forward_batch: ForwardBatch): def _replay(self, forward_batch: ForwardBatch):
self.update_attr_name = self._get_update_attr_name() self.update_attr_name = self._get_update_attr_name()
self.update_attr_type = self._get_update_attr_type() self.update_attr_type = self._get_update_attr_type()
if not is_deepseek_nsa(self.model_runner.model_config.hf_config): if not is_deepseek_dsa(self.model_runner.model_config.hf_config):
seq_lens_for_each_draft_step = [] seq_lens_for_each_draft_step = []
for speculative_step_id in range(self.speculative_num_steps - 1): for speculative_step_id in range(self.speculative_num_steps - 1):
seq_lens_cpu = forward_batch.seq_lens_cpu + speculative_step_id + 1 seq_lens_cpu = forward_batch.seq_lens_cpu + speculative_step_id + 1
python/sglang/srt/hardware_backend/npu/graph_runner/npu_graph_runner.py
+2 -2
View File
@@ -26,7 +26,7 @@ import numpy as np
import torch import torch
import sglang import sglang
from sglang.srt.configs.model_config import AttentionArch, is_deepseek_nsa from sglang.srt.configs.model_config import AttentionArch, is_deepseek_dsa
from sglang.srt.distributed.parallel_state import GroupCoordinator from sglang.srt.distributed.parallel_state import GroupCoordinator
from sglang.srt.environ import envs from sglang.srt.environ import envs
from sglang.srt.model_executor.cuda_graph_runner import CudaGraphRunner from sglang.srt.model_executor.cuda_graph_runner import CudaGraphRunner
@@ -188,7 +188,7 @@ class NPUGraphRunner(CudaGraphRunner):
self.update_attr_name = self._get_update_attr_name() self.update_attr_name = self._get_update_attr_name()
self.update_attr_type = self._get_update_attr_type() self.update_attr_type = self._get_update_attr_type()
# Replay # Replay
if not is_deepseek_nsa(self.model_runner.model_config.hf_config): if not is_deepseek_dsa(self.model_runner.model_config.hf_config):
if forward_batch.forward_mode.is_target_verify(): if forward_batch.forward_mode.is_target_verify():
seq_lens_cpu = forward_batch.seq_lens.cpu() + self.num_tokens_per_bs seq_lens_cpu = forward_batch.seq_lens.cpu() + self.num_tokens_per_bs
seq_lens = seq_lens_cpu.tolist() + [0] * (self.bs - self.raw_bs) seq_lens = seq_lens_cpu.tolist() + [0] * (self.bs - self.raw_bs)
python/sglang/srt/hardware_backend/npu/modules/deepseek_v2_attention_mla_npu.py
+9 -9
View File
@@ -11,9 +11,9 @@ from sglang.srt.hardware_backend.npu.attention.mla_preprocess import (
is_fia_nz, is_fia_nz,
is_mla_preprocess_enabled, is_mla_preprocess_enabled,
) )
from sglang.srt.layers.attention.nsa.nsa_indexer import scattered_to_tp_attn_full from sglang.srt.layers.attention.dsa.dsa_indexer import scattered_to_tp_attn_full
from sglang.srt.layers.attention.nsa.utils import ( from sglang.srt.layers.attention.dsa.utils import (
nsa_use_prefill_cp, dsa_use_prefill_cp,
) )
from sglang.srt.layers.communicator import ScatterMode, get_attn_tp_context from sglang.srt.layers.communicator import ScatterMode, get_attn_tp_context
@@ -43,9 +43,9 @@ def forward_mha_prepare_npu(
) )
) )
# NSA Indexer: cache quantized keys, auto-skip topk for sequences <= nsa_index_topk # DSA Indexer: cache quantized keys, auto-skip topk for sequences <= dsa_index_topk
if m.use_nsa: if m.use_dsa:
q_lora = m.q_a_layernorm(q) q_lora = m.q_a_layernorm(q)
q = m.q_b_proj(q_lora)[0].view(-1, m.num_local_heads, m.qk_head_dim) q = m.q_b_proj(q_lora)[0].view(-1, m.num_local_heads, m.qk_head_dim)
_ = m.indexer( _ = m.indexer(
@@ -206,7 +206,7 @@ def forward_mla_prepare_npu(
k_nope = m.kv_a_layernorm(k_nope) k_nope = m.kv_a_layernorm(k_nope)
# q_lora needed by indexer # q_lora needed by indexer
if m.use_nsa: if m.use_dsa:
q_lora = q q_lora = q
k_nope = k_nope.unsqueeze(1) k_nope = k_nope.unsqueeze(1)
@@ -226,7 +226,7 @@ def forward_mla_prepare_npu(
q_pe, k_pe = m.rotary_emb(positions, q_pe, k_pe) q_pe, k_pe = m.rotary_emb(positions, q_pe, k_pe)
if nsa_use_prefill_cp(forward_batch): if dsa_use_prefill_cp(forward_batch):
# support allgather+rerrange # support allgather+rerrange
k_nope, k_pe = m.rebuild_cp_kv_cache( k_nope, k_pe = m.rebuild_cp_kv_cache(
latent_cache, forward_batch, k_nope, k_pe latent_cache, forward_batch, k_nope, k_pe
@@ -359,7 +359,7 @@ def forward_dsa_prepare_npu(
if q_event is not None: if q_event is not None:
torch.npu.current_stream().wait_event(q_event) torch.npu.current_stream().wait_event(q_event)
else: else:
if fused_qkv_a_proj_out.shape[0] < 65535 and not nsa_use_prefill_cp( if fused_qkv_a_proj_out.shape[0] < 65535 and not dsa_use_prefill_cp(
forward_batch forward_batch
): ):
q_lora, k_nope, k_pe = fused_split_qk_norm( q_lora, k_nope, k_pe = fused_split_qk_norm(
@@ -398,7 +398,7 @@ def forward_dsa_prepare_npu(
q_pe, k_pe = m.rotary_emb(positions, q_pe, k_pe) q_pe, k_pe = m.rotary_emb(positions, q_pe, k_pe)
if nsa_use_prefill_cp(forward_batch): if dsa_use_prefill_cp(forward_batch):
# support allgather+rerrange # support allgather+rerrange
k_nope, k_pe = m.rebuild_cp_kv_cache( k_nope, k_pe = m.rebuild_cp_kv_cache(
latent_cache, forward_batch, k_nope, k_pe latent_cache, forward_batch, k_nope, k_pe
python/sglang/srt/layers/attention/attention_registry.py
+16 -4
View File
@@ -1,4 +1,5 @@
import logging import logging
import warnings
from typing import TYPE_CHECKING from typing import TYPE_CHECKING
from sglang.srt.configs.linear_attn_model_registry import ( from sglang.srt.configs.linear_attn_model_registry import (
@@ -96,11 +97,22 @@ def create_ascend_backend(runner):
return AscendAttnBackend(runner) return AscendAttnBackend(runner)
@register_attention_backend("nsa") @register_attention_backend("dsa")
def create_nsa_backend(runner): def create_dsa_backend(runner):
from sglang.srt.layers.attention.nsa_backend import NativeSparseAttnBackend from sglang.srt.layers.attention.dsa_backend import DeepseekSparseAttnBackend
return NativeSparseAttnBackend(runner) return DeepseekSparseAttnBackend(runner)
@register_attention_backend("nsa")
def _create_nsa_compat(runner):
warnings.warn(
"attention-backend='nsa' is deprecated; use 'dsa' instead. "
"The alias will be removed in a future release.",
DeprecationWarning,
stacklevel=2,
)
return create_dsa_backend(runner)
@register_attention_backend("dsv4") @register_attention_backend("dsv4")
python/sglang/srt/layers/attention/base_attn_backend.py
+1 -1
View File
@@ -9,7 +9,7 @@ from sglang.kernel_api_logging import debug_kernel_api
from sglang.srt.utils.common import is_npu from sglang.srt.utils.common import is_npu
if TYPE_CHECKING: if TYPE_CHECKING:
from sglang.srt.layers.attention.nsa.nsa_indexer import BaseIndexerMetadata from sglang.srt.layers.attention.dsa.dsa_indexer import BaseIndexerMetadata
from sglang.srt.layers.radix_attention import RadixAttention from sglang.srt.layers.radix_attention import RadixAttention
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
from sglang.srt.speculative.spec_info import SpecInput from sglang.srt.speculative.spec_info import SpecInput
python/sglang/srt/layers/attention/dsa/dequant_k_cache.py
+289
View File
@@ -0,0 +1,289 @@
import torch
import triton
import triton.language as tl
def dequantize_k_cache(quant_k_cache):
return _dequantize_k_cache_fast_wrapped(quant_k_cache)
def _dequantize_k_cache_ref(
quant_k_cache: torch.Tensor, # (num_blocks, block_size, 1, bytes_per_token)
dv: int = 512,
tile_size: int = 128,
d: int = 576,
) -> torch.Tensor:
"""
De-quantize the k-cache
"""
assert dv % tile_size == 0
original_ndim = quant_k_cache.ndim
if original_ndim == 3:
# set block_size = 1
quant_k_cache = quant_k_cache.unsqueeze(1)
num_tiles = dv // tile_size
num_blocks, block_size, h_k, _ = quant_k_cache.shape
assert h_k == 1
result = torch.empty(
(num_blocks, block_size, d), dtype=torch.bfloat16, device=quant_k_cache.device
)
quant_k_cache = quant_k_cache.view(num_blocks, block_size, -1)
input_nope = quant_k_cache[..., :dv]
input_scale = quant_k_cache[..., dv : dv + num_tiles * 4].view(torch.float32)
input_rope = quant_k_cache[..., dv + num_tiles * 4 :].view(torch.bfloat16)
result[..., dv:] = input_rope
for tile_idx in range(0, num_tiles):
cur_nope = input_nope[
..., tile_idx * tile_size : (tile_idx + 1) * tile_size
].to(torch.float32)
cur_scales = input_scale[..., tile_idx].unsqueeze(-1)
result[..., tile_idx * tile_size : (tile_idx + 1) * tile_size] = (
cur_nope * cur_scales
)
if original_ndim == 3:
return result.view(num_blocks, 1, -1)
else:
return result.view(num_blocks, block_size, 1, -1)
def _dequantize_k_cache_fast_wrapped(
quant_k_cache: torch.Tensor,
dv: int = 512,
tile_size: int = 128,
) -> torch.Tensor:
original_ndim = quant_k_cache.ndim
if original_ndim == 3:
# set block_size = 1
quant_k_cache = quant_k_cache.unsqueeze(1)
num_blocks, block_size, _, dim_quant = quant_k_cache.shape
assert dv == 512
assert dim_quant == 656
assert tile_size == 128
quant_k_cache = quant_k_cache.view((-1, dim_quant))
output = _dequantize_k_cache_fast(quant_k_cache)
if original_ndim == 3:
return output.view(num_blocks, 1, -1)
else:
return output.view(num_blocks, block_size, 1, -1)
def _dequantize_k_cache_fast(quant_k_cache, group_size: int = 128):
num_tokens, dim_quant = quant_k_cache.shape
assert quant_k_cache.dtype == torch.float8_e4m3fn
dim_nope = 512
dim_rope = 64
num_tiles = dim_nope // group_size
assert dim_quant == 656
output = torch.empty(
(num_tokens, dim_nope + dim_rope),
dtype=torch.bfloat16,
device=quant_k_cache.device,
)
num_blocks_per_token = triton.cdiv(dim_nope + dim_rope, group_size)
assert num_blocks_per_token == 5
assert dim_nope % group_size == 0
input_nope_q = quant_k_cache[:, :dim_nope]
input_nope_s = quant_k_cache[:, dim_nope : dim_nope + num_tiles * 4].view(
torch.float32
)
input_rope = quant_k_cache[:, dim_nope + num_tiles * 4 :].view(torch.bfloat16)
_dequantize_k_cache_fast_kernel[(num_tokens, num_blocks_per_token)](
output,
input_nope_q,
input_nope_s,
input_rope,
output.stride(0),
input_nope_q.stride(0),
input_nope_s.stride(0),
input_rope.stride(0),
NUM_NOPE_BLOCKS=num_tiles,
GROUP_SIZE=group_size,
DIM_NOPE=dim_nope,
DIM_ROPE=dim_rope,
)
return output
@triton.jit
def _dequantize_k_cache_fast_kernel(
output_ptr,
input_nope_q_ptr,
input_nope_s_ptr,
input_rope_ptr,
output_stride_0: int,
input_nope_q_stride_0: int,
input_nope_s_stride_0: int,
input_rope_stride_0: int,
NUM_NOPE_BLOCKS: tl.constexpr,
GROUP_SIZE: tl.constexpr,
DIM_NOPE: tl.constexpr,
DIM_ROPE: tl.constexpr,
):
token_id = tl.program_id(0)
raw_block_id = tl.program_id(1)
if raw_block_id < NUM_NOPE_BLOCKS:
# a. dequant nope
effective_block_id = raw_block_id
offs_q = effective_block_id * GROUP_SIZE + tl.arange(0, GROUP_SIZE)
mask = offs_q < DIM_NOPE
ptr_q = input_nope_q_ptr + token_id * input_nope_q_stride_0 + offs_q
ptr_s = input_nope_s_ptr + token_id * input_nope_s_stride_0 + effective_block_id
y_q = tl.load(ptr_q, mask=mask, other=0.0).to(tl.float32)
y_s = tl.load(ptr_s)
y = (y_q * y_s).to(output_ptr.dtype.element_ty)
dst_ptr = output_ptr + token_id * output_stride_0 + offs_q
tl.store(dst_ptr, y, mask=mask)
else:
# b. copy rope
effective_block_id = raw_block_id - NUM_NOPE_BLOCKS
offs = effective_block_id * GROUP_SIZE + tl.arange(0, GROUP_SIZE)
mask = offs < DIM_ROPE
src_ptr = input_rope_ptr + token_id * input_rope_stride_0 + offs
dst_ptr = output_ptr + token_id * output_stride_0 + DIM_NOPE + offs
data = tl.load(src_ptr, mask=mask).to(tl.bfloat16)
tl.store(dst_ptr, data, mask=mask)
def dequantize_k_cache_paged(
quant_k_cache: torch.Tensor,
page_table_1_flattened: torch.Tensor,
group_size: int = 128,
) -> torch.Tensor:
"""
De-quantize the k-cache with paged layout
Args:
quant_k_cache: [total_num_tokens, 1, dim_quant] or [num_blocks, block_size, 1, dim_quant], the quantized k-cache in paged layout
page_table_1_flattened: [num_tokens], the flattened page_table_1 with the page indices in each requests concatenated together
Returns:
output: [num_tokens, 1, dim_nope + dim_rope], the de-quantized k-cache
"""
dim_quant = quant_k_cache.shape[-1]
assert (
dim_quant == 656
), f"dim_quant: {dim_quant} != 656 detected in dequantize_k_cache_paged"
quant_k_cache = quant_k_cache.view((-1, dim_quant))
# num_tokens can exceed kv_cache_size due to prefix sharing (multiple seqs share same KV slots)
# Index bounds validated in dsa_backend.init_forward_metadata
num_tokens = page_table_1_flattened.shape[0]
assert quant_k_cache.dtype == torch.float8_e4m3fn
dim_nope = 512
dim_rope = 64
num_tiles = dim_nope // group_size # 512 // 128 = 4
output = torch.empty(
(num_tokens, 1, dim_nope + dim_rope),
dtype=torch.bfloat16,
device=quant_k_cache.device,
)
# cdiv(512 + 64, 128) = 5
num_blocks_per_token = triton.cdiv(dim_nope + dim_rope, group_size)
assert num_blocks_per_token == 5
assert dim_nope % group_size == 0
input_nope_q = quant_k_cache[:, :dim_nope]
# [:, 512:512+4*4] = [:, 512:528]
input_nope_s = quant_k_cache[:, dim_nope : dim_nope + num_tiles * 4].view(
torch.float32
)
# [:, 528:]
input_rope = quant_k_cache[:, dim_nope + num_tiles * 4 :].view(torch.bfloat16)
_dequantize_k_cache_paged_kernel[(num_tokens, num_blocks_per_token)](
output,
input_nope_q,
input_nope_s,
input_rope,
page_table_1_flattened,
output.stride(0),
input_nope_q.stride(0),
input_nope_s.stride(0),
input_rope.stride(0),
NUM_NOPE_BLOCKS=num_tiles,
GROUP_SIZE=group_size,
DIM_NOPE=dim_nope,
DIM_ROPE=dim_rope,
)
return output
@triton.jit
def _dequantize_k_cache_paged_kernel(
output_ptr,
input_nope_q_ptr,
input_nope_s_ptr,
input_rope_ptr,
page_table_1_ptr,
output_stride_0: int,
input_nope_q_stride_0: int,
input_nope_s_stride_0: int,
input_rope_stride_0: int,
NUM_NOPE_BLOCKS: tl.constexpr,
GROUP_SIZE: tl.constexpr,
DIM_NOPE: tl.constexpr,
DIM_ROPE: tl.constexpr,
):
token_id = tl.program_id(0)
token_id_paged = tl.load(page_table_1_ptr + token_id).to(tl.int32)
raw_block_id = tl.program_id(1)
if raw_block_id < NUM_NOPE_BLOCKS:
# a. dequant nope
effective_block_id = raw_block_id
offs_q = effective_block_id * GROUP_SIZE + tl.arange(0, GROUP_SIZE)
mask = offs_q < DIM_NOPE
ptr_q = input_nope_q_ptr + token_id_paged * input_nope_q_stride_0 + offs_q
ptr_s = (
input_nope_s_ptr
+ token_id_paged * input_nope_s_stride_0
+ effective_block_id
)
y_q = tl.load(ptr_q, mask=mask, other=0.0).to(tl.float32)
y_s = tl.load(ptr_s)
y = (y_q * y_s).to(output_ptr.dtype.element_ty)
dst_ptr = output_ptr + token_id * output_stride_0 + offs_q
tl.store(dst_ptr, y, mask=mask)
else:
# b. copy rope
effective_block_id = raw_block_id - NUM_NOPE_BLOCKS
offs = effective_block_id * GROUP_SIZE + tl.arange(0, GROUP_SIZE)
mask = offs < DIM_ROPE
src_ptr = input_rope_ptr + token_id_paged * input_rope_stride_0 + offs
dst_ptr = output_ptr + token_id * output_stride_0 + DIM_NOPE + offs
data = tl.load(src_ptr, mask=mask).to(tl.bfloat16)
tl.store(dst_ptr, data, mask=mask)
if __name__ == "__main__":
raise Exception("UT is in quant_k_cache.py")
python/sglang/srt/layers/attention/dsa/dsa_backend_mtp_precompute.py
+331
View File
@@ -0,0 +1,331 @@
"""Multi-step precompute utilities for Native Sparse Attention backend.
This module provides optimization utilities for multi-step speculative decoding
by precomputing shared metadata once and copying it to multiple backend instances.
"""
from __future__ import annotations
from dataclasses import dataclass
from typing import TYPE_CHECKING, Optional
import torch
from sglang.srt.layers.attention.dsa.utils import compute_dsa_seqlens
if TYPE_CHECKING:
from sglang.srt.model_executor.forward_batch_info import ForwardMode
from sglang.srt.speculative.spec_info import SpecInput
@dataclass
class PrecomputedMetadata:
"""Precomputed metadata shared across multiple backend instances.
Used for multi-step speculative decoding where multiple backends
need identical metadata. Precomputing once and copying N times
is much faster than computing N times.
"""
# Basic seqlens
cache_seqlens: torch.Tensor # int32, [bs]
cu_seqlens_k: torch.Tensor # int32, [bs+1]
# Page table
page_indices: torch.Tensor # int32, [bs, max_len] or [expanded_bs, max_len]
real_page_table: Optional[torch.Tensor] # int32, transformed version
# DSA seqlens
seqlens_expanded: torch.Tensor # int32, [expanded_size]
dsa_cache_seqlens: torch.Tensor # int32, [expanded_size]
dsa_cu_seqlens_k: torch.Tensor # int32, [expanded_size+1]
seqlens_expanded_size: int
# Dimensions
max_len: int # for decode/draft_extend
max_seqlen_k: int # for target_verify
# FlashMLA (optional)
flashmla_metadata: Optional[torch.Tensor] = None
def compute_cu_seqlens(seqlens: torch.Tensor) -> torch.Tensor:
"""Compute cumulative sequence lengths with padding."""
assert seqlens.dtype == torch.int32
return torch.nn.functional.pad(
torch.cumsum(seqlens, dim=0, dtype=torch.int32), (1, 0)
)
class DeepseekSparseAttnBackendMTPPrecomputeMixin:
"""Mixin class providing metadata precomputation for multi-step speculative decoding.
This mixin provides the _precompute_replay_metadata method and its helpers,
which are used to optimize CUDA graph replay in multi-step scenarios.
"""
def _precompute_replay_metadata(
self,
bs: int,
req_pool_indices: torch.Tensor,
seq_lens: torch.Tensor,
seq_lens_cpu: torch.Tensor,
forward_mode: "ForwardMode",
spec_info: Optional["SpecInput"],
) -> PrecomputedMetadata:
"""Precompute all shared metadata for multi-step backends.
This function extracts and computes all operations that are
identical across different backend instances in multi-step
speculative decoding.
Args:
bs: Batch size
req_pool_indices: Request pool indices [bs]
seq_lens: Sequence lengths [bs]
seq_lens_cpu: Sequence lengths on CPU [bs]
forward_mode: Forward mode (decode/target_verify/draft_extend)
spec_info: Speculative decoding info (for draft_extend mode)
Returns:
PrecomputedMetadata containing all shared intermediate results
"""
# Slice inputs to batch size
seq_lens = seq_lens[:bs]
seq_lens_cpu = seq_lens_cpu[:bs]
req_pool_indices = req_pool_indices[:bs]
# Dispatch to mode-specific precomputation
if forward_mode.is_decode_or_idle():
return self._precompute_decode_mode(
bs, req_pool_indices, seq_lens, seq_lens_cpu
)
elif forward_mode.is_target_verify():
return self._precompute_target_verify_mode(
bs, req_pool_indices, seq_lens, seq_lens_cpu
)
elif forward_mode.is_draft_extend():
return self._precompute_draft_extend_mode(
bs, req_pool_indices, seq_lens, seq_lens_cpu, spec_info
)
else:
raise ValueError(f"Unsupported forward mode: {forward_mode}")
def _precompute_decode_mode(
self,
bs: int,
req_pool_indices: torch.Tensor,
seq_lens: torch.Tensor,
seq_lens_cpu: torch.Tensor,
) -> PrecomputedMetadata:
"""Precompute metadata for normal decode mode."""
max_len = int(seq_lens_cpu.max().item())
# Convert to int32 and compute cumsum
cache_seqlens = seq_lens.to(torch.int32)
cu_seqlens_k = compute_cu_seqlens(cache_seqlens)
# Get page indices from cache
page_indices = self.req_to_token[req_pool_indices, :max_len].contiguous()
# Compute DSA seqlens
dsa_cache_seqlens = compute_dsa_seqlens(
cache_seqlens, dsa_index_topk=self.dsa_index_topk
)
seqlens_expanded = cache_seqlens
seqlens_expanded_size = seqlens_expanded.shape[0]
# Compute DSA cumsum
dsa_cu_seqlens_k = compute_cu_seqlens(dsa_cache_seqlens)
# Transform page table if needed
if self.real_page_size > 1:
real_page_table = self._transform_table_1_to_real(page_indices)
else:
real_page_table = None # Will use page_indices directly
# Compute FlashMLA metadata if needed
flashmla_metadata = None
if self.dsa_decode_impl == "flashmla_kv":
flashmla_metadata = self._compute_flashmla_metadata(
cache_seqlens=dsa_cache_seqlens,
seq_len_q=1,
)
return PrecomputedMetadata(
cache_seqlens=cache_seqlens,
cu_seqlens_k=cu_seqlens_k,
page_indices=page_indices,
real_page_table=real_page_table,
seqlens_expanded=seqlens_expanded,
dsa_cache_seqlens=dsa_cache_seqlens,
dsa_cu_seqlens_k=dsa_cu_seqlens_k,
seqlens_expanded_size=seqlens_expanded_size,
max_len=max_len,
max_seqlen_k=max_len,
flashmla_metadata=flashmla_metadata,
)
def _precompute_target_verify_mode(
self,
bs: int,
req_pool_indices: torch.Tensor,
seq_lens: torch.Tensor,
seq_lens_cpu: torch.Tensor,
) -> PrecomputedMetadata:
"""Precompute metadata for target verify mode."""
max_seqlen_k = int(
seq_lens_cpu.max().item() + self.speculative_num_draft_tokens
)
# Cache seqlens with draft tokens
cache_seqlens = (seq_lens + self.speculative_num_draft_tokens).to(torch.int32)
cu_seqlens_k = compute_cu_seqlens(cache_seqlens)
# Page indices (repeated for each draft token)
page_indices = self.req_to_token[req_pool_indices, :max_seqlen_k]
page_indices = torch.repeat_interleave(
page_indices, repeats=self.speculative_num_draft_tokens, dim=0
).contiguous()
# Generate expanded seqlens
extend_seq_lens_cpu = [self.speculative_num_draft_tokens] * bs
seqlens_int32_cpu = [
self.speculative_num_draft_tokens + kv_len
for kv_len in seq_lens_cpu.tolist()
]
seqlens_expanded = torch.cat(
[
torch.arange(
kv_len - qo_len + 1,
kv_len + 1,
dtype=torch.int32,
device=self.device,
)
for qo_len, kv_len in zip(
extend_seq_lens_cpu,
seqlens_int32_cpu,
strict=True,
)
]
)
# Compute DSA seqlens
dsa_cache_seqlens = compute_dsa_seqlens(seqlens_expanded, self.dsa_index_topk)
seqlens_expanded_size = seqlens_expanded.shape[0]
# DSA cumsum
dsa_cu_seqlens_k = compute_cu_seqlens(dsa_cache_seqlens)
# Transform page table
if self.real_page_size > 1:
real_page_table = self._transform_table_1_to_real(page_indices)
else:
real_page_table = None
# FlashMLA metadata
flashmla_metadata = None
if self.dsa_decode_impl == "flashmla_kv":
flashmla_metadata = self._compute_flashmla_metadata(
cache_seqlens=dsa_cache_seqlens,
seq_len_q=1,
)
return PrecomputedMetadata(
cache_seqlens=cache_seqlens,
cu_seqlens_k=cu_seqlens_k,
page_indices=page_indices,
real_page_table=real_page_table,
seqlens_expanded=seqlens_expanded,
dsa_cache_seqlens=dsa_cache_seqlens,
dsa_cu_seqlens_k=dsa_cu_seqlens_k,
seqlens_expanded_size=seqlens_expanded_size,
max_len=-1, # Not used in this mode
max_seqlen_k=max_seqlen_k,
flashmla_metadata=flashmla_metadata,
)
def _precompute_draft_extend_mode(
self,
bs: int,
req_pool_indices: torch.Tensor,
seq_lens: torch.Tensor,
seq_lens_cpu: torch.Tensor,
spec_info: "SpecInput",
) -> PrecomputedMetadata:
"""Precompute metadata for draft extend mode."""
max_seqlen_k = int(seq_lens_cpu.max().item())
# Cache seqlens
cache_seqlens = seq_lens.to(torch.int32)
cu_seqlens_k = compute_cu_seqlens(cache_seqlens)
# Extend seqlens from spec_info: num_accept_tokens already includes
# the bonus token (drafts + 1).
extend_seq_lens = spec_info.num_accept_tokens[:bs]
extend_seq_lens_cpu = extend_seq_lens.tolist()
# Page indices (repeated per accept length)
page_indices = self.req_to_token[req_pool_indices, :max_seqlen_k]
page_indices = torch.repeat_interleave(
page_indices, repeats=extend_seq_lens, dim=0
).contiguous()
# Generate expanded seqlens
seqlens_expanded = torch.cat(
[
torch.arange(
kv_len - qo_len + 1,
kv_len + 1,
dtype=torch.int32,
device=self.device,
)
for qo_len, kv_len in zip(
extend_seq_lens_cpu,
seq_lens_cpu.tolist(),
strict=True,
)
]
)
# Compute DSA seqlens
dsa_cache_seqlens = compute_dsa_seqlens(seqlens_expanded, self.dsa_index_topk)
seqlens_expanded_size = seqlens_expanded.shape[0]
# DSA cumsum
dsa_cu_seqlens_k = compute_cu_seqlens(dsa_cache_seqlens)
# Transform page table
if self.real_page_size > 1:
real_page_table = self._transform_table_1_to_real(page_indices)
else:
real_page_table = None
# FlashMLA metadata
flashmla_metadata = None
if self.dsa_decode_impl == "flashmla_kv":
flashmla_metadata = self._compute_flashmla_metadata(
cache_seqlens=dsa_cache_seqlens,
seq_len_q=1,
)
return PrecomputedMetadata(
cache_seqlens=cache_seqlens,
cu_seqlens_k=cu_seqlens_k,
page_indices=page_indices,
real_page_table=real_page_table,
seqlens_expanded=seqlens_expanded,
dsa_cache_seqlens=dsa_cache_seqlens,
dsa_cu_seqlens_k=dsa_cu_seqlens_k,
seqlens_expanded_size=seqlens_expanded_size,
max_len=max_seqlen_k,
max_seqlen_k=max_seqlen_k,
flashmla_metadata=flashmla_metadata,
)
# Backward-compat alias
DeepseekSparseAttnBackendMTPPrecomputeMixin = (
DeepseekSparseAttnBackendMTPPrecomputeMixin
)
python/sglang/srt/layers/attention/dsa/dsa_indexer.py
+1746
View File
File diff suppressed because it is too large Load Diff
python/sglang/srt/layers/attention/dsa/dsa_mtp_verification.py
+407
View File
@@ -0,0 +1,407 @@
"""
Verification utilities for DSA backend fused metadata copy operations.
This module contains verification code to ensure that fused metadata copy kernels
produce the same results as individual copy operations.
"""
import torch
def verify_single_backend_fused_metadata_copy(
metadata,
precomputed,
forward_mode,
bs,
flashmla_num_splits_src=None,
flashmla_metadata_src=None,
flashmla_num_splits_dst=None,
flashmla_metadata_dst=None,
):
"""
Verify that the fused metadata copy kernel produces the same results as individual copies.
Args:
metadata: The DSA metadata object containing destination tensors
precomputed: The precomputed metadata containing source tensors
forward_mode: The forward mode (decode, target_verify, or draft_extend)
bs: Batch size
flashmla_num_splits_src: Source FlashMLA num_splits tensor (optional)
flashmla_metadata_src: Source FlashMLA metadata tensor (optional)
flashmla_num_splits_dst: Destination FlashMLA num_splits tensor (optional)
flashmla_metadata_dst: Destination FlashMLA metadata tensor (optional)
Raises:
RuntimeError: If verification fails (tensors don't match)
"""
# Clone destination tensors to preserve fused kernel results
fused_cache_seqlens = metadata.cache_seqlens_int32.clone()
fused_cu_seqlens_k = metadata.cu_seqlens_k.clone()
fused_page_table_1 = metadata.page_table_1.clone()
fused_dsa_cache_seqlens = metadata.dsa_cache_seqlens_int32.clone()
fused_dsa_seqlens_expanded = metadata.dsa_seqlens_expanded.clone()
fused_dsa_cu_seqlens_k = metadata.dsa_cu_seqlens_k.clone()
fused_real_page_table = (
metadata.real_page_table.clone()
if precomputed.real_page_table is not None
else None
)
fused_flashmla_num_splits = None
fused_flashmla_metadata = None
if precomputed.flashmla_metadata is not None:
fused_flashmla_num_splits = flashmla_num_splits_dst.clone()
fused_flashmla_metadata = flashmla_metadata_dst.clone()
# Create reference tensors (zeroed out)
ref_cache_seqlens = torch.zeros_like(metadata.cache_seqlens_int32)
ref_cu_seqlens_k = torch.zeros_like(metadata.cu_seqlens_k)
ref_page_table_1 = torch.zeros_like(metadata.page_table_1)
ref_dsa_cache_seqlens = torch.zeros_like(metadata.dsa_cache_seqlens_int32)
ref_dsa_seqlens_expanded = torch.zeros_like(metadata.dsa_seqlens_expanded)
ref_dsa_cu_seqlens_k = torch.zeros_like(metadata.dsa_cu_seqlens_k)
ref_real_page_table = (
torch.zeros_like(metadata.real_page_table)
if precomputed.real_page_table is not None
else None
)
ref_flashmla_num_splits = None
ref_flashmla_metadata = None
if precomputed.flashmla_metadata is not None:
ref_flashmla_num_splits = torch.zeros_like(flashmla_num_splits_dst)
ref_flashmla_metadata = torch.zeros_like(flashmla_metadata_dst)
# Run individual copy operations (reference implementation)
ref_cache_seqlens.copy_(precomputed.cache_seqlens)
ref_cu_seqlens_k[1:].copy_(precomputed.cu_seqlens_k[1:])
if forward_mode.is_decode_or_idle():
# Decode mode
ref_page_table_1[:, : precomputed.max_len].copy_(precomputed.page_indices)
ref_dsa_cache_seqlens.copy_(precomputed.dsa_cache_seqlens)
elif forward_mode.is_target_verify():
# Target verify mode
ref_page_table_1[:, : precomputed.max_seqlen_k].copy_(precomputed.page_indices)
ref_dsa_seqlens_expanded.copy_(precomputed.seqlens_expanded)
ref_dsa_cache_seqlens.copy_(precomputed.dsa_cache_seqlens)
elif forward_mode.is_draft_extend():
# Draft extend mode
rows = precomputed.page_indices.shape[0]
cols = precomputed.max_seqlen_k
ref_page_table_1[:rows, :cols].copy_(precomputed.page_indices)
size = precomputed.seqlens_expanded_size
ref_dsa_seqlens_expanded[:size].copy_(precomputed.seqlens_expanded)
ref_dsa_cache_seqlens[:size].copy_(precomputed.dsa_cache_seqlens)
# Copy DSA cu_seqlens
size = precomputed.seqlens_expanded_size
ref_dsa_cu_seqlens_k[1 : 1 + size].copy_(precomputed.dsa_cu_seqlens_k[1 : 1 + size])
# Copy real page table
if precomputed.real_page_table is not None:
rows, cols = precomputed.real_page_table.shape
ref_real_page_table[:rows, :cols].copy_(precomputed.real_page_table)
# Copy FlashMLA metadata
if precomputed.flashmla_metadata is not None:
size = precomputed.seqlens_expanded_size
ref_flashmla_num_splits[: size + 1].copy_(flashmla_num_splits_src[: size + 1])
ref_flashmla_metadata.copy_(flashmla_metadata_src)
# Compare results and crash if inconsistent
def check_tensor_equal(name, fused, ref):
if not torch.equal(fused, ref):
max_diff = (fused.float() - ref.float()).abs().max().item()
mismatched_elements = (fused != ref).sum().item()
total_elements = fused.numel()
raise RuntimeError(
f"FUSED METADATA COPY VERIFICATION FAILED!\n"
f"Tensor: {name}\n"
f"Max difference: {max_diff}\n"
f"Mismatched elements: {mismatched_elements}/{total_elements}\n"
f"Fused shape: {fused.shape}, Ref shape: {ref.shape}\n"
f"Forward mode: {forward_mode}, bs={bs}\n"
f"The fused kernel produces different results than individual copies.\n"
f"This indicates a bug in the fused metadata copy kernel."
)
# Verify all tensors (only compare the slices that were actually updated)
check_tensor_equal("cache_seqlens", fused_cache_seqlens, ref_cache_seqlens)
check_tensor_equal("cu_seqlens_k", fused_cu_seqlens_k, ref_cu_seqlens_k)
# Compare page_table_1 only for the region that was updated
if forward_mode.is_decode_or_idle():
check_tensor_equal(
"page_table_1",
fused_page_table_1[:, : precomputed.max_len],
ref_page_table_1[:, : precomputed.max_len],
)
elif forward_mode.is_target_verify():
check_tensor_equal(
"page_table_1",
fused_page_table_1[:, : precomputed.max_seqlen_k],
ref_page_table_1[:, : precomputed.max_seqlen_k],
)
elif forward_mode.is_draft_extend():
rows = precomputed.page_indices.shape[0]
cols = precomputed.max_seqlen_k
check_tensor_equal(
"page_table_1",
fused_page_table_1[:rows, :cols],
ref_page_table_1[:rows, :cols],
)
# Compare dsa_cache_seqlens only for the region that was updated
if forward_mode.is_decode_or_idle():
check_tensor_equal(
"dsa_cache_seqlens",
fused_dsa_cache_seqlens,
ref_dsa_cache_seqlens,
)
else: # TARGET_VERIFY or DRAFT_EXTEND
size = precomputed.seqlens_expanded_size
check_tensor_equal(
"dsa_cache_seqlens",
fused_dsa_cache_seqlens[:size],
ref_dsa_cache_seqlens[:size],
)
# Compare dsa_seqlens_expanded only for TARGET_VERIFY and DRAFT_EXTEND
if forward_mode.is_target_verify() or forward_mode.is_draft_extend():
size = precomputed.seqlens_expanded_size
check_tensor_equal(
"dsa_seqlens_expanded",
fused_dsa_seqlens_expanded[:size],
ref_dsa_seqlens_expanded[:size],
)
# Compare dsa_cu_seqlens_k only for the region that was updated
size = precomputed.seqlens_expanded_size
check_tensor_equal(
"dsa_cu_seqlens_k",
fused_dsa_cu_seqlens_k[: 1 + size],
ref_dsa_cu_seqlens_k[: 1 + size],
)
if precomputed.real_page_table is not None:
rows, cols = precomputed.real_page_table.shape
check_tensor_equal(
"real_page_table",
fused_real_page_table[:rows, :cols],
ref_real_page_table[:rows, :cols],
)
if precomputed.flashmla_metadata is not None:
size = precomputed.seqlens_expanded_size
check_tensor_equal(
"flashmla_num_splits",
fused_flashmla_num_splits[: size + 1],
ref_flashmla_num_splits[: size + 1],
)
check_tensor_equal(
"flashmla_metadata",
fused_flashmla_metadata,
ref_flashmla_metadata,
)
def verify_multi_backend_fused_metadata_copy(
metadata0,
metadata1,
metadata2,
precomputed,
bs,
flashmla_num_splits_src=None,
flashmla_metadata_src=None,
):
"""
Verify that the multi-backend fused metadata copy kernel produces the same results
as individual copies for all three backends.
Args:
metadata0: The DSA metadata object for backend 0
metadata1: The DSA metadata object for backend 1
metadata2: The DSA metadata object for backend 2
precomputed: The precomputed metadata containing source tensors
bs: Batch size
flashmla_num_splits_src: Source FlashMLA num_splits tensor (optional)
flashmla_metadata_src: Source FlashMLA metadata tensor (optional)
Raises:
RuntimeError: If verification fails (tensors don't match)
"""
# Clone destination tensors to preserve fused kernel results
fused_results = []
for idx, metadata in enumerate([metadata0, metadata1, metadata2]):
fused_cache_seqlens = metadata.cache_seqlens_int32.clone()
fused_cu_seqlens_k = metadata.cu_seqlens_k.clone()
fused_page_table_1 = metadata.page_table_1.clone()
fused_dsa_cache_seqlens = metadata.dsa_cache_seqlens_int32.clone()
fused_dsa_cu_seqlens_k = metadata.dsa_cu_seqlens_k.clone()
fused_real_page_table = (
metadata.real_page_table.clone()
if precomputed.real_page_table is not None
else None
)
fused_flashmla_num_splits = None
fused_flashmla_metadata = None
if precomputed.flashmla_metadata is not None:
fused_flashmla_num_splits = metadata.flashmla_metadata.num_splits.clone()
fused_flashmla_metadata = (
metadata.flashmla_metadata.flashmla_metadata.clone()
)
fused_results.append(
{
"cache_seqlens": fused_cache_seqlens,
"cu_seqlens_k": fused_cu_seqlens_k,
"page_table_1": fused_page_table_1,
"dsa_cache_seqlens": fused_dsa_cache_seqlens,
"dsa_cu_seqlens_k": fused_dsa_cu_seqlens_k,
"real_page_table": fused_real_page_table,
"flashmla_num_splits": fused_flashmla_num_splits,
"flashmla_metadata": fused_flashmla_metadata,
}
)
# Run individual copy operations for each backend (reference implementation)
ref_results = []
for idx in range(3):
metadata = [metadata0, metadata1, metadata2][idx]
# Create reference tensors (zeroed out)
ref_cache_seqlens = torch.zeros_like(metadata.cache_seqlens_int32)
ref_cu_seqlens_k = torch.zeros_like(metadata.cu_seqlens_k)
ref_page_table_1 = torch.zeros_like(metadata.page_table_1)
ref_dsa_cache_seqlens = torch.zeros_like(metadata.dsa_cache_seqlens_int32)
ref_dsa_cu_seqlens_k = torch.zeros_like(metadata.dsa_cu_seqlens_k)
ref_real_page_table = (
torch.zeros_like(metadata.real_page_table)
if precomputed.real_page_table is not None
else None
)
ref_flashmla_num_splits = None
ref_flashmla_metadata = None
if precomputed.flashmla_metadata is not None:
ref_flashmla_num_splits = torch.zeros_like(
metadata.flashmla_metadata.num_splits
)
ref_flashmla_metadata = torch.zeros_like(
metadata.flashmla_metadata.flashmla_metadata
)
# Copy operations (decode mode)
ref_cache_seqlens.copy_(precomputed.cache_seqlens)
ref_cu_seqlens_k[1:].copy_(precomputed.cu_seqlens_k[1:])
ref_page_table_1[:, : precomputed.max_len].copy_(precomputed.page_indices)
ref_dsa_cache_seqlens.copy_(precomputed.dsa_cache_seqlens)
# Copy DSA cu_seqlens
size = precomputed.seqlens_expanded_size
ref_dsa_cu_seqlens_k[1 : 1 + size].copy_(
precomputed.dsa_cu_seqlens_k[1 : 1 + size]
)
# Copy real page table
if precomputed.real_page_table is not None:
rows, cols = precomputed.real_page_table.shape
ref_real_page_table[:rows, :cols].copy_(precomputed.real_page_table)
# Copy FlashMLA metadata
if precomputed.flashmla_metadata is not None:
ref_flashmla_num_splits[: size + 1].copy_(
flashmla_num_splits_src[: size + 1]
)
ref_flashmla_metadata.copy_(flashmla_metadata_src)
ref_results.append(
{
"cache_seqlens": ref_cache_seqlens,
"cu_seqlens_k": ref_cu_seqlens_k,
"page_table_1": ref_page_table_1,
"dsa_cache_seqlens": ref_dsa_cache_seqlens,
"dsa_cu_seqlens_k": ref_dsa_cu_seqlens_k,
"real_page_table": ref_real_page_table,
"flashmla_num_splits": ref_flashmla_num_splits,
"flashmla_metadata": ref_flashmla_metadata,
}
)
# Compare results for all 3 backends
def check_tensor_equal(backend_idx, name, fused, ref):
if not torch.equal(fused, ref):
max_diff = (fused.float() - ref.float()).abs().max().item()
mismatched_elements = (fused != ref).sum().item()
total_elements = fused.numel()
raise RuntimeError(
f"MULTI-BACKEND FUSED METADATA COPY VERIFICATION FAILED!\n"
f"Backend: {backend_idx}\n"
f"Tensor: {name}\n"
f"Max difference: {max_diff}\n"
f"Mismatched elements: {mismatched_elements}/{total_elements}\n"
f"Fused shape: {fused.shape}, Ref shape: {ref.shape}\n"
f"Batch size: {bs}\n"
f"The multi-backend fused kernel produces different results than individual copies.\n"
f"This indicates a bug in the fused metadata copy kernel."
)
# Verify all tensors for all 3 backends (multi-backend is DECODE mode only)
for idx in range(3):
fused = fused_results[idx]
ref = ref_results[idx]
check_tensor_equal(
idx,
"cache_seqlens",
fused["cache_seqlens"],
ref["cache_seqlens"],
)
check_tensor_equal(
idx,
"cu_seqlens_k",
fused["cu_seqlens_k"],
ref["cu_seqlens_k"],
)
# Multi-backend is DECODE mode only, so compare only [:, :max_len]
check_tensor_equal(
idx,
"page_table_1",
fused["page_table_1"][:, : precomputed.max_len],
ref["page_table_1"][:, : precomputed.max_len],
)
check_tensor_equal(
idx,
"dsa_cache_seqlens",
fused["dsa_cache_seqlens"],
ref["dsa_cache_seqlens"],
)
# DECODE mode uses bs for dsa_cu_seqlens_k size
check_tensor_equal(
idx,
"dsa_cu_seqlens_k",
fused["dsa_cu_seqlens_k"][: bs + 1],
ref["dsa_cu_seqlens_k"][: bs + 1],
)
if precomputed.real_page_table is not None:
rows, cols = precomputed.real_page_table.shape
check_tensor_equal(
idx,
"real_page_table",
fused["real_page_table"][:rows, :cols],
ref["real_page_table"][:rows, :cols],
)
if precomputed.flashmla_metadata is not None:
# DECODE mode uses bs + 1 for flashmla_num_splits
check_tensor_equal(
idx,
"flashmla_num_splits",
fused["flashmla_num_splits"][: bs + 1],
ref["flashmla_num_splits"][: bs + 1],
)
check_tensor_equal(
idx,
"flashmla_metadata",
fused["flashmla_metadata"],
ref["flashmla_metadata"],
)
python/sglang/srt/layers/attention/dsa/index_buf_accessor.py
+814
View File
@@ -0,0 +1,814 @@
from typing import TYPE_CHECKING
import torch
import triton
import triton.language as tl
from sglang.srt.layers.attention.dsa.utils import aiter_can_use_preshuffle_paged_mqa
from sglang.srt.layers.quantization.fp8_kernel import is_fp8_fnuz
from sglang.srt.utils import get_bool_env_var, is_hip
_is_hip = is_hip()
_is_fp8_fnuz = is_fp8_fnuz()
_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
# aiter cp_gather kernel with preshuffle=True is only valid when the indexer
# uses the page_size=64 preshuffle layout (i.e. when the matching MQA gluon path
# is also enabled).
_use_aiter_preshuffle = aiter_can_use_preshuffle_paged_mqa()
if _use_aiter_preshuffle:
from aiter.ops.cache import cp_gather_indexer_k_quant_cache
if TYPE_CHECKING:
from sglang.srt.mem_cache.memory_pool import DSATokenToKVPool
"""
k: data, 128 item per token, fp8
s: scale, 1 item per token, fp32
"""
class GetK:
@classmethod
def execute(cls, *args, **kwargs):
return cls.triton(*args, **kwargs)
@classmethod
def slow(
cls, pool: "DSATokenToKVPool", buf, seq_len: int, page_indices: torch.Tensor
):
num_pages = (seq_len + pool.page_size - 1) // pool.page_size
seq_len_ = num_pages * pool.page_size
index_k_fp8 = torch.empty(
(seq_len_, pool.index_head_dim),
dtype=torch.uint8,
device=pool.device,
)
for i in range(num_pages):
page_index = page_indices[i]
index_k_fp8[i * pool.page_size : (i + 1) * pool.page_size] = buf[
page_index
][: pool.page_size * pool.index_head_dim].view(-1, pool.index_head_dim)
return index_k_fp8[:seq_len]
@classmethod
def torch_fast(
cls, pool: "DSATokenToKVPool", buf, seq_len: int, page_indices: torch.Tensor
):
"""
:param page_indices: (num_pages,), int32
:return: (seq_len, index_head_dim), uint8
"""
# can handle per 128B instead of per element
# page_indices: (num_pages,), element := a page index
buf_numel_per_page = buf.shape[1]
num_k_bytes_per_page = pool.page_size * pool.index_head_dim
num_k_bytes_per_token = pool.index_head_dim
# buf: (num_pages, page_size 64 * head_dim 128 + page_size 64 * fp32_nbytes 4), uint8
# flat_buf: (whatever,), uint8
flat_buf = buf.flatten()
# flat_indices: (num_pages, num_k_bytes_per_page), int32, element := an index into flat_buf that we want to access
flat_indices = (page_indices * buf_numel_per_page)[:, None] + torch.arange(
num_k_bytes_per_page, dtype=torch.int32, device="cuda"
)[None, :]
flat_indices = flat_indices.flatten()[: seq_len * num_k_bytes_per_token]
out = flat_buf[flat_indices]
return out.view(-1, 128)
@classmethod
def triton(
cls, pool: "DSATokenToKVPool", buf, seq_len: int, page_indices: torch.Tensor
):
"""
Triton implementation for gathering K data from paged buffer.
:param page_indices: (num_pages,), int32/int64
:return: (seq_len, index_head_dim), uint8
"""
return _get_k_triton(
buf=buf,
page_indices=page_indices,
seq_len=seq_len,
page_size=pool.page_size,
index_head_dim=pool.index_head_dim,
)
class GetS:
@classmethod
def execute(cls, *args, **kwargs):
return cls.triton(*args, **kwargs)
@classmethod
def slow(
cls, pool: "DSATokenToKVPool", buf, seq_len: int, page_indices: torch.Tensor
):
num_pages = (seq_len + pool.page_size - 1) // pool.page_size
seq_len_ = num_pages * pool.page_size
assert pool.index_head_dim // pool.quant_block_size == 1
index_k_scale_fp8 = torch.empty(
(seq_len_, 4),
dtype=torch.uint8,
device=pool.device,
)
for i in range(num_pages):
page_index = page_indices[i]
index_k_scale_fp8[i * pool.page_size : (i + 1) * pool.page_size] = buf[
page_index
][pool.page_size * pool.index_head_dim :].view(-1, 4)
return index_k_scale_fp8[:seq_len]
@classmethod
def torch_fast(
cls, pool: "DSATokenToKVPool", buf, seq_len: int, page_indices: torch.Tensor
):
"""
:param page_indices: (num_pages,), int32
:return: (seq_len, index_head_dim // quant_block_size), uint8
"""
buf_numel_per_page = buf.shape[1]
num_s_bytes_per_page = buf.shape[1] - pool.page_size * pool.index_head_dim
num_s_bytes_per_token = pool.index_head_dim // pool.quant_block_size * 4
s_offset_in_page = pool.page_size * pool.index_head_dim
flat_buf = buf.flatten()
flat_indices = (
(page_indices * buf_numel_per_page)[:, None]
+ torch.arange(num_s_bytes_per_page, dtype=torch.int32, device="cuda")[
None, :
]
+ s_offset_in_page
)
flat_indices = flat_indices.flatten()[: seq_len * num_s_bytes_per_token]
out = flat_buf[flat_indices]
return out.view(-1, 4)
@classmethod
def triton(
cls, pool: "DSATokenToKVPool", buf, seq_len: int, page_indices: torch.Tensor
):
"""
Triton implementation for gathering S (scale) data from paged buffer.
:param page_indices: (num_pages,), int32/int64
:return: (seq_len, 4), uint8
"""
return _get_s_triton(
buf=buf,
page_indices=page_indices,
seq_len=seq_len,
page_size=pool.page_size,
index_head_dim=pool.index_head_dim,
)
class GetKAndS:
@classmethod
def execute(cls, *args, **kwargs):
# The aiter path uses cp_gather_indexer_k_quant_cache(preshuffle=True),
# which only matches the layout produced when the rest of the indexer
# is on the page_size=64 preshuffle path. Otherwise fall back to the
# triton implementation (which works on the page_size=1 legacy layout).
if _use_aiter_preshuffle:
return cls.aiter(*args, **kwargs)
return cls.triton(*args, **kwargs)
@classmethod
def aiter(
cls,
pool: "DSATokenToKVPool",
buf: torch.Tensor,
page_indices: torch.Tensor,
seq_len_tensor: torch.Tensor,
seq_len_sum: int,
max_seq_len: int,
):
from sglang.srt.layers.quantization.fp8_kernel import fp8_dtype
page_size = pool.page_size
index_head_dim = pool.index_head_dim
quant_block_size = pool.quant_block_size
scale_elems = index_head_dim // quant_block_size
kv_cache = buf.view(-1, page_size, index_head_dim + scale_elems * 4).view(
fp8_dtype
)
dst_k = torch.empty(
(seq_len_sum, index_head_dim), dtype=torch.uint8, device=buf.device
)
dst_scale = torch.empty(
(seq_len_sum, scale_elems * 4), dtype=torch.uint8, device=buf.device
)
cu_seq_lens = torch.zeros(
seq_len_tensor.shape[0] + 1, dtype=torch.int32, device=buf.device
)
torch.cumsum(seq_len_tensor.to(torch.int32), dim=0, out=cu_seq_lens[1:])
cp_gather_indexer_k_quant_cache(
kv_cache,
dst_k.view(fp8_dtype),
dst_scale,
page_indices.to(torch.int32),
cu_seq_lens,
preshuffle=True,
)
return dst_k, dst_scale
@classmethod
def triton(
cls,
pool: "DSATokenToKVPool",
buf: torch.Tensor,
page_indices: torch.Tensor,
seq_len_tensor: torch.Tensor,
seq_len_sum: int,
max_seq_len: int,
):
"""
Triton implementation for gathering both K and S data from paged buffer in a single call.
:param page_indices: (num_pages,), int32/int64
:param seq_len_tensor: (num_pages,), int32/int64
:param seq_len_sum: sum of all sequence len, int32
:param max_seq_len: max of all sequence len, int32
:return: tuple of (k_fp8, k_scale) where
k_fp8: (seq_len, index_head_dim), uint8
k_scale: (seq_len, 4), uint8
"""
return _get_k_and_s_triton(
buf=buf,
page_indices=page_indices,
seq_lens=seq_len_tensor,
seq_len_sum=seq_len_sum,
max_seq_len=max_seq_len,
page_size=pool.page_size,
index_head_dim=pool.index_head_dim,
)
class SetK:
@classmethod
def execute(cls, *args, buf, **kwargs):
return cls.torch_fast(*args, **kwargs, buf=buf)
@classmethod
def slow(
cls,
pool: "DSATokenToKVPool",
buf: torch.Tensor,
loc: torch.Tensor,
index_k: torch.Tensor,
):
for i in range(len(loc)):
page_index = loc[i] // pool.page_size
offset = loc[i] % pool.page_size
buf[
page_index,
offset * pool.index_head_dim : (offset + 1) * pool.index_head_dim,
] = index_k[i].view(torch.uint8)
@classmethod
def torch_fast(
cls,
pool: "DSATokenToKVPool",
buf: torch.Tensor,
loc: torch.Tensor,
index_k: torch.Tensor,
):
(num_tokens_to_write,) = loc.shape
buf_numel_per_page = buf.shape[1]
num_k_bytes_per_token = pool.index_head_dim
# loc: (num_tokens_to_write,), int32, element := the token index to write to
loc_page_index = loc // pool.page_size
loc_token_offset_in_page = loc % pool.page_size
flat_buf = buf.flatten()
flat_indices = (
(loc_page_index * buf_numel_per_page)[:, None]
+ (loc_token_offset_in_page * num_k_bytes_per_token)[:, None]
+ torch.arange(num_k_bytes_per_token, dtype=torch.int32, device="cuda")[
None, :
]
)
num_k_bytes_total = num_tokens_to_write * num_k_bytes_per_token
flat_indices = flat_indices.flatten()[:num_k_bytes_total]
flat_buf[flat_indices] = index_k.view(torch.uint8).flatten()
class SetS:
@classmethod
def execute(cls, *args, buf, **kwargs):
return cls.torch_fast(*args, **kwargs, buf=buf)
@classmethod
def slow(
cls,
pool: "DSATokenToKVPool",
buf: torch.Tensor,
loc: torch.Tensor,
index_k_scale: torch.Tensor,
):
for i in range(len(loc)):
page_index = loc[i] // pool.page_size
offset = loc[i] % pool.page_size
start = pool.page_size * pool.index_head_dim
buf[page_index, start + offset * 4 : start + (offset + 1) * 4] = (
index_k_scale[i].view(torch.uint8)
)
@classmethod
def torch_fast(
cls,
pool: "DSATokenToKVPool",
buf: torch.Tensor,
loc: torch.Tensor,
index_k_scale: torch.Tensor,
):
(num_tokens_to_write,) = loc.shape
buf_numel_per_page = buf.shape[1]
num_s_bytes_per_token = 4
s_offset_in_page = pool.page_size * pool.index_head_dim
# loc: (num_tokens_to_write,), int32, element := the token index to write to
loc_page_index = loc // pool.page_size
loc_token_offset_in_page = loc % pool.page_size
flat_buf = buf.flatten()
flat_indices = (
(loc_page_index * buf_numel_per_page)[:, None]
+ s_offset_in_page
+ (loc_token_offset_in_page * num_s_bytes_per_token)[:, None]
+ torch.arange(num_s_bytes_per_token, dtype=torch.int32, device="cuda")[
None, :
]
)
number_s_bytes_total = num_tokens_to_write * num_s_bytes_per_token
flat_indices = flat_indices.flatten()[:number_s_bytes_total]
flat_buf[flat_indices] = index_k_scale.view(torch.uint8).flatten()
class SetKAndS:
@classmethod
def execute(cls, *args, buf, **kwargs):
if 0:
# print("SetK, SetS comparison test")
buf_cloned = buf.clone()
cls.vanilla(*args, **kwargs, buf=buf)
cls.triton(*args, **kwargs, buf=buf_cloned)
def _clear_token_0(target):
target[0, :128] = target[0, 64 * 128 : 64 * 128 + 4] = 0
_clear_token_0(buf)
_clear_token_0(buf_cloned)
assert torch.all(
buf == buf_cloned
), f"{buf=} {buf_cloned=} {kwargs['loc'].to_list()=}"
return
cls.triton(*args, **kwargs, buf=buf)
@classmethod
def vanilla(cls, pool, buf, loc, index_k, index_k_scale):
SetK.execute(pool=pool, buf=buf, loc=loc, index_k=index_k)
SetS.execute(pool=pool, buf=buf, loc=loc, index_k_scale=index_k_scale)
@classmethod
def triton(cls, pool, buf, loc, index_k, index_k_scale):
loc = loc.to(torch.int64)
_set_k_and_s_triton(
buf=buf,
loc=loc,
index_k=index_k,
index_k_scale=index_k_scale,
page_size=pool.page_size,
)
def _set_k_and_s_triton(
buf: torch.Tensor,
loc: torch.Tensor,
index_k: torch.Tensor,
index_k_scale: torch.Tensor,
page_size: int,
):
"""
:param buf: (num_pages, page_size 64 * (128B data + 4B scale)), uint8
:param loc: (num_tokens_to_write,), int, element := the token index to write to
:param index_k: (num_tokens_to_write, 128 elem), fp8
:param index_k_scale: (num_tokens_to_write, 1 elem), fp32
:return:
"""
num_pages, buf_numel_per_page = buf.shape
(num_tokens_to_write,) = loc.shape
num_tokens_to_write_, index_head_dim = index_k.shape
# Handle both 1D (num_tokens,) and 2D (num_tokens, 1) shapes for index_k_scale
if index_k_scale.ndim == 1:
num_tokens_to_write__ = index_k_scale.shape[0]
scale_dim = 1
elif index_k_scale.ndim == 2:
num_tokens_to_write__, scale_dim = index_k_scale.shape
else:
raise ValueError(
f"index_k_scale must be 1D or 2D, got shape {index_k_scale.shape}"
)
assert buf_numel_per_page == page_size * (128 + 4)
assert num_tokens_to_write == num_tokens_to_write_ == num_tokens_to_write__
assert index_head_dim == 128
assert scale_dim == 1
if _is_hip:
if _use_aiter_preshuffle:
assert (
page_size % 16 == 0
), f"HIP preshuffle requires page_size to be a multiple of 16, got {page_size}"
else:
assert page_size == 64
assert buf.dtype == torch.uint8
assert loc.dtype == torch.int64, f"{loc.dtype=}" # can be int32
if _is_fp8_fnuz:
assert index_k.dtype == torch.float8_e4m3fnuz
else:
assert index_k.dtype == torch.float8_e4m3fn
assert index_k_scale.dtype == torch.float32
assert buf.is_contiguous()
assert loc.is_contiguous()
assert index_k.is_contiguous()
assert index_k_scale.is_contiguous()
if _is_fp8_fnuz:
buf_fp8 = buf.view(torch.float8_e4m3fnuz)
else:
buf_fp8 = buf.view(torch.float8_e4m3fn)
buf_fp32 = buf.view(torch.float32)
_set_k_and_s_triton_kernel[(num_tokens_to_write,)](
buf_fp8,
buf_fp32,
loc,
index_k,
index_k_scale,
index_k.stride(0),
PAGE_SIZE=page_size,
BUF_NUMEL_PER_PAGE=buf_numel_per_page,
NUM_K_ELEMS_PER_TOKEN=index_head_dim,
S_OFFSET_NBYTES_IN_PAGE=page_size * index_head_dim,
)
@triton.jit
def _set_k_and_s_triton_kernel(
buf_fp8_ptr,
buf_fp32_ptr,
loc_ptr,
index_k_ptr,
index_k_scale_ptr,
index_k_ptr_stride_0,
PAGE_SIZE: tl.constexpr,
BUF_NUMEL_PER_PAGE: tl.constexpr,
NUM_K_ELEMS_PER_TOKEN: tl.constexpr,
S_OFFSET_NBYTES_IN_PAGE: tl.constexpr,
):
token_id = tl.program_id(0)
loc = tl.load(loc_ptr + token_id)
in_k_offsets = token_id * index_k_ptr_stride_0 + tl.arange(0, NUM_K_ELEMS_PER_TOKEN)
# no need for `mask`, since we read 128B for k and 4B for scale, both pow of 2
k = tl.load(index_k_ptr + in_k_offsets)
k_scale = tl.load(index_k_scale_ptr + token_id)
loc_page_index = loc // PAGE_SIZE
loc_token_offset_in_page = loc % PAGE_SIZE
out_k_offsets = (
loc_page_index * BUF_NUMEL_PER_PAGE
+ loc_token_offset_in_page * NUM_K_ELEMS_PER_TOKEN
+ tl.arange(0, NUM_K_ELEMS_PER_TOKEN)
)
# "//4" b/c it is fp32 instead of uint8
out_s_offset = (
loc_page_index * BUF_NUMEL_PER_PAGE // 4
+ S_OFFSET_NBYTES_IN_PAGE // 4
+ loc_token_offset_in_page
)
tl.store(buf_fp8_ptr + out_k_offsets, k)
tl.store(buf_fp32_ptr + out_s_offset, k_scale)
def _get_k_triton(
buf: torch.Tensor,
page_indices: torch.Tensor,
seq_len: int,
page_size: int,
index_head_dim: int,
):
"""
Gather K (key) data from paged buffer using Triton.
:param buf: (num_pages, page_size * 128 + page_size * 4), uint8
:param page_indices: (num_pages,), int32/int64
:param seq_len: int, number of tokens to gather
:param page_size: int, typically 64
:param index_head_dim: int, typically 128
:return: (seq_len, index_head_dim), uint8
"""
num_pages, buf_numel_per_page = buf.shape
# Allocate output
out = torch.empty((seq_len, index_head_dim), dtype=torch.uint8, device=buf.device)
# Launch kernel with one thread per token
grid = (seq_len,)
_get_k_triton_kernel[grid](
buf,
page_indices,
out,
seq_len,
page_size,
buf_numel_per_page,
index_head_dim,
BLOCK_SIZE=128,
)
return out
@triton.jit
def _get_k_triton_kernel(
buf_ptr,
page_indices_ptr,
out_ptr,
seq_len: tl.constexpr,
page_size: tl.constexpr,
buf_numel_per_page: tl.constexpr,
index_head_dim: tl.constexpr,
BLOCK_SIZE: tl.constexpr,
):
"""
Each program handles one token (seq_len tokens total).
Loads 128 bytes from the appropriate page.
"""
token_id = tl.program_id(0)
# Calculate which page and offset within page
page_idx = token_id // page_size
token_offset_in_page = token_id % page_size
# Load the page index from page_indices
page_index = tl.load(page_indices_ptr + page_idx)
# Calculate source offset in buf
# buf[page_index, token_offset_in_page * index_head_dim : ...]
src_base_offset = (
page_index * buf_numel_per_page + token_offset_in_page * index_head_dim
)
# Load 128 bytes (index_head_dim elements)
offsets = tl.arange(0, BLOCK_SIZE)
mask = offsets < index_head_dim
data = tl.load(buf_ptr + src_base_offset + offsets, mask=mask)
# Store to output
dst_offset = token_id * index_head_dim
tl.store(out_ptr + dst_offset + offsets, data, mask=mask)
def _get_s_triton(
buf: torch.Tensor,
page_indices: torch.Tensor,
seq_len: int,
page_size: int,
index_head_dim: int,
):
"""
Gather S (scale) data from paged buffer using Triton.
:param buf: (num_pages, page_size * 128 + page_size * 4), uint8
:param page_indices: (num_pages,), int32/int64
:param seq_len: int, number of tokens to gather
:param page_size: int, typically 64
:param index_head_dim: int, typically 128
:return: (seq_len, 4), uint8 (representing fp32 scale)
"""
num_pages, buf_numel_per_page = buf.shape
s_offset_in_page = page_size * index_head_dim # Scales start after K data
# Allocate output
out = torch.empty((seq_len, 4), dtype=torch.uint8, device=buf.device)
# Launch kernel with one thread per token
grid = (seq_len,)
_get_s_triton_kernel[grid](
buf,
page_indices,
out,
seq_len,
page_size,
buf_numel_per_page,
s_offset_in_page,
)
return out
@triton.jit
def _get_s_triton_kernel(
buf_ptr,
page_indices_ptr,
out_ptr,
seq_len: tl.constexpr,
page_size: tl.constexpr,
buf_numel_per_page: tl.constexpr,
s_offset_in_page: tl.constexpr,
):
"""
Each program handles one token (seq_len tokens total).
Loads 4 bytes (fp32 scale) from the appropriate page.
"""
token_id = tl.program_id(0)
# Calculate which page and offset within page
page_idx = token_id // page_size
token_offset_in_page = token_id % page_size
# Load the page index from page_indices
page_index = tl.load(page_indices_ptr + page_idx)
# Calculate source offset in buf
# Scales are stored after K data: page_size * index_head_dim offset
# buf[page_index, s_offset_in_page + token_offset_in_page * 4 : ...]
src_base_offset = (
page_index * buf_numel_per_page + s_offset_in_page + token_offset_in_page * 4
)
# Load 4 bytes (fp32 scale)
offsets = tl.arange(0, 4)
data = tl.load(buf_ptr + src_base_offset + offsets)
# Store to output
dst_offset = token_id * 4
tl.store(out_ptr + dst_offset + offsets, data)
def _get_k_and_s_triton(
buf: torch.Tensor,
page_indices: torch.Tensor,
seq_lens: torch.Tensor,
seq_len_sum: int,
max_seq_len: int,
page_size: int,
index_head_dim: int,
):
"""
Fused gather of both K (key) and S (scale) data from paged buffer using Triton.
This is more efficient than calling GetK and GetS separately.
:param buf: (num_pages, page_size * 128 + page_size * 4), uint8
:param page_indices: (num_pages,), int32/int64
:param seq_lens: tensor of sequence lens, int64
:param seq_len_sum: sum of all sequence len, int32
:param max_seq_len: max of sequence len, int32
:param page_size: int, typically 64
:param index_head_dim: int, typically 128
:return: tuple of (k_out, s_out) where
k_out: (seq_len, index_head_dim), uint8
s_out: (seq_len, 4), uint8
"""
# Allocate outputs
k_out = torch.empty(
(seq_len_sum, index_head_dim), dtype=torch.uint8, device=buf.device
)
s_out = torch.empty((seq_len_sum, 4), dtype=torch.uint8, device=buf.device)
_, buf_numel_per_page = buf.shape
_, page_indice_batch_offset = page_indices.shape
s_offset_in_page = page_size * index_head_dim
# Launch kernel with one thread per token
BLOCK_SIZE = 256
BLOCK_SIZE_K = 128
num_token_blocks = (max_seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE
num_k_threads = (index_head_dim + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K
seq_num = seq_lens.shape[0]
grid = (seq_num, num_token_blocks, num_k_threads)
seq_num_pow2 = 1
while seq_num_pow2 < seq_num:
seq_num_pow2 *= 2
_get_k_and_s_triton_kernel[grid](
buf_ptr=buf,
page_indices_ptr=page_indices,
k_out_ptr=k_out,
s_out_ptr=s_out,
seq_len_ptr=seq_lens,
seq_len_num_pow=seq_num_pow2,
page_size=page_size,
buf_numel_per_page=buf_numel_per_page,
index_head_dim=index_head_dim,
s_offset_in_page=s_offset_in_page,
page_indice_batch_offset=page_indice_batch_offset,
BLOCK_SIZE=BLOCK_SIZE,
BLOCK_SIZE_K=BLOCK_SIZE_K,
)
return k_out, s_out
@triton.jit
def _get_k_and_s_triton_kernel(
buf_ptr,
page_indices_ptr,
k_out_ptr,
s_out_ptr,
seq_len_ptr,
seq_len_num_pow: tl.constexpr,
page_size: tl.constexpr,
buf_numel_per_page: tl.constexpr,
index_head_dim: tl.constexpr,
s_offset_in_page: tl.constexpr,
page_indice_batch_offset,
BLOCK_SIZE: tl.constexpr,
BLOCK_SIZE_K: tl.constexpr,
):
"""
Fused kernel that gathers both K and S data in a single pass.
Each program handles one token (seq_len tokens total).
Loads 128 bytes (K) + 4 bytes (S) from the appropriate page.
"""
batch_id = tl.program_id(0)
block_token_start = tl.program_id(1) * BLOCK_SIZE
thread_idx = tl.program_id(2)
# Define the token range within the block and the K dimension range handled by the thread.
token_ids_in_block = tl.arange(0, BLOCK_SIZE)
token_ids = block_token_start + token_ids_in_block
k_offsets = thread_idx * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K)
seq_len = tl.load(seq_len_ptr + batch_id)
token_valid_mask = token_ids < seq_len
pre_batch_idx = tl.arange(0, seq_len_num_pow)
mask_pre_batch_idx = pre_batch_idx < batch_id
prev_seq_lens = tl.load(seq_len_ptr + pre_batch_idx, mask=mask_pre_batch_idx)
batch_token_offset = tl.sum(prev_seq_lens)
# Batch calculate the page index and in-page offset of each token.
page_idx = token_ids // page_size
token_offset_in_page = token_ids % page_size
page_indices_base = batch_id * page_indice_batch_offset
page_idx_valid_mask = page_idx < page_indice_batch_offset
page_index = tl.load(
page_indices_ptr + page_idx + page_indices_base,
mask=token_valid_mask & page_idx_valid_mask,
)
# ===== Load K data =====
# The address calculation logic for K: page_index * total number of elements in a single page + K offset of the token within the page.
k_src_token_offset = token_offset_in_page * index_head_dim
k_src_base_offset = page_index * buf_numel_per_page + k_src_token_offset
k_load_addr = buf_ptr + k_src_base_offset[:, None] + k_offsets[None, :]
k_dim_mask = k_offsets[None, :] < index_head_dim
k_mask = token_valid_mask[:, None] & k_dim_mask
k_data = tl.load(k_load_addr, mask=k_mask, other=0)
# Store K to output
k_dst_token_offset = batch_token_offset + token_ids
k_dst_base_offset = k_dst_token_offset * index_head_dim
k_store_addr = k_out_ptr + k_dst_base_offset[:, None] + k_offsets[None, :]
tl.store(k_store_addr, k_data, mask=k_mask)
# ===== Load S data =====
# The address calculation logic for S: page_index * total number of elements in a single page + starting offset of S within the page + offset of token within S in the page
s_src_token_offset = s_offset_in_page + token_offset_in_page * 4
s_src_base_offset = page_index * buf_numel_per_page + s_src_token_offset
s_offsets = tl.arange(0, 4)
s_load_addr = buf_ptr + s_src_base_offset[:, None] + s_offsets[None, :]
s_mask = token_valid_mask[:, None] & (s_offsets[None, :] < 4)
s_data = tl.load(s_load_addr, mask=s_mask, other=0)
# Store S to output
s_dst_token_offset = batch_token_offset + token_ids
s_dst_base_offset = s_dst_token_offset * 4
s_store_addr = s_out_ptr + s_dst_base_offset[:, None] + s_offsets[None, :]
tl.store(s_store_addr, s_data, mask=s_mask)
python/sglang/srt/layers/attention/dsa/quant_k_cache.py
+449
View File
@@ -0,0 +1,449 @@
import torch
import triton
import triton.language as tl
def quantize_k_cache(cache_k):
return _quantize_k_cache_fast_wrapped(cache_k)
def quantize_k_cache_separate(
k_nope: torch.Tensor,
k_rope: torch.Tensor,
tile_size: int = 128,
):
"""
Quantize k_nope and k_rope separately without concat, returns two tensors.
This avoids the concat operation and enables direct reuse of set_mla_kv_buffer_triton
by returning two separate byte tensors for the nope and rope parts.
Args:
k_nope: (num_tokens, dim_nope) or (num_tokens, 1, dim_nope)
Must have dim_nope=512 for FP8 MLA quantization
k_rope: (num_tokens, dim_rope) or (num_tokens, 1, dim_rope)
Must have dim_rope=64 for FP8 MLA quantization
tile_size: quantization tile size (default 128)
Returns:
Tuple of (nope_part, rope_part) where:
- nope_part: (num_tokens, 1, 528) as uint8 view, contains [nope_fp8(512) | scales(16)]
- rope_part: (num_tokens, 1, 128) as uint8 view, contains [rope_bf16_bytes(128)]
These two tensors can be directly passed to set_mla_kv_buffer_triton(kv_buffer, loc, nope_part, rope_part)
"""
# Squeeze middle dimension if present
k_nope_2d = k_nope.squeeze(1) if k_nope.ndim == 3 else k_nope
k_rope_2d = k_rope.squeeze(1) if k_rope.ndim == 3 else k_rope
num_tokens = k_nope_2d.shape[0]
dim_nope = k_nope_2d.shape[1]
dim_rope = k_rope_2d.shape[1]
# Validate dimensions for FP8 MLA
if dim_nope != 512:
raise ValueError(f"Expected dim_nope=512 for FP8 MLA, got {dim_nope}")
if dim_rope != 64:
raise ValueError(f"Expected dim_rope=64 for FP8 MLA, got {dim_rope}")
if k_rope_2d.shape[0] != num_tokens:
raise ValueError(
f"k_nope and k_rope must have same num_tokens, got {num_tokens} vs {k_rope_2d.shape[0]}"
)
return _quantize_k_cache_fast_separate(
k_nope=k_nope_2d, k_rope=k_rope_2d, group_size=tile_size
)
# Copied from original
def _quantize_k_cache_ref(
input_k_cache: torch.Tensor, # (num_blocks, block_size, h_k, d)
dv: int = 512,
tile_size: int = 128,
) -> torch.Tensor:
"""
Quantize the k-cache
Return a tensor with shape (num_blocks, block_size, h_k, dv + 4(dv/tile_size) + t(d-dv)) of dtype uint8_t, where t = input_k_cache.element_size()
For more detail about the layout of K/V, please refer to comments in flash_mla_interface.py or README.md
"""
assert dv % tile_size == 0
num_tiles = dv // tile_size
num_blocks, block_size, h_k, d = input_k_cache.shape
assert h_k == 1
input_k_cache = input_k_cache.squeeze(2) # [num_blocks, block_size, d]
input_elem_size = input_k_cache.element_size()
result = torch.empty(
(num_blocks, block_size, dv + num_tiles * 4 + input_elem_size * (d - dv)),
dtype=torch.float8_e4m3fn,
device=input_k_cache.device,
)
result_k_nope_part = result[..., :dv]
result_k_scale_factor = result[..., dv : dv + num_tiles * 4].view(torch.float32)
result_k_rope_part = result[..., dv + num_tiles * 4 :].view(input_k_cache.dtype)
result_k_rope_part[:] = input_k_cache[..., dv:]
for tile_idx in range(0, num_tiles):
cur_scale_factors_inv = (
torch.abs(
input_k_cache[..., tile_idx * tile_size : (tile_idx + 1) * tile_size]
)
.max(dim=-1)
.values
/ 448.0
) # [num_blocks, block_size]
result_k_scale_factor[:, :, tile_idx] = cur_scale_factors_inv
cur_scale_factors_inv.unsqueeze_(-1) # [num_blocks, block_size, 1]
cur_quantized_nope = (
input_k_cache[
..., tile_idx * tile_size : (tile_idx + 1) * tile_size
].float()
/ cur_scale_factors_inv.float()
).to(torch.float8_e4m3fn)
result_k_nope_part[..., tile_idx * tile_size : (tile_idx + 1) * tile_size] = (
cur_quantized_nope
)
result = result.view(num_blocks, block_size, 1, -1)
return result
def _quantize_k_cache_fast_wrapped(
input_k_cache: torch.Tensor,
dv: int = 512,
tile_size: int = 128,
) -> torch.Tensor:
# TODO the final API may be 2D instead of 4D, thus we convert them here
num_blocks, block_size, _, dim_nope_and_rope = input_k_cache.shape
assert dv == 512
assert dim_nope_and_rope == 512 + 64
assert tile_size == 128
input_k_cache = input_k_cache.view((-1, dim_nope_and_rope))
# TODO deliberately split into two tensors, then upstream can provide the two tensors instead of concat into one
k_nope = input_k_cache[:, :dv]
k_rope = input_k_cache[:, dv:]
output = _quantize_k_cache_fast(k_nope=k_nope, k_rope=k_rope)
return output.view(num_blocks, block_size, 1, -1)
def _quantize_k_cache_fast(k_nope, k_rope, group_size: int = 128):
"""
:param k_nope: (num_tokens, dim_nope 512)
:param k_rope: (num_tokens, dim_rope 64)
"""
assert k_nope.dtype == torch.bfloat16
assert k_rope.dtype == torch.bfloat16
num_tokens, dim_nope = k_nope.shape
num_tokens_, dim_rope = k_rope.shape
assert num_tokens == num_tokens_
assert dim_nope == 512
assert dim_rope == 64
assert k_nope.dtype == k_rope.dtype
num_tiles = dim_nope // group_size
assert k_nope.stride(1) == 1
assert k_rope.stride(1) == 1
output = torch.empty(
(num_tokens, dim_nope + num_tiles * 4 + k_rope.element_size() * dim_rope),
dtype=torch.float8_e4m3fn,
device=k_nope.device,
)
output_nope_q = output[..., :dim_nope]
output_nope_s = output[..., dim_nope : dim_nope + num_tiles * 4].view(torch.float32)
output_rope = output[..., dim_nope + num_tiles * 4 :].view(torch.bfloat16)
num_blocks_per_token = triton.cdiv(dim_nope + dim_rope, group_size)
assert num_blocks_per_token == 5
assert dim_nope % group_size == 0
NUM_NOPE_BLOCKS = dim_nope // group_size
_quantize_k_cache_fast_kernel[(num_tokens, num_blocks_per_token)](
output_nope_q,
output_nope_s,
output_rope,
k_nope,
k_rope,
output_nope_q.stride(0),
output_nope_s.stride(0),
output_rope.stride(0),
k_nope.stride(0),
k_rope.stride(0),
NUM_NOPE_BLOCKS=NUM_NOPE_BLOCKS,
GROUP_SIZE=group_size,
DIM_NOPE=dim_nope,
DIM_ROPE=dim_rope,
FP8_MIN=torch.finfo(torch.float8_e4m3fn).min,
FP8_MAX=torch.finfo(torch.float8_e4m3fn).max,
)
return output
def _quantize_k_cache_fast_separate(k_nope, k_rope, group_size: int = 128):
"""
Quantize k_nope and k_rope in a single Triton kernel, directly outputting two separate tensors.
This avoids packing/unpacking and enables direct use with set_mla_kv_buffer_triton.
:param k_nope: (num_tokens, dim_nope 512) bfloat16
:param k_rope: (num_tokens, dim_rope 64) bfloat16
:param group_size: quantization tile size (default 128, kernel is tuned for this value)
:return: Tuple of (nope_part_u8, rope_part_u8)
- nope_part_u8: (num_tokens, 1, nope_part_bytes) uint8, layout [nope_fp8(dim_nope) | scales(num_tiles*4)]
- rope_part_u8: (num_tokens, 1, rope_part_bytes) uint8, layout [rope_bf16_bytes(dim_rope*2)]
"""
num_tokens, dim_nope = k_nope.shape
num_tokens_, dim_rope = k_rope.shape
assert num_tokens == num_tokens_, f"k_nope and k_rope must have same num_tokens"
# Ensure contiguous tensors for kernel
k_nope = k_nope.contiguous()
k_rope = k_rope.contiguous()
num_tiles = dim_nope // group_size
# Calculate byte sizes based on validated dimensions
# nope_part: [FP8 quantized data (dim_nope bytes)] + [FP32 scales (num_tiles * 4 bytes)]
# rope_part: [BF16 raw data (dim_rope * 2 bytes)]
nope_part_bytes = (
dim_nope + num_tiles * 4
) # e.g., 512 + 4*4 = 528 for dim_nope=512, group_size=128
rope_part_bytes = (
dim_rope * k_rope.element_size()
) # e.g., 64 * 2 = 128 for dim_rope=64, BF16
# Allocate two separate output buffers (as uint8 for direct byte-level access)
nope_part_u8 = torch.empty(
(num_tokens, nope_part_bytes), dtype=torch.uint8, device=k_nope.device
)
rope_part_u8 = torch.empty(
(num_tokens, rope_part_bytes), dtype=torch.uint8, device=k_rope.device
)
# Create typed views for the kernel to write into
# Fixed byte layout for nope_part: [nope_fp8 (dim_nope bytes) | scales_fp32 (num_tiles*4 bytes)]
# Fixed byte layout for rope_part: [rope_bf16 (dim_rope*2 bytes)]
nope_q_view = nope_part_u8[:, :dim_nope].view(torch.float8_e4m3fn)
nope_s_view = nope_part_u8[:, dim_nope:].view(torch.float32)
rope_view = rope_part_u8.view(torch.bfloat16)
# Kernel launch parameters
num_blocks_per_token = triton.cdiv(dim_nope + dim_rope, group_size)
NUM_NOPE_BLOCKS = dim_nope // group_size
# Use the same kernel as _quantize_k_cache_fast (reuse existing implementation)
_quantize_k_cache_fast_kernel[(num_tokens, num_blocks_per_token)](
nope_q_view,
nope_s_view,
rope_view,
k_nope,
k_rope,
nope_q_view.stride(0),
nope_s_view.stride(0),
rope_view.stride(0),
k_nope.stride(0),
k_rope.stride(0),
NUM_NOPE_BLOCKS=NUM_NOPE_BLOCKS,
GROUP_SIZE=group_size,
DIM_NOPE=dim_nope,
DIM_ROPE=dim_rope,
FP8_MIN=torch.finfo(torch.float8_e4m3fn).min,
FP8_MAX=torch.finfo(torch.float8_e4m3fn).max,
)
# Add middle dimension for compatibility with set_mla_kv_buffer_triton
return nope_part_u8.unsqueeze(1), rope_part_u8.unsqueeze(1)
@triton.jit
def _quantize_k_cache_fast_kernel(
output_nope_q_ptr,
output_nope_s_ptr,
output_rope_ptr,
k_nope_ptr,
k_rope_ptr,
output_nope_q_stride_0: int,
output_nope_s_stride_0: int,
output_rope_stride_0: int,
k_nope_stride_0: int,
k_rope_stride_0: int,
NUM_NOPE_BLOCKS: tl.constexpr,
GROUP_SIZE: tl.constexpr,
DIM_NOPE: tl.constexpr,
DIM_ROPE: tl.constexpr,
FP8_MIN: tl.constexpr,
FP8_MAX: tl.constexpr,
):
token_id = tl.program_id(0)
raw_block_id = tl.program_id(1)
if raw_block_id < NUM_NOPE_BLOCKS:
# a. quant nope
effective_block_id = raw_block_id
offs = effective_block_id * GROUP_SIZE + tl.arange(0, GROUP_SIZE)
mask = offs < DIM_NOPE
ptr = k_nope_ptr + token_id * k_nope_stride_0 + offs
y = tl.load(ptr, mask=mask, other=0.0).to(tl.float32)
# the ref impl do not have a `tl.maximum(... eps)`, so we remove it here
y_s = tl.max(tl.abs(y)) / FP8_MAX
y_s_inv = 1.0 / y_s
y_q = tl.clamp(y * y_s_inv, FP8_MIN, FP8_MAX).to(
output_nope_q_ptr.dtype.element_ty
)
dst_q_ptr = output_nope_q_ptr + token_id * output_nope_q_stride_0 + offs
dst_s_ptr = (
output_nope_s_ptr + token_id * output_nope_s_stride_0 + effective_block_id
)
tl.store(dst_q_ptr, y_q, mask=mask)
tl.store(dst_s_ptr, y_s)
else:
# b. copy rope
effective_block_id = raw_block_id - NUM_NOPE_BLOCKS
offs = effective_block_id * GROUP_SIZE + tl.arange(0, GROUP_SIZE)
mask = offs < DIM_ROPE
src_ptr = k_rope_ptr + token_id * k_rope_stride_0 + offs
dst_ptr = output_rope_ptr + token_id * output_rope_stride_0 + offs
data = tl.load(src_ptr, mask=mask)
tl.store(dst_ptr, data, mask=mask)
if __name__ == "__main__":
import dequant_k_cache
for num_blocks, block_size in [
(1, 1),
(10, 64),
]:
dim_nope_and_rope = 512 + 64
input_k_cache = torch.randn(
(num_blocks, block_size, 1, dim_nope_and_rope),
dtype=torch.bfloat16,
device="cuda",
)
ref_quant = _quantize_k_cache_ref(input_k_cache)
actual_quant = _quantize_k_cache_fast_wrapped(input_k_cache)
ref_ref_dequant = dequant_k_cache._dequantize_k_cache_slow(ref_quant)
ref_actual_dequant = dequant_k_cache._dequantize_k_cache_fast_wrapped(ref_quant)
actual_actual_dequant = dequant_k_cache._dequantize_k_cache_fast_wrapped(
actual_quant
)
print(f"{ref_ref_dequant=}")
print(f"{actual_actual_dequant=}")
print(f"{actual_actual_dequant - ref_ref_dequant=}")
print(f"{torch.mean(ref_ref_dequant - actual_actual_dequant)=}")
# TODO too different?
torch.testing.assert_close(
ref_ref_dequant, ref_actual_dequant, atol=0.2, rtol=0.2
)
torch.testing.assert_close(
ref_ref_dequant, actual_actual_dequant, atol=0.2, rtol=0.2
)
# test dequant_k_cache_paged
page_table_1 = torch.arange(
num_blocks * block_size, dtype=torch.int32, device="cuda"
)
actual_dequant_paged = dequant_k_cache.dequantize_k_cache_paged(
actual_quant, page_table_1
).reshape(actual_actual_dequant.shape)
print(f"{torch.mean(actual_actual_dequant - actual_dequant_paged)=}")
torch.testing.assert_close(
ref_ref_dequant, actual_dequant_paged, atol=0.2, rtol=0.2
)
print("Passed")
# Test quantize_k_cache_separate: verify output matches concat path
print("\nTesting quantize_k_cache_separate...")
for num_tokens in [64, 100]:
dim_nope = 512
dim_rope = 64
k_nope = torch.randn(
num_tokens, 1, dim_nope, dtype=torch.bfloat16, device="cuda"
)
k_rope = torch.randn(
num_tokens, 1, dim_rope, dtype=torch.bfloat16, device="cuda"
)
# Old path: concat then quantize
k_concat = torch.cat([k_nope, k_rope], dim=-1).squeeze(1) # (num_tokens, 576)
old_output = quantize_k_cache(k_concat.unsqueeze(1).unsqueeze(1)) # 4D input
old_output = old_output.squeeze(1).squeeze(1) # Back to (num_tokens, 656)
# New path: quantize separately
nope_part, rope_part = quantize_k_cache_separate(k_nope, k_rope)
new_bytes = torch.cat([nope_part.squeeze(1), rope_part.squeeze(1)], dim=-1)
# Compare byte-level equality
old_bytes = old_output.view(torch.uint8)
if old_bytes.shape != new_bytes.shape:
raise RuntimeError(
f"Shape mismatch: {old_bytes.shape} vs {new_bytes.shape}"
)
diff_bytes = (old_bytes != new_bytes).sum().item()
if diff_bytes > 0:
max_diff = (old_bytes.float() - new_bytes.float()).abs().max().item()
raise RuntimeError(
f"quantize_k_cache_separate output doesn't match concat path: "
f"{diff_bytes} differing bytes, max_diff={max_diff}"
)
print(f" num_tokens={num_tokens}: PASSED (outputs match byte-wise)")
print("quantize_k_cache_separate tests passed!")
print("\nDo benchmark...")
for num_blocks, block_size in [
(1, 64),
(64, 64),
(128, 64),
(256, 64),
(512, 64),
(1024, 64),
(2048, 64),
]:
dim_nope_and_rope = 512 + 64
input_k_cache = torch.randn(
(num_blocks, block_size, 1, dim_nope_and_rope),
dtype=torch.bfloat16,
device="cuda",
)
actual_quant = _quantize_k_cache_fast_wrapped(input_k_cache)
page_table_1 = torch.arange(
num_blocks * block_size, dtype=torch.int32, device="cuda"
)
def run_ans():
return dequant_k_cache.dequantize_k_cache_paged(actual_quant, page_table_1)
ans_time: float = triton.testing.do_bench(run_ans, warmup=10, rep=20) / 1000 # type: ignore
print(f"seq_kv: {num_blocks * block_size}, time: {ans_time * 1e6: 4.0f} us")
python/sglang/srt/layers/attention/dsa/tilelang_kernel.py
+2589
View File
File diff suppressed because it is too large Load Diff
python/sglang/srt/layers/attention/dsa/transform_index.py
+144
View File
@@ -0,0 +1,144 @@
from typing import List, Optional
import torch
import triton
import triton.language as tl
def transform_index_page_table_prefill(**kwargs):
return transform_index_page_table_prefill_ref(**kwargs)
def transform_index_page_table_decode(**kwargs):
return transform_index_page_table_decode_ref(**kwargs)
@triton.jit
def transform_index_page_table_decode_kernel(
page_table_ptr: torch.Tensor,
topk_indices_ptr: torch.Tensor,
result_ptr: torch.Tensor,
page_size: tl.constexpr,
max_seqlen_k: tl.constexpr,
):
TOPK: tl.constexpr = 2048
req_id = tl.program_id(0)
page_table_ptr = page_table_ptr + req_id * max_seqlen_k
topk_indices_ptr = topk_indices_ptr + req_id * TOPK
result_ptr = result_ptr + req_id * TOPK
offset = tl.arange(0, TOPK) # topk should be 2048
loaded_topk_indices = tl.load(topk_indices_ptr + offset)
mask = loaded_topk_indices >= 0
loaded_kv_indices = tl.load(page_table_ptr + loaded_topk_indices, mask=mask)
tl.store(result_ptr + offset, loaded_kv_indices, mask=mask)
tl.store(result_ptr + offset, -1, mask=~mask)
def transform_index_page_table_decode_fast(
page_table: torch.Tensor,
topk_indices: torch.Tensor,
result: Optional[torch.Tensor] = None,
page_size: int = 1,
) -> torch.Tensor:
"""
Transform the page table according to topk indices for sparse topk attention.
Args:
page_table: [qo_len, max_seqlen_k], the original page table
topk_indices: [qo_len, topk], the topk indices for each query position
Returns:
transformed_page_table: [qo_len, topk], the transformed page table
For out-of-bound indices in topk_indices, this should be filled with -1.
"""
assert page_size == 1
assert page_table.shape[0] == topk_indices.shape[0]
assert topk_indices.shape[1] == 2048
qo_len = topk_indices.shape[0]
max_seqlen_k = page_table.shape[1]
if result is None:
result = torch.empty_like(topk_indices, dtype=torch.int32)
# Launch triton kernel
grid = (qo_len,)
transform_index_page_table_decode_kernel[grid](
page_table,
topk_indices,
result,
page_size,
max_seqlen_k=max_seqlen_k,
)
return result
def transform_index_page_table_prefill_fast(
page_table: torch.Tensor,
topk_indices: torch.Tensor,
extend_lens_cpu: List[int],
page_size: int = 1,
) -> torch.Tensor:
# TODO(baizhou): can be implemented with another triton kernel
assert page_size == 1
result = torch.empty_like(topk_indices, dtype=torch.int32)
assert len(extend_lens_cpu) == page_table.shape[0]
offset = 0
for i, l in enumerate(extend_lens_cpu):
transform_index_page_table_decode_fast(
page_table[i].unsqueeze(0).expand(l, -1),
topk_indices[offset : offset + l],
result=result[offset : offset + l],
)
offset += l
assert offset == topk_indices.shape[0]
return result
def transform_index_page_table_decode_ref(
page_table: torch.Tensor,
topk_indices: torch.Tensor,
result: Optional[torch.Tensor] = None,
page_size: int = 1,
) -> torch.Tensor:
assert page_size == 1
assert page_table.shape[0] == topk_indices.shape[0]
if result is None:
result = torch.empty_like(topk_indices, dtype=torch.int32)
assert result.shape == topk_indices.shape
torch.gather(
page_table.to(result.dtype),
dim=1,
index=topk_indices.clamp(min=0),
out=result,
)
result[topk_indices < 0] = -1
return result
def transform_index_page_table_prefill_ref(
page_table: torch.Tensor,
topk_indices: torch.Tensor,
extend_lens_cpu: List[int],
page_size: int = 1,
) -> torch.Tensor:
assert page_size == 1
result = torch.empty_like(topk_indices, dtype=torch.int32)
assert len(extend_lens_cpu) == page_table.shape[0]
offset = 0
for i, l in enumerate(extend_lens_cpu):
transform_index_page_table_decode_ref(
page_table[i].unsqueeze(0).expand(l, -1),
topk_indices[offset : offset + l],
result=result[offset : offset + l],
)
offset += l
assert offset == topk_indices.shape[0]
return result
if __name__ == "__main__":
bs, topk, max_seqlen = 10, 2048, 3000
page_table = torch.randint(0, 100, (bs, max_seqlen), device="cuda")
topk_indices = torch.full((bs, topk), -1, device="cuda")
topk_indices[:, :1600] = torch.arange(1600).unsqueeze(0).repeat(bs, 1)
ref_result = transform_index_page_table_decode_ref(page_table, topk_indices)
result = transform_index_page_table_decode_fast(page_table, topk_indices)
assert torch.all(result == ref_result)
print("Passed")
python/sglang/srt/layers/attention/dsa/triton_kernel.py
+196
View File
@@ -0,0 +1,196 @@
from typing import Optional, Tuple
import torch
import triton
import triton.language as tl
# Triton implementation
@triton.jit
def _act_quant_kernel(
X_ptr,
Y_ptr,
S_ptr,
M,
N,
group_size: tl.constexpr,
round_scale: tl.constexpr,
BLOCK_M: tl.constexpr,
BLOCK_N: tl.constexpr,
):
"""
Triton kernel for activation quantization.
Each block processes BLOCK_M rows and group_size columns.
"""
# Get block IDs
pid_m = tl.program_id(0)
pid_n = tl.program_id(1)
# FP8 constants
fp8_min = -448.0
fp8_max = 448.0
fp8_max_inv = 1.0 / fp8_max
# Calculate row and column offsets
row_start = pid_m * BLOCK_M
col_start = pid_n * group_size
# Create offset arrays
rows = row_start + tl.arange(0, BLOCK_M)
cols = col_start + tl.arange(0, BLOCK_N)
# Mask for valid rows and columns
row_mask = rows < M
col_mask = cols < N
mask = row_mask[:, None] & col_mask[None, :]
# Load input data
x_ptrs = X_ptr + rows[:, None] * N + cols[None, :]
x = tl.load(x_ptrs, mask=mask, other=0.0).to(tl.float32)
# Compute absolute max along columns (group_size dimension) for each row
x_abs = tl.abs(x)
amax = tl.max(x_abs, axis=1) # Shape: (BLOCK_M,)
# Clamp amax to avoid division by zero
amax = tl.maximum(amax, 1e-4)
# Compute scale
if round_scale:
# Fast round scale using bit manipulation approximation
# This is a simplified version - the exact bit manipulation is harder in Triton
# Using log2 + ceil + pow2 as approximation
log_val = tl.log2(amax * fp8_max_inv)
log_ceil = tl.ceil(log_val)
scale = tl.exp2(log_ceil)
else:
scale = amax * fp8_max_inv
# Quantize: y = clamp(x / scale, fp8_min, fp8_max)
scale_broadcast = scale[:, None]
y = x / scale_broadcast
y = tl.minimum(tl.maximum(y, fp8_min), fp8_max)
# Store quantized output
y_ptrs = Y_ptr + rows[:, None] * N + cols[None, :]
tl.store(y_ptrs, y, mask=mask)
# Store scales
s_cols = pid_n
s_ptrs = S_ptr + rows * (N // group_size) + s_cols
s_mask = row_mask
tl.store(s_ptrs, scale, mask=s_mask)
def act_quant(
x: torch.Tensor, block_size: int = 128, scale_fmt: Optional[str] = None
) -> Tuple[torch.Tensor, torch.Tensor]:
"""
Quantizes the input tensor `x` using block-wise quantization with Triton.
Args:
x (torch.Tensor): The input tensor to be quantized. Must be contiguous and its last dimension size must be divisible by `block_size`.
block_size (int, optional): The size of the blocks to be used for quantization. Default is 128.
scale_fmt (Optional[str], optional): The format of the scale. Default is None.
Returns:
Tuple[torch.Tensor, torch.Tensor]: A tuple containing:
- The quantized tensor with dtype `torch.float8_e4m3fn`.
- A tensor of scaling factors with dtype `torch.float32`.
"""
assert x.is_contiguous(), "Input tensor must be contiguous"
assert (
x.size(-1) % block_size == 0
), f"Last dimension size must be divisible by block_size (block_size={block_size})"
# Flatten all dims except last
N = x.size(-1)
x_flat = x.view(-1, N)
M = x_flat.size(0)
# Allocate output tensors
y = torch.empty_like(x, dtype=torch.float8_e4m3fn)
y_flat = y.view(-1, N)
s = x.new_empty(*x.size()[:-1], N // block_size, dtype=torch.float32)
s_flat = s.view(-1, N // block_size)
# Launch kernel
BLOCK_M = 32
BLOCK_N = block_size
grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(N, block_size))
round_scale = scale_fmt is not None
_act_quant_kernel[grid](
x_flat,
y_flat,
s_flat,
M,
N,
group_size=block_size,
round_scale=round_scale,
BLOCK_M=BLOCK_M,
BLOCK_N=BLOCK_N,
num_stages=0 if round_scale else 2,
)
return y, s
@triton.jit
def _get_valid_kv_indices_kernel(
page_table_ptr, # [bs, topk]
kv_indptr_ptr, # [bs + 1]
kv_indices_ptr, # [bs * topk] output buffer
bs: tl.constexpr,
topk: tl.constexpr,
):
"""
Extract valid indices (non -1) from page_table into kv_indices.
Each program handles one batch.
"""
batch_id = tl.program_id(0)
# Get the start position for this batch in kv_indices
dst_start = tl.load(kv_indptr_ptr + batch_id)
# Load all topk indices for this batch
src_offset = batch_id * topk
offsets = tl.arange(0, topk)
indices = tl.load(page_table_ptr + src_offset + offsets)
# Count valid indices and compact them
mask = indices != -1
# Use prefix sum to compute destination positions for valid elements
# For each position, count how many valid elements are before it
prefix_sum = tl.cumsum(mask.to(tl.int32), axis=0) - 1
# Store valid indices to their compacted positions
dst_positions = dst_start + prefix_sum
tl.store(kv_indices_ptr + dst_positions, indices, mask=mask)
def get_valid_kv_indices(
page_table_1: torch.Tensor,
kv_indptr: torch.Tensor,
kv_indices: torch.Tensor,
bs: int,
):
"""
Extract valid indices from page_table_1 into kv_indices buffer.
Args:
page_table_1: [bs, topk] page table with -1 as invalid
kv_indptr: [bs + 1] cumulative count of valid indices per batch
kv_indices: [bs * topk] pre-allocated output buffer
bs: batch size
"""
topk = page_table_1.shape[1]
grid = (bs,)
_get_valid_kv_indices_kernel[grid](
page_table_1,
kv_indptr,
kv_indices,
bs,
topk,
)
python/sglang/srt/layers/attention/dsa/utils.py
+271
View File
@@ -0,0 +1,271 @@
from functools import lru_cache
from typing import TYPE_CHECKING, List, Tuple, Union
import torch
import triton
import triton.language as tl
from sglang.srt.environ import envs
from sglang.srt.layers.dp_attention import (
DpPaddingMode,
get_attention_cp_rank,
get_attention_cp_size,
get_attention_dp_rank,
)
from sglang.srt.server_args import get_global_server_args
from sglang.srt.utils import get_bool_env_var, is_hip
from sglang.srt.utils.common import ceil_align, ceil_div
@lru_cache(maxsize=1)
def aiter_can_use_preshuffle_paged_mqa() -> bool:
"""Whether aiter's preshuffle paged MQA / cache kernels can be used on this runtime.
aiter's ``deepgemm_fp8_paged_mqa_logits`` only supports ``KVBlockSize > 1`` and
``Preshuffle=True`` on its gluon kernel path. The gluon path is enabled when
Triton >= 3.5.0, OR when ``AITER_ENABLE_AOT_GLUON_PA_MQA_LOGITS=1`` is set
(which additionally requires that the AOT gluon kernel artifacts ship inside
the aiter wheel/image). Otherwise aiter asserts ``KVBlockSize == 1`` and
refuses ``Preshuffle=True``.
sglang's DSA indexer uses this single decision to pick:
* ``page_size``: 64 (preshuffle) vs 1 (legacy) on ROCm
* ``Preshuffle`` / ``preshuffle`` flags on the aiter MQA + cache kernels
* ``get_page_table_64`` vs ``get_page_table_1`` on the metadata
* whether ``GetKAndS.execute`` uses the aiter or the triton implementation
The result is cached so the cost is paid once per process.
Set ``SGLANG_DSA_HIP_DISABLE_PRESHUFFLE=1`` to force the legacy path even when
the gluon kernel would otherwise be available (useful for CI bisection).
``SGLANG_NSA_HIP_DISABLE_PRESHUFFLE`` is a deprecated alias.
"""
if not is_hip():
return False
if not get_bool_env_var("SGLANG_USE_AITER"):
return False
if envs.SGLANG_DSA_HIP_DISABLE_PRESHUFFLE.get():
return False
if get_bool_env_var("AITER_ENABLE_AOT_GLUON_PA_MQA_LOGITS"):
return True
try:
from packaging.version import Version
return Version(Version(triton.__version__).base_version) >= Version("3.5.0")
except Exception:
return False
if TYPE_CHECKING:
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
def compute_dsa_seqlens(original_seq_lens, dsa_index_topk: int):
return original_seq_lens.clamp(max=dsa_index_topk)
def is_dsa_enable_prefill_cp():
return get_global_server_args().enable_dsa_prefill_context_parallel
def is_dsa_prefill_cp_in_seq_split():
return (
is_dsa_enable_prefill_cp()
and get_global_server_args().dsa_prefill_cp_mode == "in-seq-split"
)
def is_dsa_prefill_cp_round_robin_split():
return (
is_dsa_enable_prefill_cp()
and get_global_server_args().dsa_prefill_cp_mode == "round-robin-split"
)
def can_dsa_prefill_cp_round_robin_split(forward_batch: "ForwardBatch"):
if not forward_batch.forward_mode.is_context_parallel_extend():
return False
cp_size = get_attention_cp_size()
seq_len = sum(forward_batch.extend_seq_lens_cpu)
return (
is_dsa_prefill_cp_round_robin_split()
and seq_len > 0
and seq_len >= cp_size
and cp_size > 1
)
def dsa_cp_round_robin_split_data(input_: Union[torch.Tensor, List]):
"""
# for round-robin-split, split the tokens evenly according to the rule of token_idx % cp_size.
| +-----------before split------------+|
| token0, token1, token2, token3, token4, token5, token6, token7, ...
|
| +--------------result-------------------+
| dp_atten_tp0: token0, token4, token8, token12, token16, ... |
| dp_atten_tp1: token1, token5, token9, token13, token17, ... |
| dp_atten_tp2: token2, token6, token10, token14, token18, ... |
| dp_atten_tp3: token3, token7, token11, token15, token19, ... |
| +-------------------------+
"""
cp_size = get_attention_cp_size()
cp_rank = get_attention_cp_rank()
if isinstance(input_, (tuple, list)):
indices = range(cp_rank, len(input_), cp_size)
return input_[indices]
tokens = len(input_)
if tokens % cp_size != 0:
cur_len = tokens // cp_size + (tokens % cp_size > cp_rank)
if cur_len == 0:
return input_.new_empty(0, *input_.shape[1:])
indices = torch.arange(cp_rank, tokens, cp_size, device=input_.device)
return input_[indices]
# for torch device tensor
return input_.view(-1, cp_size, *input_.shape[1:])[:, cp_rank].contiguous()
def cal_padded_tokens(forward_batch: "ForwardBatch"):
# Consistent with the padding calculation logic in ForwardBatch.prepare_mlp_sync_batch,
# calculate the actual token length after padding when attn_tp_size > 1 or in the MAX_LEN padding mode.
global_num_tokens = forward_batch.global_num_tokens_cpu.copy()
sync_group_size = len(global_num_tokens)
attn_cp_size = get_attention_cp_size()
for i in range(sync_group_size):
global_num_tokens[i] = ceil_align(global_num_tokens[i], attn_cp_size)
dp_padding_mode = DpPaddingMode.get_dp_padding_mode(
forward_batch.is_extend_in_batch, global_num_tokens
)
if dp_padding_mode.is_max_len():
tokens = max(global_num_tokens)
elif len(global_num_tokens) > 1:
tokens = global_num_tokens[get_attention_dp_rank()]
else:
tokens = global_num_tokens[0]
if can_dsa_prefill_cp_round_robin_split(forward_batch):
tokens = ceil_div(tokens, attn_cp_size)
return tokens
def pad_dsa_cache_seqlens(forward_batch: "ForwardBatch", dsa_cache_seqlens):
attn_cp_size = get_attention_cp_size()
needs_cp_pad = attn_cp_size > 1 and can_dsa_prefill_cp_round_robin_split(
forward_batch
)
needs_dp_pad = forward_batch.global_num_tokens_cpu is not None
if not needs_cp_pad and not needs_dp_pad:
return dsa_cache_seqlens
tokens = cal_padded_tokens(forward_batch)
pad_len = tokens - dsa_cache_seqlens.shape[0]
if pad_len > 0:
dsa_cache_seqlens = torch.cat(
[
dsa_cache_seqlens,
dsa_cache_seqlens.new_zeros(pad_len, *dsa_cache_seqlens.shape[1:]),
]
)
return dsa_cache_seqlens
def can_dsa_cp_split(seq_len: int, cp_size: int, use_dsa: bool, forward_batch):
if is_dsa_prefill_cp_round_robin_split():
cur_cp_seq_len = seq_len // cp_size
assert (
seq_len % cp_size == 0
), f"seq_len {seq_len} is not divisible by cp_size {cp_size} when dsa_prefill_cp_mode is round-robin-split"
else:
# TODO current just support prefill batch=1 and len(input_ids) > self.cp_size * 2
# Note: (self.cp_size * 2) To achieve load balancing for seq computation,
# the seq data needs to be divided and recombined at twice the size of cp_size.
cur_cp_seq_len = seq_len // (cp_size * 2)
if (
cur_cp_seq_len != 0
and cp_size > 1
and use_dsa
and forward_batch.forward_mode.is_context_parallel_extend()
and is_dsa_enable_prefill_cp()
and sum(forward_batch.extend_seq_lens_cpu) >= cp_size
):
return True
else:
return False
@triton.jit
def dsa_cp_round_robin_split_q_seqs_kernel(
in_seqs_ptr,
out_seqs_ptr,
bs_idx_ptr,
tokens: tl.constexpr,
cp_size: tl.constexpr,
cp_rank: tl.constexpr,
):
extra_seq = 0
bs_idx = 0
for bs in range(tokens):
cur_len = tl.load(in_seqs_ptr + bs)
cur_len += extra_seq
cur_seq = cur_len // cp_size + (cur_len % cp_size > cp_rank)
if cur_seq > 0:
tl.store(bs_idx_ptr + bs_idx, bs)
tl.store(out_seqs_ptr + bs_idx, cur_seq)
bs_idx += 1
extra_seq = cur_len - cur_seq * cp_size
def dsa_cp_round_robin_split_q_seqs_cpu(extend_seqs):
cp_size = get_attention_cp_size()
cp_rank = get_attention_cp_rank()
extra_seq = 0
q_seqs = []
for bs, cur_len in enumerate(extend_seqs):
cur_len += extra_seq
cur_seq = cur_len // cp_size + int(cur_len % cp_size > cp_rank)
q_seqs.append(cur_seq)
extra_seq = cur_len - cur_seq * cp_size
bs_idx = list([i for i, x in enumerate(q_seqs) if x > 0])
q_seqs = [q_len for q_len in q_seqs if q_len > 0]
return q_seqs, bs_idx
def dsa_cp_round_robin_split_q_seqs(
extend_seqs_cpu, extend_seqs
) -> Tuple[List, torch.Tensor, List, torch.Tensor]:
"""
round-robin-split distributes tokens across ranks based on token_idx % cp_size.
Return:
ret_q_lens_cpu(List) and ret_q_lens(torch.Tensor): the partitioned length (excluding zeros) on the current cp rank
for each sequence after distribution across cp ranks.
bs_idx_cpu(List) and bs_idx(torch.Tensor): marks which sequences are ultimately selected,
i.e., those with a partitioned length greater than zero.
"""
cp_size = get_attention_cp_size()
cp_rank = get_attention_cp_rank()
# len(ret_q_lens_cpu) == len(bs_idx_cpu)
ret_q_lens_cpu, bs_idx_cpu = dsa_cp_round_robin_split_q_seqs_cpu(extend_seqs_cpu)
ret_q_lens = torch.empty(
(len(bs_idx_cpu),), device=extend_seqs.device, dtype=extend_seqs.dtype
)
bs_idx = torch.empty(
(len(bs_idx_cpu),), device=extend_seqs.device, dtype=torch.int32
)
grid = (1,)
dsa_cp_round_robin_split_q_seqs_kernel[grid](
extend_seqs, ret_q_lens, bs_idx, len(extend_seqs), cp_size, cp_rank
)
return ret_q_lens_cpu, ret_q_lens, bs_idx_cpu, bs_idx
def dsa_use_prefill_cp(forward_batch, dsa_enable_prefill_cp=None):
if dsa_enable_prefill_cp is None:
dsa_enable_prefill_cp = is_dsa_enable_prefill_cp()
if (
forward_batch.attn_cp_metadata is not None
and dsa_enable_prefill_cp
and forward_batch.forward_mode.is_context_parallel_extend()
):
return True
else:
return False
python/sglang/srt/layers/attention/dsa_backend.py
+2528
View File
File diff suppressed because it is too large Load Diff
python/sglang/srt/layers/attention/dsv4/compress_hip.py
+1 -1
View File
@@ -10,8 +10,8 @@ import triton
import triton.language as tl import triton.language as tl
from sglang.srt.environ import envs from sglang.srt.environ import envs
from sglang.srt.layers.attention.dsa.dsa_indexer import rotate_activation
from sglang.srt.layers.attention.dsv4.compressor import Compressor as _CompressorBase from sglang.srt.layers.attention.dsv4.compressor import Compressor as _CompressorBase
from sglang.srt.layers.attention.nsa.nsa_indexer import rotate_activation
from sglang.srt.layers.deepseek_v4_rope import ( from sglang.srt.layers.deepseek_v4_rope import (
apply_rotary_emb_triton, apply_rotary_emb_triton,
fused_norm_rope_inplace_triton, fused_norm_rope_inplace_triton,
python/sglang/srt/layers/attention/dsv4/compressor.py
+4 -4
View File
@@ -15,11 +15,11 @@ from sglang.jit_kernel.deepseek_v4 import (
) )
from sglang.srt.configs.deepseek_v4 import DeepSeekV4Config from sglang.srt.configs.deepseek_v4 import DeepSeekV4Config
from sglang.srt.environ import envs from sglang.srt.environ import envs
from sglang.srt.layers.attention.dsa.triton_kernel import act_quant
from sglang.srt.layers.attention.dsa.utils import dsa_use_prefill_cp
from sglang.srt.layers.attention.dsv4.quant_k_cache import ( from sglang.srt.layers.attention.dsv4.quant_k_cache import (
quant_to_nope_fp8_rope_bf16_pack_triton, quant_to_nope_fp8_rope_bf16_pack_triton,
) )
from sglang.srt.layers.attention.nsa.triton_kernel import act_quant
from sglang.srt.layers.attention.nsa.utils import nsa_use_prefill_cp
from sglang.srt.layers.dp_attention import get_attention_cp_size from sglang.srt.layers.dp_attention import get_attention_cp_size
from sglang.srt.layers.layernorm import RMSNorm from sglang.srt.layers.layernorm import RMSNorm
from sglang.srt.layers.linear import ReplicatedLinear from sglang.srt.layers.linear import ReplicatedLinear
@@ -71,7 +71,7 @@ class CompressorBackendMixin:
compress_ratio: int, compress_ratio: int,
is_paged: bool = False, is_paged: bool = False,
) -> torch.Tensor: ) -> torch.Tensor:
from sglang.srt.layers.attention.nsa.nsa_indexer import rotate_activation from sglang.srt.layers.attention.dsa.dsa_indexer import rotate_activation
assert compress_ratio in ( assert compress_ratio in (
4, 4,
@@ -358,7 +358,7 @@ class Compressor(nn.Module):
kv_score = linear_bf16_fp32(x, self.wkv_gate.weight) kv_score = linear_bf16_fp32(x, self.wkv_gate.weight)
# CUDA path: delegate to backend # CUDA path: delegate to backend
if nsa_use_prefill_cp(forward_batch): if dsa_use_prefill_cp(forward_batch):
kv_score = cp_all_gather_rerange_output( kv_score = cp_all_gather_rerange_output(
kv_score, kv_score,
get_attention_cp_size(), get_attention_cp_size(),
python/sglang/srt/layers/attention/dsv4/indexer.py
+1 -1
View File
@@ -368,7 +368,7 @@ class C4IndexerBackendMixin:
assert len(weights.shape) == 3 assert len(weights.shape) == 3
weights = weights.squeeze(2) weights = weights.squeeze(2)
if envs.SGLANG_OPT_USE_TILELANG_INDEXER.get(): if envs.SGLANG_OPT_USE_TILELANG_INDEXER.get():
from sglang.srt.layers.attention.nsa.tilelang_kernel import ( from sglang.srt.layers.attention.dsa.tilelang_kernel import (
tilelang_fp8_paged_mqa_logits as fn, tilelang_fp8_paged_mqa_logits as fn,
) )
elif envs.SGLANG_FP8_PAGED_MQA_LOGITS_TORCH.get(): elif envs.SGLANG_FP8_PAGED_MQA_LOGITS_TORCH.get():
python/sglang/srt/layers/attention/hip_flash_mla.py
+1 -1
View File
@@ -12,7 +12,7 @@ def flash_mla_with_kvcache_entrypoint(backend: str, **kwargs):
if is_hip(): if is_hip():
import os import os
from sglang.srt.layers.attention.nsa.tilelang_kernel import ( from sglang.srt.layers.attention.dsa.tilelang_kernel import (
dpsk_v4_fp8_attention_fwd, dpsk_v4_fp8_attention_fwd,
) )
python/sglang/srt/layers/attention/hybrid_attn_backend.py
+1 -1
View File
@@ -3,7 +3,7 @@ from typing import Optional
import torch import torch
from sglang.srt.layers.attention.base_attn_backend import AttentionBackend from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
from sglang.srt.layers.attention.nsa.nsa_indexer import BaseIndexerMetadata from sglang.srt.layers.attention.dsa.dsa_indexer import BaseIndexerMetadata
from sglang.srt.layers.radix_attention import RadixAttention from sglang.srt.layers.radix_attention import RadixAttention
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
from sglang.srt.model_executor.model_runner import ModelRunner from sglang.srt.model_executor.model_runner import ModelRunner
python/sglang/srt/layers/attention/nsa/__init__.py
+11
View File
@@ -0,0 +1,11 @@
# [Deprecated] attention/nsa/ is a thin re-export shim for backward compatibility.
# Use attention/dsa/ instead. This directory will be removed in a future release.
import warnings
warnings.warn(
"sglang.srt.layers.attention.nsa is deprecated; "
"use sglang.srt.layers.attention.dsa instead.",
DeprecationWarning,
stacklevel=2,
)
from sglang.srt.layers.attention.dsa import * # noqa: F401, F403
python/sglang/srt/layers/attention/nsa/dequant_k_cache.py
+8 -287
View File
@@ -1,289 +1,10 @@
import torch # [Deprecated] Re-export shim for backward compatibility. Use dsa.dequant_k_cache instead.
import triton import warnings
import triton.language as tl
warnings.warn(
def dequantize_k_cache(quant_k_cache): "sglang.srt.layers.attention.nsa.dequant_k_cache is deprecated; "
return _dequantize_k_cache_fast_wrapped(quant_k_cache) "use sglang.srt.layers.attention.dsa.dequant_k_cache instead.",
DeprecationWarning,
stacklevel=2,
def _dequantize_k_cache_ref(
quant_k_cache: torch.Tensor, # (num_blocks, block_size, 1, bytes_per_token)
dv: int = 512,
tile_size: int = 128,
d: int = 576,
) -> torch.Tensor:
"""
De-quantize the k-cache
"""
assert dv % tile_size == 0
original_ndim = quant_k_cache.ndim
if original_ndim == 3:
# set block_size = 1
quant_k_cache = quant_k_cache.unsqueeze(1)
num_tiles = dv // tile_size
num_blocks, block_size, h_k, _ = quant_k_cache.shape
assert h_k == 1
result = torch.empty(
(num_blocks, block_size, d), dtype=torch.bfloat16, device=quant_k_cache.device
) )
from sglang.srt.layers.attention.dsa.dequant_k_cache import * # noqa: F401, F403
quant_k_cache = quant_k_cache.view(num_blocks, block_size, -1)
input_nope = quant_k_cache[..., :dv]
input_scale = quant_k_cache[..., dv : dv + num_tiles * 4].view(torch.float32)
input_rope = quant_k_cache[..., dv + num_tiles * 4 :].view(torch.bfloat16)
result[..., dv:] = input_rope
for tile_idx in range(0, num_tiles):
cur_nope = input_nope[
..., tile_idx * tile_size : (tile_idx + 1) * tile_size
].to(torch.float32)
cur_scales = input_scale[..., tile_idx].unsqueeze(-1)
result[..., tile_idx * tile_size : (tile_idx + 1) * tile_size] = (
cur_nope * cur_scales
)
if original_ndim == 3:
return result.view(num_blocks, 1, -1)
else:
return result.view(num_blocks, block_size, 1, -1)
def _dequantize_k_cache_fast_wrapped(
quant_k_cache: torch.Tensor,
dv: int = 512,
tile_size: int = 128,
) -> torch.Tensor:
original_ndim = quant_k_cache.ndim
if original_ndim == 3:
# set block_size = 1
quant_k_cache = quant_k_cache.unsqueeze(1)
num_blocks, block_size, _, dim_quant = quant_k_cache.shape
assert dv == 512
assert dim_quant == 656
assert tile_size == 128
quant_k_cache = quant_k_cache.view((-1, dim_quant))
output = _dequantize_k_cache_fast(quant_k_cache)
if original_ndim == 3:
return output.view(num_blocks, 1, -1)
else:
return output.view(num_blocks, block_size, 1, -1)
def _dequantize_k_cache_fast(quant_k_cache, group_size: int = 128):
num_tokens, dim_quant = quant_k_cache.shape
assert quant_k_cache.dtype == torch.float8_e4m3fn
dim_nope = 512
dim_rope = 64
num_tiles = dim_nope // group_size
assert dim_quant == 656
output = torch.empty(
(num_tokens, dim_nope + dim_rope),
dtype=torch.bfloat16,
device=quant_k_cache.device,
)
num_blocks_per_token = triton.cdiv(dim_nope + dim_rope, group_size)
assert num_blocks_per_token == 5
assert dim_nope % group_size == 0
input_nope_q = quant_k_cache[:, :dim_nope]
input_nope_s = quant_k_cache[:, dim_nope : dim_nope + num_tiles * 4].view(
torch.float32
)
input_rope = quant_k_cache[:, dim_nope + num_tiles * 4 :].view(torch.bfloat16)
_dequantize_k_cache_fast_kernel[(num_tokens, num_blocks_per_token)](
output,
input_nope_q,
input_nope_s,
input_rope,
output.stride(0),
input_nope_q.stride(0),
input_nope_s.stride(0),
input_rope.stride(0),
NUM_NOPE_BLOCKS=num_tiles,
GROUP_SIZE=group_size,
DIM_NOPE=dim_nope,
DIM_ROPE=dim_rope,
)
return output
@triton.jit
def _dequantize_k_cache_fast_kernel(
output_ptr,
input_nope_q_ptr,
input_nope_s_ptr,
input_rope_ptr,
output_stride_0: int,
input_nope_q_stride_0: int,
input_nope_s_stride_0: int,
input_rope_stride_0: int,
NUM_NOPE_BLOCKS: tl.constexpr,
GROUP_SIZE: tl.constexpr,
DIM_NOPE: tl.constexpr,
DIM_ROPE: tl.constexpr,
):
token_id = tl.program_id(0)
raw_block_id = tl.program_id(1)
if raw_block_id < NUM_NOPE_BLOCKS:
# a. dequant nope
effective_block_id = raw_block_id
offs_q = effective_block_id * GROUP_SIZE + tl.arange(0, GROUP_SIZE)
mask = offs_q < DIM_NOPE
ptr_q = input_nope_q_ptr + token_id * input_nope_q_stride_0 + offs_q
ptr_s = input_nope_s_ptr + token_id * input_nope_s_stride_0 + effective_block_id
y_q = tl.load(ptr_q, mask=mask, other=0.0).to(tl.float32)
y_s = tl.load(ptr_s)
y = (y_q * y_s).to(output_ptr.dtype.element_ty)
dst_ptr = output_ptr + token_id * output_stride_0 + offs_q
tl.store(dst_ptr, y, mask=mask)
else:
# b. copy rope
effective_block_id = raw_block_id - NUM_NOPE_BLOCKS
offs = effective_block_id * GROUP_SIZE + tl.arange(0, GROUP_SIZE)
mask = offs < DIM_ROPE
src_ptr = input_rope_ptr + token_id * input_rope_stride_0 + offs
dst_ptr = output_ptr + token_id * output_stride_0 + DIM_NOPE + offs
data = tl.load(src_ptr, mask=mask).to(tl.bfloat16)
tl.store(dst_ptr, data, mask=mask)
def dequantize_k_cache_paged(
quant_k_cache: torch.Tensor,
page_table_1_flattened: torch.Tensor,
group_size: int = 128,
) -> torch.Tensor:
"""
De-quantize the k-cache with paged layout
Args:
quant_k_cache: [total_num_tokens, 1, dim_quant] or [num_blocks, block_size, 1, dim_quant], the quantized k-cache in paged layout
page_table_1_flattened: [num_tokens], the flattened page_table_1 with the page indices in each requests concatenated together
Returns:
output: [num_tokens, 1, dim_nope + dim_rope], the de-quantized k-cache
"""
dim_quant = quant_k_cache.shape[-1]
assert (
dim_quant == 656
), f"dim_quant: {dim_quant} != 656 detected in dequantize_k_cache_paged"
quant_k_cache = quant_k_cache.view((-1, dim_quant))
# num_tokens can exceed kv_cache_size due to prefix sharing (multiple seqs share same KV slots)
# Index bounds validated in nsa_backend.init_forward_metadata
num_tokens = page_table_1_flattened.shape[0]
assert quant_k_cache.dtype == torch.float8_e4m3fn
dim_nope = 512
dim_rope = 64
num_tiles = dim_nope // group_size # 512 // 128 = 4
output = torch.empty(
(num_tokens, 1, dim_nope + dim_rope),
dtype=torch.bfloat16,
device=quant_k_cache.device,
)
# cdiv(512 + 64, 128) = 5
num_blocks_per_token = triton.cdiv(dim_nope + dim_rope, group_size)
assert num_blocks_per_token == 5
assert dim_nope % group_size == 0
input_nope_q = quant_k_cache[:, :dim_nope]
# [:, 512:512+4*4] = [:, 512:528]
input_nope_s = quant_k_cache[:, dim_nope : dim_nope + num_tiles * 4].view(
torch.float32
)
# [:, 528:]
input_rope = quant_k_cache[:, dim_nope + num_tiles * 4 :].view(torch.bfloat16)
_dequantize_k_cache_paged_kernel[(num_tokens, num_blocks_per_token)](
output,
input_nope_q,
input_nope_s,
input_rope,
page_table_1_flattened,
output.stride(0),
input_nope_q.stride(0),
input_nope_s.stride(0),
input_rope.stride(0),
NUM_NOPE_BLOCKS=num_tiles,
GROUP_SIZE=group_size,
DIM_NOPE=dim_nope,
DIM_ROPE=dim_rope,
)
return output
@triton.jit
def _dequantize_k_cache_paged_kernel(
output_ptr,
input_nope_q_ptr,
input_nope_s_ptr,
input_rope_ptr,
page_table_1_ptr,
output_stride_0: int,
input_nope_q_stride_0: int,
input_nope_s_stride_0: int,
input_rope_stride_0: int,
NUM_NOPE_BLOCKS: tl.constexpr,
GROUP_SIZE: tl.constexpr,
DIM_NOPE: tl.constexpr,
DIM_ROPE: tl.constexpr,
):
token_id = tl.program_id(0)
token_id_paged = tl.load(page_table_1_ptr + token_id).to(tl.int32)
raw_block_id = tl.program_id(1)
if raw_block_id < NUM_NOPE_BLOCKS:
# a. dequant nope
effective_block_id = raw_block_id
offs_q = effective_block_id * GROUP_SIZE + tl.arange(0, GROUP_SIZE)
mask = offs_q < DIM_NOPE
ptr_q = input_nope_q_ptr + token_id_paged * input_nope_q_stride_0 + offs_q
ptr_s = (
input_nope_s_ptr
+ token_id_paged * input_nope_s_stride_0
+ effective_block_id
)
y_q = tl.load(ptr_q, mask=mask, other=0.0).to(tl.float32)
y_s = tl.load(ptr_s)
y = (y_q * y_s).to(output_ptr.dtype.element_ty)
dst_ptr = output_ptr + token_id * output_stride_0 + offs_q
tl.store(dst_ptr, y, mask=mask)
else:
# b. copy rope
effective_block_id = raw_block_id - NUM_NOPE_BLOCKS
offs = effective_block_id * GROUP_SIZE + tl.arange(0, GROUP_SIZE)
mask = offs < DIM_ROPE
src_ptr = input_rope_ptr + token_id_paged * input_rope_stride_0 + offs
dst_ptr = output_ptr + token_id * output_stride_0 + DIM_NOPE + offs
data = tl.load(src_ptr, mask=mask).to(tl.bfloat16)
tl.store(dst_ptr, data, mask=mask)
if __name__ == "__main__":
raise Exception("UT is in quant_k_cache.py")
python/sglang/srt/layers/attention/nsa/index_buf_accessor.py
+8 -812
View File
@@ -1,814 +1,10 @@
from typing import TYPE_CHECKING # [Deprecated] Re-export shim for backward compatibility. Use dsa.index_buf_accessor instead.
import warnings
import torch warnings.warn(
import triton "sglang.srt.layers.attention.nsa.index_buf_accessor is deprecated; "
import triton.language as tl "use sglang.srt.layers.attention.dsa.index_buf_accessor instead.",
DeprecationWarning,
from sglang.srt.layers.attention.nsa.utils import aiter_can_use_preshuffle_paged_mqa stacklevel=2,
from sglang.srt.layers.quantization.fp8_kernel import is_fp8_fnuz
from sglang.srt.utils import get_bool_env_var, is_hip
_is_hip = is_hip()
_is_fp8_fnuz = is_fp8_fnuz()
_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
# aiter cp_gather kernel with preshuffle=True is only valid when the indexer
# uses the page_size=64 preshuffle layout (i.e. when the matching MQA gluon path
# is also enabled).
_use_aiter_preshuffle = aiter_can_use_preshuffle_paged_mqa()
if _use_aiter_preshuffle:
from aiter.ops.cache import cp_gather_indexer_k_quant_cache
if TYPE_CHECKING:
from sglang.srt.mem_cache.memory_pool import NSATokenToKVPool
"""
k: data, 128 item per token, fp8
s: scale, 1 item per token, fp32
"""
class GetK:
@classmethod
def execute(cls, *args, **kwargs):
return cls.triton(*args, **kwargs)
@classmethod
def slow(
cls, pool: "NSATokenToKVPool", buf, seq_len: int, page_indices: torch.Tensor
):
num_pages = (seq_len + pool.page_size - 1) // pool.page_size
seq_len_ = num_pages * pool.page_size
index_k_fp8 = torch.empty(
(seq_len_, pool.index_head_dim),
dtype=torch.uint8,
device=pool.device,
) )
for i in range(num_pages): from sglang.srt.layers.attention.dsa.index_buf_accessor import * # noqa: F401, F403
page_index = page_indices[i]
index_k_fp8[i * pool.page_size : (i + 1) * pool.page_size] = buf[
page_index
][: pool.page_size * pool.index_head_dim].view(-1, pool.index_head_dim)
return index_k_fp8[:seq_len]
@classmethod
def torch_fast(
cls, pool: "NSATokenToKVPool", buf, seq_len: int, page_indices: torch.Tensor
):
"""
:param page_indices: (num_pages,), int32
:return: (seq_len, index_head_dim), uint8
"""
# can handle per 128B instead of per element
# page_indices: (num_pages,), element := a page index
buf_numel_per_page = buf.shape[1]
num_k_bytes_per_page = pool.page_size * pool.index_head_dim
num_k_bytes_per_token = pool.index_head_dim
# buf: (num_pages, page_size 64 * head_dim 128 + page_size 64 * fp32_nbytes 4), uint8
# flat_buf: (whatever,), uint8
flat_buf = buf.flatten()
# flat_indices: (num_pages, num_k_bytes_per_page), int32, element := an index into flat_buf that we want to access
flat_indices = (page_indices * buf_numel_per_page)[:, None] + torch.arange(
num_k_bytes_per_page, dtype=torch.int32, device="cuda"
)[None, :]
flat_indices = flat_indices.flatten()[: seq_len * num_k_bytes_per_token]
out = flat_buf[flat_indices]
return out.view(-1, 128)
@classmethod
def triton(
cls, pool: "NSATokenToKVPool", buf, seq_len: int, page_indices: torch.Tensor
):
"""
Triton implementation for gathering K data from paged buffer.
:param page_indices: (num_pages,), int32/int64
:return: (seq_len, index_head_dim), uint8
"""
return _get_k_triton(
buf=buf,
page_indices=page_indices,
seq_len=seq_len,
page_size=pool.page_size,
index_head_dim=pool.index_head_dim,
)
class GetS:
@classmethod
def execute(cls, *args, **kwargs):
return cls.triton(*args, **kwargs)
@classmethod
def slow(
cls, pool: "NSATokenToKVPool", buf, seq_len: int, page_indices: torch.Tensor
):
num_pages = (seq_len + pool.page_size - 1) // pool.page_size
seq_len_ = num_pages * pool.page_size
assert pool.index_head_dim // pool.quant_block_size == 1
index_k_scale_fp8 = torch.empty(
(seq_len_, 4),
dtype=torch.uint8,
device=pool.device,
)
for i in range(num_pages):
page_index = page_indices[i]
index_k_scale_fp8[i * pool.page_size : (i + 1) * pool.page_size] = buf[
page_index
][pool.page_size * pool.index_head_dim :].view(-1, 4)
return index_k_scale_fp8[:seq_len]
@classmethod
def torch_fast(
cls, pool: "NSATokenToKVPool", buf, seq_len: int, page_indices: torch.Tensor
):
"""
:param page_indices: (num_pages,), int32
:return: (seq_len, index_head_dim // quant_block_size), uint8
"""
buf_numel_per_page = buf.shape[1]
num_s_bytes_per_page = buf.shape[1] - pool.page_size * pool.index_head_dim
num_s_bytes_per_token = pool.index_head_dim // pool.quant_block_size * 4
s_offset_in_page = pool.page_size * pool.index_head_dim
flat_buf = buf.flatten()
flat_indices = (
(page_indices * buf_numel_per_page)[:, None]
+ torch.arange(num_s_bytes_per_page, dtype=torch.int32, device="cuda")[
None, :
]
+ s_offset_in_page
)
flat_indices = flat_indices.flatten()[: seq_len * num_s_bytes_per_token]
out = flat_buf[flat_indices]
return out.view(-1, 4)
@classmethod
def triton(
cls, pool: "NSATokenToKVPool", buf, seq_len: int, page_indices: torch.Tensor
):
"""
Triton implementation for gathering S (scale) data from paged buffer.
:param page_indices: (num_pages,), int32/int64
:return: (seq_len, 4), uint8
"""
return _get_s_triton(
buf=buf,
page_indices=page_indices,
seq_len=seq_len,
page_size=pool.page_size,
index_head_dim=pool.index_head_dim,
)
class GetKAndS:
@classmethod
def execute(cls, *args, **kwargs):
# The aiter path uses cp_gather_indexer_k_quant_cache(preshuffle=True),
# which only matches the layout produced when the rest of the indexer
# is on the page_size=64 preshuffle path. Otherwise fall back to the
# triton implementation (which works on the page_size=1 legacy layout).
if _use_aiter_preshuffle:
return cls.aiter(*args, **kwargs)
return cls.triton(*args, **kwargs)
@classmethod
def aiter(
cls,
pool: "NSATokenToKVPool",
buf: torch.Tensor,
page_indices: torch.Tensor,
seq_len_tensor: torch.Tensor,
seq_len_sum: int,
max_seq_len: int,
):
from sglang.srt.layers.quantization.fp8_kernel import fp8_dtype
page_size = pool.page_size
index_head_dim = pool.index_head_dim
quant_block_size = pool.quant_block_size
scale_elems = index_head_dim // quant_block_size
kv_cache = buf.view(-1, page_size, index_head_dim + scale_elems * 4).view(
fp8_dtype
)
dst_k = torch.empty(
(seq_len_sum, index_head_dim), dtype=torch.uint8, device=buf.device
)
dst_scale = torch.empty(
(seq_len_sum, scale_elems * 4), dtype=torch.uint8, device=buf.device
)
cu_seq_lens = torch.zeros(
seq_len_tensor.shape[0] + 1, dtype=torch.int32, device=buf.device
)
torch.cumsum(seq_len_tensor.to(torch.int32), dim=0, out=cu_seq_lens[1:])
cp_gather_indexer_k_quant_cache(
kv_cache,
dst_k.view(fp8_dtype),
dst_scale,
page_indices.to(torch.int32),
cu_seq_lens,
preshuffle=True,
)
return dst_k, dst_scale
@classmethod
def triton(
cls,
pool: "NSATokenToKVPool",
buf: torch.Tensor,
page_indices: torch.Tensor,
seq_len_tensor: torch.Tensor,
seq_len_sum: int,
max_seq_len: int,
):
"""
Triton implementation for gathering both K and S data from paged buffer in a single call.
:param page_indices: (num_pages,), int32/int64
:param seq_len_tensor: (num_pages,), int32/int64
:param seq_len_sum: sum of all sequence len, int32
:param max_seq_len: max of all sequence len, int32
:return: tuple of (k_fp8, k_scale) where
k_fp8: (seq_len, index_head_dim), uint8
k_scale: (seq_len, 4), uint8
"""
return _get_k_and_s_triton(
buf=buf,
page_indices=page_indices,
seq_lens=seq_len_tensor,
seq_len_sum=seq_len_sum,
max_seq_len=max_seq_len,
page_size=pool.page_size,
index_head_dim=pool.index_head_dim,
)
class SetK:
@classmethod
def execute(cls, *args, buf, **kwargs):
return cls.torch_fast(*args, **kwargs, buf=buf)
@classmethod
def slow(
cls,
pool: "NSATokenToKVPool",
buf: torch.Tensor,
loc: torch.Tensor,
index_k: torch.Tensor,
):
for i in range(len(loc)):
page_index = loc[i] // pool.page_size
offset = loc[i] % pool.page_size
buf[
page_index,
offset * pool.index_head_dim : (offset + 1) * pool.index_head_dim,
] = index_k[i].view(torch.uint8)
@classmethod
def torch_fast(
cls,
pool: "NSATokenToKVPool",
buf: torch.Tensor,
loc: torch.Tensor,
index_k: torch.Tensor,
):
(num_tokens_to_write,) = loc.shape
buf_numel_per_page = buf.shape[1]
num_k_bytes_per_token = pool.index_head_dim
# loc: (num_tokens_to_write,), int32, element := the token index to write to
loc_page_index = loc // pool.page_size
loc_token_offset_in_page = loc % pool.page_size
flat_buf = buf.flatten()
flat_indices = (
(loc_page_index * buf_numel_per_page)[:, None]
+ (loc_token_offset_in_page * num_k_bytes_per_token)[:, None]
+ torch.arange(num_k_bytes_per_token, dtype=torch.int32, device="cuda")[
None, :
]
)
num_k_bytes_total = num_tokens_to_write * num_k_bytes_per_token
flat_indices = flat_indices.flatten()[:num_k_bytes_total]
flat_buf[flat_indices] = index_k.view(torch.uint8).flatten()
class SetS:
@classmethod
def execute(cls, *args, buf, **kwargs):
return cls.torch_fast(*args, **kwargs, buf=buf)
@classmethod
def slow(
cls,
pool: "NSATokenToKVPool",
buf: torch.Tensor,
loc: torch.Tensor,
index_k_scale: torch.Tensor,
):
for i in range(len(loc)):
page_index = loc[i] // pool.page_size
offset = loc[i] % pool.page_size
start = pool.page_size * pool.index_head_dim
buf[page_index, start + offset * 4 : start + (offset + 1) * 4] = (
index_k_scale[i].view(torch.uint8)
)
@classmethod
def torch_fast(
cls,
pool: "NSATokenToKVPool",
buf: torch.Tensor,
loc: torch.Tensor,
index_k_scale: torch.Tensor,
):
(num_tokens_to_write,) = loc.shape
buf_numel_per_page = buf.shape[1]
num_s_bytes_per_token = 4
s_offset_in_page = pool.page_size * pool.index_head_dim
# loc: (num_tokens_to_write,), int32, element := the token index to write to
loc_page_index = loc // pool.page_size
loc_token_offset_in_page = loc % pool.page_size
flat_buf = buf.flatten()
flat_indices = (
(loc_page_index * buf_numel_per_page)[:, None]
+ s_offset_in_page
+ (loc_token_offset_in_page * num_s_bytes_per_token)[:, None]
+ torch.arange(num_s_bytes_per_token, dtype=torch.int32, device="cuda")[
None, :
]
)
number_s_bytes_total = num_tokens_to_write * num_s_bytes_per_token
flat_indices = flat_indices.flatten()[:number_s_bytes_total]
flat_buf[flat_indices] = index_k_scale.view(torch.uint8).flatten()
class SetKAndS:
@classmethod
def execute(cls, *args, buf, **kwargs):
if 0:
# print("SetK, SetS comparison test")
buf_cloned = buf.clone()
cls.vanilla(*args, **kwargs, buf=buf)
cls.triton(*args, **kwargs, buf=buf_cloned)
def _clear_token_0(target):
target[0, :128] = target[0, 64 * 128 : 64 * 128 + 4] = 0
_clear_token_0(buf)
_clear_token_0(buf_cloned)
assert torch.all(
buf == buf_cloned
), f"{buf=} {buf_cloned=} {kwargs['loc'].to_list()=}"
return
cls.triton(*args, **kwargs, buf=buf)
@classmethod
def vanilla(cls, pool, buf, loc, index_k, index_k_scale):
SetK.execute(pool=pool, buf=buf, loc=loc, index_k=index_k)
SetS.execute(pool=pool, buf=buf, loc=loc, index_k_scale=index_k_scale)
@classmethod
def triton(cls, pool, buf, loc, index_k, index_k_scale):
loc = loc.to(torch.int64)
_set_k_and_s_triton(
buf=buf,
loc=loc,
index_k=index_k,
index_k_scale=index_k_scale,
page_size=pool.page_size,
)
def _set_k_and_s_triton(
buf: torch.Tensor,
loc: torch.Tensor,
index_k: torch.Tensor,
index_k_scale: torch.Tensor,
page_size: int,
):
"""
:param buf: (num_pages, page_size 64 * (128B data + 4B scale)), uint8
:param loc: (num_tokens_to_write,), int, element := the token index to write to
:param index_k: (num_tokens_to_write, 128 elem), fp8
:param index_k_scale: (num_tokens_to_write, 1 elem), fp32
:return:
"""
num_pages, buf_numel_per_page = buf.shape
(num_tokens_to_write,) = loc.shape
num_tokens_to_write_, index_head_dim = index_k.shape
# Handle both 1D (num_tokens,) and 2D (num_tokens, 1) shapes for index_k_scale
if index_k_scale.ndim == 1:
num_tokens_to_write__ = index_k_scale.shape[0]
scale_dim = 1
elif index_k_scale.ndim == 2:
num_tokens_to_write__, scale_dim = index_k_scale.shape
else:
raise ValueError(
f"index_k_scale must be 1D or 2D, got shape {index_k_scale.shape}"
)
assert buf_numel_per_page == page_size * (128 + 4)
assert num_tokens_to_write == num_tokens_to_write_ == num_tokens_to_write__
assert index_head_dim == 128
assert scale_dim == 1
if _is_hip:
if _use_aiter_preshuffle:
assert (
page_size % 16 == 0
), f"HIP preshuffle requires page_size to be a multiple of 16, got {page_size}"
else:
assert page_size == 64
assert buf.dtype == torch.uint8
assert loc.dtype == torch.int64, f"{loc.dtype=}" # can be int32
if _is_fp8_fnuz:
assert index_k.dtype == torch.float8_e4m3fnuz
else:
assert index_k.dtype == torch.float8_e4m3fn
assert index_k_scale.dtype == torch.float32
assert buf.is_contiguous()
assert loc.is_contiguous()
assert index_k.is_contiguous()
assert index_k_scale.is_contiguous()
if _is_fp8_fnuz:
buf_fp8 = buf.view(torch.float8_e4m3fnuz)
else:
buf_fp8 = buf.view(torch.float8_e4m3fn)
buf_fp32 = buf.view(torch.float32)
_set_k_and_s_triton_kernel[(num_tokens_to_write,)](
buf_fp8,
buf_fp32,
loc,
index_k,
index_k_scale,
index_k.stride(0),
PAGE_SIZE=page_size,
BUF_NUMEL_PER_PAGE=buf_numel_per_page,
NUM_K_ELEMS_PER_TOKEN=index_head_dim,
S_OFFSET_NBYTES_IN_PAGE=page_size * index_head_dim,
)
@triton.jit
def _set_k_and_s_triton_kernel(
buf_fp8_ptr,
buf_fp32_ptr,
loc_ptr,
index_k_ptr,
index_k_scale_ptr,
index_k_ptr_stride_0,
PAGE_SIZE: tl.constexpr,
BUF_NUMEL_PER_PAGE: tl.constexpr,
NUM_K_ELEMS_PER_TOKEN: tl.constexpr,
S_OFFSET_NBYTES_IN_PAGE: tl.constexpr,
):
token_id = tl.program_id(0)
loc = tl.load(loc_ptr + token_id)
in_k_offsets = token_id * index_k_ptr_stride_0 + tl.arange(0, NUM_K_ELEMS_PER_TOKEN)
# no need for `mask`, since we read 128B for k and 4B for scale, both pow of 2
k = tl.load(index_k_ptr + in_k_offsets)
k_scale = tl.load(index_k_scale_ptr + token_id)
loc_page_index = loc // PAGE_SIZE
loc_token_offset_in_page = loc % PAGE_SIZE
out_k_offsets = (
loc_page_index * BUF_NUMEL_PER_PAGE
+ loc_token_offset_in_page * NUM_K_ELEMS_PER_TOKEN
+ tl.arange(0, NUM_K_ELEMS_PER_TOKEN)
)
# "//4" b/c it is fp32 instead of uint8
out_s_offset = (
loc_page_index * BUF_NUMEL_PER_PAGE // 4
+ S_OFFSET_NBYTES_IN_PAGE // 4
+ loc_token_offset_in_page
)
tl.store(buf_fp8_ptr + out_k_offsets, k)
tl.store(buf_fp32_ptr + out_s_offset, k_scale)
def _get_k_triton(
buf: torch.Tensor,
page_indices: torch.Tensor,
seq_len: int,
page_size: int,
index_head_dim: int,
):
"""
Gather K (key) data from paged buffer using Triton.
:param buf: (num_pages, page_size * 128 + page_size * 4), uint8
:param page_indices: (num_pages,), int32/int64
:param seq_len: int, number of tokens to gather
:param page_size: int, typically 64
:param index_head_dim: int, typically 128
:return: (seq_len, index_head_dim), uint8
"""
num_pages, buf_numel_per_page = buf.shape
# Allocate output
out = torch.empty((seq_len, index_head_dim), dtype=torch.uint8, device=buf.device)
# Launch kernel with one thread per token
grid = (seq_len,)
_get_k_triton_kernel[grid](
buf,
page_indices,
out,
seq_len,
page_size,
buf_numel_per_page,
index_head_dim,
BLOCK_SIZE=128,
)
return out
@triton.jit
def _get_k_triton_kernel(
buf_ptr,
page_indices_ptr,
out_ptr,
seq_len: tl.constexpr,
page_size: tl.constexpr,
buf_numel_per_page: tl.constexpr,
index_head_dim: tl.constexpr,
BLOCK_SIZE: tl.constexpr,
):
"""
Each program handles one token (seq_len tokens total).
Loads 128 bytes from the appropriate page.
"""
token_id = tl.program_id(0)
# Calculate which page and offset within page
page_idx = token_id // page_size
token_offset_in_page = token_id % page_size
# Load the page index from page_indices
page_index = tl.load(page_indices_ptr + page_idx)
# Calculate source offset in buf
# buf[page_index, token_offset_in_page * index_head_dim : ...]
src_base_offset = (
page_index * buf_numel_per_page + token_offset_in_page * index_head_dim
)
# Load 128 bytes (index_head_dim elements)
offsets = tl.arange(0, BLOCK_SIZE)
mask = offsets < index_head_dim
data = tl.load(buf_ptr + src_base_offset + offsets, mask=mask)
# Store to output
dst_offset = token_id * index_head_dim
tl.store(out_ptr + dst_offset + offsets, data, mask=mask)
def _get_s_triton(
buf: torch.Tensor,
page_indices: torch.Tensor,
seq_len: int,
page_size: int,
index_head_dim: int,
):
"""
Gather S (scale) data from paged buffer using Triton.
:param buf: (num_pages, page_size * 128 + page_size * 4), uint8
:param page_indices: (num_pages,), int32/int64
:param seq_len: int, number of tokens to gather
:param page_size: int, typically 64
:param index_head_dim: int, typically 128
:return: (seq_len, 4), uint8 (representing fp32 scale)
"""
num_pages, buf_numel_per_page = buf.shape
s_offset_in_page = page_size * index_head_dim # Scales start after K data
# Allocate output
out = torch.empty((seq_len, 4), dtype=torch.uint8, device=buf.device)
# Launch kernel with one thread per token
grid = (seq_len,)
_get_s_triton_kernel[grid](
buf,
page_indices,
out,
seq_len,
page_size,
buf_numel_per_page,
s_offset_in_page,
)
return out
@triton.jit
def _get_s_triton_kernel(
buf_ptr,
page_indices_ptr,
out_ptr,
seq_len: tl.constexpr,
page_size: tl.constexpr,
buf_numel_per_page: tl.constexpr,
s_offset_in_page: tl.constexpr,
):
"""
Each program handles one token (seq_len tokens total).
Loads 4 bytes (fp32 scale) from the appropriate page.
"""
token_id = tl.program_id(0)
# Calculate which page and offset within page
page_idx = token_id // page_size
token_offset_in_page = token_id % page_size
# Load the page index from page_indices
page_index = tl.load(page_indices_ptr + page_idx)
# Calculate source offset in buf
# Scales are stored after K data: page_size * index_head_dim offset
# buf[page_index, s_offset_in_page + token_offset_in_page * 4 : ...]
src_base_offset = (
page_index * buf_numel_per_page + s_offset_in_page + token_offset_in_page * 4
)
# Load 4 bytes (fp32 scale)
offsets = tl.arange(0, 4)
data = tl.load(buf_ptr + src_base_offset + offsets)
# Store to output
dst_offset = token_id * 4
tl.store(out_ptr + dst_offset + offsets, data)
def _get_k_and_s_triton(
buf: torch.Tensor,
page_indices: torch.Tensor,
seq_lens: torch.Tensor,
seq_len_sum: int,
max_seq_len: int,
page_size: int,
index_head_dim: int,
):
"""
Fused gather of both K (key) and S (scale) data from paged buffer using Triton.
This is more efficient than calling GetK and GetS separately.
:param buf: (num_pages, page_size * 128 + page_size * 4), uint8
:param page_indices: (num_pages,), int32/int64
:param seq_lens: tensor of sequence lens, int64
:param seq_len_sum: sum of all sequence len, int32
:param max_seq_len: max of sequence len, int32
:param page_size: int, typically 64
:param index_head_dim: int, typically 128
:return: tuple of (k_out, s_out) where
k_out: (seq_len, index_head_dim), uint8
s_out: (seq_len, 4), uint8
"""
# Allocate outputs
k_out = torch.empty(
(seq_len_sum, index_head_dim), dtype=torch.uint8, device=buf.device
)
s_out = torch.empty((seq_len_sum, 4), dtype=torch.uint8, device=buf.device)
_, buf_numel_per_page = buf.shape
_, page_indice_batch_offset = page_indices.shape
s_offset_in_page = page_size * index_head_dim
# Launch kernel with one thread per token
BLOCK_SIZE = 256
BLOCK_SIZE_K = 128
num_token_blocks = (max_seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE
num_k_threads = (index_head_dim + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K
seq_num = seq_lens.shape[0]
grid = (seq_num, num_token_blocks, num_k_threads)
seq_num_pow2 = 1
while seq_num_pow2 < seq_num:
seq_num_pow2 *= 2
_get_k_and_s_triton_kernel[grid](
buf_ptr=buf,
page_indices_ptr=page_indices,
k_out_ptr=k_out,
s_out_ptr=s_out,
seq_len_ptr=seq_lens,
seq_len_num_pow=seq_num_pow2,
page_size=page_size,
buf_numel_per_page=buf_numel_per_page,
index_head_dim=index_head_dim,
s_offset_in_page=s_offset_in_page,
page_indice_batch_offset=page_indice_batch_offset,
BLOCK_SIZE=BLOCK_SIZE,
BLOCK_SIZE_K=BLOCK_SIZE_K,
)
return k_out, s_out
@triton.jit
def _get_k_and_s_triton_kernel(
buf_ptr,
page_indices_ptr,
k_out_ptr,
s_out_ptr,
seq_len_ptr,
seq_len_num_pow: tl.constexpr,
page_size: tl.constexpr,
buf_numel_per_page: tl.constexpr,
index_head_dim: tl.constexpr,
s_offset_in_page: tl.constexpr,
page_indice_batch_offset,
BLOCK_SIZE: tl.constexpr,
BLOCK_SIZE_K: tl.constexpr,
):
"""
Fused kernel that gathers both K and S data in a single pass.
Each program handles one token (seq_len tokens total).
Loads 128 bytes (K) + 4 bytes (S) from the appropriate page.
"""
batch_id = tl.program_id(0)
block_token_start = tl.program_id(1) * BLOCK_SIZE
thread_idx = tl.program_id(2)
# Define the token range within the block and the K dimension range handled by the thread.
token_ids_in_block = tl.arange(0, BLOCK_SIZE)
token_ids = block_token_start + token_ids_in_block
k_offsets = thread_idx * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K)
seq_len = tl.load(seq_len_ptr + batch_id)
token_valid_mask = token_ids < seq_len
pre_batch_idx = tl.arange(0, seq_len_num_pow)
mask_pre_batch_idx = pre_batch_idx < batch_id
prev_seq_lens = tl.load(seq_len_ptr + pre_batch_idx, mask=mask_pre_batch_idx)
batch_token_offset = tl.sum(prev_seq_lens)
# Batch calculate the page index and in-page offset of each token.
page_idx = token_ids // page_size
token_offset_in_page = token_ids % page_size
page_indices_base = batch_id * page_indice_batch_offset
page_idx_valid_mask = page_idx < page_indice_batch_offset
page_index = tl.load(
page_indices_ptr + page_idx + page_indices_base,
mask=token_valid_mask & page_idx_valid_mask,
)
# ===== Load K data =====
# The address calculation logic for K: page_index * total number of elements in a single page + K offset of the token within the page.
k_src_token_offset = token_offset_in_page * index_head_dim
k_src_base_offset = page_index * buf_numel_per_page + k_src_token_offset
k_load_addr = buf_ptr + k_src_base_offset[:, None] + k_offsets[None, :]
k_dim_mask = k_offsets[None, :] < index_head_dim
k_mask = token_valid_mask[:, None] & k_dim_mask
k_data = tl.load(k_load_addr, mask=k_mask, other=0)
# Store K to output
k_dst_token_offset = batch_token_offset + token_ids
k_dst_base_offset = k_dst_token_offset * index_head_dim
k_store_addr = k_out_ptr + k_dst_base_offset[:, None] + k_offsets[None, :]
tl.store(k_store_addr, k_data, mask=k_mask)
# ===== Load S data =====
# The address calculation logic for S: page_index * total number of elements in a single page + starting offset of S within the page + offset of token within S in the page
s_src_token_offset = s_offset_in_page + token_offset_in_page * 4
s_src_base_offset = page_index * buf_numel_per_page + s_src_token_offset
s_offsets = tl.arange(0, 4)
s_load_addr = buf_ptr + s_src_base_offset[:, None] + s_offsets[None, :]
s_mask = token_valid_mask[:, None] & (s_offsets[None, :] < 4)
s_data = tl.load(s_load_addr, mask=s_mask, other=0)
# Store S to output
s_dst_token_offset = batch_token_offset + token_ids
s_dst_base_offset = s_dst_token_offset * 4
s_store_addr = s_out_ptr + s_dst_base_offset[:, None] + s_offsets[None, :]
tl.store(s_store_addr, s_data, mask=s_mask)
python/sglang/srt/layers/attention/nsa/nsa_backend_mtp_precompute.py
+8 -323
View File
@@ -1,325 +1,10 @@
"""Multi-step precompute utilities for Native Sparse Attention backend. # [Deprecated] Re-export shim for backward compatibility. Use dsa.dsa_backend_mtp_precompute instead.
import warnings
This module provides optimization utilities for multi-step speculative decoding warnings.warn(
by precomputing shared metadata once and copying it to multiple backend instances. "sglang.srt.layers.attention.nsa.nsa_backend_mtp_precompute is deprecated; "
""" "use sglang.srt.layers.attention.dsa.dsa_backend_mtp_precompute instead.",
DeprecationWarning,
from __future__ import annotations stacklevel=2,
from dataclasses import dataclass
from typing import TYPE_CHECKING, Optional
import torch
from sglang.srt.layers.attention.nsa.utils import compute_nsa_seqlens
if TYPE_CHECKING:
from sglang.srt.model_executor.forward_batch_info import ForwardMode
from sglang.srt.speculative.spec_info import SpecInput
@dataclass
class PrecomputedMetadata:
"""Precomputed metadata shared across multiple backend instances.
Used for multi-step speculative decoding where multiple backends
need identical metadata. Precomputing once and copying N times
is much faster than computing N times.
"""
# Basic seqlens
cache_seqlens: torch.Tensor # int32, [bs]
cu_seqlens_k: torch.Tensor # int32, [bs+1]
# Page table
page_indices: torch.Tensor # int32, [bs, max_len] or [expanded_bs, max_len]
real_page_table: Optional[torch.Tensor] # int32, transformed version
# NSA seqlens
seqlens_expanded: torch.Tensor # int32, [expanded_size]
nsa_cache_seqlens: torch.Tensor # int32, [expanded_size]
nsa_cu_seqlens_k: torch.Tensor # int32, [expanded_size+1]
seqlens_expanded_size: int
# Dimensions
max_len: int # for decode/draft_extend
max_seqlen_k: int # for target_verify
# FlashMLA (optional)
flashmla_metadata: Optional[torch.Tensor] = None
def compute_cu_seqlens(seqlens: torch.Tensor) -> torch.Tensor:
"""Compute cumulative sequence lengths with padding."""
assert seqlens.dtype == torch.int32
return torch.nn.functional.pad(
torch.cumsum(seqlens, dim=0, dtype=torch.int32), (1, 0)
)
class NativeSparseAttnBackendMTPPrecomputeMixin:
"""Mixin class providing metadata precomputation for multi-step speculative decoding.
This mixin provides the _precompute_replay_metadata method and its helpers,
which are used to optimize CUDA graph replay in multi-step scenarios.
"""
def _precompute_replay_metadata(
self,
bs: int,
req_pool_indices: torch.Tensor,
seq_lens: torch.Tensor,
seq_lens_cpu: torch.Tensor,
forward_mode: "ForwardMode",
spec_info: Optional["SpecInput"],
) -> PrecomputedMetadata:
"""Precompute all shared metadata for multi-step backends.
This function extracts and computes all operations that are
identical across different backend instances in multi-step
speculative decoding.
Args:
bs: Batch size
req_pool_indices: Request pool indices [bs]
seq_lens: Sequence lengths [bs]
seq_lens_cpu: Sequence lengths on CPU [bs]
forward_mode: Forward mode (decode/target_verify/draft_extend)
spec_info: Speculative decoding info (for draft_extend mode)
Returns:
PrecomputedMetadata containing all shared intermediate results
"""
# Slice inputs to batch size
seq_lens = seq_lens[:bs]
seq_lens_cpu = seq_lens_cpu[:bs]
req_pool_indices = req_pool_indices[:bs]
# Dispatch to mode-specific precomputation
if forward_mode.is_decode_or_idle():
return self._precompute_decode_mode(
bs, req_pool_indices, seq_lens, seq_lens_cpu
)
elif forward_mode.is_target_verify():
return self._precompute_target_verify_mode(
bs, req_pool_indices, seq_lens, seq_lens_cpu
)
elif forward_mode.is_draft_extend():
return self._precompute_draft_extend_mode(
bs, req_pool_indices, seq_lens, seq_lens_cpu, spec_info
)
else:
raise ValueError(f"Unsupported forward mode: {forward_mode}")
def _precompute_decode_mode(
self,
bs: int,
req_pool_indices: torch.Tensor,
seq_lens: torch.Tensor,
seq_lens_cpu: torch.Tensor,
) -> PrecomputedMetadata:
"""Precompute metadata for normal decode mode."""
max_len = int(seq_lens_cpu.max().item())
# Convert to int32 and compute cumsum
cache_seqlens = seq_lens.to(torch.int32)
cu_seqlens_k = compute_cu_seqlens(cache_seqlens)
# Get page indices from cache
page_indices = self.req_to_token[req_pool_indices, :max_len].contiguous()
# Compute NSA seqlens
nsa_cache_seqlens = compute_nsa_seqlens(
cache_seqlens, nsa_index_topk=self.nsa_index_topk
)
seqlens_expanded = cache_seqlens
seqlens_expanded_size = seqlens_expanded.shape[0]
# Compute NSA cumsum
nsa_cu_seqlens_k = compute_cu_seqlens(nsa_cache_seqlens)
# Transform page table if needed
if self.real_page_size > 1:
real_page_table = self._transform_table_1_to_real(page_indices)
else:
real_page_table = None # Will use page_indices directly
# Compute FlashMLA metadata if needed
flashmla_metadata = None
if self.nsa_decode_impl == "flashmla_kv":
flashmla_metadata = self._compute_flashmla_metadata(
cache_seqlens=nsa_cache_seqlens,
seq_len_q=1,
)
return PrecomputedMetadata(
cache_seqlens=cache_seqlens,
cu_seqlens_k=cu_seqlens_k,
page_indices=page_indices,
real_page_table=real_page_table,
seqlens_expanded=seqlens_expanded,
nsa_cache_seqlens=nsa_cache_seqlens,
nsa_cu_seqlens_k=nsa_cu_seqlens_k,
seqlens_expanded_size=seqlens_expanded_size,
max_len=max_len,
max_seqlen_k=max_len,
flashmla_metadata=flashmla_metadata,
)
def _precompute_target_verify_mode(
self,
bs: int,
req_pool_indices: torch.Tensor,
seq_lens: torch.Tensor,
seq_lens_cpu: torch.Tensor,
) -> PrecomputedMetadata:
"""Precompute metadata for target verify mode."""
max_seqlen_k = int(
seq_lens_cpu.max().item() + self.speculative_num_draft_tokens
)
# Cache seqlens with draft tokens
cache_seqlens = (seq_lens + self.speculative_num_draft_tokens).to(torch.int32)
cu_seqlens_k = compute_cu_seqlens(cache_seqlens)
# Page indices (repeated for each draft token)
page_indices = self.req_to_token[req_pool_indices, :max_seqlen_k]
page_indices = torch.repeat_interleave(
page_indices, repeats=self.speculative_num_draft_tokens, dim=0
).contiguous()
# Generate expanded seqlens
extend_seq_lens_cpu = [self.speculative_num_draft_tokens] * bs
seqlens_int32_cpu = [
self.speculative_num_draft_tokens + kv_len
for kv_len in seq_lens_cpu.tolist()
]
seqlens_expanded = torch.cat(
[
torch.arange(
kv_len - qo_len + 1,
kv_len + 1,
dtype=torch.int32,
device=self.device,
)
for qo_len, kv_len in zip(
extend_seq_lens_cpu,
seqlens_int32_cpu,
strict=True,
)
]
)
# Compute NSA seqlens
nsa_cache_seqlens = compute_nsa_seqlens(seqlens_expanded, self.nsa_index_topk)
seqlens_expanded_size = seqlens_expanded.shape[0]
# NSA cumsum
nsa_cu_seqlens_k = compute_cu_seqlens(nsa_cache_seqlens)
# Transform page table
if self.real_page_size > 1:
real_page_table = self._transform_table_1_to_real(page_indices)
else:
real_page_table = None
# FlashMLA metadata
flashmla_metadata = None
if self.nsa_decode_impl == "flashmla_kv":
flashmla_metadata = self._compute_flashmla_metadata(
cache_seqlens=nsa_cache_seqlens,
seq_len_q=1,
)
return PrecomputedMetadata(
cache_seqlens=cache_seqlens,
cu_seqlens_k=cu_seqlens_k,
page_indices=page_indices,
real_page_table=real_page_table,
seqlens_expanded=seqlens_expanded,
nsa_cache_seqlens=nsa_cache_seqlens,
nsa_cu_seqlens_k=nsa_cu_seqlens_k,
seqlens_expanded_size=seqlens_expanded_size,
max_len=-1, # Not used in this mode
max_seqlen_k=max_seqlen_k,
flashmla_metadata=flashmla_metadata,
)
def _precompute_draft_extend_mode(
self,
bs: int,
req_pool_indices: torch.Tensor,
seq_lens: torch.Tensor,
seq_lens_cpu: torch.Tensor,
spec_info: "SpecInput",
) -> PrecomputedMetadata:
"""Precompute metadata for draft extend mode."""
max_seqlen_k = int(seq_lens_cpu.max().item())
# Cache seqlens
cache_seqlens = seq_lens.to(torch.int32)
cu_seqlens_k = compute_cu_seqlens(cache_seqlens)
# Extend seqlens from spec_info: num_accept_tokens already includes
# the bonus token (drafts + 1).
extend_seq_lens = spec_info.num_accept_tokens[:bs]
extend_seq_lens_cpu = extend_seq_lens.tolist()
# Page indices (repeated per accept length)
page_indices = self.req_to_token[req_pool_indices, :max_seqlen_k]
page_indices = torch.repeat_interleave(
page_indices, repeats=extend_seq_lens, dim=0
).contiguous()
# Generate expanded seqlens
seqlens_expanded = torch.cat(
[
torch.arange(
kv_len - qo_len + 1,
kv_len + 1,
dtype=torch.int32,
device=self.device,
)
for qo_len, kv_len in zip(
extend_seq_lens_cpu,
seq_lens_cpu.tolist(),
strict=True,
)
]
)
# Compute NSA seqlens
nsa_cache_seqlens = compute_nsa_seqlens(seqlens_expanded, self.nsa_index_topk)
seqlens_expanded_size = seqlens_expanded.shape[0]
# NSA cumsum
nsa_cu_seqlens_k = compute_cu_seqlens(nsa_cache_seqlens)
# Transform page table
if self.real_page_size > 1:
real_page_table = self._transform_table_1_to_real(page_indices)
else:
real_page_table = None
# FlashMLA metadata
flashmla_metadata = None
if self.nsa_decode_impl == "flashmla_kv":
flashmla_metadata = self._compute_flashmla_metadata(
cache_seqlens=nsa_cache_seqlens,
seq_len_q=1,
)
return PrecomputedMetadata(
cache_seqlens=cache_seqlens,
cu_seqlens_k=cu_seqlens_k,
page_indices=page_indices,
real_page_table=real_page_table,
seqlens_expanded=seqlens_expanded,
nsa_cache_seqlens=nsa_cache_seqlens,
nsa_cu_seqlens_k=nsa_cu_seqlens_k,
seqlens_expanded_size=seqlens_expanded_size,
max_len=max_seqlen_k,
max_seqlen_k=max_seqlen_k,
flashmla_metadata=flashmla_metadata,
) )
from sglang.srt.layers.attention.dsa.dsa_backend_mtp_precompute import * # noqa: F401, F403
python/sglang/srt/layers/attention/nsa/nsa_indexer.py
+8 -1744
View File
File diff suppressed because it is too large Load Diff
python/sglang/srt/layers/attention/nsa/nsa_mtp_verification.py
+8 -405
View File
@@ -1,407 +1,10 @@
""" # [Deprecated] Re-export shim for backward compatibility. Use dsa.dsa_mtp_verification instead.
Verification utilities for NSA backend fused metadata copy operations. import warnings
This module contains verification code to ensure that fused metadata copy kernels warnings.warn(
produce the same results as individual copy operations. "sglang.srt.layers.attention.nsa.nsa_mtp_verification is deprecated; "
""" "use sglang.srt.layers.attention.dsa.dsa_mtp_verification instead.",
DeprecationWarning,
import torch stacklevel=2,
def verify_single_backend_fused_metadata_copy(
metadata,
precomputed,
forward_mode,
bs,
flashmla_num_splits_src=None,
flashmla_metadata_src=None,
flashmla_num_splits_dst=None,
flashmla_metadata_dst=None,
):
"""
Verify that the fused metadata copy kernel produces the same results as individual copies.
Args:
metadata: The NSA metadata object containing destination tensors
precomputed: The precomputed metadata containing source tensors
forward_mode: The forward mode (decode, target_verify, or draft_extend)
bs: Batch size
flashmla_num_splits_src: Source FlashMLA num_splits tensor (optional)
flashmla_metadata_src: Source FlashMLA metadata tensor (optional)
flashmla_num_splits_dst: Destination FlashMLA num_splits tensor (optional)
flashmla_metadata_dst: Destination FlashMLA metadata tensor (optional)
Raises:
RuntimeError: If verification fails (tensors don't match)
"""
# Clone destination tensors to preserve fused kernel results
fused_cache_seqlens = metadata.cache_seqlens_int32.clone()
fused_cu_seqlens_k = metadata.cu_seqlens_k.clone()
fused_page_table_1 = metadata.page_table_1.clone()
fused_nsa_cache_seqlens = metadata.nsa_cache_seqlens_int32.clone()
fused_nsa_seqlens_expanded = metadata.nsa_seqlens_expanded.clone()
fused_nsa_cu_seqlens_k = metadata.nsa_cu_seqlens_k.clone()
fused_real_page_table = (
metadata.real_page_table.clone()
if precomputed.real_page_table is not None
else None
)
fused_flashmla_num_splits = None
fused_flashmla_metadata = None
if precomputed.flashmla_metadata is not None:
fused_flashmla_num_splits = flashmla_num_splits_dst.clone()
fused_flashmla_metadata = flashmla_metadata_dst.clone()
# Create reference tensors (zeroed out)
ref_cache_seqlens = torch.zeros_like(metadata.cache_seqlens_int32)
ref_cu_seqlens_k = torch.zeros_like(metadata.cu_seqlens_k)
ref_page_table_1 = torch.zeros_like(metadata.page_table_1)
ref_nsa_cache_seqlens = torch.zeros_like(metadata.nsa_cache_seqlens_int32)
ref_nsa_seqlens_expanded = torch.zeros_like(metadata.nsa_seqlens_expanded)
ref_nsa_cu_seqlens_k = torch.zeros_like(metadata.nsa_cu_seqlens_k)
ref_real_page_table = (
torch.zeros_like(metadata.real_page_table)
if precomputed.real_page_table is not None
else None
)
ref_flashmla_num_splits = None
ref_flashmla_metadata = None
if precomputed.flashmla_metadata is not None:
ref_flashmla_num_splits = torch.zeros_like(flashmla_num_splits_dst)
ref_flashmla_metadata = torch.zeros_like(flashmla_metadata_dst)
# Run individual copy operations (reference implementation)
ref_cache_seqlens.copy_(precomputed.cache_seqlens)
ref_cu_seqlens_k[1:].copy_(precomputed.cu_seqlens_k[1:])
if forward_mode.is_decode_or_idle():
# Decode mode
ref_page_table_1[:, : precomputed.max_len].copy_(precomputed.page_indices)
ref_nsa_cache_seqlens.copy_(precomputed.nsa_cache_seqlens)
elif forward_mode.is_target_verify():
# Target verify mode
ref_page_table_1[:, : precomputed.max_seqlen_k].copy_(precomputed.page_indices)
ref_nsa_seqlens_expanded.copy_(precomputed.seqlens_expanded)
ref_nsa_cache_seqlens.copy_(precomputed.nsa_cache_seqlens)
elif forward_mode.is_draft_extend():
# Draft extend mode
rows = precomputed.page_indices.shape[0]
cols = precomputed.max_seqlen_k
ref_page_table_1[:rows, :cols].copy_(precomputed.page_indices)
size = precomputed.seqlens_expanded_size
ref_nsa_seqlens_expanded[:size].copy_(precomputed.seqlens_expanded)
ref_nsa_cache_seqlens[:size].copy_(precomputed.nsa_cache_seqlens)
# Copy NSA cu_seqlens
size = precomputed.seqlens_expanded_size
ref_nsa_cu_seqlens_k[1 : 1 + size].copy_(precomputed.nsa_cu_seqlens_k[1 : 1 + size])
# Copy real page table
if precomputed.real_page_table is not None:
rows, cols = precomputed.real_page_table.shape
ref_real_page_table[:rows, :cols].copy_(precomputed.real_page_table)
# Copy FlashMLA metadata
if precomputed.flashmla_metadata is not None:
size = precomputed.seqlens_expanded_size
ref_flashmla_num_splits[: size + 1].copy_(flashmla_num_splits_src[: size + 1])
ref_flashmla_metadata.copy_(flashmla_metadata_src)
# Compare results and crash if inconsistent
def check_tensor_equal(name, fused, ref):
if not torch.equal(fused, ref):
max_diff = (fused.float() - ref.float()).abs().max().item()
mismatched_elements = (fused != ref).sum().item()
total_elements = fused.numel()
raise RuntimeError(
f"FUSED METADATA COPY VERIFICATION FAILED!\n"
f"Tensor: {name}\n"
f"Max difference: {max_diff}\n"
f"Mismatched elements: {mismatched_elements}/{total_elements}\n"
f"Fused shape: {fused.shape}, Ref shape: {ref.shape}\n"
f"Forward mode: {forward_mode}, bs={bs}\n"
f"The fused kernel produces different results than individual copies.\n"
f"This indicates a bug in the fused metadata copy kernel."
)
# Verify all tensors (only compare the slices that were actually updated)
check_tensor_equal("cache_seqlens", fused_cache_seqlens, ref_cache_seqlens)
check_tensor_equal("cu_seqlens_k", fused_cu_seqlens_k, ref_cu_seqlens_k)
# Compare page_table_1 only for the region that was updated
if forward_mode.is_decode_or_idle():
check_tensor_equal(
"page_table_1",
fused_page_table_1[:, : precomputed.max_len],
ref_page_table_1[:, : precomputed.max_len],
)
elif forward_mode.is_target_verify():
check_tensor_equal(
"page_table_1",
fused_page_table_1[:, : precomputed.max_seqlen_k],
ref_page_table_1[:, : precomputed.max_seqlen_k],
)
elif forward_mode.is_draft_extend():
rows = precomputed.page_indices.shape[0]
cols = precomputed.max_seqlen_k
check_tensor_equal(
"page_table_1",
fused_page_table_1[:rows, :cols],
ref_page_table_1[:rows, :cols],
)
# Compare nsa_cache_seqlens only for the region that was updated
if forward_mode.is_decode_or_idle():
check_tensor_equal(
"nsa_cache_seqlens",
fused_nsa_cache_seqlens,
ref_nsa_cache_seqlens,
)
else: # TARGET_VERIFY or DRAFT_EXTEND
size = precomputed.seqlens_expanded_size
check_tensor_equal(
"nsa_cache_seqlens",
fused_nsa_cache_seqlens[:size],
ref_nsa_cache_seqlens[:size],
)
# Compare nsa_seqlens_expanded only for TARGET_VERIFY and DRAFT_EXTEND
if forward_mode.is_target_verify() or forward_mode.is_draft_extend():
size = precomputed.seqlens_expanded_size
check_tensor_equal(
"nsa_seqlens_expanded",
fused_nsa_seqlens_expanded[:size],
ref_nsa_seqlens_expanded[:size],
)
# Compare nsa_cu_seqlens_k only for the region that was updated
size = precomputed.seqlens_expanded_size
check_tensor_equal(
"nsa_cu_seqlens_k",
fused_nsa_cu_seqlens_k[: 1 + size],
ref_nsa_cu_seqlens_k[: 1 + size],
)
if precomputed.real_page_table is not None:
rows, cols = precomputed.real_page_table.shape
check_tensor_equal(
"real_page_table",
fused_real_page_table[:rows, :cols],
ref_real_page_table[:rows, :cols],
)
if precomputed.flashmla_metadata is not None:
size = precomputed.seqlens_expanded_size
check_tensor_equal(
"flashmla_num_splits",
fused_flashmla_num_splits[: size + 1],
ref_flashmla_num_splits[: size + 1],
)
check_tensor_equal(
"flashmla_metadata",
fused_flashmla_metadata,
ref_flashmla_metadata,
)
def verify_multi_backend_fused_metadata_copy(
metadata0,
metadata1,
metadata2,
precomputed,
bs,
flashmla_num_splits_src=None,
flashmla_metadata_src=None,
):
"""
Verify that the multi-backend fused metadata copy kernel produces the same results
as individual copies for all three backends.
Args:
metadata0: The NSA metadata object for backend 0
metadata1: The NSA metadata object for backend 1
metadata2: The NSA metadata object for backend 2
precomputed: The precomputed metadata containing source tensors
bs: Batch size
flashmla_num_splits_src: Source FlashMLA num_splits tensor (optional)
flashmla_metadata_src: Source FlashMLA metadata tensor (optional)
Raises:
RuntimeError: If verification fails (tensors don't match)
"""
# Clone destination tensors to preserve fused kernel results
fused_results = []
for idx, metadata in enumerate([metadata0, metadata1, metadata2]):
fused_cache_seqlens = metadata.cache_seqlens_int32.clone()
fused_cu_seqlens_k = metadata.cu_seqlens_k.clone()
fused_page_table_1 = metadata.page_table_1.clone()
fused_nsa_cache_seqlens = metadata.nsa_cache_seqlens_int32.clone()
fused_nsa_cu_seqlens_k = metadata.nsa_cu_seqlens_k.clone()
fused_real_page_table = (
metadata.real_page_table.clone()
if precomputed.real_page_table is not None
else None
)
fused_flashmla_num_splits = None
fused_flashmla_metadata = None
if precomputed.flashmla_metadata is not None:
fused_flashmla_num_splits = metadata.flashmla_metadata.num_splits.clone()
fused_flashmla_metadata = (
metadata.flashmla_metadata.flashmla_metadata.clone()
)
fused_results.append(
{
"cache_seqlens": fused_cache_seqlens,
"cu_seqlens_k": fused_cu_seqlens_k,
"page_table_1": fused_page_table_1,
"nsa_cache_seqlens": fused_nsa_cache_seqlens,
"nsa_cu_seqlens_k": fused_nsa_cu_seqlens_k,
"real_page_table": fused_real_page_table,
"flashmla_num_splits": fused_flashmla_num_splits,
"flashmla_metadata": fused_flashmla_metadata,
}
)
# Run individual copy operations for each backend (reference implementation)
ref_results = []
for idx in range(3):
metadata = [metadata0, metadata1, metadata2][idx]
# Create reference tensors (zeroed out)
ref_cache_seqlens = torch.zeros_like(metadata.cache_seqlens_int32)
ref_cu_seqlens_k = torch.zeros_like(metadata.cu_seqlens_k)
ref_page_table_1 = torch.zeros_like(metadata.page_table_1)
ref_nsa_cache_seqlens = torch.zeros_like(metadata.nsa_cache_seqlens_int32)
ref_nsa_cu_seqlens_k = torch.zeros_like(metadata.nsa_cu_seqlens_k)
ref_real_page_table = (
torch.zeros_like(metadata.real_page_table)
if precomputed.real_page_table is not None
else None
)
ref_flashmla_num_splits = None
ref_flashmla_metadata = None
if precomputed.flashmla_metadata is not None:
ref_flashmla_num_splits = torch.zeros_like(
metadata.flashmla_metadata.num_splits
)
ref_flashmla_metadata = torch.zeros_like(
metadata.flashmla_metadata.flashmla_metadata
)
# Copy operations (decode mode)
ref_cache_seqlens.copy_(precomputed.cache_seqlens)
ref_cu_seqlens_k[1:].copy_(precomputed.cu_seqlens_k[1:])
ref_page_table_1[:, : precomputed.max_len].copy_(precomputed.page_indices)
ref_nsa_cache_seqlens.copy_(precomputed.nsa_cache_seqlens)
# Copy NSA cu_seqlens
size = precomputed.seqlens_expanded_size
ref_nsa_cu_seqlens_k[1 : 1 + size].copy_(
precomputed.nsa_cu_seqlens_k[1 : 1 + size]
)
# Copy real page table
if precomputed.real_page_table is not None:
rows, cols = precomputed.real_page_table.shape
ref_real_page_table[:rows, :cols].copy_(precomputed.real_page_table)
# Copy FlashMLA metadata
if precomputed.flashmla_metadata is not None:
ref_flashmla_num_splits[: size + 1].copy_(
flashmla_num_splits_src[: size + 1]
)
ref_flashmla_metadata.copy_(flashmla_metadata_src)
ref_results.append(
{
"cache_seqlens": ref_cache_seqlens,
"cu_seqlens_k": ref_cu_seqlens_k,
"page_table_1": ref_page_table_1,
"nsa_cache_seqlens": ref_nsa_cache_seqlens,
"nsa_cu_seqlens_k": ref_nsa_cu_seqlens_k,
"real_page_table": ref_real_page_table,
"flashmla_num_splits": ref_flashmla_num_splits,
"flashmla_metadata": ref_flashmla_metadata,
}
)
# Compare results for all 3 backends
def check_tensor_equal(backend_idx, name, fused, ref):
if not torch.equal(fused, ref):
max_diff = (fused.float() - ref.float()).abs().max().item()
mismatched_elements = (fused != ref).sum().item()
total_elements = fused.numel()
raise RuntimeError(
f"MULTI-BACKEND FUSED METADATA COPY VERIFICATION FAILED!\n"
f"Backend: {backend_idx}\n"
f"Tensor: {name}\n"
f"Max difference: {max_diff}\n"
f"Mismatched elements: {mismatched_elements}/{total_elements}\n"
f"Fused shape: {fused.shape}, Ref shape: {ref.shape}\n"
f"Batch size: {bs}\n"
f"The multi-backend fused kernel produces different results than individual copies.\n"
f"This indicates a bug in the fused metadata copy kernel."
)
# Verify all tensors for all 3 backends (multi-backend is DECODE mode only)
for idx in range(3):
fused = fused_results[idx]
ref = ref_results[idx]
check_tensor_equal(
idx,
"cache_seqlens",
fused["cache_seqlens"],
ref["cache_seqlens"],
)
check_tensor_equal(
idx,
"cu_seqlens_k",
fused["cu_seqlens_k"],
ref["cu_seqlens_k"],
)
# Multi-backend is DECODE mode only, so compare only [:, :max_len]
check_tensor_equal(
idx,
"page_table_1",
fused["page_table_1"][:, : precomputed.max_len],
ref["page_table_1"][:, : precomputed.max_len],
)
check_tensor_equal(
idx,
"nsa_cache_seqlens",
fused["nsa_cache_seqlens"],
ref["nsa_cache_seqlens"],
)
# DECODE mode uses bs for nsa_cu_seqlens_k size
check_tensor_equal(
idx,
"nsa_cu_seqlens_k",
fused["nsa_cu_seqlens_k"][: bs + 1],
ref["nsa_cu_seqlens_k"][: bs + 1],
)
if precomputed.real_page_table is not None:
rows, cols = precomputed.real_page_table.shape
check_tensor_equal(
idx,
"real_page_table",
fused["real_page_table"][:rows, :cols],
ref["real_page_table"][:rows, :cols],
)
if precomputed.flashmla_metadata is not None:
# DECODE mode uses bs + 1 for flashmla_num_splits
check_tensor_equal(
idx,
"flashmla_num_splits",
fused["flashmla_num_splits"][: bs + 1],
ref["flashmla_num_splits"][: bs + 1],
)
check_tensor_equal(
idx,
"flashmla_metadata",
fused["flashmla_metadata"],
ref["flashmla_metadata"],
) )
from sglang.srt.layers.attention.dsa.dsa_mtp_verification import * # noqa: F401, F403
python/sglang/srt/layers/attention/nsa/quant_k_cache.py
+8 -447
View File
@@ -1,449 +1,10 @@
import torch # [Deprecated] Re-export shim for backward compatibility. Use dsa.quant_k_cache instead.
import triton import warnings
import triton.language as tl
warnings.warn(
def quantize_k_cache(cache_k): "sglang.srt.layers.attention.nsa.quant_k_cache is deprecated; "
return _quantize_k_cache_fast_wrapped(cache_k) "use sglang.srt.layers.attention.dsa.quant_k_cache instead.",
DeprecationWarning,
stacklevel=2,
def quantize_k_cache_separate(
k_nope: torch.Tensor,
k_rope: torch.Tensor,
tile_size: int = 128,
):
"""
Quantize k_nope and k_rope separately without concat, returns two tensors.
This avoids the concat operation and enables direct reuse of set_mla_kv_buffer_triton
by returning two separate byte tensors for the nope and rope parts.
Args:
k_nope: (num_tokens, dim_nope) or (num_tokens, 1, dim_nope)
Must have dim_nope=512 for FP8 MLA quantization
k_rope: (num_tokens, dim_rope) or (num_tokens, 1, dim_rope)
Must have dim_rope=64 for FP8 MLA quantization
tile_size: quantization tile size (default 128)
Returns:
Tuple of (nope_part, rope_part) where:
- nope_part: (num_tokens, 1, 528) as uint8 view, contains [nope_fp8(512) | scales(16)]
- rope_part: (num_tokens, 1, 128) as uint8 view, contains [rope_bf16_bytes(128)]
These two tensors can be directly passed to set_mla_kv_buffer_triton(kv_buffer, loc, nope_part, rope_part)
"""
# Squeeze middle dimension if present
k_nope_2d = k_nope.squeeze(1) if k_nope.ndim == 3 else k_nope
k_rope_2d = k_rope.squeeze(1) if k_rope.ndim == 3 else k_rope
num_tokens = k_nope_2d.shape[0]
dim_nope = k_nope_2d.shape[1]
dim_rope = k_rope_2d.shape[1]
# Validate dimensions for FP8 MLA
if dim_nope != 512:
raise ValueError(f"Expected dim_nope=512 for FP8 MLA, got {dim_nope}")
if dim_rope != 64:
raise ValueError(f"Expected dim_rope=64 for FP8 MLA, got {dim_rope}")
if k_rope_2d.shape[0] != num_tokens:
raise ValueError(
f"k_nope and k_rope must have same num_tokens, got {num_tokens} vs {k_rope_2d.shape[0]}"
) )
from sglang.srt.layers.attention.dsa.quant_k_cache import * # noqa: F401, F403
return _quantize_k_cache_fast_separate(
k_nope=k_nope_2d, k_rope=k_rope_2d, group_size=tile_size
)
# Copied from original
def _quantize_k_cache_ref(
input_k_cache: torch.Tensor, # (num_blocks, block_size, h_k, d)
dv: int = 512,
tile_size: int = 128,
) -> torch.Tensor:
"""
Quantize the k-cache
Return a tensor with shape (num_blocks, block_size, h_k, dv + 4(dv/tile_size) + t(d-dv)) of dtype uint8_t, where t = input_k_cache.element_size()
For more detail about the layout of K/V, please refer to comments in flash_mla_interface.py or README.md
"""
assert dv % tile_size == 0
num_tiles = dv // tile_size
num_blocks, block_size, h_k, d = input_k_cache.shape
assert h_k == 1
input_k_cache = input_k_cache.squeeze(2) # [num_blocks, block_size, d]
input_elem_size = input_k_cache.element_size()
result = torch.empty(
(num_blocks, block_size, dv + num_tiles * 4 + input_elem_size * (d - dv)),
dtype=torch.float8_e4m3fn,
device=input_k_cache.device,
)
result_k_nope_part = result[..., :dv]
result_k_scale_factor = result[..., dv : dv + num_tiles * 4].view(torch.float32)
result_k_rope_part = result[..., dv + num_tiles * 4 :].view(input_k_cache.dtype)
result_k_rope_part[:] = input_k_cache[..., dv:]
for tile_idx in range(0, num_tiles):
cur_scale_factors_inv = (
torch.abs(
input_k_cache[..., tile_idx * tile_size : (tile_idx + 1) * tile_size]
)
.max(dim=-1)
.values
/ 448.0
) # [num_blocks, block_size]
result_k_scale_factor[:, :, tile_idx] = cur_scale_factors_inv
cur_scale_factors_inv.unsqueeze_(-1) # [num_blocks, block_size, 1]
cur_quantized_nope = (
input_k_cache[
..., tile_idx * tile_size : (tile_idx + 1) * tile_size
].float()
/ cur_scale_factors_inv.float()
).to(torch.float8_e4m3fn)
result_k_nope_part[..., tile_idx * tile_size : (tile_idx + 1) * tile_size] = (
cur_quantized_nope
)
result = result.view(num_blocks, block_size, 1, -1)
return result
def _quantize_k_cache_fast_wrapped(
input_k_cache: torch.Tensor,
dv: int = 512,
tile_size: int = 128,
) -> torch.Tensor:
# TODO the final API may be 2D instead of 4D, thus we convert them here
num_blocks, block_size, _, dim_nope_and_rope = input_k_cache.shape
assert dv == 512
assert dim_nope_and_rope == 512 + 64
assert tile_size == 128
input_k_cache = input_k_cache.view((-1, dim_nope_and_rope))
# TODO deliberately split into two tensors, then upstream can provide the two tensors instead of concat into one
k_nope = input_k_cache[:, :dv]
k_rope = input_k_cache[:, dv:]
output = _quantize_k_cache_fast(k_nope=k_nope, k_rope=k_rope)
return output.view(num_blocks, block_size, 1, -1)
def _quantize_k_cache_fast(k_nope, k_rope, group_size: int = 128):
"""
:param k_nope: (num_tokens, dim_nope 512)
:param k_rope: (num_tokens, dim_rope 64)
"""
assert k_nope.dtype == torch.bfloat16
assert k_rope.dtype == torch.bfloat16
num_tokens, dim_nope = k_nope.shape
num_tokens_, dim_rope = k_rope.shape
assert num_tokens == num_tokens_
assert dim_nope == 512
assert dim_rope == 64
assert k_nope.dtype == k_rope.dtype
num_tiles = dim_nope // group_size
assert k_nope.stride(1) == 1
assert k_rope.stride(1) == 1
output = torch.empty(
(num_tokens, dim_nope + num_tiles * 4 + k_rope.element_size() * dim_rope),
dtype=torch.float8_e4m3fn,
device=k_nope.device,
)
output_nope_q = output[..., :dim_nope]
output_nope_s = output[..., dim_nope : dim_nope + num_tiles * 4].view(torch.float32)
output_rope = output[..., dim_nope + num_tiles * 4 :].view(torch.bfloat16)
num_blocks_per_token = triton.cdiv(dim_nope + dim_rope, group_size)
assert num_blocks_per_token == 5
assert dim_nope % group_size == 0
NUM_NOPE_BLOCKS = dim_nope // group_size
_quantize_k_cache_fast_kernel[(num_tokens, num_blocks_per_token)](
output_nope_q,
output_nope_s,
output_rope,
k_nope,
k_rope,
output_nope_q.stride(0),
output_nope_s.stride(0),
output_rope.stride(0),
k_nope.stride(0),
k_rope.stride(0),
NUM_NOPE_BLOCKS=NUM_NOPE_BLOCKS,
GROUP_SIZE=group_size,
DIM_NOPE=dim_nope,
DIM_ROPE=dim_rope,
FP8_MIN=torch.finfo(torch.float8_e4m3fn).min,
FP8_MAX=torch.finfo(torch.float8_e4m3fn).max,
)
return output
def _quantize_k_cache_fast_separate(k_nope, k_rope, group_size: int = 128):
"""
Quantize k_nope and k_rope in a single Triton kernel, directly outputting two separate tensors.
This avoids packing/unpacking and enables direct use with set_mla_kv_buffer_triton.
:param k_nope: (num_tokens, dim_nope 512) bfloat16
:param k_rope: (num_tokens, dim_rope 64) bfloat16
:param group_size: quantization tile size (default 128, kernel is tuned for this value)
:return: Tuple of (nope_part_u8, rope_part_u8)
- nope_part_u8: (num_tokens, 1, nope_part_bytes) uint8, layout [nope_fp8(dim_nope) | scales(num_tiles*4)]
- rope_part_u8: (num_tokens, 1, rope_part_bytes) uint8, layout [rope_bf16_bytes(dim_rope*2)]
"""
num_tokens, dim_nope = k_nope.shape
num_tokens_, dim_rope = k_rope.shape
assert num_tokens == num_tokens_, f"k_nope and k_rope must have same num_tokens"
# Ensure contiguous tensors for kernel
k_nope = k_nope.contiguous()
k_rope = k_rope.contiguous()
num_tiles = dim_nope // group_size
# Calculate byte sizes based on validated dimensions
# nope_part: [FP8 quantized data (dim_nope bytes)] + [FP32 scales (num_tiles * 4 bytes)]
# rope_part: [BF16 raw data (dim_rope * 2 bytes)]
nope_part_bytes = (
dim_nope + num_tiles * 4
) # e.g., 512 + 4*4 = 528 for dim_nope=512, group_size=128
rope_part_bytes = (
dim_rope * k_rope.element_size()
) # e.g., 64 * 2 = 128 for dim_rope=64, BF16
# Allocate two separate output buffers (as uint8 for direct byte-level access)
nope_part_u8 = torch.empty(
(num_tokens, nope_part_bytes), dtype=torch.uint8, device=k_nope.device
)
rope_part_u8 = torch.empty(
(num_tokens, rope_part_bytes), dtype=torch.uint8, device=k_rope.device
)
# Create typed views for the kernel to write into
# Fixed byte layout for nope_part: [nope_fp8 (dim_nope bytes) | scales_fp32 (num_tiles*4 bytes)]
# Fixed byte layout for rope_part: [rope_bf16 (dim_rope*2 bytes)]
nope_q_view = nope_part_u8[:, :dim_nope].view(torch.float8_e4m3fn)
nope_s_view = nope_part_u8[:, dim_nope:].view(torch.float32)
rope_view = rope_part_u8.view(torch.bfloat16)
# Kernel launch parameters
num_blocks_per_token = triton.cdiv(dim_nope + dim_rope, group_size)
NUM_NOPE_BLOCKS = dim_nope // group_size
# Use the same kernel as _quantize_k_cache_fast (reuse existing implementation)
_quantize_k_cache_fast_kernel[(num_tokens, num_blocks_per_token)](
nope_q_view,
nope_s_view,
rope_view,
k_nope,
k_rope,
nope_q_view.stride(0),
nope_s_view.stride(0),
rope_view.stride(0),
k_nope.stride(0),
k_rope.stride(0),
NUM_NOPE_BLOCKS=NUM_NOPE_BLOCKS,
GROUP_SIZE=group_size,
DIM_NOPE=dim_nope,
DIM_ROPE=dim_rope,
FP8_MIN=torch.finfo(torch.float8_e4m3fn).min,
FP8_MAX=torch.finfo(torch.float8_e4m3fn).max,
)
# Add middle dimension for compatibility with set_mla_kv_buffer_triton
return nope_part_u8.unsqueeze(1), rope_part_u8.unsqueeze(1)
@triton.jit
def _quantize_k_cache_fast_kernel(
output_nope_q_ptr,
output_nope_s_ptr,
output_rope_ptr,
k_nope_ptr,
k_rope_ptr,
output_nope_q_stride_0: int,
output_nope_s_stride_0: int,
output_rope_stride_0: int,
k_nope_stride_0: int,
k_rope_stride_0: int,
NUM_NOPE_BLOCKS: tl.constexpr,
GROUP_SIZE: tl.constexpr,
DIM_NOPE: tl.constexpr,
DIM_ROPE: tl.constexpr,
FP8_MIN: tl.constexpr,
FP8_MAX: tl.constexpr,
):
token_id = tl.program_id(0)
raw_block_id = tl.program_id(1)
if raw_block_id < NUM_NOPE_BLOCKS:
# a. quant nope
effective_block_id = raw_block_id
offs = effective_block_id * GROUP_SIZE + tl.arange(0, GROUP_SIZE)
mask = offs < DIM_NOPE
ptr = k_nope_ptr + token_id * k_nope_stride_0 + offs
y = tl.load(ptr, mask=mask, other=0.0).to(tl.float32)
# the ref impl do not have a `tl.maximum(... eps)`, so we remove it here
y_s = tl.max(tl.abs(y)) / FP8_MAX
y_s_inv = 1.0 / y_s
y_q = tl.clamp(y * y_s_inv, FP8_MIN, FP8_MAX).to(
output_nope_q_ptr.dtype.element_ty
)
dst_q_ptr = output_nope_q_ptr + token_id * output_nope_q_stride_0 + offs
dst_s_ptr = (
output_nope_s_ptr + token_id * output_nope_s_stride_0 + effective_block_id
)
tl.store(dst_q_ptr, y_q, mask=mask)
tl.store(dst_s_ptr, y_s)
else:
# b. copy rope
effective_block_id = raw_block_id - NUM_NOPE_BLOCKS
offs = effective_block_id * GROUP_SIZE + tl.arange(0, GROUP_SIZE)
mask = offs < DIM_ROPE
src_ptr = k_rope_ptr + token_id * k_rope_stride_0 + offs
dst_ptr = output_rope_ptr + token_id * output_rope_stride_0 + offs
data = tl.load(src_ptr, mask=mask)
tl.store(dst_ptr, data, mask=mask)
if __name__ == "__main__":
import dequant_k_cache
for num_blocks, block_size in [
(1, 1),
(10, 64),
]:
dim_nope_and_rope = 512 + 64
input_k_cache = torch.randn(
(num_blocks, block_size, 1, dim_nope_and_rope),
dtype=torch.bfloat16,
device="cuda",
)
ref_quant = _quantize_k_cache_ref(input_k_cache)
actual_quant = _quantize_k_cache_fast_wrapped(input_k_cache)
ref_ref_dequant = dequant_k_cache._dequantize_k_cache_slow(ref_quant)
ref_actual_dequant = dequant_k_cache._dequantize_k_cache_fast_wrapped(ref_quant)
actual_actual_dequant = dequant_k_cache._dequantize_k_cache_fast_wrapped(
actual_quant
)
print(f"{ref_ref_dequant=}")
print(f"{actual_actual_dequant=}")
print(f"{actual_actual_dequant - ref_ref_dequant=}")
print(f"{torch.mean(ref_ref_dequant - actual_actual_dequant)=}")
# TODO too different?
torch.testing.assert_close(
ref_ref_dequant, ref_actual_dequant, atol=0.2, rtol=0.2
)
torch.testing.assert_close(
ref_ref_dequant, actual_actual_dequant, atol=0.2, rtol=0.2
)
# test dequant_k_cache_paged
page_table_1 = torch.arange(
num_blocks * block_size, dtype=torch.int32, device="cuda"
)
actual_dequant_paged = dequant_k_cache.dequantize_k_cache_paged(
actual_quant, page_table_1
).reshape(actual_actual_dequant.shape)
print(f"{torch.mean(actual_actual_dequant - actual_dequant_paged)=}")
torch.testing.assert_close(
ref_ref_dequant, actual_dequant_paged, atol=0.2, rtol=0.2
)
print("Passed")
# Test quantize_k_cache_separate: verify output matches concat path
print("\nTesting quantize_k_cache_separate...")
for num_tokens in [64, 100]:
dim_nope = 512
dim_rope = 64
k_nope = torch.randn(
num_tokens, 1, dim_nope, dtype=torch.bfloat16, device="cuda"
)
k_rope = torch.randn(
num_tokens, 1, dim_rope, dtype=torch.bfloat16, device="cuda"
)
# Old path: concat then quantize
k_concat = torch.cat([k_nope, k_rope], dim=-1).squeeze(1) # (num_tokens, 576)
old_output = quantize_k_cache(k_concat.unsqueeze(1).unsqueeze(1)) # 4D input
old_output = old_output.squeeze(1).squeeze(1) # Back to (num_tokens, 656)
# New path: quantize separately
nope_part, rope_part = quantize_k_cache_separate(k_nope, k_rope)
new_bytes = torch.cat([nope_part.squeeze(1), rope_part.squeeze(1)], dim=-1)
# Compare byte-level equality
old_bytes = old_output.view(torch.uint8)
if old_bytes.shape != new_bytes.shape:
raise RuntimeError(
f"Shape mismatch: {old_bytes.shape} vs {new_bytes.shape}"
)
diff_bytes = (old_bytes != new_bytes).sum().item()
if diff_bytes > 0:
max_diff = (old_bytes.float() - new_bytes.float()).abs().max().item()
raise RuntimeError(
f"quantize_k_cache_separate output doesn't match concat path: "
f"{diff_bytes} differing bytes, max_diff={max_diff}"
)
print(f" num_tokens={num_tokens}: PASSED (outputs match byte-wise)")
print("quantize_k_cache_separate tests passed!")
print("\nDo benchmark...")
for num_blocks, block_size in [
(1, 64),
(64, 64),
(128, 64),
(256, 64),
(512, 64),
(1024, 64),
(2048, 64),
]:
dim_nope_and_rope = 512 + 64
input_k_cache = torch.randn(
(num_blocks, block_size, 1, dim_nope_and_rope),
dtype=torch.bfloat16,
device="cuda",
)
actual_quant = _quantize_k_cache_fast_wrapped(input_k_cache)
page_table_1 = torch.arange(
num_blocks * block_size, dtype=torch.int32, device="cuda"
)
def run_ans():
return dequant_k_cache.dequantize_k_cache_paged(actual_quant, page_table_1)
ans_time: float = triton.testing.do_bench(run_ans, warmup=10, rep=20) / 1000 # type: ignore
print(f"seq_kv: {num_blocks * block_size}, time: {ans_time * 1e6: 4.0f} us")
python/sglang/srt/layers/attention/nsa/tilelang_kernel.py
+8 -2587
View File
File diff suppressed because it is too large Load Diff
python/sglang/srt/layers/attention/nsa/transform_index.py
+8 -142
View File
@@ -1,144 +1,10 @@
from typing import List, Optional # [Deprecated] Re-export shim for backward compatibility. Use dsa.transform_index instead.
import warnings
import torch warnings.warn(
import triton "sglang.srt.layers.attention.nsa.transform_index is deprecated; "
import triton.language as tl "use sglang.srt.layers.attention.dsa.transform_index instead.",
DeprecationWarning,
stacklevel=2,
def transform_index_page_table_prefill(**kwargs):
return transform_index_page_table_prefill_ref(**kwargs)
def transform_index_page_table_decode(**kwargs):
return transform_index_page_table_decode_ref(**kwargs)
@triton.jit
def transform_index_page_table_decode_kernel(
page_table_ptr: torch.Tensor,
topk_indices_ptr: torch.Tensor,
result_ptr: torch.Tensor,
page_size: tl.constexpr,
max_seqlen_k: tl.constexpr,
):
TOPK: tl.constexpr = 2048
req_id = tl.program_id(0)
page_table_ptr = page_table_ptr + req_id * max_seqlen_k
topk_indices_ptr = topk_indices_ptr + req_id * TOPK
result_ptr = result_ptr + req_id * TOPK
offset = tl.arange(0, TOPK) # topk should be 2048
loaded_topk_indices = tl.load(topk_indices_ptr + offset)
mask = loaded_topk_indices >= 0
loaded_kv_indices = tl.load(page_table_ptr + loaded_topk_indices, mask=mask)
tl.store(result_ptr + offset, loaded_kv_indices, mask=mask)
tl.store(result_ptr + offset, -1, mask=~mask)
def transform_index_page_table_decode_fast(
page_table: torch.Tensor,
topk_indices: torch.Tensor,
result: Optional[torch.Tensor] = None,
page_size: int = 1,
) -> torch.Tensor:
"""
Transform the page table according to topk indices for sparse topk attention.
Args:
page_table: [qo_len, max_seqlen_k], the original page table
topk_indices: [qo_len, topk], the topk indices for each query position
Returns:
transformed_page_table: [qo_len, topk], the transformed page table
For out-of-bound indices in topk_indices, this should be filled with -1.
"""
assert page_size == 1
assert page_table.shape[0] == topk_indices.shape[0]
assert topk_indices.shape[1] == 2048
qo_len = topk_indices.shape[0]
max_seqlen_k = page_table.shape[1]
if result is None:
result = torch.empty_like(topk_indices, dtype=torch.int32)
# Launch triton kernel
grid = (qo_len,)
transform_index_page_table_decode_kernel[grid](
page_table,
topk_indices,
result,
page_size,
max_seqlen_k=max_seqlen_k,
) )
return result from sglang.srt.layers.attention.dsa.transform_index import * # noqa: F401, F403
def transform_index_page_table_prefill_fast(
page_table: torch.Tensor,
topk_indices: torch.Tensor,
extend_lens_cpu: List[int],
page_size: int = 1,
) -> torch.Tensor:
# TODO(baizhou): can be implemented with another triton kernel
assert page_size == 1
result = torch.empty_like(topk_indices, dtype=torch.int32)
assert len(extend_lens_cpu) == page_table.shape[0]
offset = 0
for i, l in enumerate(extend_lens_cpu):
transform_index_page_table_decode_fast(
page_table[i].unsqueeze(0).expand(l, -1),
topk_indices[offset : offset + l],
result=result[offset : offset + l],
)
offset += l
assert offset == topk_indices.shape[0]
return result
def transform_index_page_table_decode_ref(
page_table: torch.Tensor,
topk_indices: torch.Tensor,
result: Optional[torch.Tensor] = None,
page_size: int = 1,
) -> torch.Tensor:
assert page_size == 1
assert page_table.shape[0] == topk_indices.shape[0]
if result is None:
result = torch.empty_like(topk_indices, dtype=torch.int32)
assert result.shape == topk_indices.shape
torch.gather(
page_table.to(result.dtype),
dim=1,
index=topk_indices.clamp(min=0),
out=result,
)
result[topk_indices < 0] = -1
return result
def transform_index_page_table_prefill_ref(
page_table: torch.Tensor,
topk_indices: torch.Tensor,
extend_lens_cpu: List[int],
page_size: int = 1,
) -> torch.Tensor:
assert page_size == 1
result = torch.empty_like(topk_indices, dtype=torch.int32)
assert len(extend_lens_cpu) == page_table.shape[0]
offset = 0
for i, l in enumerate(extend_lens_cpu):
transform_index_page_table_decode_ref(
page_table[i].unsqueeze(0).expand(l, -1),
topk_indices[offset : offset + l],
result=result[offset : offset + l],
)
offset += l
assert offset == topk_indices.shape[0]
return result
if __name__ == "__main__":
bs, topk, max_seqlen = 10, 2048, 3000
page_table = torch.randint(0, 100, (bs, max_seqlen), device="cuda")
topk_indices = torch.full((bs, topk), -1, device="cuda")
topk_indices[:, :1600] = torch.arange(1600).unsqueeze(0).repeat(bs, 1)
ref_result = transform_index_page_table_decode_ref(page_table, topk_indices)
result = transform_index_page_table_decode_fast(page_table, topk_indices)
assert torch.all(result == ref_result)
print("Passed")
python/sglang/srt/layers/attention/nsa/triton_kernel.py
+8 -194
View File
@@ -1,196 +1,10 @@
from typing import Optional, Tuple # [Deprecated] Re-export shim for backward compatibility. Use dsa.triton_kernel instead.
import warnings
import torch warnings.warn(
import triton "sglang.srt.layers.attention.nsa.triton_kernel is deprecated; "
import triton.language as tl "use sglang.srt.layers.attention.dsa.triton_kernel instead.",
DeprecationWarning,
stacklevel=2,
# Triton implementation
@triton.jit
def _act_quant_kernel(
X_ptr,
Y_ptr,
S_ptr,
M,
N,
group_size: tl.constexpr,
round_scale: tl.constexpr,
BLOCK_M: tl.constexpr,
BLOCK_N: tl.constexpr,
):
"""
Triton kernel for activation quantization.
Each block processes BLOCK_M rows and group_size columns.
"""
# Get block IDs
pid_m = tl.program_id(0)
pid_n = tl.program_id(1)
# FP8 constants
fp8_min = -448.0
fp8_max = 448.0
fp8_max_inv = 1.0 / fp8_max
# Calculate row and column offsets
row_start = pid_m * BLOCK_M
col_start = pid_n * group_size
# Create offset arrays
rows = row_start + tl.arange(0, BLOCK_M)
cols = col_start + tl.arange(0, BLOCK_N)
# Mask for valid rows and columns
row_mask = rows < M
col_mask = cols < N
mask = row_mask[:, None] & col_mask[None, :]
# Load input data
x_ptrs = X_ptr + rows[:, None] * N + cols[None, :]
x = tl.load(x_ptrs, mask=mask, other=0.0).to(tl.float32)
# Compute absolute max along columns (group_size dimension) for each row
x_abs = tl.abs(x)
amax = tl.max(x_abs, axis=1) # Shape: (BLOCK_M,)
# Clamp amax to avoid division by zero
amax = tl.maximum(amax, 1e-4)
# Compute scale
if round_scale:
# Fast round scale using bit manipulation approximation
# This is a simplified version - the exact bit manipulation is harder in Triton
# Using log2 + ceil + pow2 as approximation
log_val = tl.log2(amax * fp8_max_inv)
log_ceil = tl.ceil(log_val)
scale = tl.exp2(log_ceil)
else:
scale = amax * fp8_max_inv
# Quantize: y = clamp(x / scale, fp8_min, fp8_max)
scale_broadcast = scale[:, None]
y = x / scale_broadcast
y = tl.minimum(tl.maximum(y, fp8_min), fp8_max)
# Store quantized output
y_ptrs = Y_ptr + rows[:, None] * N + cols[None, :]
tl.store(y_ptrs, y, mask=mask)
# Store scales
s_cols = pid_n
s_ptrs = S_ptr + rows * (N // group_size) + s_cols
s_mask = row_mask
tl.store(s_ptrs, scale, mask=s_mask)
def act_quant(
x: torch.Tensor, block_size: int = 128, scale_fmt: Optional[str] = None
) -> Tuple[torch.Tensor, torch.Tensor]:
"""
Quantizes the input tensor `x` using block-wise quantization with Triton.
Args:
x (torch.Tensor): The input tensor to be quantized. Must be contiguous and its last dimension size must be divisible by `block_size`.
block_size (int, optional): The size of the blocks to be used for quantization. Default is 128.
scale_fmt (Optional[str], optional): The format of the scale. Default is None.
Returns:
Tuple[torch.Tensor, torch.Tensor]: A tuple containing:
- The quantized tensor with dtype `torch.float8_e4m3fn`.
- A tensor of scaling factors with dtype `torch.float32`.
"""
assert x.is_contiguous(), "Input tensor must be contiguous"
assert (
x.size(-1) % block_size == 0
), f"Last dimension size must be divisible by block_size (block_size={block_size})"
# Flatten all dims except last
N = x.size(-1)
x_flat = x.view(-1, N)
M = x_flat.size(0)
# Allocate output tensors
y = torch.empty_like(x, dtype=torch.float8_e4m3fn)
y_flat = y.view(-1, N)
s = x.new_empty(*x.size()[:-1], N // block_size, dtype=torch.float32)
s_flat = s.view(-1, N // block_size)
# Launch kernel
BLOCK_M = 32
BLOCK_N = block_size
grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(N, block_size))
round_scale = scale_fmt is not None
_act_quant_kernel[grid](
x_flat,
y_flat,
s_flat,
M,
N,
group_size=block_size,
round_scale=round_scale,
BLOCK_M=BLOCK_M,
BLOCK_N=BLOCK_N,
num_stages=0 if round_scale else 2,
)
return y, s
@triton.jit
def _get_valid_kv_indices_kernel(
page_table_ptr, # [bs, topk]
kv_indptr_ptr, # [bs + 1]
kv_indices_ptr, # [bs * topk] output buffer
bs: tl.constexpr,
topk: tl.constexpr,
):
"""
Extract valid indices (non -1) from page_table into kv_indices.
Each program handles one batch.
"""
batch_id = tl.program_id(0)
# Get the start position for this batch in kv_indices
dst_start = tl.load(kv_indptr_ptr + batch_id)
# Load all topk indices for this batch
src_offset = batch_id * topk
offsets = tl.arange(0, topk)
indices = tl.load(page_table_ptr + src_offset + offsets)
# Count valid indices and compact them
mask = indices != -1
# Use prefix sum to compute destination positions for valid elements
# For each position, count how many valid elements are before it
prefix_sum = tl.cumsum(mask.to(tl.int32), axis=0) - 1
# Store valid indices to their compacted positions
dst_positions = dst_start + prefix_sum
tl.store(kv_indices_ptr + dst_positions, indices, mask=mask)
def get_valid_kv_indices(
page_table_1: torch.Tensor,
kv_indptr: torch.Tensor,
kv_indices: torch.Tensor,
bs: int,
):
"""
Extract valid indices from page_table_1 into kv_indices buffer.
Args:
page_table_1: [bs, topk] page table with -1 as invalid
kv_indptr: [bs + 1] cumulative count of valid indices per batch
kv_indices: [bs * topk] pre-allocated output buffer
bs: batch size
"""
topk = page_table_1.shape[1]
grid = (bs,)
_get_valid_kv_indices_kernel[grid](
page_table_1,
kv_indptr,
kv_indices,
bs,
topk,
) )
from sglang.srt.layers.attention.dsa.triton_kernel import * # noqa: F401, F403
python/sglang/srt/layers/attention/nsa/utils.py
+8 -267
View File
@@ -1,269 +1,10 @@
from functools import lru_cache # [Deprecated] Re-export shim for backward compatibility. Use dsa.utils instead.
from typing import TYPE_CHECKING, List, Tuple, Union import warnings
import torch warnings.warn(
import triton "sglang.srt.layers.attention.nsa.utils is deprecated; "
import triton.language as tl "use sglang.srt.layers.attention.dsa.utils instead.",
DeprecationWarning,
from sglang.srt.layers.dp_attention import ( stacklevel=2,
DpPaddingMode,
get_attention_cp_rank,
get_attention_cp_size,
get_attention_dp_rank,
) )
from sglang.srt.server_args import get_global_server_args from sglang.srt.layers.attention.dsa.utils import * # noqa: F401, F403
from sglang.srt.utils import get_bool_env_var, is_hip
from sglang.srt.utils.common import ceil_align, ceil_div
@lru_cache(maxsize=1)
def aiter_can_use_preshuffle_paged_mqa() -> bool:
"""Whether aiter's preshuffle paged MQA / cache kernels can be used on this runtime.
aiter's ``deepgemm_fp8_paged_mqa_logits`` only supports ``KVBlockSize > 1`` and
``Preshuffle=True`` on its gluon kernel path. The gluon path is enabled when
Triton >= 3.5.0, OR when ``AITER_ENABLE_AOT_GLUON_PA_MQA_LOGITS=1`` is set
(which additionally requires that the AOT gluon kernel artifacts ship inside
the aiter wheel/image). Otherwise aiter asserts ``KVBlockSize == 1`` and
refuses ``Preshuffle=True``.
sglang's NSA indexer uses this single decision to pick:
* ``page_size``: 64 (preshuffle) vs 1 (legacy) on ROCm
* ``Preshuffle`` / ``preshuffle`` flags on the aiter MQA + cache kernels
* ``get_page_table_64`` vs ``get_page_table_1`` on the metadata
* whether ``GetKAndS.execute`` uses the aiter or the triton implementation
The result is cached so the cost is paid once per process.
Set ``SGLANG_NSA_HIP_DISABLE_PRESHUFFLE=1`` to force the legacy path even when
the gluon kernel would otherwise be available (useful for CI bisection).
"""
if not is_hip():
return False
if not get_bool_env_var("SGLANG_USE_AITER"):
return False
if get_bool_env_var("SGLANG_NSA_HIP_DISABLE_PRESHUFFLE"):
return False
if get_bool_env_var("AITER_ENABLE_AOT_GLUON_PA_MQA_LOGITS"):
return True
try:
from packaging.version import Version
return Version(Version(triton.__version__).base_version) >= Version("3.5.0")
except Exception:
return False
if TYPE_CHECKING:
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
def compute_nsa_seqlens(original_seq_lens, nsa_index_topk: int):
return original_seq_lens.clamp(max=nsa_index_topk)
def is_nsa_enable_prefill_cp():
return get_global_server_args().enable_nsa_prefill_context_parallel
def is_nsa_prefill_cp_in_seq_split():
return (
is_nsa_enable_prefill_cp()
and get_global_server_args().nsa_prefill_cp_mode == "in-seq-split"
)
def is_nsa_prefill_cp_round_robin_split():
return (
is_nsa_enable_prefill_cp()
and get_global_server_args().nsa_prefill_cp_mode == "round-robin-split"
)
def can_nsa_prefill_cp_round_robin_split(forward_batch: "ForwardBatch"):
if not forward_batch.forward_mode.is_context_parallel_extend():
return False
cp_size = get_attention_cp_size()
seq_len = sum(forward_batch.extend_seq_lens_cpu)
return (
is_nsa_prefill_cp_round_robin_split()
and seq_len > 0
and seq_len >= cp_size
and cp_size > 1
)
def nsa_cp_round_robin_split_data(input_: Union[torch.Tensor, List]):
"""
# for round-robin-split, split the tokens evenly according to the rule of token_idx % cp_size.
| +-----------before split------------+|
| token0, token1, token2, token3, token4, token5, token6, token7, ...
|
| +--------------result-------------------+
| dp_atten_tp0: token0, token4, token8, token12, token16, ... |
| dp_atten_tp1: token1, token5, token9, token13, token17, ... |
| dp_atten_tp2: token2, token6, token10, token14, token18, ... |
| dp_atten_tp3: token3, token7, token11, token15, token19, ... |
| +-------------------------+
"""
cp_size = get_attention_cp_size()
cp_rank = get_attention_cp_rank()
if isinstance(input_, (tuple, list)):
indices = range(cp_rank, len(input_), cp_size)
return input_[indices]
tokens = len(input_)
if tokens % cp_size != 0:
cur_len = tokens // cp_size + (tokens % cp_size > cp_rank)
if cur_len == 0:
return input_.new_empty(0, *input_.shape[1:])
indices = torch.arange(cp_rank, tokens, cp_size, device=input_.device)
return input_[indices]
# for torch device tensor
return input_.view(-1, cp_size, *input_.shape[1:])[:, cp_rank].contiguous()
def cal_padded_tokens(forward_batch: "ForwardBatch"):
# Consistent with the padding calculation logic in ForwardBatch.prepare_mlp_sync_batch,
# calculate the actual token length after padding when attn_tp_size > 1 or in the MAX_LEN padding mode.
global_num_tokens = forward_batch.global_num_tokens_cpu.copy()
sync_group_size = len(global_num_tokens)
attn_cp_size = get_attention_cp_size()
for i in range(sync_group_size):
global_num_tokens[i] = ceil_align(global_num_tokens[i], attn_cp_size)
dp_padding_mode = DpPaddingMode.get_dp_padding_mode(
forward_batch.is_extend_in_batch, global_num_tokens
)
if dp_padding_mode.is_max_len():
tokens = max(global_num_tokens)
elif len(global_num_tokens) > 1:
tokens = global_num_tokens[get_attention_dp_rank()]
else:
tokens = global_num_tokens[0]
if can_nsa_prefill_cp_round_robin_split(forward_batch):
tokens = ceil_div(tokens, attn_cp_size)
return tokens
def pad_nsa_cache_seqlens(forward_batch: "ForwardBatch", nsa_cache_seqlens):
attn_cp_size = get_attention_cp_size()
needs_cp_pad = attn_cp_size > 1 and can_nsa_prefill_cp_round_robin_split(
forward_batch
)
needs_dp_pad = forward_batch.global_num_tokens_cpu is not None
if not needs_cp_pad and not needs_dp_pad:
return nsa_cache_seqlens
tokens = cal_padded_tokens(forward_batch)
pad_len = tokens - nsa_cache_seqlens.shape[0]
if pad_len > 0:
nsa_cache_seqlens = torch.cat(
[
nsa_cache_seqlens,
nsa_cache_seqlens.new_zeros(pad_len, *nsa_cache_seqlens.shape[1:]),
]
)
return nsa_cache_seqlens
def can_nsa_cp_split(seq_len: int, cp_size: int, use_nsa: bool, forward_batch):
if is_nsa_prefill_cp_round_robin_split():
cur_cp_seq_len = seq_len // cp_size
assert (
seq_len % cp_size == 0
), f"seq_len {seq_len} is not divisible by cp_size {cp_size} when nsa_prefill_cp_mode is round-robin-split"
else:
# TODO current just support prefill batch=1 and len(input_ids) > self.cp_size * 2
# Note: (self.cp_size * 2) To achieve load balancing for seq computation,
# the seq data needs to be divided and recombined at twice the size of cp_size.
cur_cp_seq_len = seq_len // (cp_size * 2)
if (
cur_cp_seq_len != 0
and cp_size > 1
and use_nsa
and forward_batch.forward_mode.is_context_parallel_extend()
and is_nsa_enable_prefill_cp()
and sum(forward_batch.extend_seq_lens_cpu) >= cp_size
):
return True
else:
return False
@triton.jit
def nsa_cp_round_robin_split_q_seqs_kernel(
in_seqs_ptr,
out_seqs_ptr,
bs_idx_ptr,
tokens: tl.constexpr,
cp_size: tl.constexpr,
cp_rank: tl.constexpr,
):
extra_seq = 0
bs_idx = 0
for bs in range(tokens):
cur_len = tl.load(in_seqs_ptr + bs)
cur_len += extra_seq
cur_seq = cur_len // cp_size + (cur_len % cp_size > cp_rank)
if cur_seq > 0:
tl.store(bs_idx_ptr + bs_idx, bs)
tl.store(out_seqs_ptr + bs_idx, cur_seq)
bs_idx += 1
extra_seq = cur_len - cur_seq * cp_size
def nsa_cp_round_robin_split_q_seqs_cpu(extend_seqs):
cp_size = get_attention_cp_size()
cp_rank = get_attention_cp_rank()
extra_seq = 0
q_seqs = []
for bs, cur_len in enumerate(extend_seqs):
cur_len += extra_seq
cur_seq = cur_len // cp_size + int(cur_len % cp_size > cp_rank)
q_seqs.append(cur_seq)
extra_seq = cur_len - cur_seq * cp_size
bs_idx = list([i for i, x in enumerate(q_seqs) if x > 0])
q_seqs = [q_len for q_len in q_seqs if q_len > 0]
return q_seqs, bs_idx
def nsa_cp_round_robin_split_q_seqs(
extend_seqs_cpu, extend_seqs
) -> Tuple[List, torch.Tensor, List, torch.Tensor]:
"""
round-robin-split distributes tokens across ranks based on token_idx % cp_size.
Return:
ret_q_lens_cpu(List) and ret_q_lens(torch.Tensor): the partitioned length (excluding zeros) on the current cp rank
for each sequence after distribution across cp ranks.
bs_idx_cpu(List) and bs_idx(torch.Tensor): marks which sequences are ultimately selected,
i.e., those with a partitioned length greater than zero.
"""
cp_size = get_attention_cp_size()
cp_rank = get_attention_cp_rank()
# len(ret_q_lens_cpu) == len(bs_idx_cpu)
ret_q_lens_cpu, bs_idx_cpu = nsa_cp_round_robin_split_q_seqs_cpu(extend_seqs_cpu)
ret_q_lens = torch.empty(
(len(bs_idx_cpu),), device=extend_seqs.device, dtype=extend_seqs.dtype
)
bs_idx = torch.empty(
(len(bs_idx_cpu),), device=extend_seqs.device, dtype=torch.int32
)
grid = (1,)
nsa_cp_round_robin_split_q_seqs_kernel[grid](
extend_seqs, ret_q_lens, bs_idx, len(extend_seqs), cp_size, cp_rank
)
return ret_q_lens_cpu, ret_q_lens, bs_idx_cpu, bs_idx
def nsa_use_prefill_cp(forward_batch, nsa_enable_prefill_cp=None):
if nsa_enable_prefill_cp is None:
nsa_enable_prefill_cp = is_nsa_enable_prefill_cp()
if (
forward_batch.attn_cp_metadata is not None
and nsa_enable_prefill_cp
and forward_batch.forward_mode.is_context_parallel_extend()
):
return True
else:
return False
python/sglang/srt/layers/attention/nsa_backend.py
+22 -2519
View File
File diff suppressed because it is too large Load Diff
python/sglang/srt/layers/communicator.py
+20 -20
View File
@@ -33,9 +33,9 @@ from sglang.srt.distributed.device_communicators.pynccl_allocator import (
use_symmetric_memory, use_symmetric_memory,
) )
from sglang.srt.environ import envs from sglang.srt.environ import envs
from sglang.srt.layers.attention.nsa.utils import ( from sglang.srt.layers.attention.dsa.utils import (
is_nsa_enable_prefill_cp, dsa_use_prefill_cp,
nsa_use_prefill_cp, is_dsa_enable_prefill_cp,
) )
from sglang.srt.layers.dp_attention import ( from sglang.srt.layers.dp_attention import (
attn_tp_all_gather_into_tensor, attn_tp_all_gather_into_tensor,
@@ -202,7 +202,7 @@ class ScatterMode(Enum):
@staticmethod @staticmethod
def model_input_output(): def model_input_output():
"""The scatter mode for model forward pass input and output data""" """The scatter mode for model forward pass input and output data"""
if is_nsa_enable_prefill_cp(): if is_dsa_enable_prefill_cp():
return ScatterMode.SCATTERED return ScatterMode.SCATTERED
return ScatterMode.TP_ATTN_FULL return ScatterMode.TP_ATTN_FULL
@@ -256,15 +256,15 @@ class AttnTpContext:
self.allow_input_scattered = False self.allow_input_scattered = False
self.input_scattered_ = False self.input_scattered_ = False
self.attn_inputs_: Optional[AttentionInputs] = None self.attn_inputs_: Optional[AttentionInputs] = None
self.is_nsa = False self.is_dsa = False
def init_context(self, q_lora_rank, is_nsa): def init_context(self, q_lora_rank, is_dsa):
self.is_nsa = is_nsa self.is_dsa = is_dsa
self.allow_input_scattered = ( self.allow_input_scattered = (
get_global_server_args().enable_attn_tp_input_scattered get_global_server_args().enable_attn_tp_input_scattered
and (_is_cuda or _is_npu) and (_is_cuda or _is_npu)
and q_lora_rank is not None and q_lora_rank is not None
and not is_nsa and not is_dsa
and get_tensor_model_parallel_world_size() > 1 and get_tensor_model_parallel_world_size() > 1
and not is_dp_attention_enabled() and not is_dp_attention_enabled()
and get_moe_a2a_backend().is_none() and get_moe_a2a_backend().is_none()
@@ -379,8 +379,8 @@ class LayerScatterModes:
or should_use_flashinfer_cutlass_moe_fp4_allgather() or should_use_flashinfer_cutlass_moe_fp4_allgather()
): ):
return ScatterMode.SCATTERED return ScatterMode.SCATTERED
# NSA CP doesn't support MOE_FULL yet; fall back to FULL # DSA CP doesn't support MOE_FULL yet; fall back to FULL
if is_enable_moe_cp_allgather() and not is_nsa_enable_prefill_cp(): if is_enable_moe_cp_allgather() and not is_dsa_enable_prefill_cp():
return ScatterMode.MOE_FULL return ScatterMode.MOE_FULL
return ScatterMode.FULL return ScatterMode.FULL
else: else:
@@ -551,10 +551,10 @@ class LayerCommunicator:
) )
elif _use_aiter and _is_gfx95_supported and (quant_format == "fp8"): elif _use_aiter and _is_gfx95_supported and (quant_format == "fp8"):
# aiter (ROCm gfx95) fused RMSNorm + FP8 group quant. # aiter (ROCm gfx95) fused RMSNorm + FP8 group quant.
# When NSA is active, also preserve the unquantized bf16 # When DSA is active, also preserve the unquantized bf16
# output as a 3-tuple (fp8, scale, bf16) so the NSA # output as a 3-tuple (fp8, scale, bf16) so the DSA
# indexer can skip redundant FP8 dequantization. # indexer can skip redundant FP8 dequantization.
_nsa_needs_bf16 = get_attn_tp_context().is_nsa _dsa_needs_bf16 = get_attn_tp_context().is_dsa
hidden_states, _unq_bf16, _, _res = fused_rms_fp8_group_quant( hidden_states, _unq_bf16, _, _res = fused_rms_fp8_group_quant(
hidden_states, hidden_states,
self.input_layernorm.weight, self.input_layernorm.weight,
@@ -565,9 +565,9 @@ class LayerCommunicator:
group_size=128, group_size=128,
dtype_quant=torch.float8_e4m3fn, dtype_quant=torch.float8_e4m3fn,
res1=None, res1=None,
output_unquantized_inp1=_nsa_needs_bf16, output_unquantized_inp1=_dsa_needs_bf16,
) )
if _nsa_needs_bf16: if _dsa_needs_bf16:
hidden_states = ( hidden_states = (
hidden_states[0], hidden_states[0],
hidden_states[1], hidden_states[1],
@@ -596,9 +596,9 @@ class LayerCommunicator:
) )
elif _use_aiter and _is_gfx95_supported and (quant_format == "fp8"): elif _use_aiter and _is_gfx95_supported and (quant_format == "fp8"):
# aiter (ROCm gfx95) fused RMSNorm + FP8 group quant # aiter (ROCm gfx95) fused RMSNorm + FP8 group quant
# with residual addition. When NSA is active, pack # with residual addition. When DSA is active, pack
# the unquantized bf16 as a 3-tuple (fp8, scale, bf16). # the unquantized bf16 as a 3-tuple (fp8, scale, bf16).
_nsa_needs_bf16 = get_attn_tp_context().is_nsa _dsa_needs_bf16 = get_attn_tp_context().is_dsa
hidden_states, _unq_bf16, _, residual = ( hidden_states, _unq_bf16, _, residual = (
fused_rms_fp8_group_quant( fused_rms_fp8_group_quant(
hidden_states, hidden_states,
@@ -610,10 +610,10 @@ class LayerCommunicator:
group_size=128, group_size=128,
dtype_quant=torch.float8_e4m3fn, dtype_quant=torch.float8_e4m3fn,
res1=residual, res1=residual,
output_unquantized_inp1=_nsa_needs_bf16, output_unquantized_inp1=_dsa_needs_bf16,
) )
) )
if _nsa_needs_bf16: if _dsa_needs_bf16:
hidden_states = ( hidden_states = (
hidden_states[0], hidden_states[0],
hidden_states[1], hidden_states[1],
@@ -709,7 +709,7 @@ class LayerCommunicator:
return True return True
if forward_batch.dp_padding_mode.is_max_len(): if forward_batch.dp_padding_mode.is_max_len():
return True return True
if nsa_use_prefill_cp(forward_batch): if dsa_use_prefill_cp(forward_batch):
return True return True
if get_attn_tp_context().input_scattered and not self.is_last_layer: if get_attn_tp_context().input_scattered and not self.is_last_layer:
return True return True
python/sglang/srt/layers/communicator_nsa_cp.py → python/sglang/srt/layers/communicator_dsa_cp.py
+19 -19
View File
@@ -18,9 +18,9 @@ from typing import Callable, Optional
import torch import torch
from sglang.srt.layers.attention.nsa.utils import ( from sglang.srt.layers.attention.dsa.utils import (
is_nsa_enable_prefill_cp, dsa_use_prefill_cp,
nsa_use_prefill_cp, is_dsa_enable_prefill_cp,
) )
from sglang.srt.layers.communicator import ( from sglang.srt.layers.communicator import (
CommunicateContext, CommunicateContext,
@@ -40,14 +40,14 @@ from sglang.srt.layers.dp_attention import (
from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.model_executor.forward_batch_info import ForwardBatch
def nsa_enable_prefill_cp(): def dsa_enable_prefill_cp():
# After using cp, the communication mode of this part changes. # After using cp, the communication mode of this part changes.
# The three parts of prepare_attn, prepare_mlp, and postprocess_layer # The three parts of prepare_attn, prepare_mlp, and postprocess_layer
# no longer require additional communication for reduce, scatter, etc. # no longer require additional communication for reduce, scatter, etc.
return is_nsa_enable_prefill_cp() return is_dsa_enable_prefill_cp()
class NSACPLayerCommunicator(LayerCommunicator): class DSACPLayerCommunicator(LayerCommunicator):
def __init__( def __init__(
self, self,
layer_scatter_modes: LayerScatterModes, layer_scatter_modes: LayerScatterModes,
@@ -73,19 +73,19 @@ class NSACPLayerCommunicator(LayerCommunicator):
assert ( assert (
self._context.attn_dp_size == 1 self._context.attn_dp_size == 1
), f"dp_size should be 1 when moe_runner_backend is none" ), f"dp_size should be 1 when moe_runner_backend is none"
self._communicate_simple_fn = NSACPCommunicateSimpleFn.get_fn( self._communicate_simple_fn = DSACPCommunicateSimpleFn.get_fn(
input_mode=ScatterMode.SCATTERED, input_mode=ScatterMode.SCATTERED,
output_mode=ScatterMode.SCATTERED, output_mode=ScatterMode.SCATTERED,
context=self._context, context=self._context,
) )
self._communicate_with_all_reduce_and_layer_norm_fn = NSACPCommunicateWithAllReduceAndLayerNormFn.get_fn( self._communicate_with_all_reduce_and_layer_norm_fn = DSACPCommunicateWithAllReduceAndLayerNormFn.get_fn(
hidden_states_input_mode=ScatterMode.SCATTERED, hidden_states_input_mode=ScatterMode.SCATTERED,
residual_input_mode=ScatterMode.SCATTERED, residual_input_mode=ScatterMode.SCATTERED,
hidden_states_output_mode=self.layer_scatter_modes.mlp_mode, # SCATTERED, FULL hidden_states_output_mode=self.layer_scatter_modes.mlp_mode, # SCATTERED, FULL
residual_output_mode=ScatterMode.SCATTERED, residual_output_mode=ScatterMode.SCATTERED,
context=self._context, context=self._context,
) )
self._communicate_summable_tensor_pair_fn = NSACPCommunicateSummableTensorPairFn.get_fn( self._communicate_summable_tensor_pair_fn = DSACPCommunicateSummableTensorPairFn.get_fn(
hidden_states_input_mode=self.layer_scatter_modes.mlp_mode, # SCATTERED, FULL hidden_states_input_mode=self.layer_scatter_modes.mlp_mode, # SCATTERED, FULL
residual_input_mode=ScatterMode.SCATTERED, residual_input_mode=ScatterMode.SCATTERED,
output_mode=ScatterMode.SCATTERED, output_mode=ScatterMode.SCATTERED,
@@ -93,7 +93,7 @@ class NSACPLayerCommunicator(LayerCommunicator):
) )
class NSACPCommunicateSimpleFn(CommunicateSimpleFn): class DSACPCommunicateSimpleFn(CommunicateSimpleFn):
@staticmethod @staticmethod
def get_fn( def get_fn(
input_mode: ScatterMode, input_mode: ScatterMode,
@@ -101,12 +101,12 @@ class NSACPCommunicateSimpleFn(CommunicateSimpleFn):
context: CommunicateContext, context: CommunicateContext,
): ):
if context.is_same_group_size(input_mode, output_mode): if context.is_same_group_size(input_mode, output_mode):
return NSACPCommunicateSimpleFn._trivial return DSACPCommunicateSimpleFn._trivial
raise NotImplementedError(f"{input_mode=} {output_mode=}") raise NotImplementedError(f"{input_mode=} {output_mode=}")
class NSACPCommunicateWithAllReduceAndLayerNormFn( class DSACPCommunicateWithAllReduceAndLayerNormFn(
CommunicateWithAllReduceAndLayerNormFn CommunicateWithAllReduceAndLayerNormFn
): ):
"""Besides communication, needs to """Besides communication, needs to
@@ -126,11 +126,11 @@ class NSACPCommunicateWithAllReduceAndLayerNormFn(
assert residual_input_mode == ScatterMode.SCATTERED assert residual_input_mode == ScatterMode.SCATTERED
assert residual_output_mode == ScatterMode.SCATTERED assert residual_output_mode == ScatterMode.SCATTERED
if hidden_states_output_mode == ScatterMode.SCATTERED: if hidden_states_output_mode == ScatterMode.SCATTERED:
return NSACPCommunicateWithAllReduceAndLayerNormFn._simple return DSACPCommunicateWithAllReduceAndLayerNormFn._simple
if hidden_states_output_mode == ScatterMode.FULL: if hidden_states_output_mode == ScatterMode.FULL:
return partial( return partial(
NSACPCommunicateWithAllReduceAndLayerNormFn._gather_hidden_states_and_residual, DSACPCommunicateWithAllReduceAndLayerNormFn._gather_hidden_states_and_residual,
residual_input_mode=residual_input_mode, residual_input_mode=residual_input_mode,
) )
@@ -152,7 +152,7 @@ class NSACPCommunicateWithAllReduceAndLayerNormFn(
hidden_states, residual = layernorm(hidden_states, residual) hidden_states, residual = layernorm(hidden_states, residual)
# for prefill: attn tp scattered -> full # for prefill: attn tp scattered -> full
# for decode: attn tp full -> full # for decode: attn tp full -> full
if nsa_use_prefill_cp(forward_batch): if dsa_use_prefill_cp(forward_batch):
assert context.attn_dp_size == 1 assert context.attn_dp_size == 1
hidden_states, local_hidden_states = ( hidden_states, local_hidden_states = (
get_local_dp_buffer(get_attention_cp_group()), get_local_dp_buffer(get_attention_cp_group()),
@@ -165,7 +165,7 @@ class NSACPCommunicateWithAllReduceAndLayerNormFn(
return hidden_states, residual return hidden_states, residual
class NSACPCommunicateSummableTensorPairFn(CommunicateSummableTensorPairFn): class DSACPCommunicateSummableTensorPairFn(CommunicateSummableTensorPairFn):
"""It is allowed to make (hidden_states, residual) := (hidden_states + residual, None) if needed.""" """It is allowed to make (hidden_states, residual) := (hidden_states + residual, None) if needed."""
@staticmethod @staticmethod
@@ -184,12 +184,12 @@ class NSACPCommunicateSummableTensorPairFn(CommunicateSummableTensorPairFn):
and (residual_input_mode == ScatterMode.SCATTERED) and (residual_input_mode == ScatterMode.SCATTERED)
and (output_mode == ScatterMode.SCATTERED) and (output_mode == ScatterMode.SCATTERED)
): ):
return NSACPCommunicateSummableTensorPairFn._scatter_hidden_states return DSACPCommunicateSummableTensorPairFn._scatter_hidden_states
if context.is_same_group_size( if context.is_same_group_size(
hidden_states_input_mode, output_mode hidden_states_input_mode, output_mode
) and context.is_same_group_size(residual_input_mode, output_mode): ) and context.is_same_group_size(residual_input_mode, output_mode):
return NSACPCommunicateSummableTensorPairFn._trivial return DSACPCommunicateSummableTensorPairFn._trivial
raise NotImplementedError( raise NotImplementedError(
f"{hidden_states_input_mode=} {residual_input_mode=} {output_mode=}" f"{hidden_states_input_mode=} {residual_input_mode=} {output_mode=}"
@@ -205,7 +205,7 @@ class NSACPCommunicateSummableTensorPairFn(CommunicateSummableTensorPairFn):
): ):
# for prefill: full -> attn tp scattered # for prefill: full -> attn tp scattered
# for decode: full -> attn tp full # for decode: full -> attn tp full
if nsa_use_prefill_cp(forward_batch): if dsa_use_prefill_cp(forward_batch):
assert context.attn_dp_size == 1 assert context.attn_dp_size == 1
input_hidden_states = hidden_states input_hidden_states = hidden_states
hidden_states = hidden_states.tensor_split(context.attn_cp_size)[ hidden_states = hidden_states.tensor_split(context.attn_cp_size)[
python/sglang/srt/layers/dp_attention.py
+1 -1
View File
@@ -72,7 +72,7 @@ class DpPaddingMode(IntEnum):
# When is_extend_in_batch and dp_size > 1, use SUM_LEN to avoid padding # When is_extend_in_batch and dp_size > 1, use SUM_LEN to avoid padding
# overhead from uneven token distribution. # overhead from uneven token distribution.
# For dp_size=1, max_len equals sum_len, so prefer MAX_LEN mode # For dp_size=1, max_len equals sum_len, so prefer MAX_LEN mode
# to enable symmetric memory optimization (needed for NSA CP, etc.). # to enable symmetric memory optimization (needed for DSA CP, etc.).
if is_extend_in_batch and dp_size > 1: if is_extend_in_batch and dp_size > 1:
return DpPaddingMode.SUM_LEN return DpPaddingMode.SUM_LEN
python/sglang/srt/layers/mhc.py
+2 -2
View File
@@ -8,7 +8,7 @@ import torch
from sglang.jit_kernel.utils import is_arch_support_pdl from sglang.jit_kernel.utils import is_arch_support_pdl
from sglang.srt.environ import envs from sglang.srt.environ import envs
from sglang.srt.layers.attention.nsa.utils import is_nsa_prefill_cp_round_robin_split from sglang.srt.layers.attention.dsa.utils import is_dsa_prefill_cp_round_robin_split
from sglang.srt.layers.utils.common import strict_contiguous from sglang.srt.layers.utils.common import strict_contiguous
tilelang.set_log_level("WARNING") tilelang.set_log_level("WARNING")
@@ -880,7 +880,7 @@ def mhc_post(
post_layer_mix: torch.Tensor, post_layer_mix: torch.Tensor,
comb_res_mix: torch.Tensor, comb_res_mix: torch.Tensor,
) -> torch.Tensor: ) -> torch.Tensor:
if is_nsa_prefill_cp_round_robin_split(): if is_dsa_prefill_cp_round_robin_split():
x = strict_contiguous(x) x = strict_contiguous(x)
residual = strict_contiguous(residual) residual = strict_contiguous(residual)
post_layer_mix = strict_contiguous(post_layer_mix) post_layer_mix = strict_contiguous(post_layer_mix)
python/sglang/srt/layers/utils/cp_utils.py
+20 -20
View File
@@ -66,17 +66,17 @@ def can_cp_split(seq_len: int, cp_size: int, forward_batch):
def cp_split_and_rebuild_data(forward_batch, input_: torch.Tensor): def cp_split_and_rebuild_data(forward_batch, input_: torch.Tensor):
from sglang.srt.layers.attention.nsa.utils import ( from sglang.srt.layers.attention.dsa.utils import (
is_nsa_prefill_cp_round_robin_split, dsa_cp_round_robin_split_data,
nsa_cp_round_robin_split_data, is_dsa_prefill_cp_round_robin_split,
) )
if is_nsa_prefill_cp_round_robin_split(): if is_dsa_prefill_cp_round_robin_split():
cp_size = get_attention_cp_size() cp_size = get_attention_cp_size()
assert ( assert (
input_.shape[0] % cp_size == 0 input_.shape[0] % cp_size == 0
), f"Expect input shape 0 can divided by cp size, but got input shape {input_.shape}, cp size {cp_size}" ), f"Expect input shape 0 can divided by cp size, but got input shape {input_.shape}, cp size {cp_size}"
return nsa_cp_round_robin_split_data(input_) return dsa_cp_round_robin_split_data(input_)
input_list = list( input_list = list(
torch.split(input_, forward_batch.attn_cp_metadata.split_list, dim=0) torch.split(input_, forward_batch.attn_cp_metadata.split_list, dim=0)
@@ -88,18 +88,18 @@ def cp_split_and_rebuild_data(forward_batch, input_: torch.Tensor):
def cp_split_and_rebuild_position(forward_batch, positions: torch.Tensor): def cp_split_and_rebuild_position(forward_batch, positions: torch.Tensor):
from sglang.srt.layers.attention.nsa.utils import ( from sglang.srt.layers.attention.dsa.utils import (
is_nsa_prefill_cp_round_robin_split, dsa_cp_round_robin_split_data,
nsa_cp_round_robin_split_data, is_dsa_prefill_cp_round_robin_split,
) )
if is_nsa_prefill_cp_round_robin_split(): if is_dsa_prefill_cp_round_robin_split():
cp_size = get_attention_cp_size() cp_size = get_attention_cp_size()
assert positions.shape[0] % cp_size == 0, ( assert positions.shape[0] % cp_size == 0, (
f"Expect positions shape 0 can divided by cp size, but got positions shape {positions.shape}, " f"Expect positions shape 0 can divided by cp size, but got positions shape {positions.shape}, "
f"cp size {cp_size}" f"cp size {cp_size}"
) )
return nsa_cp_round_robin_split_data(positions) return dsa_cp_round_robin_split_data(positions)
position_id_list = list( position_id_list = list(
torch.split(positions, forward_batch.attn_cp_metadata.split_list, dim=-1) torch.split(positions, forward_batch.attn_cp_metadata.split_list, dim=-1)
@@ -238,11 +238,11 @@ def cp_all_gather_rerange_output(input_tensor, cp_size, forward_batch, stream):
| token0, token1, token2, token3, token4, token5, token6, token7, ... | token0, token1, token2, token3, token4, token5, token6, token7, ...
| +-------------------------+ | +-------------------------+
""" """
from sglang.srt.layers.attention.nsa.utils import ( from sglang.srt.layers.attention.dsa.utils import (
is_nsa_prefill_cp_round_robin_split, is_dsa_prefill_cp_round_robin_split,
) )
if is_nsa_prefill_cp_round_robin_split(): if is_dsa_prefill_cp_round_robin_split():
with use_symmetric_memory( with use_symmetric_memory(
get_attention_cp_group(), disabled=not is_allocation_symmetric() get_attention_cp_group(), disabled=not is_allocation_symmetric()
): ):
@@ -395,11 +395,11 @@ def prepare_context_parallel_metadata(
cp_size, cp_size,
seqs_len, seqs_len,
): ):
from sglang.srt.layers.attention.nsa.utils import ( from sglang.srt.layers.attention.dsa.utils import (
is_nsa_prefill_cp_round_robin_split, is_dsa_prefill_cp_round_robin_split,
) )
if is_nsa_prefill_cp_round_robin_split(): if is_dsa_prefill_cp_round_robin_split():
return ContextParallelMetadata() return ContextParallelMetadata()
"""prepare_input_dp_with_cp_dsa-zigzag index """prepare_input_dp_with_cp_dsa-zigzag index
@@ -505,16 +505,16 @@ def prepare_context_parallel_metadata(
# TODO Support multi-batch-cp-split, multi-batch-cp support has accuracy issues # TODO Support multi-batch-cp-split, multi-batch-cp support has accuracy issues
# Prefix offset is critical when radix cache hits (prefix_len > 0). # Prefix offset is critical when radix cache hits (prefix_len > 0).
# For non-NSA CP (e.g. qwen3-moe), consumers use these values directly as # For non-DSA CP (e.g. qwen3-moe), consumers use these values directly as
# FlashAttention cache_seqlens, so the prefix must be baked in here. # FlashAttention cache_seqlens, so the prefix must be baked in here.
# For NSA CP, `_get_topk_ragged_with_cp` re-adds the cached-prefix offset # For DSA CP, `_get_topk_ragged_with_cp` re-adds the cached-prefix offset
# from (seq_lens_cpu - extend_seq_lens_cpu); baking prefix_len in here # from (seq_lens_cpu - extend_seq_lens_cpu); baking prefix_len in here
# would silently drop it whenever the scheduler packs multiple requests # would silently drop it whenever the scheduler packs multiple requests
# into a single CP extend (len(seqs_len) != 1 -> prefix_len falls back # into a single CP extend (len(seqs_len) != 1 -> prefix_len falls back
# to 0), corrupting the indexer's ke_offset on prefix-cache hits. # to 0), corrupting the indexer's ke_offset on prefix-cache hits.
from sglang.srt.layers.attention.nsa.utils import is_nsa_enable_prefill_cp from sglang.srt.layers.attention.dsa.utils import is_dsa_enable_prefill_cp
if is_nsa_enable_prefill_cp(): if is_dsa_enable_prefill_cp():
kv_len_prev = prefix_sum_list[cp_rank] kv_len_prev = prefix_sum_list[cp_rank]
kv_len_next = prefix_sum_list[cp_size * 2 - cp_rank - 1] kv_len_next = prefix_sum_list[cp_size * 2 - cp_rank - 1]
else: else:
python/sglang/srt/managers/hisparse_coordinator.py
+2 -2
View File
@@ -9,7 +9,7 @@ from sglang.srt.managers.schedule_batch import Req
from sglang.srt.mem_cache.hisparse_memory_pool import ( from sglang.srt.mem_cache.hisparse_memory_pool import (
DeepSeekV4HiSparseTokenToKVPoolAllocator, DeepSeekV4HiSparseTokenToKVPoolAllocator,
DeepSeekV4SingleKVPoolHost, DeepSeekV4SingleKVPoolHost,
HiSparseNSATokenToKVPool, HiSparseDSATokenToKVPool,
HiSparseTokenToKVPoolAllocator, HiSparseTokenToKVPoolAllocator,
) )
from sglang.srt.mem_cache.memory_pool_host import MLATokenToKVPoolHost from sglang.srt.mem_cache.memory_pool_host import MLATokenToKVPoolHost
@@ -79,7 +79,7 @@ class HiSparseCoordinator:
assert isinstance( assert isinstance(
self.token_to_kv_pool_allocator, HiSparseTokenToKVPoolAllocator self.token_to_kv_pool_allocator, HiSparseTokenToKVPoolAllocator
) )
self.mem_pool_device: HiSparseNSATokenToKVPool = ( self.mem_pool_device: HiSparseDSATokenToKVPool = (
self.token_to_kv_pool_allocator.get_kvcache() self.token_to_kv_pool_allocator.get_kvcache()
) )
self.mem_pool_host = MLATokenToKVPoolHost( self.mem_pool_host = MLATokenToKVPoolHost(
python/sglang/srt/managers/schedule_policy.py
+3 -3
View File
@@ -34,7 +34,7 @@ from typing import TYPE_CHECKING, Dict, List, Optional, Set, Union
import torch import torch
from sglang.srt.dllm.config import DllmConfig from sglang.srt.dllm.config import DllmConfig
from sglang.srt.layers.attention.nsa.utils import is_nsa_prefill_cp_in_seq_split from sglang.srt.layers.attention.dsa.utils import is_dsa_prefill_cp_in_seq_split
from sglang.srt.layers.utils.cp_utils import is_prefill_context_parallel_enabled from sglang.srt.layers.utils.cp_utils import is_prefill_context_parallel_enabled
from sglang.srt.managers.schedule_batch import Req, ScheduleBatch from sglang.srt.managers.schedule_batch import Req, ScheduleBatch
from sglang.srt.mem_cache.base_prefix_cache import ( from sglang.srt.mem_cache.base_prefix_cache import (
@@ -467,7 +467,7 @@ class PrefillAdder:
self.priority_scheduling_preemption_threshold = ( self.priority_scheduling_preemption_threshold = (
priority_scheduling_preemption_threshold priority_scheduling_preemption_threshold
) )
self.nsa_prefill_cp_in_seq_split = is_nsa_prefill_cp_in_seq_split() self.dsa_prefill_cp_in_seq_split = is_dsa_prefill_cp_in_seq_split()
self.max_running_requests = max_running_requests self.max_running_requests = max_running_requests
self.prefill_context_parallel_enabled = is_prefill_context_parallel_enabled() self.prefill_context_parallel_enabled = is_prefill_context_parallel_enabled()
self.prefill_max_requests = prefill_max_requests self.prefill_max_requests = prefill_max_requests
@@ -826,7 +826,7 @@ class PrefillAdder:
# Enabling context parallelism currently presents precision issues; # Enabling context parallelism currently presents precision issues;
# therefore, the prefill-batch setting is temporarily set to 1. # therefore, the prefill-batch setting is temporarily set to 1.
if ( if (
self.nsa_prefill_cp_in_seq_split or self.prefill_context_parallel_enabled self.dsa_prefill_cp_in_seq_split or self.prefill_context_parallel_enabled
) and len(self.can_run_list) >= 1: ) and len(self.can_run_list) >= 1:
return AddReqResult.OTHER return AddReqResult.OTHER
python/sglang/srt/managers/scheduler_pp_mixin.py
+1 -1
View File
@@ -523,7 +523,7 @@ class SchedulerPPMixin:
self.pp_loop_size: int = self.ps.pp_size + self.server_args.pp_async_batch_depth self.pp_loop_size: int = self.ps.pp_size + self.server_args.pp_async_batch_depth
# In CP mode, attention weights are duplicated, eliminating the need for the attention TP all-gather operation. # In CP mode, attention weights are duplicated, eliminating the need for the attention TP all-gather operation.
self.require_attn_tp_allgather = ( self.require_attn_tp_allgather = (
not self.server_args.enable_nsa_prefill_context_parallel not self.server_args.enable_dsa_prefill_context_parallel
) )
self.mbs = [None] * self.pp_loop_size self.mbs = [None] * self.pp_loop_size
self.last_mbs = [None] * self.pp_loop_size self.last_mbs = [None] * self.pp_loop_size
python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py
+1 -1
View File
@@ -9,11 +9,11 @@ import torch
from sglang.jit_kernel.deepseek_v4 import fused_k_norm_rope_flashmla, fused_store_cache from sglang.jit_kernel.deepseek_v4 import fused_k_norm_rope_flashmla, fused_store_cache
from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE
from sglang.srt.environ import envs from sglang.srt.environ import envs
from sglang.srt.layers.attention.dsa import index_buf_accessor
from sglang.srt.layers.attention.dsv4 import ( from sglang.srt.layers.attention.dsv4 import (
index_buf_accessor as dsv4_index_buf_accessor, index_buf_accessor as dsv4_index_buf_accessor,
) )
from sglang.srt.layers.attention.dsv4.index_buf_accessor import NopeFp8RopeBf16Pack from sglang.srt.layers.attention.dsv4.index_buf_accessor import NopeFp8RopeBf16Pack
from sglang.srt.layers.attention.nsa import index_buf_accessor
from sglang.srt.mem_cache.base_swa_memory_pool import BaseSWAKVPool from sglang.srt.mem_cache.base_swa_memory_pool import BaseSWAKVPool
from sglang.srt.mem_cache.deepseek_v4_compress_state import CompressStatePool from sglang.srt.mem_cache.deepseek_v4_compress_state import CompressStatePool
from sglang.srt.mem_cache.memory_pool import KVCache from sglang.srt.mem_cache.memory_pool import KVCache
python/sglang/srt/mem_cache/hiradix_cache.py
+8 -10
View File
@@ -36,12 +36,12 @@ from sglang.srt.mem_cache.hybrid_cache.hybrid_cache_controller import (
HybridCacheController, HybridCacheController,
) )
from sglang.srt.mem_cache.hybrid_cache.hybrid_pool_assembler import ( from sglang.srt.mem_cache.hybrid_cache.hybrid_pool_assembler import (
attach_hybrid_nsa_pool_to_hiradix_cache, attach_hybrid_dsa_pool_to_hiradix_cache,
) )
from sglang.srt.mem_cache.memory_pool import ( from sglang.srt.mem_cache.memory_pool import (
DSATokenToKVPool,
MHATokenToKVPool, MHATokenToKVPool,
MLATokenToKVPool, MLATokenToKVPool,
NSATokenToKVPool,
) )
from sglang.srt.mem_cache.memory_pool_host import ( from sglang.srt.mem_cache.memory_pool_host import (
MHATokenToKVPoolHost, MHATokenToKVPoolHost,
@@ -82,8 +82,8 @@ class HiRadixCache(RadixCache):
server_args.hicache_mem_layout, server_args.hicache_mem_layout,
allocator_type=server_args.hicache_storage_backend, allocator_type=server_args.hicache_storage_backend,
) )
elif isinstance(self.kv_cache, NSATokenToKVPool): elif isinstance(self.kv_cache, DSATokenToKVPool):
# Filled by attach_hybrid_nsa_pool_to_hiradix_cache after storage extra_config is parsed. # Filled by attach_hybrid_dsa_pool_to_hiradix_cache after storage extra_config is parsed.
self.token_to_kv_pool_host = None self.token_to_kv_pool_host = None
elif isinstance(self.kv_cache, MLATokenToKVPool): elif isinstance(self.kv_cache, MLATokenToKVPool):
self.token_to_kv_pool_host = MLATokenToKVPoolHost( self.token_to_kv_pool_host = MLATokenToKVPoolHost(
@@ -95,9 +95,7 @@ class HiRadixCache(RadixCache):
allocator_type=server_args.hicache_storage_backend, allocator_type=server_args.hicache_storage_backend,
) )
else: else:
raise ValueError( raise ValueError("HiRadixCache only supports MHA, MLA, and DSA models")
"HiRadixCache only supports MHA, MLA, and NSA (DSA) models"
)
self.tp_group = params.tp_cache_group self.tp_group = params.tp_cache_group
self.attn_cp_group = params.attn_cp_cache_group self.attn_cp_group = params.attn_cp_cache_group
@@ -122,8 +120,8 @@ class HiRadixCache(RadixCache):
self.prefetch_stop_policy = server_args.hicache_storage_prefetch_policy self.prefetch_stop_policy = server_args.hicache_storage_prefetch_policy
self.load_cache_event = threading.Event() self.load_cache_event = threading.Event()
if isinstance(self.kv_cache, NSATokenToKVPool): if isinstance(self.kv_cache, DSATokenToKVPool):
attach_hybrid_nsa_pool_to_hiradix_cache( attach_hybrid_dsa_pool_to_hiradix_cache(
self, self,
params, params,
server_args, server_args,
@@ -643,7 +641,7 @@ class HiRadixCache(RadixCache):
def _get_extra_pools(self) -> dict: def _get_extra_pools(self) -> dict:
if not isinstance(self.cache_controller, HybridCacheController): if not isinstance(self.cache_controller, HybridCacheController):
return {} return {}
if isinstance(self.kv_cache, NSATokenToKVPool): if isinstance(self.kv_cache, DSATokenToKVPool):
pool = PoolTransfer( pool = PoolTransfer(
name=PoolName.INDEXER, name=PoolName.INDEXER,
hit_policy=PoolHitPolicy.ALL_PAGES, hit_policy=PoolHitPolicy.ALL_PAGES,
python/sglang/srt/mem_cache/hisparse_memory_pool.py
+3 -3
View File
@@ -16,7 +16,7 @@ from sglang.srt.mem_cache.deepseek_v4_memory_pool import (
DeepSeekV4TokenToKVPool, DeepSeekV4TokenToKVPool,
HiSparseC4DevicePool, HiSparseC4DevicePool,
) )
from sglang.srt.mem_cache.memory_pool import NSATokenToKVPool from sglang.srt.mem_cache.memory_pool import DSATokenToKVPool
from sglang.srt.mem_cache.memory_pool_host import HiSparseHostPoolMixin from sglang.srt.mem_cache.memory_pool_host import HiSparseHostPoolMixin
from sglang.srt.utils import is_cuda, is_hip from sglang.srt.utils import is_cuda, is_hip
from sglang.srt.utils.common import get_num_new_pages from sglang.srt.utils.common import get_num_new_pages
@@ -37,7 +37,7 @@ else:
) )
class HiSparseNSATokenToKVPool(NSATokenToKVPool): class HiSparseDSATokenToKVPool(DSATokenToKVPool):
def __init__( def __init__(
self, self,
size: int, size: int,
@@ -143,7 +143,7 @@ class HiSparseTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
page_size: int, page_size: int,
dtype: torch.dtype, dtype: torch.dtype,
device: torch.device, device: torch.device,
kvcache: HiSparseNSATokenToKVPool, kvcache: HiSparseDSATokenToKVPool,
need_sort: bool, need_sort: bool,
host_to_device_ratio: int = 2, host_to_device_ratio: int = 2,
): ):
python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py
+11 -11
View File
@@ -10,12 +10,12 @@ from sglang.srt.mem_cache.hybrid_cache.hybrid_cache_controller import (
from sglang.srt.mem_cache.memory_pool_host import ( from sglang.srt.mem_cache.memory_pool_host import (
DeepSeekV4PagedHostPool, DeepSeekV4PagedHostPool,
DeepSeekV4StateHostPool, DeepSeekV4StateHostPool,
DSAIndexerPoolHost,
HostPoolGroup, HostPoolGroup,
LogicalHostPool, LogicalHostPool,
MambaPoolHost, MambaPoolHost,
MHATokenToKVPoolHost, MHATokenToKVPoolHost,
MLATokenToKVPoolHost, MLATokenToKVPoolHost,
NSAIndexerPoolHost,
PoolEntry, PoolEntry,
) )
@@ -656,9 +656,9 @@ def attach_hybrid_pool_to_unified_cache(
from sglang.srt.mem_cache.base_prefix_cache import EvictParams from sglang.srt.mem_cache.base_prefix_cache import EvictParams
from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool
from sglang.srt.mem_cache.memory_pool import ( from sglang.srt.mem_cache.memory_pool import (
DSATokenToKVPool,
HybridLinearKVPool, HybridLinearKVPool,
MLATokenToKVPool, MLATokenToKVPool,
NSATokenToKVPool,
) )
from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
from sglang.srt.mem_cache.unified_cache_components import ComponentType from sglang.srt.mem_cache.unified_cache_components import ComponentType
@@ -667,7 +667,7 @@ def attach_hybrid_pool_to_unified_cache(
kvcache = params.token_to_kv_pool_allocator.get_kvcache() kvcache = params.token_to_kv_pool_allocator.get_kvcache()
swa_stack = isinstance(kvcache, SWAKVPool) swa_stack = isinstance(kvcache, SWAKVPool)
mamba_stack = isinstance(kvcache, HybridLinearKVPool) mamba_stack = isinstance(kvcache, HybridLinearKVPool)
nsa_stack = isinstance(kvcache, NSATokenToKVPool) dsa_stack = isinstance(kvcache, DSATokenToKVPool)
deepseek_v4_stack = isinstance(kvcache, DeepSeekV4TokenToKVPool) deepseek_v4_stack = isinstance(kvcache, DeepSeekV4TokenToKVPool)
if deepseek_v4_stack: if deepseek_v4_stack:
@@ -820,7 +820,7 @@ def attach_hybrid_pool_to_unified_cache(
cache.swa_kv_pool_host cache.swa_kv_pool_host
) )
transfer_layer_num = len(full_layer_mapping | swa_layer_mapping) transfer_layer_num = len(full_layer_mapping | swa_layer_mapping)
elif nsa_stack: elif dsa_stack:
full_layer_mapping = { full_layer_mapping = {
layer_id: layer_id for layer_id in range(full_kv_pool.layer_num) layer_id: layer_id for layer_id in range(full_kv_pool.layer_num)
} }
@@ -838,7 +838,7 @@ def attach_hybrid_pool_to_unified_cache(
storage_backend=None, storage_backend=None,
use_mla=use_mla, use_mla=use_mla,
override_kv_cache_dim=full_kv_pool.kv_cache_dim, override_kv_cache_dim=full_kv_pool.kv_cache_dim,
sidecar_host_pool_factory=lambda kv_host_pool: NSAIndexerPoolHost( sidecar_host_pool_factory=lambda kv_host_pool: DSAIndexerPoolHost(
full_kv_pool, full_kv_pool,
kv_host_pool, kv_host_pool,
server_args.hicache_mem_layout, server_args.hicache_mem_layout,
@@ -897,7 +897,7 @@ def attach_hybrid_pool_to_unified_cache(
pools_desc = "KV + MAMBA" pools_desc = "KV + MAMBA"
elif swa_stack: elif swa_stack:
pools_desc = "KV + SWA" pools_desc = "KV + SWA"
elif nsa_stack: elif dsa_stack:
pools_desc = "KV + INDEXER" pools_desc = "KV + INDEXER"
else: else:
pools_desc = "KV" pools_desc = "KV"
@@ -911,7 +911,7 @@ def attach_hybrid_pool_to_unified_cache(
raise raise
def attach_hybrid_nsa_pool_to_hiradix_cache( def attach_hybrid_dsa_pool_to_hiradix_cache(
radix_cache: HiRadixCache, radix_cache: HiRadixCache,
params: CacheInitParams, params: CacheInitParams,
server_args: ServerArgs, server_args: ServerArgs,
@@ -925,7 +925,7 @@ def attach_hybrid_nsa_pool_to_hiradix_cache(
) -> None: ) -> None:
"""Attach HostPoolGroup (KV + indexer) + HybridCacheController for HiRadixCache. """Attach HostPoolGroup (KV + indexer) + HybridCacheController for HiRadixCache.
This entrypoint is currently intended only for HiRadixCache's NSA path. This entrypoint is currently intended only for HiRadixCache's DSA path.
""" """
try: try:
kv = radix_cache.kv_cache kv = radix_cache.kv_cache
@@ -945,7 +945,7 @@ def attach_hybrid_nsa_pool_to_hiradix_cache(
use_mla=True, use_mla=True,
override_kv_cache_dim=kv.kv_cache_dim, override_kv_cache_dim=kv.kv_cache_dim,
prefetch_threshold=prefetch_threshold, prefetch_threshold=prefetch_threshold,
sidecar_host_pool_factory=lambda kv_host_pool: NSAIndexerPoolHost( sidecar_host_pool_factory=lambda kv_host_pool: DSAIndexerPoolHost(
kv, kv,
kv_host_pool, kv_host_pool,
server_args.hicache_mem_layout, server_args.hicache_mem_layout,
@@ -961,12 +961,12 @@ def attach_hybrid_nsa_pool_to_hiradix_cache(
radix_cache.token_to_kv_pool_host = host_pool_group radix_cache.token_to_kv_pool_host = host_pool_group
radix_cache.cache_controller = cache_controller radix_cache.cache_controller = cache_controller
logger.info( logger.info(
"Attached hybrid NSA pool stack to HiRadixCache: pools=KV + INDEXER, " "Attached hybrid DSA pool stack to HiRadixCache: pools=KV + INDEXER, "
"transfer_layer_num=%s", "transfer_layer_num=%s",
len(layer_mapping), len(layer_mapping),
) )
except Exception: except Exception:
logger.exception("attach_hybrid_nsa_pool_to_hiradix_cache failed") logger.exception("attach_hybrid_dsa_pool_to_hiradix_cache failed")
raise raise
python/sglang/srt/mem_cache/memory_pool.py
+22 -22
View File
@@ -40,12 +40,12 @@ from sglang.jit_kernel.kvcache import can_use_store_cache, store_cache
from sglang.srt.configs.mamba_utils import BaseLinearStateParams from sglang.srt.configs.mamba_utils import BaseLinearStateParams
from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE
from sglang.srt.environ import envs from sglang.srt.environ import envs
from sglang.srt.layers.attention.nsa import index_buf_accessor from sglang.srt.layers.attention.dsa import index_buf_accessor
from sglang.srt.layers.attention.nsa.quant_k_cache import ( from sglang.srt.layers.attention.dsa.quant_k_cache import (
quantize_k_cache, quantize_k_cache,
quantize_k_cache_separate, quantize_k_cache_separate,
) )
from sglang.srt.layers.attention.nsa.utils import aiter_can_use_preshuffle_paged_mqa from sglang.srt.layers.attention.dsa.utils import aiter_can_use_preshuffle_paged_mqa
from sglang.srt.layers.quantization.fp8_kernel import fp8_dtype, is_fp8_fnuz from sglang.srt.layers.quantization.fp8_kernel import fp8_dtype, is_fp8_fnuz
from sglang.srt.layers.radix_attention import RadixAttention from sglang.srt.layers.radix_attention import RadixAttention
from sglang.srt.mem_cache.utils import ( from sglang.srt.mem_cache.utils import (
@@ -1618,7 +1618,7 @@ class MLATokenToKVPool(KVCache):
enable_memory_saver: bool, enable_memory_saver: bool,
start_layer: Optional[int] = None, start_layer: Optional[int] = None,
end_layer: Optional[int] = None, end_layer: Optional[int] = None,
use_nsa: bool = False, use_dsa: bool = False,
override_kv_cache_dim: Optional[int] = None, override_kv_cache_dim: Optional[int] = None,
): ):
super().__init__( super().__init__(
@@ -1634,17 +1634,17 @@ class MLATokenToKVPool(KVCache):
self.kv_lora_rank = kv_lora_rank self.kv_lora_rank = kv_lora_rank
self.qk_rope_head_dim = qk_rope_head_dim self.qk_rope_head_dim = qk_rope_head_dim
self.use_nsa = use_nsa self.use_dsa = use_dsa
self.nsa_kv_cache_store_fp8 = ( self.dsa_kv_cache_store_fp8 = (
use_nsa use_dsa
and dtype == torch.float8_e4m3fn and dtype == torch.float8_e4m3fn
and override_kv_cache_dim is not None and override_kv_cache_dim is not None
) )
# When override_kv_cache_dim is provided with nsa model, we assume the # When override_kv_cache_dim is provided with dsa model, we assume the
# override kv cache dim is correct and use it directly. # override kv cache dim is correct and use it directly.
self.kv_cache_dim = ( self.kv_cache_dim = (
override_kv_cache_dim override_kv_cache_dim
if self.nsa_kv_cache_store_fp8 if self.dsa_kv_cache_store_fp8
else (kv_lora_rank + qk_rope_head_dim) else (kv_lora_rank + qk_rope_head_dim)
) )
@@ -1655,8 +1655,8 @@ class MLATokenToKVPool(KVCache):
dtype=torch.uint64, dtype=torch.uint64,
device=self.device, device=self.device,
) )
if not use_nsa: if not use_dsa:
# NSA will allocate indexer KV cache later and then log the total size # DSA will allocate indexer KV cache later and then log the total size
self._finalize_allocation_log(size) self._finalize_allocation_log(size)
def _create_buffers(self): def _create_buffers(self):
@@ -1726,7 +1726,7 @@ class MLATokenToKVPool(KVCache):
cache_v: torch.Tensor, cache_v: torch.Tensor,
): ):
layer_id = layer.layer_id layer_id = layer.layer_id
assert not self.nsa_kv_cache_store_fp8 assert not self.dsa_kv_cache_store_fp8
if cache_k.dtype != self.dtype: if cache_k.dtype != self.dtype:
cache_k = cache_k.to(self.dtype) cache_k = cache_k.to(self.dtype)
@@ -1746,7 +1746,7 @@ class MLATokenToKVPool(KVCache):
): ):
layer_id = layer.layer_id layer_id = layer.layer_id
if _is_hip and self.use_nsa and self.dtype == fp8_dtype: if _is_hip and self.use_dsa and self.dtype == fp8_dtype:
# HIP FP8 path uses raw MLA KV layout (nope + rope) without per-block scales. # HIP FP8 path uses raw MLA KV layout (nope + rope) without per-block scales.
# Fuse BF16/FP16 -> FP8 cast with paged KV write. # Fuse BF16/FP16 -> FP8 cast with paged KV write.
set_mla_kv_buffer_triton_fp8_quant( set_mla_kv_buffer_triton_fp8_quant(
@@ -1756,7 +1756,7 @@ class MLATokenToKVPool(KVCache):
cache_k_rope, cache_k_rope,
fp8_dtype, fp8_dtype,
) )
elif self.nsa_kv_cache_store_fp8: elif self.dsa_kv_cache_store_fp8:
# OPTIMIZATION: Quantize k_nope and k_rope separately to avoid concat overhead # OPTIMIZATION: Quantize k_nope and k_rope separately to avoid concat overhead
# This also enables reuse of set_mla_kv_buffer_triton two-tensor write path # This also enables reuse of set_mla_kv_buffer_triton two-tensor write path
# quantize_k_cache_separate returns (nope_part, rope_part) as uint8 bytes # quantize_k_cache_separate returns (nope_part, rope_part) as uint8 bytes
@@ -1905,7 +1905,7 @@ class MLATokenToKVPoolFP4(MLATokenToKVPool):
cache_v: torch.Tensor, cache_v: torch.Tensor,
): ):
layer_id = layer.layer_id layer_id = layer.layer_id
assert not self.nsa_kv_cache_store_fp8 assert not self.dsa_kv_cache_store_fp8
if cache_k.dtype != self.dtype: if cache_k.dtype != self.dtype:
from sglang.srt.layers.quantization.kvfp4_tensor import KVFP4QuantizeUtil from sglang.srt.layers.quantization.kvfp4_tensor import KVFP4QuantizeUtil
@@ -1930,7 +1930,7 @@ class MLATokenToKVPoolFP4(MLATokenToKVPool):
): ):
layer_id = layer.layer_id layer_id = layer.layer_id
if self.nsa_kv_cache_store_fp8: if self.dsa_kv_cache_store_fp8:
# original cache_k: (num_tokens, num_heads 1, hidden 576); we unsqueeze the page_size=1 dim here # original cache_k: (num_tokens, num_heads 1, hidden 576); we unsqueeze the page_size=1 dim here
# TODO no need to cat # TODO no need to cat
cache_k = torch.cat([cache_k_nope, cache_k_rope], dim=-1) cache_k = torch.cat([cache_k_nope, cache_k_rope], dim=-1)
@@ -1968,7 +1968,7 @@ class MLATokenToKVPoolFP4(MLATokenToKVPool):
) )
class NSATokenToKVPool(MLATokenToKVPool): class DSATokenToKVPool(MLATokenToKVPool):
quant_block_size = 128 quant_block_size = 128
index_k_with_scale_buffer_dtype = torch.uint8 index_k_with_scale_buffer_dtype = torch.uint8
rope_storage_dtype = torch.bfloat16 # rope is always stored in bf16 rope_storage_dtype = torch.bfloat16 # rope is always stored in bf16
@@ -2005,7 +2005,7 @@ class NSATokenToKVPool(MLATokenToKVPool):
enable_memory_saver, enable_memory_saver,
start_layer, start_layer,
end_layer, end_layer,
use_nsa=True, use_dsa=True,
override_kv_cache_dim=override_dim, override_kv_cache_dim=override_dim,
) )
# self.index_k_dtype = torch.float8_e4m3fn # self.index_k_dtype = torch.float8_e4m3fn
@@ -2013,7 +2013,7 @@ class NSATokenToKVPool(MLATokenToKVPool):
self.index_head_dim = index_head_dim self.index_head_dim = index_head_dim
if index_buf_size is None: if index_buf_size is None:
index_buf_size = size index_buf_size = size
# num head == 1 and head dim == 128 for index_k in NSA # num head == 1 and head dim == 128 for index_k in DSA
assert index_head_dim == 128 assert index_head_dim == 128
if _is_hip: if _is_hip:
@@ -2024,7 +2024,7 @@ class NSATokenToKVPool(MLATokenToKVPool):
else: else:
assert ( assert (
self.page_size == 1 self.page_size == 1
), f"HIP legacy NSA path requires page_size == 1, got {self.page_size}" ), f"HIP legacy DSA path requires page_size == 1, got {self.page_size}"
else: else:
assert self.page_size == 64 assert self.page_size == 64
with ( with (
@@ -2133,11 +2133,11 @@ class NSATokenToKVPool(MLATokenToKVPool):
) )
def get_cpu_copy(self, indices): def get_cpu_copy(self, indices):
# NSA keeps a page-indexed index_k_with_scale_buffer alongside kv_buffer. # DSA keeps a page-indexed index_k_with_scale_buffer alongside kv_buffer.
# Retract frees the slots/pages and they get reused by other reqs' # Retract frees the slots/pages and they get reused by other reqs'
# set_index_k_scale_buffer, so we must offload it here too -- otherwise # set_index_k_scale_buffer, so we must offload it here too -- otherwise
# resume restores kv_buffer but leaves foreign index/scale in place and # resume restores kv_buffer but leaves foreign index/scale in place and
# NSA attention reads garbage at those token positions. # DSA attention reads garbage at those token positions.
kv_cache_cpu = super().get_cpu_copy(indices) kv_cache_cpu = super().get_cpu_copy(indices)
page_indices = indices[:: self.page_size] // self.page_size page_indices = indices[:: self.page_size] // self.page_size
python/sglang/srt/mem_cache/memory_pool_host.py
+9 -9
View File
@@ -31,11 +31,11 @@ from sglang.jit_kernel.hicache import (
transfer_hicache_one_layer_mla as jit_transfer_hicache_one_layer_mla, transfer_hicache_one_layer_mla as jit_transfer_hicache_one_layer_mla,
) )
from sglang.srt.mem_cache.memory_pool import ( from sglang.srt.mem_cache.memory_pool import (
DSATokenToKVPool,
KVCache, KVCache,
MambaPool, MambaPool,
MHATokenToKVPool, MHATokenToKVPool,
MLATokenToKVPool, MLATokenToKVPool,
NSATokenToKVPool,
) )
from sglang.srt.utils import is_cuda, is_hip, is_mps, is_npu, is_xpu from sglang.srt.utils import is_cuda, is_hip, is_mps, is_npu, is_xpu
@@ -2608,14 +2608,14 @@ class HostPoolGroup:
) )
class NSAIndexerPoolHost(HostKVCache): class DSAIndexerPoolHost(HostKVCache):
"""Host-side NSA index buffers only. Slot layout matches the anchor MLA host pool.""" """Host-side DSA index buffers only. Slot layout matches the anchor MLA host pool."""
device_pool: NSATokenToKVPool device_pool: DSATokenToKVPool
def __init__( def __init__(
self, self,
device_pool: NSATokenToKVPool, device_pool: DSATokenToKVPool,
anchor_host: MLATokenToKVPoolHost, anchor_host: MLATokenToKVPoolHost,
layout: str, layout: str,
pin_memory: bool = True, pin_memory: bool = True,
@@ -2635,7 +2635,7 @@ class NSAIndexerPoolHost(HostKVCache):
self.index_head_dim = device_pool.index_head_dim self.index_head_dim = device_pool.index_head_dim
self.indexer_quant_block_size = device_pool.quant_block_size self.indexer_quant_block_size = device_pool.quant_block_size
self.indexer_dtype = NSATokenToKVPool.index_k_with_scale_buffer_dtype self.indexer_dtype = DSATokenToKVPool.index_k_with_scale_buffer_dtype
self.indexer_size_per_token = ( self.indexer_size_per_token = (
self.index_head_dim self.index_head_dim
+ self.index_head_dim // self.indexer_quant_block_size * 4 + self.index_head_dim // self.indexer_quant_block_size * 4
@@ -2658,12 +2658,12 @@ class NSAIndexerPoolHost(HostKVCache):
available_bytes = host_mem.available - HICACHE_HOST_MEMORY_RESERVE_BYTES available_bytes = host_mem.available - HICACHE_HOST_MEMORY_RESERVE_BYTES
if requested_bytes > available_bytes: if requested_bytes > available_bytes:
raise ValueError( raise ValueError(
f"Not enough host memory for NSA indexer hierarchical cache. " f"Not enough host memory for DSA indexer hierarchical cache. "
f"Requesting {requested_bytes / 1e9:.2f} GB but only have " f"Requesting {requested_bytes / 1e9:.2f} GB but only have "
f"{available_bytes / 1e9:.2f} GB free." f"{available_bytes / 1e9:.2f} GB free."
) )
logger.info( logger.info(
"Allocating %.2f GB host memory for NSA indexer (layout=%s).", "Allocating %.2f GB host memory for DSA indexer (layout=%s).",
requested_bytes / 1e9, requested_bytes / 1e9,
layout, layout,
) )
@@ -2726,7 +2726,7 @@ class NSAIndexerPoolHost(HostKVCache):
return host_indices, device_indices return host_indices, device_indices
if host_indices.numel() % self.page_size != 0: if host_indices.numel() % self.page_size != 0:
raise ValueError( raise ValueError(
"Index buffer transfer expects page-aligned indices for NSA." "Index buffer transfer expects page-aligned indices for DSA."
) )
host_page_indices = ( host_page_indices = (
host_indices.reshape(-1, self.page_size)[:, 0] // self.page_size host_indices.reshape(-1, self.page_size)[:, 0] // self.page_size
python/sglang/srt/mem_cache/sparsity/__init__.py
+2 -2
View File
@@ -1,7 +1,7 @@
from sglang.srt.mem_cache.sparsity.algorithms import ( from sglang.srt.mem_cache.sparsity.algorithms import (
BaseSparseAlgorithm, BaseSparseAlgorithm,
BaseSparseAlgorithmImpl, BaseSparseAlgorithmImpl,
DeepSeekNSAAlgorithm, DeepSeekDSAAlgorithm,
QuestAlgorithm, QuestAlgorithm,
) )
from sglang.srt.mem_cache.sparsity.backend import BackendAdaptor, FlashAttentionAdaptor from sglang.srt.mem_cache.sparsity.backend import BackendAdaptor, FlashAttentionAdaptor
@@ -17,7 +17,7 @@ __all__ = [
"BaseSparseAlgorithm", "BaseSparseAlgorithm",
"BaseSparseAlgorithmImpl", "BaseSparseAlgorithmImpl",
"QuestAlgorithm", "QuestAlgorithm",
"DeepSeekNSAAlgorithm", "DeepSeekDSAAlgorithm",
"BackendAdaptor", "BackendAdaptor",
"FlashAttentionAdaptor", "FlashAttentionAdaptor",
"SparseConfig", "SparseConfig",

Some files were not shown because too many files have changed in this diff Show More