diff --git a/.claude/skills/llm-torch-profiler-analysis/references/fuse-overlap-catalog.md b/.claude/skills/llm-torch-profiler-analysis/references/fuse-overlap-catalog.md index 917de36b6..ed05c517f 100644 --- a/.claude/skills/llm-torch-profiler-analysis/references/fuse-overlap-catalog.md +++ b/.claude/skills/llm-torch-profiler-analysis/references/fuse-overlap-catalog.md @@ -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 [`vllm-torch-compile-fusions.md`](vllm-torch-compile-fusions.md). Stable 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 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*`
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`
RoPE followed by KV-store, DtoD, or cache-write kernels | `python/sglang/jit_kernel/rope.py`
`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`
`cache_seqlens_int32`
`cu_seqlens_k`
`page_table`
`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`
`fused_metadata_copy_multi`
`fused_nsa_cache_seqlens`
`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`
`fused_metadata_copy_multi`
`fused_dsa_cache_seqlens`
`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`
`fused_qkv_a_proj_with_mqa`
`forward_absorb_fused_mla_rope*` | `python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_fused_rope_cpu.py`
`python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_fused_rope_rocm.py`
`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`
`set_mla_kv_buffer` | `python/sglang/srt/layers/rocm_linear_utils.py`
`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`
`mrope`
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`
`python/sglang/srt/models/qwen3.py`
`python/sglang/srt/models/qwen3_moe.py`
`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`
`fp8 kv cache write`
`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*`
`set_mla_kv_buffer_triton_fp8_quant` | `python/sglang/srt/mem_cache/utils.py`
`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*`
`set_mla_kv_buffer_triton_fp8_quant` | `python/sglang/srt/mem_cache/utils.py`
`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`
`fused_moe_router*`
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`
`moe_fused_gate`
`aiter_fused_topk`
`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`
`fused_append_shared_experts_with_weights`
`num_fused_shared_experts` | `python/sglang/srt/models/qwen2_moe.py`
`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`
`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*`
`npu_dequant_swiglu_quant`
`swiglu_quant` | `python/sglang/srt/layers/moe/ep_moe/kernels.py`
`python/sglang/jit_kernel/nvfp4.py`
`python/sglang/srt/layers/moe/cutlass_w4a8_moe.py`
`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`
`fused_rms_mxfp4_quant`
`fused_flatten_fp8_group_quant`
`fused_flatten_mxfp4_quant` | `python/sglang/srt/layers/communicator.py`
`python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py`
`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`
`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. | -| NSA fused quantize + indexed K-cache store | `fused_store_index_k_cache`
`act_quant`
`index_k_with_scale_buffer` | `python/sglang/jit_kernel/fused_store_index_cache.py`
`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 top-k transform / page-table build | `fast_topk_transform_fused`
`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. | +| DSA fused quantize + indexed K-cache store | `fused_store_index_k_cache`
`act_quant`
`index_k_with_scale_buffer` | `python/sglang/jit_kernel/fused_store_index_cache.py`
`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`
`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`
`final_logit_softcapping` | `python/sglang/srt/layers/elementwise.py`
`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*`
`qkvz_proj`
`ba_proj`
`qkvabz_proj`
`fused_qkvbfg_a_proj` | `python/sglang/jit_kernel/triton/gdn_fused_proj.py`
`python/sglang/srt/models/qwen3_next.py`
`python/sglang/srt/models/qwen3_5.py`
`python/sglang/srt/models/kimi_linear.py`
`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. | | 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. | -| 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`
`_comm_stream`
`dispatch`
`combine`
`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`
`_scatter_stream`
`staging` | `python/sglang/srt/disaggregation/common/staging_handler.py`
`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`
`qwen3_moe.py`
`glm4_moe.py`
`bailing_moe.py`
`llada2.py`
`grok.py`
`olmo2.py`
`step3p5.py`
`longcat_flash.py`
`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 | | --- | --- | --- | --- | --- | | PR `#21877` fused grouped down-GEMM + combine | `grouped_gemm_nt_masked`
`combine`
`fused grouped gemm combine` | `PR #21877`
`python/sglang/srt/layers/moe/ep_moe/flashinfer_cutedsl_moe.py`
`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`
`fp4 kv cache` | `PR #21889`
`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 FP4 paged dequant to FP8 + page-table remap | `_dequant_fp4_to_fp8_paged_kernel`
`WRITE_PT`
`dequant_fp4_paged_decode` | `PR #21889`
`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 BF16 to FP4 quant + paged KV write | `set_mla_kv_buffer_fp4_quant_kernel`
`fp4 kv cache` | `PR #21889`
`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`
`WRITE_PT`
`dequant_fp4_paged_decode` | `PR #21889`
`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`
`trtllm_fp8_block_scale_moe` | `PR #21491`
`python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py`
`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`
`per_token_quant_fp8` | `PR #22005`
`python/sglang/jit_kernel/csrc/elementwise/fused_add_rmsnorm_per_token_quant.cuh`
`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`
`fused_qk_norm_mrope_3d_cache_pts_quant_shuffle`
`rotary_dim` | `PR #20667`
`python/sglang/srt/models/qwen3_5.py`
`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`
`fp8_scaled_mm`
`nvjet`
`cudaMemsetAsync` | `PR #22392`
`sgl-kernel/python/sgl_kernel/gemm.py`
`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`
`nvfp4 expert quant`
`cutlass moe` | `PR #18612`
`python/sglang/srt/layers/moe/cutlass_w4a8_moe.py`
`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`
`trtllm_fp4_block_scale_moe`
`FlashInfer MoE` | `PR #22918`
`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`
`flashinfer_topk`
`pytorch_topk`
`fast_topk_transform` | `PR #22851`
`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 `#24125` GLM5 NSA decode CatArrayBatchedCopy removal | `CatArrayBatchedCopy`
`GLM-5`
`NSA`
`TileLang decode` | `PR #24125`
`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 `#22851` DSA top-k backend and FlashInfer / PyTorch top-k split | `dsa topk`
`flashinfer_topk`
`pytorch_topk`
`fast_topk_transform` | `PR #22851`
`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 DSA decode CatArrayBatchedCopy removal | `CatArrayBatchedCopy`
`GLM-5`
`DSA`
`TileLang decode` | `PR #24125`
`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`
`virtual experts`
`MoE LoRA`
`fused_moe_lora` | `PR #24007`
`python/sglang/srt/layers/lora_backend.py`
`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`
`local compile`
`decode compile`
`torchinductor` | `PR #24150`
`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_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_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_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. | diff --git a/.claude/skills/llm-torch-profiler-analysis/references/overlap-catalog.md b/.claude/skills/llm-torch-profiler-analysis/references/overlap-catalog.md index 5a38cb204..6647869b8 100644 --- a/.claude/skills/llm-torch-profiler-analysis/references/overlap-catalog.md +++ b/.claude/skills/llm-torch-profiler-analysis/references/overlap-catalog.md @@ -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. | | 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. | -| 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`
`_comm_stream`
`dispatch`
`combine`
`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`
`qwen3_moe.py`
`glm4_moe.py`
`bailing_moe.py`
`llada2.py`
`grok.py`
`olmo2.py`
`step3p5.py`
`longcat_flash.py`
`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. | diff --git a/.claude/skills/llm-torch-profiler-analysis/scripts/triage_kernel_helpers.py b/.claude/skills/llm-torch-profiler-analysis/scripts/triage_kernel_helpers.py index ca1d47d18..8352f47bf 100644 --- a/.claude/skills/llm-torch-profiler-analysis/scripts/triage_kernel_helpers.py +++ b/.claude/skills/llm-torch-profiler-analysis/scripts/triage_kernel_helpers.py @@ -530,16 +530,16 @@ FUSION_PATTERN_REGISTRY: Tuple[FusionPatternSpec, ...] = ( likely_share=0.5, ), 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", active_keywords=( "fused_metadata_copy", "fused_metadata_copy_multi", - "fused_nsa_cache_seqlens", + "fused_dsa_cache_seqlens", "fused_flashmla_metadata", ), 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, likely_share=0.2, @@ -744,23 +744,23 @@ FUSION_PATTERN_REGISTRY: Tuple[FusionPatternSpec, ...] = ( likely_share=1.5, ), FusionPatternSpec( - pattern="NSA fused top-k transform / page-table build", - candidate_path="python/sglang/srt/layers/attention/nsa_backend.py", + pattern="DSA fused top-k transform / page-table build", + candidate_path="python/sglang/srt/layers/attention/dsa_backend.py", active_keywords=( "fast_topk_transform_fused", "fast_topk_transform_ragged_fused", ), 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, likely_share=0.3, ), FusionPatternSpec( - pattern="NSA fused quantize + indexed K-cache store", + pattern="DSA fused quantize + indexed K-cache store", candidate_path=( "python/sglang/jit_kernel/fused_store_index_cache.py" - "
python/sglang/srt/layers/attention/nsa/nsa_indexer.py" + "
python/sglang/srt/layers/attention/dsa/dsa_indexer.py" ), active_keywords=("fused_store_index_k_cache",), split_groups=( @@ -768,7 +768,7 @@ FUSION_PATTERN_REGISTRY: Tuple[FusionPatternSpec, ...] = ( ("index_k", "cache", "store"), ), 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, likely_share=1.0, diff --git a/.github/workflows/nightly-test-amd-rocm720.yml b/.github/workflows/nightly-test-amd-rocm720.yml index 9f5a6fa8f..cab07292b 100644 --- a/.github/workflows/nightly-test-amd-rocm720.yml +++ b/.github/workflows/nightly-test-amd-rocm720.yml @@ -742,7 +742,7 @@ jobs: 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 - - 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 run: | > 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 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 run: | > github_summary.md # Clear summary file diff --git a/.github/workflows/nightly-test-amd.yml b/.github/workflows/nightly-test-amd.yml index 201cd13cc..01543b668 100644 --- a/.github/workflows/nightly-test-amd.yml +++ b/.github/workflows/nightly-test-amd.yml @@ -744,7 +744,7 @@ jobs: 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 - - name: Accuracy Test (8-GPU GLM-5.1 NSA) + - name: Accuracy Test (8-GPU GLM-5.1 DSA) timeout-minutes: 120 run: | > 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 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 run: | > github_summary.md # Clear summary file diff --git a/docs/advanced_features/attention_backend.md b/docs/advanced_features/attention_backend.md index 98d07d31a..b6b604c28 100644 --- a/docs/advanced_features/attention_backend.md +++ b/docs/advanced_features/attention_backend.md @@ -50,7 +50,7 @@ Multimodal attention is selected by `--mm-attention-backend`. The "MultiModal" c ```{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. -- 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} @@ -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. ``` -### 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** | |-----------------------|-------------|------------|-----------------------------------------------| diff --git a/docs/advanced_features/hisparse_guide.md b/docs/advanced_features/hisparse_guide.md index 57aa5e7c2..941318c86 100644 --- a/docs/advanced_features/hisparse_guide.md +++ b/docs/advanced_features/hisparse_guide.md @@ -90,7 +90,7 @@ python3 -m sglang.launch_server \ --tp-size 8 --dp-size 8 --enable-dp-attention \ --mem-fraction-static 0.85 \ --kv-cache-dtype bfloat16 \ - --nsa-decode-backend flashmla_sparse \ + --dsa-decode-backend flashmla_sparse \ --disaggregation-mode decode \ --disaggregation-ib-device mlx5_0,mlx5_1,mlx5_2,mlx5_3 \ --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. - 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). - - `--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. - `--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: diff --git a/docs/advanced_features/server_arguments.md b/docs/advanced_features/server_arguments.md index 072679b5e..a880f518c 100644 --- a/docs/advanced_features/server_arguments.md +++ b/docs/advanced_features/server_arguments.md @@ -265,14 +265,14 @@ Please consult the documentation below and [server_args.py](https://github.com/s ## Kernel Backends (Attention, Sampling, Grammar, GEMM) | 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` | -| `--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` | -| `--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` | +| `--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`, `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`, `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` | | `--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` | -| `--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` | -| `--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-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` | +| `--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` | | `--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) | @@ -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) | | `--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-nsa-prefill-context-parallel` | Enable context parallelism used in the long sequence prefill phase of DeepSeek v3.2. | `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` | +| `--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) | +| `--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-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 | | `--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 | -| `--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-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-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` | 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` | diff --git a/docs/basic_usage/deepseek_v32.md b/docs/basic_usage/deepseek_v32.md index 095060a7f..7bd0a447b 100644 --- a/docs/basic_usage/deepseek_v32.md +++ b/docs/basic_usage/deepseek_v32.md @@ -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 # 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`. @@ -61,9 +61,9 @@ To serve GLM-5, just replace the `--model` argument with `zai-org/GLM-5-FP8`. ### Configuration Tips - **DP Attention**: To enable [DP Attention](../advanced_features/dp_dpa_smg_guide.md), please include `--enable-dp-attention --dp ` 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. -- **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. -- **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. -- **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: +- **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_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 `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_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. @@ -319,11 +319,11 @@ DeepSeek-V3.2-Speciale: **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 -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: - 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: ```bash # 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) -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. @@ -349,7 +349,7 @@ For more details, please refer to PR https://github.com/sgl-project/sglang/pull/ Example usage: ```bash # 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) @@ -372,9 +372,9 @@ python3 -m sglang.launch_server \ --dist-init-addr :62001 \ --tp 8 --pp-size 2 \ --dp-size 1 --moe-dense-tp-size 1 \ - --enable-nsa-prefill-context-parallel \ + --enable-dsa-prefill-context-parallel \ --attn-cp-size 8 \ - --nsa-prefill-cp-mode round-robin-split \ + --dsa-prefill-cp-mode round-robin-split \ --trust-remote-code \ --disable-radix-cache \ --mem-fraction-static 0.8 \ @@ -396,9 +396,9 @@ python3 -m sglang.launch_server \ --dist-init-addr :62001 \ --tp 8 --pp-size 2 \ --dp-size 1 --moe-dense-tp-size 1 \ - --enable-nsa-prefill-context-parallel \ + --enable-dsa-prefill-context-parallel \ --attn-cp-size 8 \ - --nsa-prefill-cp-mode round-robin-split \ + --dsa-prefill-cp-mode round-robin-split \ --trust-remote-code \ --disable-radix-cache \ --mem-fraction-static 0.8 \ @@ -424,9 +424,9 @@ python -m sglang.launch_server \ --dist-init-addr :20102 \ --tp 8 --pp-size 2 \ --dp-size 1 --moe-dense-tp-size 1 \ - --enable-nsa-prefill-context-parallel \ + --enable-dsa-prefill-context-parallel \ --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 \ --trust-remote-code \ --disable-radix-cache \ @@ -450,9 +450,9 @@ python -m sglang.launch_server \ --dist-init-addr :20102 \ --tp 8 --pp-size 2 \ --dp-size 1 --moe-dense-tp-size 1 \ - --enable-nsa-prefill-context-parallel \ + --enable-dsa-prefill-context-parallel \ --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 \ --trust-remote-code \ --disable-radix-cache \ diff --git a/docs/platforms/ascend/ascend_npu_best_practice.md b/docs/platforms/ascend/ascend_npu_best_practice.md index 4ec6d3a8c..927941af8 100644 --- a/docs/platforms/ascend/ascend_npu_best_practice.md +++ b/docs/platforms/ascend/ascend_npu_best_practice.md @@ -1113,8 +1113,8 @@ do --nnodes 2 --node-rank $i \ --disaggregation-bootstrap-port 8995 \ --moe-dense-tp-size 1 \ - --enable-nsa-prefill-context-parallel \ - --nsa-prefill-cp-mode in-seq-split \ + --enable-dsa-prefill-context-parallel \ + --dsa-prefill-cp-mode in-seq-split \ --attn-cp-size 32 \ --speculative-algorithm NEXTN --speculative-num-steps 1 --speculative-eagle-topk 1 --speculative-num-draft-tokens 2 \ --dist-init-addr ${P_IP[0]}:10000 diff --git a/docs/platforms/ascend/ascend_npu_environment_variables.md b/docs/platforms/ascend/ascend_npu_environment_variables.md index cfb1bb0b8..c638298b5 100644 --- a/docs/platforms/ascend/ascend_npu_environment_variables.md +++ b/docs/platforms/ascend/ascend_npu_environment_variables.md @@ -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
preprocessing stage of the MLA model. | `false` | | `SGLANG_USE_FIA_NZ` | Reshapes KV Cache for FIA NZ format.
`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
and routing experts in DeepSeek models.
Enable dual-stream computation in DeepSeek NSA Indexer. | `false` | +| `SGLANG_NPU_USE_MULTI_STREAM` | Enable dual-stream computation of shared experts
and routing experts in DeepSeek models.
Enable dual-stream computation in DeepSeek DSA Indexer. | `false` | | `SGLANG_NPU_DISABLE_ACL_FORMAT_WEIGHT` | Disable cast model weight tensor to a specific NPU
ACL format. | `false` | | `SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK` | The maximum number of dispatched tokens on each rank. | `128` | diff --git a/docs/platforms/ascend/ascend_npu_support_features.md b/docs/platforms/ascend/ascend_npu_support_features.md index 729702ed6..a28c51add 100644 --- a/docs/platforms/ascend/ascend_npu_support_features.md +++ b/docs/platforms/ascend/ascend_npu_support_features.md @@ -204,8 +204,8 @@ click [Server Arguments](https://docs.sglang.io/advanced_features/server_argumen | `--sampling-backend` | `None` | `pytorch`,
`ascend` | A2, A3 | | `--grammar-backend` | `None` | `xgrammar` | A2, A3 | | `--mm-attention-backend` | `None` | `ascend_attn` | A2, A3 | -| `--nsa-prefill-backend` | `flashmla_sparse` | `flashmla_sparse`,
`flashmla_decode`,
`fa3`,
`tilelang`,
`aiter` | Special for GPU | -| `--nsa-decode-backend` | `fa3` | `flashmla_prefill`,
`flashmla_kv`,
`fa3`,
`tilelang`,
`aiter` | Special for GPU | +| `--dsa-prefill-backend` | `flashmla_sparse` | `flashmla_sparse`,
`flashmla_decode`,
`fa3`,
`tilelang`,
`aiter` | Special for GPU | +| `--dsa-decode-backend` | `fa3` | `flashmla_prefill`,
`flashmla_kv`,
`fa3`,
`tilelang`,
`aiter` | Special for GPU | | `--fp8-gemm-backend` | `auto` | `auto`,
`deep_gemm`,
`flashinfer_trtllm`,
`flashinfer_cutlass`,
`flashinfer_deepgemm`,
`cutlass`,
`triton`,
`aiter` | Special for GPU | | `--disable-flashinfer-`
`autotune` | `False` | bool flag
(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 | | `--enable-layerwise-`
`nvtx-marker` | `False` | bool flag
(set to enable) | Special for GPU | | `--enable-attn-tp-`
`input-scattered` | `False` | bool flag
(set to enable) | Experimental | -| `--enable-nsa-prefill-`
`context-parallel` | `False` | bool flag
(set to enable) | A2, A3 | +| `--enable-dsa-prefill-`
`context-parallel` | `False` | bool flag
(set to enable) | A2, A3 | | `--enable-fused-qk-`
`norm-rope` | `False` | bool flag
(set to enable) | Special for GPU | ## Dynamic batch tokenizer diff --git a/docs/platforms/plugin.md b/docs/platforms/plugin.md index 8a4c4ee1c..a956ebffa 100644 --- a/docs/platforms/plugin.md +++ b/docs/platforms/plugin.md @@ -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_mha_kv_pool_cls()` | `raise NotImplementedError` | MHA 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_piecewise_backend_cls()` | `raise NotImplementedError` | Piecewise compilation backend class | | `get_compile_backend(mode)` | `"inductor"` | Compilation backend string | diff --git a/docs/references/environment_variables.md b/docs/references/environment_variables.md index 99f65c7c5..87e085880 100644 --- a/docs/references/environment_variables.md +++ b/docs/references/environment_variables.md @@ -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_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 | Description | Default Value | | --- | --- | --- | -| `SGLANG_NSA_FUSE_TOPK` | Fuse the operation of picking topk logits and picking topk indices from page table | `true` | -| `SGLANG_NSA_ENABLE_MTP_PRECOMPUTE_METADATA` | Precompute metadata that can be shared among different draft steps when MTP is enabled | `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_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_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 diff --git a/docs_new/cookbook/autoregressive/DeepSeek/DeepSeek-V3_2.mdx b/docs_new/cookbook/autoregressive/DeepSeek/DeepSeek-V3_2.mdx index e5c44234a..8ec0ca505 100644 --- a/docs_new/cookbook/autoregressive/DeepSeek/DeepSeek-V3_2.mdx +++ b/docs_new/cookbook/autoregressive/DeepSeek/DeepSeek-V3_2.mdx @@ -349,11 +349,11 @@ sglang_args=$(echo serve \ --moe-a2a-backend deepep --ep-size 16 \ --page-size 128 \ --chunked-prefill-size 16384 \ - --attention-backend nsa \ - --nsa-prefill-backend flashmla_sparse \ - --nsa-decode-backend flashmla_sparse \ - --enable-nsa-prefill-context-parallel \ - --nsa-prefill-cp-mode round-robin-split \ + --attention-backend dsa \ + --dsa-prefill-backend flashmla_sparse \ + --dsa-decode-backend flashmla_sparse \ + --enable-dsa-prefill-context-parallel \ + --dsa-prefill-cp-mode round-robin-split \ --cuda-graph-max-bs 128 \ --max-running-requests 128 \ --trust-remote-code --host "0.0.0.0" --port 30000 \ @@ -361,7 +361,7 @@ sglang_args=$(echo serve \ --context-length 65536 \ --allow-auto-truncate --enable-metrics \ --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) @@ -382,20 +382,20 @@ dp_config=" \ " cp_config=" \ - --enable-nsa-prefill-context-parallel \ + --enable-dsa-prefill-context-parallel \ " if [ "$dp" -eq 1 ]; then cp_config=" \ $cp_config \ - --nsa-prefill-cp-mode round-robin-split \ + --dsa-prefill-cp-mode round-robin-split \ " else cp_config=" \ $cp_config \ - --nsa-prefill-cp-mode in-seq-split \ + --dsa-prefill-cp-mode in-seq-split \ " fi diff --git a/docs_new/cookbook/autoregressive/GLM/GLM-5.1.mdx b/docs_new/cookbook/autoregressive/GLM/GLM-5.1.mdx index cd4083561..f16407106 100644 --- a/docs_new/cookbook/autoregressive/GLM/GLM-5.1.mdx +++ b/docs_new/cookbook/autoregressive/GLM/GLM-5.1.mdx @@ -79,7 +79,7 @@ import { GLM51Deployment } from '/src/snippets/autoregressive/glm-51-deployment. -- **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`. - 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. @@ -116,8 +116,8 @@ sglang serve \ --trust-remote-code \ --tool-call-parser glm47 \ --reasoning-parser glm45 \ - --nsa-prefill-backend tilelang \ - --nsa-decode-backend tilelang \ + --dsa-prefill-backend tilelang \ + --dsa-decode-backend tilelang \ --chunked-prefill-size 131072 \ --mem-fraction-static 0.80 \ --watchdog-timeout 1200 \ @@ -132,8 +132,8 @@ sglang serve \ --model-path zai-org/GLM-5.1 \ --tp 8 \ --trust-remote-code \ - --nsa-prefill-backend tilelang \ - --nsa-decode-backend tilelang \ + --dsa-prefill-backend tilelang \ + --dsa-decode-backend tilelang \ --chunked-prefill-size 131072 \ --mem-fraction-static 0.80 \ --watchdog-timeout 1200 \ @@ -627,7 +627,7 @@ Average accuracy: 0.877 #### 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 python3 benchmark/gsm8k/bench_sglang.py --num-questions 200 diff --git a/docs_new/cookbook/autoregressive/GLM/GLM-5.mdx b/docs_new/cookbook/autoregressive/GLM/GLM-5.mdx index b937d71a8..e1d1f2d3c 100644 --- a/docs_new/cookbook/autoregressive/GLM/GLM-5.mdx +++ b/docs_new/cookbook/autoregressive/GLM/GLM-5.mdx @@ -85,9 +85,9 @@ import { GLM5Deployment } from '/src/snippets/autoregressive/glm-5-deployment.js -- **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...). - 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 \ --tp 8 \ --trust-remote-code \ - --nsa-prefill-backend tilelang \ - --nsa-decode-backend tilelang \ + --dsa-prefill-backend tilelang \ + --dsa-decode-backend tilelang \ --chunked-prefill-size 131072 \ --mem-fraction-static 0.80 \ --watchdog-timeout 1200 \ @@ -653,7 +653,7 @@ Average accuracy: 0.877 #### 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 python3 benchmark/gsm8k/bench_sglang.py --num-questions 200 diff --git a/docs_new/docs/advanced_features/attention_backend.mdx b/docs_new/docs/advanced_features/attention_backend.mdx index 9ad998f51..d9fea017c 100644 --- a/docs_new/docs/advanced_features/attention_backend.mdx +++ b/docs_new/docs/advanced_features/attention_backend.mdx @@ -280,7 +280,7 @@ Multimodal attention is selected by `--mm-attention-backend`. The "MultiModal" c - 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. @@ -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. -### 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`: diff --git a/docs_new/docs/advanced_features/hisparse_guide.mdx b/docs_new/docs/advanced_features/hisparse_guide.mdx index 9ec2e082b..2c62b0b0a 100644 --- a/docs_new/docs/advanced_features/hisparse_guide.mdx +++ b/docs_new/docs/advanced_features/hisparse_guide.mdx @@ -142,7 +142,7 @@ python3 -m sglang.launch_server \ --tp-size 8 --dp-size 8 --enable-dp-attention \ --mem-fraction-static 0.85 \ --kv-cache-dtype bfloat16 \ - --nsa-decode-backend flashmla_sparse \ + --dsa-decode-backend flashmla_sparse \ --disaggregation-mode decode \ --disaggregation-ib-device mlx5_0,mlx5_1,mlx5_2,mlx5_3 \ --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. - 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). - - `--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. - `--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: diff --git a/docs_new/docs/advanced_features/server_arguments.mdx b/docs_new/docs/advanced_features/server_arguments.mdx index 134a99123..d40f0dce0 100644 --- a/docs_new/docs/advanced_features/server_arguments.mdx +++ b/docs_new/docs/advanced_features/server_arguments.mdx @@ -1162,19 +1162,19 @@ Please consult the documentation below and [server_args.py](https://github.com/s - + - + - + @@ -1195,14 +1195,14 @@ Please consult the documentation below and [server_args.py](https://github.com/s - - + + - - + + @@ -2278,13 +2278,13 @@ Please consult the documentation below and [server_args.py](https://github.com/s - + - + @@ -2841,13 +2841,13 @@ Please consult the documentation below and [server_args.py](https://github.com/s - + - + diff --git a/docs_new/docs/basic_usage/deepseek_v32.mdx b/docs_new/docs/basic_usage/deepseek_v32.mdx index 1077a9956..a1cbf0d9c 100644 --- a/docs_new/docs/basic_usage/deepseek_v32.mdx +++ b/docs_new/docs/basic_usage/deepseek_v32.mdx @@ -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 # 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`. @@ -64,9 +64,9 @@ To serve GLM-5, just replace the `--model` argument with `zai-org/GLM-5-FP8`. ### Configuration Tips - **DP Attention**: To enable [DP Attention](../advanced_features/dp_dpa_smg_guide), please include `--enable-dp-attention --dp ` 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. -- **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. -- **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. -- **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: +- **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_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 `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_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. @@ -444,11 +444,11 @@ DeepSeek-V3.2-Speciale: **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 -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: - 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: ```bash Command # 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) -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. @@ -474,7 +474,7 @@ For more details, please refer to PR https://github.com/sgl-project/sglang/pull/ Example usage: ```bash Command # 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) @@ -497,9 +497,9 @@ python3 -m sglang.launch_server \ --dist-init-addr :62001 \ --tp 8 --pp-size 2 \ --dp-size 1 --moe-dense-tp-size 1 \ - --enable-nsa-prefill-context-parallel \ + --enable-dsa-prefill-context-parallel \ --attn-cp-size 8 \ - --nsa-prefill-cp-mode round-robin-split \ + --dsa-prefill-cp-mode round-robin-split \ --trust-remote-code \ --disable-radix-cache \ --mem-fraction-static 0.8 \ @@ -521,9 +521,9 @@ python3 -m sglang.launch_server \ --dist-init-addr :62001 \ --tp 8 --pp-size 2 \ --dp-size 1 --moe-dense-tp-size 1 \ - --enable-nsa-prefill-context-parallel \ + --enable-dsa-prefill-context-parallel \ --attn-cp-size 8 \ - --nsa-prefill-cp-mode round-robin-split \ + --dsa-prefill-cp-mode round-robin-split \ --trust-remote-code \ --disable-radix-cache \ --mem-fraction-static 0.8 \ @@ -549,9 +549,9 @@ python -m sglang.launch_server \ --dist-init-addr :20102 \ --tp 8 --pp-size 2 \ --dp-size 1 --moe-dense-tp-size 1 \ - --enable-nsa-prefill-context-parallel \ + --enable-dsa-prefill-context-parallel \ --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 \ --trust-remote-code \ --disable-radix-cache \ @@ -575,9 +575,9 @@ python -m sglang.launch_server \ --dist-init-addr :20102 \ --tp 8 --pp-size 2 \ --dp-size 1 --moe-dense-tp-size 1 \ - --enable-nsa-prefill-context-parallel \ + --enable-dsa-prefill-context-parallel \ --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 \ --trust-remote-code \ --disable-radix-cache \ diff --git a/docs_new/docs/hardware-platforms/ascend-npus/ascend_npu_best_practice.mdx b/docs_new/docs/hardware-platforms/ascend-npus/ascend_npu_best_practice.mdx index 73e2377a1..2617008cb 100644 --- a/docs_new/docs/hardware-platforms/ascend-npus/ascend_npu_best_practice.mdx +++ b/docs_new/docs/hardware-platforms/ascend-npus/ascend_npu_best_practice.mdx @@ -1598,8 +1598,8 @@ do --nnodes 2 --node-rank $i \ --disaggregation-bootstrap-port 8995 \ --moe-dense-tp-size 1 \ - --enable-nsa-prefill-context-parallel \ - --nsa-prefill-cp-mode in-seq-split \ + --enable-dsa-prefill-context-parallel \ + --dsa-prefill-cp-mode in-seq-split \ --attn-cp-size 32 \ --speculative-algorithm NEXTN --speculative-num-steps 1 --speculative-eagle-topk 1 --speculative-num-draft-tokens 2 \ --dist-init-addr ${P_IP[0]}:10000 diff --git a/docs_new/docs/hardware-platforms/ascend-npus/ascend_npu_environment_variables.mdx b/docs_new/docs/hardware-platforms/ascend-npus/ascend_npu_environment_variables.mdx index f403f19b1..4a88e7052 100644 --- a/docs_new/docs/hardware-platforms/ascend-npus/ascend_npu_environment_variables.mdx +++ b/docs_new/docs/hardware-platforms/ascend-npus/ascend_npu_environment_variables.mdx @@ -29,7 +29,7 @@ This document provides a list of commonly used environment variables and aims to - + diff --git a/docs_new/docs/hardware-platforms/ascend-npus/ascend_npu_optimization.mdx b/docs_new/docs/hardware-platforms/ascend-npus/ascend_npu_optimization.mdx index 2e6a3e3b6..2f5bdd0f1 100644 --- a/docs_new/docs/hardware-platforms/ascend-npus/ascend_npu_optimization.mdx +++ b/docs_new/docs/hardware-platforms/ascend-npus/ascend_npu_optimization.mdx @@ -187,18 +187,18 @@ These arguments and environment variables are critical for tuning prefill perfor - - + + - - + + - + @@ -360,7 +360,7 @@ The following environment variables are used in other best practice configuratio - + diff --git a/docs_new/docs/hardware-platforms/ascend-npus/ascend_npu_support_features.mdx b/docs_new/docs/hardware-platforms/ascend-npus/ascend_npu_support_features.mdx index e4b3143d1..701ad2688 100644 --- a/docs_new/docs/hardware-platforms/ascend-npus/ascend_npu_support_features.mdx +++ b/docs_new/docs/hardware-platforms/ascend-npus/ascend_npu_support_features.mdx @@ -1225,13 +1225,13 @@ click [Server Arguments](../../advanced_features/server_arguments). - + - + @@ -2226,7 +2226,7 @@ click [Server Arguments](../../advanced_features/server_arguments). - + diff --git a/docs_new/docs/hardware-platforms/plugin.mdx b/docs_new/docs/hardware-platforms/plugin.mdx index 676eb0450..3376ef953 100644 --- a/docs_new/docs/hardware-platforms/plugin.mdx +++ b/docs_new/docs/hardware-platforms/plugin.mdx @@ -512,9 +512,9 @@ python -c "from sglang.srt.platforms import current_platform; print(current_plat - + - + diff --git a/docs_new/docs/references/environment_variables.mdx b/docs_new/docs/references/environment_variables.mdx index af52bd385..61dae0e07 100644 --- a/docs_new/docs/references/environment_variables.mdx +++ b/docs_new/docs/references/environment_variables.mdx @@ -393,7 +393,7 @@ SGLang supports various environment variables that can be used to configure its
`--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, ascendtriton, torch_native, flex_attention, dsa, 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, ascendtriton, torch_native, flex_attention, dsa, 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, ascendtriton, torch_native, flex_attention, dsa, cutlass_mla, fa3, fa4, flashinfer, flashmla, trtllm_mla, trtllm_mha, dual_chunk_flash_attn, aiter, wave, intel_amx, ascend
`--sampling-backend`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).`--dsa-prefill-backend`DSA backend for the prefill stage (overrides `--attention-backend` when running DeepSeek DSA-style attention). `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.`--dsa-decode-backend`DSA backend for the decode stage when running DeepSeek DSA-style attention. Overrides `--attention-backend` for decoding. `fa3` flashmla_sparse, flashmla_kv, fa3, tilelang, aiter, trtllm
bool flag (set to enable)
--enable-nsa-prefill-context-parallel--enable-dsa-prefill-context-parallel Enable context parallelism used in the long sequence prefill phase of DeepSeek v3.2. False bool flag (set to enable)
--nsa-prefill-cp-mode--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. in-seq-split in-seq-split, round-robin-split
`--nsa-prefill`Choose the NSA backend for the prefill stage (overrides `--attention-backend` when running DeepSeek NSA-style attention).Deprecated alias for `--dsa-prefill-backend`. `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.Deprecated alias for `--dsa-decode-backend`. `flashmla_kv` `flashmla_prefill`, `flashmla_kv`, `fa3`, `tilelang`, `aiter`
SGLANG_NPU_USE_MULTI_STREAMEnable dual-stream computation of shared experts
and routing experts in DeepSeek models.
Enable dual-stream computation in DeepSeek NSA Indexer.		Enable dual-stream computation of shared experts
and routing experts in DeepSeek models.
Enable dual-stream computation in DeepSeek DSA Indexer.		 false
true
`--enable-nsa-prefill-context-parallel`(DeepSeek V3.2 NSA-specific) 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.`--enable-dsa-prefill-context-parallel`(DeepSeek V3.2 DSA-specific) 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. Enabled
`--nsa-prefill-cp-mode`(DeepSeek V3.2 NSA-specific) 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.`--dsa-prefill-cp-mode`(DeepSeek V3.2 DSA-specific) 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. `in-seq-split`
`--attn-cp-size`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.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. `32`
`HCCL_OP_EXPANSION_MODE=AIV` Configures the HCCL communication algorithm scheduling to use AIV (Ascend Intelligent Vision) expansion mode, which can improve communication efficiency for certain collective operations.Used in Qwen MoE and R1 non-NSA configurationsUsed in Qwen MoE and R1 non-DSA configurations
`SGLANG_NPU_FUSED_MOE_MODE`A2, A3
`--nsa-prefill-backend``--dsa-prefill-backend` `flashmla_sparse` flashmla_sparse,
flashmla_decode,
fa3,
tilelang,
aiter		 Special for GPU
`--nsa-decode-backend``--dsa-decode-backend` `fa3` flashmla_prefill,
flashmla_kv,
fa3,
tilelang,
aiter		 Special for GPUExperimental
`--enable-nsa-prefill-`
`context-parallel`		`--enable-dsa-prefill-`
`context-parallel`		 `False` bool flag
(set to enable)		 A2, A3MLA KV cache pool class
get_nsa_kv_pool_cls()get_dsa_kv_pool_cls() raise NotImplementedErrorNSA KV cache pool class (DeepSeek V3.2)DSA KV cache pool class (DeepSeek V3.2)
get_paged_allocator_cls()
-## 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 */} @@ -412,13 +412,13 @@ SGLang supports various environment variables that can be used to configure its - SGLANG_NSA_FUSE_TOPK - Fuse the operation of picking topk logits and picking topk indices from page table + 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 + 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 @@ -427,8 +427,8 @@ SGLang supports various environment variables that can be used to configure its 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) + 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 diff --git a/docs_new/src/snippets/autoregressive/deepseek-v32-deployment.jsx b/docs_new/src/snippets/autoregressive/deepseek-v32-deployment.jsx index 641afd385..ea687f5a7 100644 --- a/docs_new/src/snippets/autoregressive/deepseek-v32-deployment.jsx +++ b/docs_new/src/snippets/autoregressive/deepseek-v32-deployment.jsx @@ -196,8 +196,8 @@ export const DeepSeekV32Deployment = () => { // Hardware platform specific parameters if (isAMD) { cmd += ' \\\n --trust-remote-code'; - cmd += ' \\\n --nsa-prefill-backend tilelang'; - cmd += ' \\\n --nsa-decode-backend tilelang'; + cmd += ' \\\n --dsa-prefill-backend tilelang'; + cmd += ' \\\n --dsa-decode-backend tilelang'; cmd += ' \\\n --cuda-graph-max-bs 64'; } diff --git a/docs_new/src/snippets/autoregressive/deepseek-v4-deployment.jsx b/docs_new/src/snippets/autoregressive/deepseek-v4-deployment.jsx index 6d1a07e6e..455fcd14a 100644 --- a/docs_new/src/snippets/autoregressive/deepseek-v4-deployment.jsx +++ b/docs_new/src/snippets/autoregressive/deepseek-v4-deployment.jsx @@ -650,8 +650,8 @@ export const DeepSeekV4Deployment = () => { flags.push(` --tp ${tp}`); if (multinode) flags.push(...multiNodeFlags(nnodes)); flags.push(" --moe-a2a-backend deepep"); - flags.push(" --enable-nsa-prefill-context-parallel"); - flags.push(" --nsa-prefill-cp-mode round-robin-split"); + flags.push(" --enable-dsa-prefill-context-parallel"); + flags.push(" --dsa-prefill-cp-mode round-robin-split"); flags.push(" --chunked-prefill-size 16384"); // 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 diff --git a/docs_new/src/snippets/autoregressive/glm-5-deployment.jsx b/docs_new/src/snippets/autoregressive/glm-5-deployment.jsx index e7f8c9550..16188cfba 100644 --- a/docs_new/src/snippets/autoregressive/glm-5-deployment.jsx +++ b/docs_new/src/snippets/autoregressive/glm-5-deployment.jsx @@ -155,13 +155,13 @@ export const GLM5Deployment = () => { cmd += ` --model-path ${modelName}`; 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) { cmd += ' \\\n --trust-remote-code'; cmd += ' \\\n --quantization modelopt_fp4'; cmd += ' \\\n --kv-cache-dtype fp8_e4m3'; - cmd += ' \\\n --nsa-decode-backend trtllm'; - cmd += ' \\\n --nsa-prefill-backend trtllm'; + cmd += ' \\\n --dsa-decode-backend trtllm'; + cmd += ' \\\n --dsa-prefill-backend trtllm'; cmd += ' \\\n --moe-runner-backend flashinfer_trtllm'; cmd += ' \\\n --enable-flashinfer-allreduce-fusion'; cmd += ' \\\n --enable-dp-lm-head'; @@ -174,11 +174,11 @@ export const GLM5Deployment = () => { return cmd; } - // AMD-specific: NSA tilelang backend. + // AMD-specific: DSA tilelang backend. if (isAMD) { cmd += ' \\\n --trust-remote-code'; - cmd += ' \\\n --nsa-prefill-backend tilelang'; - cmd += ' \\\n --nsa-decode-backend tilelang'; + cmd += ' \\\n --dsa-prefill-backend tilelang'; + cmd += ' \\\n --dsa-decode-backend tilelang'; cmd += ' \\\n --chunked-prefill-size 131072'; cmd += ' \\\n --watchdog-timeout 1200'; } @@ -199,9 +199,9 @@ export const GLM5Deployment = () => { if (hardware === 'b200' && effectiveQuant === 'fp8') { cmd += ' \\\n --ep 1'; cmd += ' \\\n --quantization fp8'; - cmd += ' \\\n --attention-backend nsa'; - cmd += ' \\\n --nsa-decode-backend trtllm'; - cmd += ' \\\n --nsa-prefill-backend trtllm'; + cmd += ' \\\n --attention-backend dsa'; + cmd += ' \\\n --dsa-decode-backend trtllm'; + cmd += ' \\\n --dsa-prefill-backend trtllm'; cmd += ' \\\n --moe-runner-backend flashinfer_trtllm'; cmd += ' \\\n --enable-flashinfer-allreduce-fusion'; } diff --git a/docs_new/src/snippets/autoregressive/glm-51-deployment.jsx b/docs_new/src/snippets/autoregressive/glm-51-deployment.jsx index fd2734362..314318ba5 100644 --- a/docs_new/src/snippets/autoregressive/glm-51-deployment.jsx +++ b/docs_new/src/snippets/autoregressive/glm-51-deployment.jsx @@ -152,8 +152,8 @@ export const GLM51Deployment = () => { if (isAMD) { cmd += ' \\\n --trust-remote-code'; - cmd += ' \\\n --nsa-prefill-backend tilelang'; - cmd += ' \\\n --nsa-decode-backend tilelang'; + cmd += ' \\\n --dsa-prefill-backend tilelang'; + cmd += ' \\\n --dsa-decode-backend tilelang'; cmd += ' \\\n --chunked-prefill-size 131072'; cmd += ' \\\n --watchdog-timeout 1200'; } diff --git a/python/sglang/jit_kernel/csrc/nsa/fused_store_index_cache.cuh b/python/sglang/jit_kernel/csrc/dsa/fused_store_index_cache.cuh similarity index 100% rename from python/sglang/jit_kernel/csrc/nsa/fused_store_index_cache.cuh rename to python/sglang/jit_kernel/csrc/dsa/fused_store_index_cache.cuh diff --git a/python/sglang/jit_kernel/csrc/elementwise/fused_metadata_copy.cuh b/python/sglang/jit_kernel/csrc/elementwise/fused_metadata_copy.cuh index c996f6f1b..7e1f91dcf 100644 --- a/python/sglang/jit_kernel/csrc/elementwise/fused_metadata_copy.cuh +++ b/python/sglang/jit_kernel/csrc/elementwise/fused_metadata_copy.cuh @@ -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. * * OVERVIEW: * 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 * graph replay performance during inference. * @@ -37,7 +37,7 @@ #include // for std::min #include -// 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 }; /** @@ -49,9 +49,9 @@ struct SourcePointers { 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__ 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__ 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__ flashmla_num_splits; // optional FlashMLA split counts 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__ cu_seqlens_k; // [bs+1] cumulative sequence lengths 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__ 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__ flashmla_num_splits; // optional FlashMLA split counts 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 #pragma unroll 8 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 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 #pragma unroll 4 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 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; } - // Copy nsa_cache_seqlens to all 3 backends + // Copy dsa_cache_seqlens to all 3 backends #pragma unroll 8 for (int i = tid; i < bs; i += total_threads) { - int32_t val = src.nsa_cache_seqlens[i]; - dst0.nsa_cache_seqlens[i] = val; - dst1.nsa_cache_seqlens[i] = val; - dst2.nsa_cache_seqlens[i] = val; + int32_t val = src.dsa_cache_seqlens[i]; + dst0.dsa_cache_seqlens[i] = val; + dst1.dsa_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 for (int i = tid; i < bs; i += total_threads) { - int32_t val = src.nsa_cu_seqlens_k[i + 1]; - dst0.nsa_cu_seqlens_k[i + 1] = val; - dst1.nsa_cu_seqlens_k[i + 1] = val; - dst2.nsa_cu_seqlens_k[i + 1] = val; + int32_t val = src.dsa_cu_seqlens_k[i + 1]; + dst0.dsa_cu_seqlens_k[i + 1] = val; + dst1.dsa_cu_seqlens_k[i + 1] = val; + dst2.dsa_cu_seqlens_k[i + 1] = val; } // Copy real page table to all 3 backends @@ -493,18 +493,18 @@ struct FusedMetadataCopyKernel { run(const tvm::ffi::TensorView cache_seqlens_src, const tvm::ffi::TensorView cu_seqlens_k_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 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 real_page_table_src, const tvm::ffi::Optional flashmla_num_splits_src, const tvm::ffi::Optional flashmla_metadata_src, const tvm::ffi::TensorView cache_seqlens_dst, const tvm::ffi::TensorView cu_seqlens_k_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 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 real_page_table_dst, const tvm::ffi::Optional flashmla_num_splits_dst, const tvm::ffi::Optional flashmla_metadata_dst, @@ -522,9 +522,9 @@ struct FusedMetadataCopyKernel { .cache_seqlens = unwrap_data_ptr(cache_seqlens_src, "cache_seqlens_src"), .cu_seqlens_k = unwrap_data_ptr(cu_seqlens_k_src, "cu_seqlens_k_src"), .page_indices = unwrap_data_ptr(page_indices_src, "page_indices_src"), - .nsa_cache_seqlens = unwrap_data_ptr(nsa_cache_seqlens_src, "nsa_cache_seqlens_src"), + .dsa_cache_seqlens = unwrap_data_ptr(dsa_cache_seqlens_src, "dsa_cache_seqlens_src"), .seqlens_expanded = unwrap_optional_data_ptr(seqlens_expanded_src, "seqlens_expanded_src"), - .nsa_cu_seqlens_k = unwrap_data_ptr(nsa_cu_seqlens_k_src, "nsa_cu_seqlens_k_src"), + .dsa_cu_seqlens_k = unwrap_data_ptr(dsa_cu_seqlens_k_src, "dsa_cu_seqlens_k_src"), .real_page_table = unwrap_optional_data_ptr(real_page_table_src, "real_page_table_src"), .flashmla_num_splits = unwrap_optional_data_ptr(flashmla_num_splits_src, "flashmla_num_splits_src"), @@ -535,9 +535,9 @@ struct FusedMetadataCopyKernel { .cache_seqlens = unwrap_data_ptr_mut(cache_seqlens_dst, "cache_seqlens_dst"), .cu_seqlens_k = unwrap_data_ptr_mut(cu_seqlens_k_dst, "cu_seqlens_k_dst"), .page_table_1 = unwrap_data_ptr_mut(page_table_1_dst, "page_table_1_dst"), - .nsa_cache_seqlens = unwrap_data_ptr_mut(nsa_cache_seqlens_dst, "nsa_cache_seqlens_dst"), + .dsa_cache_seqlens = unwrap_data_ptr_mut(dsa_cache_seqlens_dst, "dsa_cache_seqlens_dst"), .seqlens_expanded = unwrap_optional_data_ptr_mut(seqlens_expanded_dst, "seqlens_expanded_dst"), - .nsa_cu_seqlens_k = unwrap_data_ptr_mut(nsa_cu_seqlens_k_dst, "nsa_cu_seqlens_k_dst"), + .dsa_cu_seqlens_k = unwrap_data_ptr_mut(dsa_cu_seqlens_k_dst, "dsa_cu_seqlens_k_dst"), .real_page_table = unwrap_optional_data_ptr_mut(real_page_table_dst, "real_page_table_dst"), .flashmla_num_splits = unwrap_optional_data_ptr_mut(flashmla_num_splits_dst, "flashmla_num_splits_dst"), @@ -605,32 +605,32 @@ struct FusedMetadataCopyMultiKernel { run(const tvm::ffi::TensorView cache_seqlens_src, const tvm::ffi::TensorView cu_seqlens_k_src, const tvm::ffi::TensorView page_indices_src, - const tvm::ffi::TensorView nsa_cache_seqlens_src, - const tvm::ffi::TensorView nsa_cu_seqlens_k_src, + const tvm::ffi::TensorView dsa_cache_seqlens_src, + const tvm::ffi::TensorView dsa_cu_seqlens_k_src, const tvm::ffi::Optional real_page_table_src, const tvm::ffi::Optional flashmla_num_splits_src, const tvm::ffi::Optional flashmla_metadata_src, const tvm::ffi::TensorView cache_seqlens_dst0, const tvm::ffi::TensorView cu_seqlens_k_dst0, const tvm::ffi::TensorView page_table_1_dst0, - const tvm::ffi::TensorView nsa_cache_seqlens_dst0, - const tvm::ffi::TensorView nsa_cu_seqlens_k_dst0, + const tvm::ffi::TensorView dsa_cache_seqlens_dst0, + const tvm::ffi::TensorView dsa_cu_seqlens_k_dst0, const tvm::ffi::Optional real_page_table_dst0, const tvm::ffi::Optional flashmla_num_splits_dst0, const tvm::ffi::Optional flashmla_metadata_dst0, const tvm::ffi::TensorView cache_seqlens_dst1, const tvm::ffi::TensorView cu_seqlens_k_dst1, const tvm::ffi::TensorView page_table_1_dst1, - const tvm::ffi::TensorView nsa_cache_seqlens_dst1, - const tvm::ffi::TensorView nsa_cu_seqlens_k_dst1, + const tvm::ffi::TensorView dsa_cache_seqlens_dst1, + const tvm::ffi::TensorView dsa_cu_seqlens_k_dst1, const tvm::ffi::Optional real_page_table_dst1, const tvm::ffi::Optional flashmla_num_splits_dst1, const tvm::ffi::Optional flashmla_metadata_dst1, const tvm::ffi::TensorView cache_seqlens_dst2, const tvm::ffi::TensorView cu_seqlens_k_dst2, const tvm::ffi::TensorView page_table_1_dst2, - const tvm::ffi::TensorView nsa_cache_seqlens_dst2, - const tvm::ffi::TensorView nsa_cu_seqlens_k_dst2, + const tvm::ffi::TensorView dsa_cache_seqlens_dst2, + const tvm::ffi::TensorView dsa_cu_seqlens_k_dst2, const tvm::ffi::Optional real_page_table_dst2, const tvm::ffi::Optional flashmla_num_splits_dst2, const tvm::ffi::Optional flashmla_metadata_dst2, @@ -647,9 +647,9 @@ struct FusedMetadataCopyMultiKernel { .cache_seqlens = unwrap_data_ptr(cache_seqlens_src, "cache_seqlens_src"), .cu_seqlens_k = unwrap_data_ptr(cu_seqlens_k_src, "cu_seqlens_k_src"), .page_indices = unwrap_data_ptr(page_indices_src, "page_indices_src"), - .nsa_cache_seqlens = unwrap_data_ptr(nsa_cache_seqlens_src, "nsa_cache_seqlens_src"), + .dsa_cache_seqlens = unwrap_data_ptr(dsa_cache_seqlens_src, "dsa_cache_seqlens_src"), .seqlens_expanded = nullptr, // Not used in multi-backend DECODE mode - .nsa_cu_seqlens_k = unwrap_data_ptr(nsa_cu_seqlens_k_src, "nsa_cu_seqlens_k_src"), + .dsa_cu_seqlens_k = unwrap_data_ptr(dsa_cu_seqlens_k_src, "dsa_cu_seqlens_k_src"), .real_page_table = unwrap_optional_data_ptr(real_page_table_src, "real_page_table_src"), .flashmla_num_splits = unwrap_optional_data_ptr(flashmla_num_splits_src, "flashmla_num_splits_src"), @@ -660,9 +660,9 @@ struct FusedMetadataCopyMultiKernel { .cache_seqlens = unwrap_data_ptr_mut(cache_seqlens_dst0, "cache_seqlens_dst0"), .cu_seqlens_k = unwrap_data_ptr_mut(cu_seqlens_k_dst0, "cu_seqlens_k_dst0"), .page_table_1 = unwrap_data_ptr_mut(page_table_1_dst0, "page_table_1_dst0"), - .nsa_cache_seqlens = unwrap_data_ptr_mut(nsa_cache_seqlens_dst0, "nsa_cache_seqlens_dst0"), + .dsa_cache_seqlens = unwrap_data_ptr_mut(dsa_cache_seqlens_dst0, "dsa_cache_seqlens_dst0"), .seqlens_expanded = nullptr, - .nsa_cu_seqlens_k = unwrap_data_ptr_mut(nsa_cu_seqlens_k_dst0, "nsa_cu_seqlens_k_dst0"), + .dsa_cu_seqlens_k = unwrap_data_ptr_mut(dsa_cu_seqlens_k_dst0, "dsa_cu_seqlens_k_dst0"), .real_page_table = unwrap_optional_data_ptr_mut(real_page_table_dst0, "real_page_table_dst0"), .flashmla_num_splits = unwrap_optional_data_ptr_mut(flashmla_num_splits_dst0, "flashmla_num_splits_dst0"), @@ -674,9 +674,9 @@ struct FusedMetadataCopyMultiKernel { .cache_seqlens = unwrap_data_ptr_mut(cache_seqlens_dst1, "cache_seqlens_dst1"), .cu_seqlens_k = unwrap_data_ptr_mut(cu_seqlens_k_dst1, "cu_seqlens_k_dst1"), .page_table_1 = unwrap_data_ptr_mut(page_table_1_dst1, "page_table_1_dst1"), - .nsa_cache_seqlens = unwrap_data_ptr_mut(nsa_cache_seqlens_dst1, "nsa_cache_seqlens_dst1"), + .dsa_cache_seqlens = unwrap_data_ptr_mut(dsa_cache_seqlens_dst1, "dsa_cache_seqlens_dst1"), .seqlens_expanded = nullptr, - .nsa_cu_seqlens_k = unwrap_data_ptr_mut(nsa_cu_seqlens_k_dst1, "nsa_cu_seqlens_k_dst1"), + .dsa_cu_seqlens_k = unwrap_data_ptr_mut(dsa_cu_seqlens_k_dst1, "dsa_cu_seqlens_k_dst1"), .real_page_table = unwrap_optional_data_ptr_mut(real_page_table_dst1, "real_page_table_dst1"), .flashmla_num_splits = unwrap_optional_data_ptr_mut(flashmla_num_splits_dst1, "flashmla_num_splits_dst1"), @@ -688,9 +688,9 @@ struct FusedMetadataCopyMultiKernel { .cache_seqlens = unwrap_data_ptr_mut(cache_seqlens_dst2, "cache_seqlens_dst2"), .cu_seqlens_k = unwrap_data_ptr_mut(cu_seqlens_k_dst2, "cu_seqlens_k_dst2"), .page_table_1 = unwrap_data_ptr_mut(page_table_1_dst2, "page_table_1_dst2"), - .nsa_cache_seqlens = unwrap_data_ptr_mut(nsa_cache_seqlens_dst2, "nsa_cache_seqlens_dst2"), + .dsa_cache_seqlens = unwrap_data_ptr_mut(dsa_cache_seqlens_dst2, "dsa_cache_seqlens_dst2"), .seqlens_expanded = nullptr, - .nsa_cu_seqlens_k = unwrap_data_ptr_mut(nsa_cu_seqlens_k_dst2, "nsa_cu_seqlens_k_dst2"), + .dsa_cu_seqlens_k = unwrap_data_ptr_mut(dsa_cu_seqlens_k_dst2, "dsa_cu_seqlens_k_dst2"), .real_page_table = unwrap_optional_data_ptr_mut(real_page_table_dst2, "real_page_table_dst2"), .flashmla_num_splits = unwrap_optional_data_ptr_mut(flashmla_num_splits_dst2, "flashmla_num_splits_dst2"), diff --git a/python/sglang/jit_kernel/fused_metadata_copy.py b/python/sglang/jit_kernel/fused_metadata_copy.py index b4d347f6a..68d0f9227 100644 --- a/python/sglang/jit_kernel/fused_metadata_copy.py +++ b/python/sglang/jit_kernel/fused_metadata_copy.py @@ -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 copy operations into single kernel launches, reducing kernel launch overhead @@ -98,18 +98,18 @@ def fused_metadata_copy_cuda( cache_seqlens_src: torch.Tensor, cu_seqlens_k_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], - nsa_cu_seqlens_k_src: torch.Tensor, + dsa_cu_seqlens_k_src: torch.Tensor, real_page_table_src: Optional[torch.Tensor], flashmla_num_splits_src: Optional[torch.Tensor], flashmla_metadata_src: Optional[torch.Tensor], cache_seqlens_dst: torch.Tensor, cu_seqlens_k_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], - nsa_cu_seqlens_k_dst: torch.Tensor, + dsa_cu_seqlens_k_dst: torch.Tensor, real_page_table_dst: Optional[torch.Tensor], flashmla_num_splits_dst: Optional[torch.Tensor], flashmla_metadata_dst: Optional[torch.Tensor], @@ -120,7 +120,7 @@ def fused_metadata_copy_cuda( seqlens_expanded_size: int, ) -> 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, reducing kernel launch overhead and improving performance. @@ -129,18 +129,18 @@ def fused_metadata_copy_cuda( cache_seqlens_src: Source cache sequence lengths [bs] cu_seqlens_k_src: Source cumulative sequence lengths [bs+1] 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) - 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] flashmla_num_splits_src: Optional source FlashMLA num_splits [size+1] flashmla_metadata_src: Optional source FlashMLA metadata tensor cache_seqlens_dst: Destination cache sequence lengths [bs] cu_seqlens_k_dst: Destination cumulative sequence lengths [bs+1] 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) - 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] flashmla_num_splits_dst: Optional destination FlashMLA num_splits [size+1] flashmla_metadata_dst: Optional destination FlashMLA metadata tensor @@ -164,28 +164,28 @@ def fused_metadata_copy_cuda( cache_seqlens_src = cache_seqlens_src.contiguous() cu_seqlens_k_src = cu_seqlens_k_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: 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) module.fused_metadata_copy( cache_seqlens_src, cu_seqlens_k_src, page_indices_src, - nsa_cache_seqlens_src, + dsa_cache_seqlens_src, seqlens_expanded_src, - nsa_cu_seqlens_k_src, + dsa_cu_seqlens_k_src, real_page_table_src, flashmla_num_splits_src, flashmla_metadata_src, cache_seqlens_dst, cu_seqlens_k_dst, page_table_1_dst, - nsa_cache_seqlens_dst, + dsa_cache_seqlens_dst, seqlens_expanded_dst, - nsa_cu_seqlens_k_dst, + dsa_cu_seqlens_k_dst, real_page_table_dst, flashmla_num_splits_dst, flashmla_metadata_dst, @@ -200,32 +200,32 @@ def fused_metadata_copy_multi_cuda( cache_seqlens_src: torch.Tensor, cu_seqlens_k_src: torch.Tensor, page_indices_src: torch.Tensor, - nsa_cache_seqlens_src: torch.Tensor, - nsa_cu_seqlens_k_src: torch.Tensor, + dsa_cache_seqlens_src: torch.Tensor, + dsa_cu_seqlens_k_src: torch.Tensor, real_page_table_src: Optional[torch.Tensor], flashmla_num_splits_src: Optional[torch.Tensor], flashmla_metadata_src: Optional[torch.Tensor], cache_seqlens_dst0: torch.Tensor, cu_seqlens_k_dst0: torch.Tensor, page_table_1_dst0: torch.Tensor, - nsa_cache_seqlens_dst0: torch.Tensor, - nsa_cu_seqlens_k_dst0: torch.Tensor, + dsa_cache_seqlens_dst0: torch.Tensor, + dsa_cu_seqlens_k_dst0: torch.Tensor, real_page_table_dst0: Optional[torch.Tensor], flashmla_num_splits_dst0: Optional[torch.Tensor], flashmla_metadata_dst0: Optional[torch.Tensor], cache_seqlens_dst1: torch.Tensor, cu_seqlens_k_dst1: torch.Tensor, page_table_1_dst1: torch.Tensor, - nsa_cache_seqlens_dst1: torch.Tensor, - nsa_cu_seqlens_k_dst1: torch.Tensor, + dsa_cache_seqlens_dst1: torch.Tensor, + dsa_cu_seqlens_k_dst1: torch.Tensor, real_page_table_dst1: Optional[torch.Tensor], flashmla_num_splits_dst1: Optional[torch.Tensor], flashmla_metadata_dst1: Optional[torch.Tensor], cache_seqlens_dst2: torch.Tensor, cu_seqlens_k_dst2: torch.Tensor, page_table_1_dst2: torch.Tensor, - nsa_cache_seqlens_dst2: torch.Tensor, - nsa_cu_seqlens_k_dst2: torch.Tensor, + dsa_cache_seqlens_dst2: torch.Tensor, + dsa_cu_seqlens_k_dst2: torch.Tensor, real_page_table_dst2: Optional[torch.Tensor], flashmla_num_splits_dst2: Optional[torch.Tensor], flashmla_metadata_dst2: Optional[torch.Tensor], @@ -234,7 +234,7 @@ def fused_metadata_copy_multi_cuda( seqlens_expanded_size: int, ) -> 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 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] cu_seqlens_k_src: Source cumulative sequence lengths [bs+1] page_indices_src: Source page indices [bs, max_len] - nsa_cache_seqlens_src: Source NSA cache sequence lengths [bs] - nsa_cu_seqlens_k_src: Source NSA cumulative sequence lengths [bs+1] + dsa_cache_seqlens_src: Source DSA cache sequence lengths [bs] + dsa_cu_seqlens_k_src: Source DSA cumulative sequence lengths [bs+1] real_page_table_src: Optional source real page table [bs, cols] flashmla_num_splits_src: Optional source FlashMLA num_splits [bs+1] flashmla_metadata_src: Optional source FlashMLA metadata tensor 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 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 - nsa_cu_seqlens_k_dst0-2: Destination NSA cumulative sequence lengths for backends 0-2 + dsa_cache_seqlens_dst0-2: Destination DSA cache 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 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 @@ -273,40 +273,40 @@ def fused_metadata_copy_multi_cuda( cache_seqlens_src = cache_seqlens_src.contiguous() cu_seqlens_k_src = cu_seqlens_k_src.contiguous() page_indices_src = page_indices_src.contiguous() - nsa_cache_seqlens_src = nsa_cache_seqlens_src.contiguous() - nsa_cu_seqlens_k_src = nsa_cu_seqlens_k_src.contiguous() + dsa_cache_seqlens_src = dsa_cache_seqlens_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) module.fused_metadata_copy_multi( cache_seqlens_src, cu_seqlens_k_src, page_indices_src, - nsa_cache_seqlens_src, - nsa_cu_seqlens_k_src, + dsa_cache_seqlens_src, + dsa_cu_seqlens_k_src, real_page_table_src, flashmla_num_splits_src, flashmla_metadata_src, cache_seqlens_dst0, cu_seqlens_k_dst0, page_table_1_dst0, - nsa_cache_seqlens_dst0, - nsa_cu_seqlens_k_dst0, + dsa_cache_seqlens_dst0, + dsa_cu_seqlens_k_dst0, real_page_table_dst0, flashmla_num_splits_dst0, flashmla_metadata_dst0, cache_seqlens_dst1, cu_seqlens_k_dst1, page_table_1_dst1, - nsa_cache_seqlens_dst1, - nsa_cu_seqlens_k_dst1, + dsa_cache_seqlens_dst1, + dsa_cu_seqlens_k_dst1, real_page_table_dst1, flashmla_num_splits_dst1, flashmla_metadata_dst1, cache_seqlens_dst2, cu_seqlens_k_dst2, page_table_1_dst2, - nsa_cache_seqlens_dst2, - nsa_cu_seqlens_k_dst2, + dsa_cache_seqlens_dst2, + dsa_cu_seqlens_k_dst2, real_page_table_dst2, flashmla_num_splits_dst2, flashmla_metadata_dst2, diff --git a/python/sglang/jit_kernel/fused_store_index_cache.py b/python/sglang/jit_kernel/fused_store_index_cache.py index dc50e21b5..f8b3b1432 100644 --- a/python/sglang/jit_kernel/fused_store_index_cache.py +++ b/python/sglang/jit_kernel/fused_store_index_cache.py @@ -28,7 +28,7 @@ logger = logging.getLogger(__name__) @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 ) -> Module: """ @@ -39,13 +39,13 @@ def _jit_nsa_fused_store_module( return load_jit( "fused_store_index_k_cache", *args, - cuda_files=["nsa/fused_store_index_cache.cuh"], + cuda_files=["dsa/fused_store_index_cache.cuh"], cuda_wrappers=[ ( "fused_store_index_k_cache", # - Float = bf16_t (sgl_kernel/type.cuh) # - IndicesT = int64_t (out_cache_loc is int64 in SGLang SetKAndS) - # - kPageSize = 64 (CUDA NSA) + # - kPageSize = 64 (CUDA DSA) f"FusedStoreCacheIndexerKernel<{args}>::run", ) ], @@ -53,15 +53,15 @@ def _jit_nsa_fused_store_module( @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 ) -> bool: logger = logging.getLogger(__name__) 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 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 @@ -73,7 +73,7 @@ def fused_store_index_k_cache( page_size: int = 64, ) -> 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) 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(): 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) diff --git a/python/sglang/jit_kernel/tests/test_fused_metadata_copy.py b/python/sglang/jit_kernel/tests/test_fused_metadata_copy.py index 4daeae01c..08be9df14 100644 --- a/python/sglang/jit_kernel/tests/test_fused_metadata_copy.py +++ b/python/sglang/jit_kernel/tests/test_fused_metadata_copy.py @@ -33,7 +33,7 @@ def create_test_metadata( has_flashmla: bool = False, device: str = "cuda", ): - """Create test metadata tensors matching NSA backend structure.""" + """Create test metadata tensors matching DSA backend structure.""" # Basic tensors (always present) cache_seqlens_src = torch.randint( 1, max_len, (bs,), dtype=torch.int32, device=device @@ -44,28 +44,28 @@ def create_test_metadata( page_indices_src = torch.randint( 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 ) seqlens_expanded_src = torch.randint( 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 ) - 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 cache_seqlens_dst = torch.zeros(bs, 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) - nsa_cache_seqlens_dst = torch.zeros( + dsa_cache_seqlens_dst = torch.zeros( 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 ) - nsa_cu_seqlens_k_dst = torch.zeros( + dsa_cu_seqlens_k_dst = torch.zeros( seqlens_expanded_size + 1, dtype=torch.int32, device=device ) @@ -107,9 +107,9 @@ def create_test_metadata( "cache_seqlens": cache_seqlens_src, "cu_seqlens_k": cu_seqlens_k_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, - "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, "flashmla_num_splits": flashmla_num_splits_src, "flashmla_metadata": flashmla_metadata_src, @@ -118,9 +118,9 @@ def create_test_metadata( "cache_seqlens": cache_seqlens_dst, "cu_seqlens_k": cu_seqlens_k_dst, "page_table_1": page_table_1_dst, - "nsa_cache_seqlens": nsa_cache_seqlens_dst, - "nsa_seqlens_expanded": nsa_seqlens_expanded_dst, - "nsa_cu_seqlens_k": nsa_cu_seqlens_k_dst, + "dsa_cache_seqlens": dsa_cache_seqlens_dst, + "dsa_seqlens_expanded": dsa_seqlens_expanded_dst, + "dsa_cu_seqlens_k": dsa_cu_seqlens_k_dst, "real_page_table": real_page_table_dst, "flashmla_num_splits": flashmla_num_splits_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["cu_seqlens_k"][1:].copy_(src["cu_seqlens_k"][1:]) dst["page_table_1"][:, :max_len].copy_(src["page_indices"]) - dst["nsa_cache_seqlens"].copy_(src["nsa_cache_seqlens"]) - dst["nsa_cu_seqlens_k"][1 : bs + 1].copy_(src["nsa_cu_seqlens_k"][1 : bs + 1]) + dst["dsa_cache_seqlens"].copy_(src["dsa_cache_seqlens"]) + 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: 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 dst["page_table_1"][:rows, :cols].copy_(src["page_indices"]) - dst["nsa_seqlens_expanded"][:seqlens_expanded_size].copy_(src["seqlens_expanded"]) - dst["nsa_cache_seqlens"][:seqlens_expanded_size].copy_(src["nsa_cache_seqlens"]) - dst["nsa_cu_seqlens_k"][1 : seqlens_expanded_size + 1].copy_( - src["nsa_cu_seqlens_k"][1 : seqlens_expanded_size + 1] + dst["dsa_seqlens_expanded"][:seqlens_expanded_size].copy_(src["seqlens_expanded"]) + dst["dsa_cache_seqlens"][:seqlens_expanded_size].copy_(src["dsa_cache_seqlens"]) + dst["dsa_cu_seqlens_k"][1 : seqlens_expanded_size + 1].copy_( + src["dsa_cu_seqlens_k"][1 : seqlens_expanded_size + 1] ) 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 dst["page_table_1"][:rows, :cols].copy_(src["page_indices"]) - dst["nsa_seqlens_expanded"][:seqlens_expanded_size].copy_(src["seqlens_expanded"]) - dst["nsa_cache_seqlens"][:seqlens_expanded_size].copy_(src["nsa_cache_seqlens"]) - dst["nsa_cu_seqlens_k"][1 : seqlens_expanded_size + 1].copy_( - src["nsa_cu_seqlens_k"][1 : seqlens_expanded_size + 1] + dst["dsa_seqlens_expanded"][:seqlens_expanded_size].copy_(src["seqlens_expanded"]) + dst["dsa_cache_seqlens"][:seqlens_expanded_size].copy_(src["dsa_cache_seqlens"]) + dst["dsa_cu_seqlens_k"][1 : seqlens_expanded_size + 1].copy_( + src["dsa_cu_seqlens_k"][1 : seqlens_expanded_size + 1] ) if src["real_page_table"] is not None: @@ -233,13 +233,13 @@ def test_fused_metadata_copy_dtype_validation(): page_indices_src = torch.randint( 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 ) seqlens_expanded_src = torch.randint( 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 ) @@ -247,13 +247,13 @@ def test_fused_metadata_copy_dtype_validation(): cache_seqlens_dst = torch.zeros(bs, 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) - nsa_cache_seqlens_dst = torch.zeros( + dsa_cache_seqlens_dst = torch.zeros( 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 ) - nsa_cu_seqlens_k_dst = torch.zeros( + dsa_cu_seqlens_k_dst = torch.zeros( 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 cu_seqlens_k_src, page_indices_src, - nsa_cache_seqlens_src, + dsa_cache_seqlens_src, seqlens_expanded_src, - nsa_cu_seqlens_k_src, + dsa_cu_seqlens_k_src, None, # real_page_table_src None, # flashmla_num_splits_src None, # flashmla_metadata_src cache_seqlens_dst, cu_seqlens_k_dst, page_table_1_dst, - nsa_cache_seqlens_dst, - nsa_seqlens_expanded_dst, - nsa_cu_seqlens_k_dst, + dsa_cache_seqlens_dst, + dsa_seqlens_expanded_dst, + dsa_cu_seqlens_k_dst, None, # real_page_table_dst None, # flashmla_num_splits_dst None, # flashmla_metadata_dst @@ -296,18 +296,18 @@ def test_fused_metadata_copy_dtype_validation(): cache_seqlens_src, cu_seqlens_k_src, page_indices_src, - nsa_cache_seqlens_src, + dsa_cache_seqlens_src, seqlens_expanded_src, - nsa_cu_seqlens_k_src, + dsa_cu_seqlens_k_src, None, None, None, cache_seqlens_dst_wrong, # Wrong dtype: int64 cu_seqlens_k_dst, page_table_1_dst, - nsa_cache_seqlens_dst, - nsa_seqlens_expanded_dst, - nsa_cu_seqlens_k_dst, + dsa_cache_seqlens_dst, + dsa_seqlens_expanded_dst, + dsa_cu_seqlens_k_dst, 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"]["cu_seqlens_k"], data["src"]["page_indices"], - data["src"]["nsa_cache_seqlens"], + data["src"]["dsa_cache_seqlens"], data["src"]["seqlens_expanded"], - data["src"]["nsa_cu_seqlens_k"], + data["src"]["dsa_cu_seqlens_k"], data["src"]["real_page_table"], data["src"]["flashmla_num_splits"], data["src"]["flashmla_metadata"], dst_fused["cache_seqlens"], dst_fused["cu_seqlens_k"], dst_fused["page_table_1"], - dst_fused["nsa_cache_seqlens"], - dst_fused["nsa_seqlens_expanded"], - dst_fused["nsa_cu_seqlens_k"], + dst_fused["dsa_cache_seqlens"], + dst_fused["dsa_seqlens_expanded"], + dst_fused["dsa_cu_seqlens_k"], dst_fused["real_page_table"], dst_fused["flashmla_num_splits"], 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"] ), "page_table_1 mismatch" assert torch.equal( - dst_ref["nsa_cache_seqlens"], dst_fused["nsa_cache_seqlens"] - ), "nsa_cache_seqlens mismatch" + dst_ref["dsa_cache_seqlens"], dst_fused["dsa_cache_seqlens"] + ), "dsa_cache_seqlens mismatch" assert torch.equal( - dst_ref["nsa_seqlens_expanded"], dst_fused["nsa_seqlens_expanded"] - ), "nsa_seqlens_expanded mismatch" + dst_ref["dsa_seqlens_expanded"], dst_fused["dsa_seqlens_expanded"] + ), "dsa_seqlens_expanded mismatch" assert torch.equal( - dst_ref["nsa_cu_seqlens_k"], dst_fused["nsa_cu_seqlens_k"] - ), "nsa_cu_seqlens_k mismatch" + dst_ref["dsa_cu_seqlens_k"], dst_fused["dsa_cu_seqlens_k"] + ), "dsa_cu_seqlens_k mismatch" if has_real_page_table: assert torch.equal( @@ -458,18 +458,18 @@ def test_fused_metadata_copy_large_batch(bs): data["src"]["cache_seqlens"], data["src"]["cu_seqlens_k"], data["src"]["page_indices"], - data["src"]["nsa_cache_seqlens"], + data["src"]["dsa_cache_seqlens"], data["src"]["seqlens_expanded"], - data["src"]["nsa_cu_seqlens_k"], + data["src"]["dsa_cu_seqlens_k"], data["src"]["real_page_table"], data["src"]["flashmla_num_splits"], data["src"]["flashmla_metadata"], dst_fused["cache_seqlens"], dst_fused["cu_seqlens_k"], dst_fused["page_table_1"], - dst_fused["nsa_cache_seqlens"], - dst_fused["nsa_seqlens_expanded"], - dst_fused["nsa_cu_seqlens_k"], + dst_fused["dsa_cache_seqlens"], + dst_fused["dsa_seqlens_expanded"], + dst_fused["dsa_cu_seqlens_k"], dst_fused["real_page_table"], dst_fused["flashmla_num_splits"], dst_fused["flashmla_metadata"], @@ -510,13 +510,13 @@ def create_test_metadata_multi( page_indices_src = torch.randint( 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 ) - nsa_cu_seqlens_k_src = torch.zeros( + dsa_cu_seqlens_k_src = torch.zeros( 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 real_page_table_src = None @@ -544,10 +544,10 @@ def create_test_metadata_multi( 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 ) - nsa_cu_seqlens_k_dst = torch.zeros( + dsa_cu_seqlens_k_dst = torch.zeros( seqlens_expanded_size + 1, dtype=torch.int32, device=device ) @@ -573,8 +573,8 @@ def create_test_metadata_multi( "cache_seqlens_int32": cache_seqlens_dst, "cu_seqlens_k": cu_seqlens_k_dst, "page_table_1": page_table_1_dst, - "nsa_cache_seqlens_int32": nsa_cache_seqlens_dst, - "nsa_cu_seqlens_k": nsa_cu_seqlens_k_dst, + "dsa_cache_seqlens_int32": dsa_cache_seqlens_dst, + "dsa_cu_seqlens_k": dsa_cu_seqlens_k_dst, "real_page_table": real_page_table_dst, "flashmla_num_splits": flashmla_num_splits_dst, "flashmla_metadata": flashmla_metadata_dst, @@ -585,8 +585,8 @@ def create_test_metadata_multi( "cache_seqlens": cache_seqlens_src, "cu_seqlens_k": cu_seqlens_k_src, "page_indices": page_indices_src, - "nsa_cache_seqlens": nsa_cache_seqlens_src, - "nsa_cu_seqlens_k": nsa_cu_seqlens_k_src, + "dsa_cache_seqlens": dsa_cache_seqlens_src, + "dsa_cu_seqlens_k": dsa_cu_seqlens_k_src, "real_page_table": real_page_table_src, "flashmla_num_splits": flashmla_num_splits_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["cu_seqlens_k"][1:].copy_(src["cu_seqlens_k"][1:]) dst["page_table_1"][:, :max_len].copy_(src["page_indices"]) - dst["nsa_cache_seqlens_int32"].copy_(src["nsa_cache_seqlens"]) - dst["nsa_cu_seqlens_k"][1 : bs + 1].copy_(src["nsa_cu_seqlens_k"][1 : bs + 1]) + dst["dsa_cache_seqlens_int32"].copy_(src["dsa_cache_seqlens"]) + 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: rows, cols = src["real_page_table"].shape @@ -641,10 +641,10 @@ def test_fused_metadata_copy_multi_dtype_validation(): page_indices_src = torch.randint( 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 ) - nsa_cu_seqlens_k_src = torch.zeros( + dsa_cu_seqlens_k_src = torch.zeros( 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( (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 ), - "nsa_cu_seqlens_k": torch.zeros( + "dsa_cu_seqlens_k": torch.zeros( 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 cu_seqlens_k_src, page_indices_src, - nsa_cache_seqlens_src, - nsa_cu_seqlens_k_src, + dsa_cache_seqlens_src, + dsa_cu_seqlens_k_src, None, # real_page_table_src None, # flashmla_num_splits_src None, # flashmla_metadata_src @@ -683,8 +683,8 @@ def test_fused_metadata_copy_multi_dtype_validation(): dst0["cache_seqlens"], dst0["cu_seqlens_k"], dst0["page_table_1"], - dst0["nsa_cache_seqlens"], - dst0["nsa_cu_seqlens_k"], + dst0["dsa_cache_seqlens"], + dst0["dsa_cu_seqlens_k"], None, None, None, @@ -692,8 +692,8 @@ def test_fused_metadata_copy_multi_dtype_validation(): dst1["cache_seqlens"], dst1["cu_seqlens_k"], dst1["page_table_1"], - dst1["nsa_cache_seqlens"], - dst1["nsa_cu_seqlens_k"], + dst1["dsa_cache_seqlens"], + dst1["dsa_cu_seqlens_k"], None, None, None, @@ -701,8 +701,8 @@ def test_fused_metadata_copy_multi_dtype_validation(): dst2["cache_seqlens"], dst2["cu_seqlens_k"], dst2["page_table_1"], - dst2["nsa_cache_seqlens"], - dst2["nsa_cu_seqlens_k"], + dst2["dsa_cache_seqlens"], + dst2["dsa_cu_seqlens_k"], 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"]["cu_seqlens_k"], data["src"]["page_indices"], - data["src"]["nsa_cache_seqlens"], - data["src"]["nsa_cu_seqlens_k"], + data["src"]["dsa_cache_seqlens"], + data["src"]["dsa_cu_seqlens_k"], data["src"]["real_page_table"], data["src"]["flashmla_num_splits"], 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["cu_seqlens_k"], dst_fused_0["page_table_1"], - dst_fused_0["nsa_cache_seqlens_int32"], - dst_fused_0["nsa_cu_seqlens_k"], + dst_fused_0["dsa_cache_seqlens_int32"], + dst_fused_0["dsa_cu_seqlens_k"], dst_fused_0["real_page_table"], dst_fused_0["flashmla_num_splits"], 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["cu_seqlens_k"], dst_fused_1["page_table_1"], - dst_fused_1["nsa_cache_seqlens_int32"], - dst_fused_1["nsa_cu_seqlens_k"], + dst_fused_1["dsa_cache_seqlens_int32"], + dst_fused_1["dsa_cu_seqlens_k"], dst_fused_1["real_page_table"], dst_fused_1["flashmla_num_splits"], 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["cu_seqlens_k"], dst_fused_2["page_table_1"], - dst_fused_2["nsa_cache_seqlens_int32"], - dst_fused_2["nsa_cu_seqlens_k"], + dst_fused_2["dsa_cache_seqlens_int32"], + dst_fused_2["dsa_cu_seqlens_k"], dst_fused_2["real_page_table"], dst_fused_2["flashmla_num_splits"], 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", "cu_seqlens_k", "page_table_1", - "nsa_cache_seqlens_int32", - "nsa_cu_seqlens_k", + "dsa_cache_seqlens_int32", + "dsa_cu_seqlens_k", ]: if not torch.equal(dst_ref[key], dst_fused[key]): diff = ( @@ -965,32 +965,32 @@ def test_fused_metadata_copy_multi_large_batch(bs): data["src"]["cache_seqlens"], data["src"]["cu_seqlens_k"], data["src"]["page_indices"], - data["src"]["nsa_cache_seqlens"], - data["src"]["nsa_cu_seqlens_k"], + data["src"]["dsa_cache_seqlens"], + data["src"]["dsa_cu_seqlens_k"], data["src"]["real_page_table"], data["src"]["flashmla_num_splits"], data["src"]["flashmla_metadata"], dst_fused_0["cache_seqlens_int32"], dst_fused_0["cu_seqlens_k"], dst_fused_0["page_table_1"], - dst_fused_0["nsa_cache_seqlens_int32"], - dst_fused_0["nsa_cu_seqlens_k"], + dst_fused_0["dsa_cache_seqlens_int32"], + dst_fused_0["dsa_cu_seqlens_k"], dst_fused_0["real_page_table"], dst_fused_0["flashmla_num_splits"], dst_fused_0["flashmla_metadata"], dst_fused_1["cache_seqlens_int32"], dst_fused_1["cu_seqlens_k"], dst_fused_1["page_table_1"], - dst_fused_1["nsa_cache_seqlens_int32"], - dst_fused_1["nsa_cu_seqlens_k"], + dst_fused_1["dsa_cache_seqlens_int32"], + dst_fused_1["dsa_cu_seqlens_k"], dst_fused_1["real_page_table"], dst_fused_1["flashmla_num_splits"], dst_fused_1["flashmla_metadata"], dst_fused_2["cache_seqlens_int32"], dst_fused_2["cu_seqlens_k"], dst_fused_2["page_table_1"], - dst_fused_2["nsa_cache_seqlens_int32"], - dst_fused_2["nsa_cu_seqlens_k"], + dst_fused_2["dsa_cache_seqlens_int32"], + dst_fused_2["dsa_cu_seqlens_k"], dst_fused_2["real_page_table"], dst_fused_2["flashmla_num_splits"], dst_fused_2["flashmla_metadata"], @@ -1013,32 +1013,32 @@ def test_fused_metadata_copy_multi_large_batch(bs): data["src"]["cache_seqlens"], data["src"]["cu_seqlens_k"], data["src"]["page_indices"], - data["src"]["nsa_cache_seqlens"], - data["src"]["nsa_cu_seqlens_k"], + data["src"]["dsa_cache_seqlens"], + data["src"]["dsa_cu_seqlens_k"], data["src"]["real_page_table"], data["src"]["flashmla_num_splits"], data["src"]["flashmla_metadata"], dst_fused_0["cache_seqlens_int32"], dst_fused_0["cu_seqlens_k"], dst_fused_0["page_table_1"], - dst_fused_0["nsa_cache_seqlens_int32"], - dst_fused_0["nsa_cu_seqlens_k"], + dst_fused_0["dsa_cache_seqlens_int32"], + dst_fused_0["dsa_cu_seqlens_k"], dst_fused_0["real_page_table"], dst_fused_0["flashmla_num_splits"], dst_fused_0["flashmla_metadata"], dst_fused_1["cache_seqlens_int32"], dst_fused_1["cu_seqlens_k"], dst_fused_1["page_table_1"], - dst_fused_1["nsa_cache_seqlens_int32"], - dst_fused_1["nsa_cu_seqlens_k"], + dst_fused_1["dsa_cache_seqlens_int32"], + dst_fused_1["dsa_cu_seqlens_k"], dst_fused_1["real_page_table"], dst_fused_1["flashmla_num_splits"], dst_fused_1["flashmla_metadata"], dst_fused_2["cache_seqlens_int32"], dst_fused_2["cu_seqlens_k"], dst_fused_2["page_table_1"], - dst_fused_2["nsa_cache_seqlens_int32"], - dst_fused_2["nsa_cu_seqlens_k"], + dst_fused_2["dsa_cache_seqlens_int32"], + dst_fused_2["dsa_cu_seqlens_k"], dst_fused_2["real_page_table"], dst_fused_2["flashmla_num_splits"], dst_fused_2["flashmla_metadata"], diff --git a/python/sglang/jit_kernel/tests/test_fused_store_index_cache.py b/python/sglang/jit_kernel/tests/test_fused_store_index_cache.py index 6dc382166..1766e8d0c 100644 --- a/python/sglang/jit_kernel/tests/test_fused_store_index_cache.py +++ b/python/sglang/jit_kernel/tests/test_fused_store_index_cache.py @@ -26,7 +26,7 @@ from sglang.test.ci.ci_register import register_cuda_ci try: from sglang.jit_kernel.fused_store_index_cache import ( - can_use_nsa_fused_store, + can_use_dsa_fused_store, 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") if not HAS_FUSED: 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") @@ -187,7 +187,7 @@ def _reference_quantize_and_store( def _import_act_quant(): 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 except Exception: diff --git a/python/sglang/jit_kernel/tests/test_set_mla_kv_buffer.py b/python/sglang/jit_kernel/tests/test_set_mla_kv_buffer.py index f9aa08724..d7f967647 100644 --- a/python/sglang/jit_kernel/tests/test_set_mla_kv_buffer.py +++ b/python/sglang/jit_kernel/tests/test_set_mla_kv_buffer.py @@ -75,7 +75,7 @@ def test_set_mla_kv_buffer_loc_dtypes(loc_dtype): 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.""" nope_bytes, rope_bytes = 528, 128 batch_size = 64 diff --git a/python/sglang/srt/arg_groups/deepseek_v4_hook.py b/python/sglang/srt/arg_groups/deepseek_v4_hook.py index 235b3415a..3dd5d31c2 100644 --- a/python/sglang/srt/arg_groups/deepseek_v4_hook.py +++ b/python/sglang/srt/arg_groups/deepseek_v4_hook.py @@ -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: """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 - if server_args.nsa_prefill_cp_mode != "round-robin-split": + if server_args.dsa_prefill_cp_mode != "round-robin-split": raise ValueError( 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 diff --git a/python/sglang/srt/arg_groups/hisparse_hook.py b/python/sglang/srt/arg_groups/hisparse_hook.py index 379d76dd7..f9f1197ae 100644 --- a/python/sglang/srt/arg_groups/hisparse_hook.py +++ b/python/sglang/srt/arg_groups/hisparse_hook.py @@ -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" -def apply_hisparse_nsa_backend_defaults( +def apply_hisparse_dsa_backend_defaults( server_args: "ServerArgs", user_set_prefill: bool, user_set_decode: bool, kv_cache_dtype: str, ) -> 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 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) if not user_set_prefill: - server_args.nsa_prefill_backend = backend + server_args.dsa_prefill_backend = backend if not user_set_decode: - server_args.nsa_decode_backend = backend + server_args.dsa_decode_backend = backend logger.warning( - f"HiSparse enabled ({kv_cache_dtype}): using NSA backends " - f"prefill={server_args.nsa_prefill_backend}, decode={server_args.nsa_decode_backend}." + f"HiSparse enabled ({kv_cache_dtype}): using DSA backends " + f"prefill={server_args.dsa_prefill_backend}, decode={server_args.dsa_decode_backend}." ) return True 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: return from sglang.srt.configs.model_config import ( - is_deepseek_nsa, + is_deepseek_dsa, is_deepseek_v4, ) hf_config = server_args.get_model_config().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) " "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"} ) for attr, label in [ - ("nsa_prefill_backend", "prefill"), - ("nsa_decode_backend", "decode"), + ("dsa_prefill_backend", "prefill"), + ("dsa_decode_backend", "decode"), ]: backend = getattr(server_args, attr) if backend is not None and backend not in allowed_backends: raise ValueError( 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}." ) diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index 111145ef6..6f89a6902 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -99,7 +99,7 @@ def _hf_attr(config, name): return getattr(config, name, None) -def is_deepseek_nsa(config) -> bool: +def is_deepseek_dsa(config) -> bool: return ( _hf_arch(config) in ( @@ -121,31 +121,31 @@ def is_deepseek_v4(config) -> bool: ) -def get_nsa_index_head_dim(config: PretrainedConfig) -> int: - assert is_deepseek_nsa(config) or is_deepseek_v4(config) +def get_dsa_index_head_dim(config: PretrainedConfig) -> int: + assert is_deepseek_dsa(config) or is_deepseek_v4(config) return config.index_head_dim -def get_nsa_index_topk(config: PretrainedConfig) -> int: - assert is_deepseek_nsa(config) +def get_dsa_index_topk(config: PretrainedConfig) -> int: + assert is_deepseek_dsa(config) return config.index_topk -def get_nsa_index_n_heads(config: PretrainedConfig) -> int: - assert is_deepseek_nsa(config) +def get_dsa_index_n_heads(config: PretrainedConfig) -> int: + assert is_deepseek_dsa(config) return config.index_n_heads def get_num_indexer_layers(config) -> int: """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 get a slot (mirrored at the MLA call site). DSv4 has C4 indexers only on layers whose compress_ratio == 4. Other architectures: set 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 if is_deepseek_v4(config): 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.is_piecewise_cuda_graph_disabled_model = ( 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) @@ -622,8 +622,8 @@ class ModelConfig: 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.index_head_dim = ( - get_nsa_index_head_dim(self.hf_text_config) - if is_deepseek_nsa(self.hf_text_config) + get_dsa_index_head_dim(self.hf_text_config) + if is_deepseek_dsa(self.hf_text_config) else None ) # Handle rope scaling diff --git a/python/sglang/srt/disaggregation/base/conn.py b/python/sglang/srt/disaggregation/base/conn.py index 097a84159..a7bf9904a 100644 --- a/python/sglang/srt/disaggregation/base/conn.py +++ b/python/sglang/srt/disaggregation/base/conn.py @@ -17,7 +17,7 @@ if TYPE_CHECKING: class StateType(str, enum.Enum): MAMBA = "mamba" SWA = "swa" - NSA = "nsa" + DSA = "dsa" @dataclasses.dataclass diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 3bf4797f3..65f2e4cb8 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -954,7 +954,7 @@ class DecodePreallocQueue: 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[ decode_req.req.req_pool_idx, :seq_len ] @@ -971,8 +971,8 @@ class DecodePreallocQueue: state_indices.append(_mamba_payload()) elif st == StateType.SWA: state_indices.append(_swa_payload()) - elif st == StateType.NSA: - state_indices.append(_nsa_payload()) + elif st == StateType.DSA: + state_indices.append(_dsa_payload()) else: state_indices.append(None) diff --git a/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py b/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py index ebc660bbf..27b33846f 100644 --- a/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py +++ b/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py @@ -238,7 +238,7 @@ class DecodeKVCacheOffloadManager: kv_committed_len = req.pop_committed_kv_cache() start = start_offset 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] self.token_to_kv_pool_allocator.free(kv_indices) diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index 18a455301..ef7405624 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py @@ -1012,7 +1012,7 @@ class MooncakeKVManager(CommonKVManager): ) or rc ) - elif st in (StateType.SWA, StateType.NSA): + elif st in (StateType.SWA, StateType.DSA): if ( target_rank_registration_info is not None and not self.is_mla_backend diff --git a/python/sglang/srt/disaggregation/mori/conn.py b/python/sglang/srt/disaggregation/mori/conn.py index 5f1925af6..45bdb6b50 100644 --- a/python/sglang/srt/disaggregation/mori/conn.py +++ b/python/sglang/srt/disaggregation/mori/conn.py @@ -960,8 +960,8 @@ class MoriKVManager(CommonKVManager): return self._send_mamba_state( peer_info, src_state_indices, dst_state_indices ) - elif state_type in ("swa", "nsa"): - return self._send_swa_nsa_state( + elif state_type in ("swa", "dsa"): + return self._send_swa_dsa_state( peer_info, src_state_indices, dst_state_indices, state_type ) else: @@ -1056,7 +1056,7 @@ class MoriKVManager(CommonKVManager): return statuses - def _send_swa_nsa_state( + def _send_swa_dsa_state( self, peer_info: KVArgsRegisterInfo, src_state_indices: npt.NDArray[np.int32], diff --git a/python/sglang/srt/disaggregation/nixl/conn.py b/python/sglang/srt/disaggregation/nixl/conn.py index 173308cf0..6dc92555d 100644 --- a/python/sglang/srt/disaggregation/nixl/conn.py +++ b/python/sglang/srt/disaggregation/nixl/conn.py @@ -1541,7 +1541,7 @@ class NixlKVManager(CommonKVManager): dst_gpu_id, 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: raise RuntimeError( f"PD Disaggregation does NOT support PD different TP sizes for non-MLA {st.upper()} hybrid models yet." diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index fd267a802..bb07f4012 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -818,7 +818,7 @@ class SchedulerDisaggregationPrefillMixin: 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[ req.req_pool_idx, :seq_len ] @@ -833,8 +833,8 @@ class SchedulerDisaggregationPrefillMixin: state_indices.append(_mamba_payload()) elif st == StateType.SWA: state_indices.append(_swa_payload()) - elif st == StateType.NSA: - state_indices.append(_nsa_payload()) + elif st == StateType.DSA: + state_indices.append(_dsa_payload()) else: state_indices.append(None) diff --git a/python/sglang/srt/disaggregation/utils.py b/python/sglang/srt/disaggregation/utils.py index ad6904a0a..030b0cb7a 100644 --- a/python/sglang/srt/disaggregation/utils.py +++ b/python/sglang/srt/disaggregation/utils.py @@ -567,7 +567,7 @@ def setup_state_kv_args( from sglang.srt.disaggregation.base.conn import StateType 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.memory_pool import HybridLinearKVPool, NSATokenToKVPool + from sglang.srt.mem_cache.memory_pool import DSATokenToKVPool, HybridLinearKVPool kv_args.state_types = [] kv_args.state_data_ptrs = [] @@ -593,9 +593,9 @@ def setup_state_kv_args( append_state_component( 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( - draft_token_to_kv_pool, NSATokenToKVPool + draft_token_to_kv_pool, DSATokenToKVPool ): ( draft_data_ptrs, @@ -612,7 +612,7 @@ def setup_state_kv_args( kv_args.total_kv_layers = total_kv_layers else: append_state_component( - kv_args, StateType.NSA, data_ptrs, data_lens, item_lens + kv_args, StateType.DSA, data_ptrs, data_lens, item_lens ) if ( diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 1e65a50cb..238d8d176 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -134,6 +134,41 @@ class EnvInt(EnvField): 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): def parse(self, value: str) -> float: try: @@ -428,11 +463,18 @@ class Envs: SGLANG_NIXL_EP_BF16_DISPATCH = EnvBool(False) SGLANG_NIXL_EP_NUM_MAX_DISPATCH_TOKENS_PER_RANK = EnvInt(128) - # NSA Backend - SGLANG_NSA_FUSE_TOPK = EnvBool(True) - SGLANG_NSA_ENABLE_MTP_PRECOMPUTE_METADATA = EnvBool(True) + # DSA Backend (canonical names; fall back to SGLANG_NSA_* with deprecation warning) + SGLANG_DSA_FUSE_TOPK = EnvBoolWithAlias(True, deprecated_name="SGLANG_NSA_FUSE_TOPK") + 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_NSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD = EnvInt(2048) # sgl-kernel SGLANG_SKIP_SGL_KERNEL_VERSION_CHECK = EnvBool(False) diff --git a/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py b/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py index c20fdf679..1035a96eb 100644 --- a/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py +++ b/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py @@ -20,7 +20,7 @@ from sglang.srt.hardware_backend.npu.attention.mla_preprocess import ( is_mla_preprocess_enabled, ) 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.radix_attention import AttentionType from sglang.srt.layers.utils.cp_utils import cp_all_gather_rerange_kv_cache @@ -932,7 +932,7 @@ class AscendAttnBackend(AttentionBackend): if ( is_prefill - and is_nsa_enable_prefill_cp() + and is_dsa_enable_prefill_cp() and forward_batch.attn_cp_metadata is not None ): attn_out = self.do_cp_balance_attn( diff --git a/python/sglang/srt/hardware_backend/npu/graph_runner/eagle_draft_extend_npu_graph_runner.py b/python/sglang/srt/hardware_backend/npu/graph_runner/eagle_draft_extend_npu_graph_runner.py index 92308ca46..5415f4160 100644 --- a/python/sglang/srt/hardware_backend/npu/graph_runner/eagle_draft_extend_npu_graph_runner.py +++ b/python/sglang/srt/hardware_backend/npu/graph_runner/eagle_draft_extend_npu_graph_runner.py @@ -20,7 +20,7 @@ from typing import TYPE_CHECKING 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.speculative.eagle_draft_extend_cuda_graph_runner import ( EAGLEDraftExtendCudaGraphRunner, @@ -59,7 +59,7 @@ class EAGLEDraftExtendNpuGraphRunner(EAGLEDraftExtendCudaGraphRunner): ) 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] * ( self.bs - self.raw_bs ) diff --git a/python/sglang/srt/hardware_backend/npu/graph_runner/eagle_draft_npu_graph_runner.py b/python/sglang/srt/hardware_backend/npu/graph_runner/eagle_draft_npu_graph_runner.py index 77c5d4f24..16bc1bf37 100644 --- a/python/sglang/srt/hardware_backend/npu/graph_runner/eagle_draft_npu_graph_runner.py +++ b/python/sglang/srt/hardware_backend/npu/graph_runner/eagle_draft_npu_graph_runner.py @@ -22,7 +22,7 @@ from typing import TYPE_CHECKING, Dict, Union import numpy as np 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.speculative.eagle_draft_cuda_graph_runner import ( EAGLEDraftCudaGraphRunner, @@ -96,7 +96,7 @@ class EAGLEDraftNpuGraphRunner(EAGLEDraftCudaGraphRunner): def _replay(self, forward_batch: ForwardBatch): self.update_attr_name = self._get_update_attr_name() 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 = [] for speculative_step_id in range(self.speculative_num_steps - 1): seq_lens_cpu = forward_batch.seq_lens_cpu + speculative_step_id + 1 diff --git a/python/sglang/srt/hardware_backend/npu/graph_runner/npu_graph_runner.py b/python/sglang/srt/hardware_backend/npu/graph_runner/npu_graph_runner.py index f1a4ebf09..c383331c1 100644 --- a/python/sglang/srt/hardware_backend/npu/graph_runner/npu_graph_runner.py +++ b/python/sglang/srt/hardware_backend/npu/graph_runner/npu_graph_runner.py @@ -26,7 +26,7 @@ import numpy as np import torch 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.environ import envs 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_type = self._get_update_attr_type() # 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(): 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) diff --git a/python/sglang/srt/hardware_backend/npu/modules/deepseek_v2_attention_mla_npu.py b/python/sglang/srt/hardware_backend/npu/modules/deepseek_v2_attention_mla_npu.py index 6726f8589..68f23f1ac 100644 --- a/python/sglang/srt/hardware_backend/npu/modules/deepseek_v2_attention_mla_npu.py +++ b/python/sglang/srt/hardware_backend/npu/modules/deepseek_v2_attention_mla_npu.py @@ -11,9 +11,9 @@ from sglang.srt.hardware_backend.npu.attention.mla_preprocess import ( is_fia_nz, is_mla_preprocess_enabled, ) -from sglang.srt.layers.attention.nsa.nsa_indexer import scattered_to_tp_attn_full -from sglang.srt.layers.attention.nsa.utils import ( - nsa_use_prefill_cp, +from sglang.srt.layers.attention.dsa.dsa_indexer import scattered_to_tp_attn_full +from sglang.srt.layers.attention.dsa.utils import ( + dsa_use_prefill_cp, ) 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 = m.q_b_proj(q_lora)[0].view(-1, m.num_local_heads, m.qk_head_dim) _ = m.indexer( @@ -206,7 +206,7 @@ def forward_mla_prepare_npu( k_nope = m.kv_a_layernorm(k_nope) # q_lora needed by indexer - if m.use_nsa: + if m.use_dsa: q_lora = q 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) - if nsa_use_prefill_cp(forward_batch): + if dsa_use_prefill_cp(forward_batch): # support allgather+rerrange k_nope, k_pe = m.rebuild_cp_kv_cache( latent_cache, forward_batch, k_nope, k_pe @@ -359,7 +359,7 @@ def forward_dsa_prepare_npu( if q_event is not None: torch.npu.current_stream().wait_event(q_event) 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 ): 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) - if nsa_use_prefill_cp(forward_batch): + if dsa_use_prefill_cp(forward_batch): # support allgather+rerrange k_nope, k_pe = m.rebuild_cp_kv_cache( latent_cache, forward_batch, k_nope, k_pe diff --git a/python/sglang/srt/layers/attention/attention_registry.py b/python/sglang/srt/layers/attention/attention_registry.py index 6e0f48bc6..6d86c0189 100644 --- a/python/sglang/srt/layers/attention/attention_registry.py +++ b/python/sglang/srt/layers/attention/attention_registry.py @@ -1,4 +1,5 @@ import logging +import warnings from typing import TYPE_CHECKING from sglang.srt.configs.linear_attn_model_registry import ( @@ -96,11 +97,22 @@ def create_ascend_backend(runner): return AscendAttnBackend(runner) -@register_attention_backend("nsa") -def create_nsa_backend(runner): - from sglang.srt.layers.attention.nsa_backend import NativeSparseAttnBackend +@register_attention_backend("dsa") +def create_dsa_backend(runner): + 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") diff --git a/python/sglang/srt/layers/attention/base_attn_backend.py b/python/sglang/srt/layers/attention/base_attn_backend.py index 7a46a9ba2..37a5e62d0 100644 --- a/python/sglang/srt/layers/attention/base_attn_backend.py +++ b/python/sglang/srt/layers/attention/base_attn_backend.py @@ -9,7 +9,7 @@ from sglang.kernel_api_logging import debug_kernel_api from sglang.srt.utils.common import is_npu 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.model_executor.forward_batch_info import ForwardBatch, ForwardMode from sglang.srt.speculative.spec_info import SpecInput diff --git a/python/sglang/srt/layers/attention/dsa/dequant_k_cache.py b/python/sglang/srt/layers/attention/dsa/dequant_k_cache.py new file mode 100644 index 000000000..5ffca4a39 --- /dev/null +++ b/python/sglang/srt/layers/attention/dsa/dequant_k_cache.py @@ -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") diff --git a/python/sglang/srt/layers/attention/dsa/dsa_backend_mtp_precompute.py b/python/sglang/srt/layers/attention/dsa/dsa_backend_mtp_precompute.py new file mode 100644 index 000000000..1f25bebed --- /dev/null +++ b/python/sglang/srt/layers/attention/dsa/dsa_backend_mtp_precompute.py @@ -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 +) diff --git a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py new file mode 100644 index 000000000..5d2fc69ea --- /dev/null +++ b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py @@ -0,0 +1,1746 @@ +from __future__ import annotations + +import contextlib +import logging +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union + +import torch +from einops import rearrange + +from sglang.jit_kernel.fused_store_index_cache import ( + can_use_dsa_fused_store, + fused_store_index_k_cache, +) +from sglang.srt.environ import envs +from sglang.srt.layers.attention.dsa.utils import ( + aiter_can_use_preshuffle_paged_mqa, + is_dsa_enable_prefill_cp, + is_dsa_prefill_cp_in_seq_split, +) +from sglang.srt.layers.dp_attention import attn_tp_all_gather_into_tensor +from sglang.srt.layers.layernorm import LayerNorm +from sglang.srt.layers.quantization.fp8_kernel import fp8_dtype, is_fp8_fnuz +from sglang.srt.layers.utils import MultiPlatformOp +from sglang.srt.state_capturer.indexer_topk import ( + maybe_capture_indexer_topk, +) +from sglang.srt.utils import ( + add_prefix, + ceil_align, + get_bool_env_var, + is_cuda, + is_gfx95_supported, + is_hip, + is_npu, +) + +logger = logging.getLogger(__name__) + +global _use_multi_stream +_is_cuda = is_cuda() +_is_hip = is_hip() +_is_npu = is_npu() +_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip +_is_fp8_fnuz = is_fp8_fnuz() +_is_gfx95_supported = is_gfx95_supported() +# Whether the aiter preshuffle paged-MQA path (page_size=64 + Preshuffle=True + +# KVBlockSize=64) can be used. Falls back to the legacy page_size=1 / KVBlockSize=1 +# path when the gluon kernel is unavailable (Triton<3.5 and no AOT bundle). +_use_aiter_preshuffle = aiter_can_use_preshuffle_paged_mqa() +if _use_aiter and not _use_aiter_preshuffle: + logger.warning( + "ROCm DSA indexer: aiter preshuffle paged-MQA path is unavailable " + "(needs Triton>=3.5.0 or AITER_ENABLE_AOT_GLUON_PA_MQA_LOGITS=1); " + "falling back to legacy page_size=1 / KVBlockSize=1 path." + ) +if _is_cuda: + try: + import deep_gemm + except ImportError as e: + deep_gemm = e + +if _use_aiter: + from aiter.ops.cache import indexer_k_quant_and_cache + +if is_npu(): + import torch_npu + from sglang.srt.hardware_backend.npu.utils import get_indexer_weight_stream + +from sglang.srt.distributed import ( + get_attn_context_model_parallel_rank, + get_attn_context_model_parallel_world_size, +) +from sglang.srt.distributed.parallel_state import get_pp_group +from sglang.srt.layers import deep_gemm_wrapper +from sglang.srt.layers.communicator import ScatterMode +from sglang.srt.layers.linear import ReplicatedLinear +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.layers.rotary_embedding import get_rope_wrapper +from sglang.srt.layers.utils.cp_utils import cp_all_gather_rerange_output +from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode +from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.server_args import get_global_server_args + +_use_ag_after_qlora = envs.SGLANG_USE_AG_AFTER_QLORA.get() +if TYPE_CHECKING: + from sglang.srt.mem_cache.memory_pool import DSATokenToKVPool + + +DUAL_STREAM_TOKEN_THRESHOLD = 1024 if _is_cuda else 0 + + +class BaseIndexerMetadata(ABC): + @abstractmethod + def get_seqlens_int32(self) -> torch.Tensor: + """ + Return: (batch_size,) int32 tensor + """ + + @abstractmethod + def get_page_table_64(self) -> torch.Tensor: + """ + Return: (batch_size, num_blocks) int32, page table. + The page size of the table is 64. + """ + + @abstractmethod + def get_page_table_1(self) -> torch.Tensor: + """ + Return: (batch_size, num_blocks) int32, page table. + The page size of the table is 1. + """ + + @abstractmethod + def get_seqlens_expanded(self) -> torch.Tensor: + """ + Return: (sum_extend_seq_len,) int32 tensor + """ + + def get_indexer_kvcache_range(self) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Return: (tokens, ), (tokens, ) int32, k_start and k_end in kv cache(token,xxx) for each token. + """ + + def get_indexer_seq_len_cpu(self) -> torch.Tensor: + """ + Return: seq lens for each batch. + """ + + def get_indexer_seq_len(self) -> torch.Tensor: + """ + Return: seq lens for each batch. + """ + + def get_dsa_extend_len_cpu(self) -> List[int]: + """ + Return: extend seq lens for each batch. + """ + + def get_token_to_batch_idx(self) -> torch.Tensor: + """ + Return: batch idx for each token. + """ + + @abstractmethod + def topk_transform( + self, + logits: torch.Tensor, + topk: int, + ) -> torch.Tensor: + """ + Perform topk selection on the logits and possibly transform the result. + + NOTE that attention backend may override this function to do some + transformation, which means the result of this topk_transform may not + be the topk indices of the input logits. + + Return: Anything, since it will be passed to the attention backend + for further processing on sparse attention computation. + Don't assume it is the topk indices of the input logits. + """ + + +def rotate_activation(x: torch.Tensor) -> torch.Tensor: + # from sgl_kernel import hadamard_transform + if _is_hip: + from fast_hadamard_transform import hadamard_transform + else: + from sglang.jit_kernel.hadamard import hadamard_transform + + hidden_size = x.size(-1) + assert ( + hidden_size & (hidden_size - 1) + ) == 0, "Hidden size must be a power of 2 for Hadamard transform." + return hadamard_transform(x, scale=hidden_size**-0.5) + + +class Indexer(MultiPlatformOp): + _MQA_LOGITS_BYTES_PER_ELEM = 4 + _MQA_LOGITS_STATIC_SKIP_ELEMS = 8_000_000 + _MQA_LOGITS_FREE_MEM_FRACTION = 0.5 + _MQA_LOGITS_TOTAL_MEM_FRACTION = 0.3 + _mqa_logits_budget_bytes: Dict[int, int] = {} + + def __init__( + self, + hidden_size: int, + index_n_heads: int, + index_head_dim: int, + rope_head_dim: int, + index_topk: int, + q_lora_rank: int, + max_position_embeddings: int, + rope_theta: float, + layer_id: int, + scale_fmt: Optional[str], + block_size: int = 128, + rope_scaling: Optional[Dict[str, Any]] = None, + is_neox_style: bool = True, + prefix: str = "", + quant_config: Optional[QuantizationConfig] = None, + alt_stream: Optional[torch.cuda.Stream] = None, + ): + super().__init__() + self.hidden_size = hidden_size + self.n_heads = index_n_heads + self.head_dim = index_head_dim + self.rope_head_dim = rope_head_dim + self.index_topk = index_topk + self.q_lora_rank = q_lora_rank + self.layer_id = layer_id + self.alt_stream = alt_stream + self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp() + if self.dsa_enable_prefill_cp: + self.cp_size = get_attn_context_model_parallel_world_size() + self.cp_rank = get_attn_context_model_parallel_rank() + else: + self.cp_size = None + self.cp_rank = None + if _is_cuda: + self.sm_count = deep_gemm.get_num_sms() + self.half_device_sm_count = ceil_align(self.sm_count // 2, 8) + pp_size = get_global_server_args().pp_size + self.logits_with_pp_recv = pp_size > 1 and not get_pp_group().is_last_rank + else: + self.logits_with_pp_recv = False + + self.wq_b = ReplicatedLinear( + self.q_lora_rank, + self.n_heads * self.head_dim, + bias=False, + quant_config=quant_config, + prefix=add_prefix("wq_b", prefix), + ) + + self.wk = ReplicatedLinear( + self.hidden_size, + self.head_dim, + bias=False, + quant_config=quant_config, + prefix=add_prefix("wk", prefix), + ) + self.weights_proj = ReplicatedLinear( + self.hidden_size, + self.n_heads, + bias=False, + params_dtype=torch.bfloat16, + prefix=add_prefix("weights_proj", prefix), + ) + self.k_norm = LayerNorm( + self.head_dim, dtype=torch.bfloat16 if _use_aiter else torch.float32 + ) + self.rotary_emb = get_rope_wrapper( + rope_head_dim, + rotary_dim=rope_head_dim, + max_position=max_position_embeddings, + base=rope_theta, # type: ignore + rope_scaling=rope_scaling, + is_neox_style=is_neox_style, + device=get_global_server_args().device, + ) + self.block_size = block_size + self.scale_fmt = scale_fmt + self.softmax_scale = self.head_dim**-0.5 + + @contextlib.contextmanager + def _with_real_sm_count(self): + # When pipeline parallelism is enabled, each PP rank initiates a recv operation after the _pp_launch_batch + # request to receive the PP proxy tensor or output from the previous stage, occupying one SM resource. + # Model execution runs in parallel with the recv operation, so the SMs available to the indexer must be reduced + # by 1. Currently, the last rank starts the send result + recv request only after waiting for execution results. + if self.logits_with_pp_recv: + pp_recv_sm_count = 1 + with deep_gemm_wrapper.configure_deep_gemm_num_sms( + self.sm_count - pp_recv_sm_count + ): + yield + else: + yield + + def _weights_proj_bf16_in_fp32_out( + self, x: Union[torch.Tensor, Tuple[torch.Tensor, ...]] + ) -> torch.Tensor: + # aiter (ROCm gfx95): extract the passthrough bf16 tensor from the + # 3-tuple (fp8, scale, bf16) produced by fused_rms_fp8_group_quant, + # avoiding an expensive FP8-to-bf16 dequantization. + if _use_aiter and _is_gfx95_supported and isinstance(x, tuple) and len(x) == 3: + x = x[2] + if _is_cuda: + return torch.mm(x, self.weights_proj.weight.t(), out_dtype=torch.float32) + + weights, _ = self.weights_proj(x) + if _is_hip: + # Return bf16; multiplying with q_scale promotes back to fp32. + return weights + return weights.float() + + @torch.compile(dynamic=True) + def _project_and_scale_head_gates( + self, x: Union[torch.Tensor, Tuple[torch.Tensor, ...]] + ): + weights = self._weights_proj_bf16_in_fp32_out(x) + weights = weights * self.n_heads**-0.5 + return weights + + @torch.compile(dynamic=True) + def _get_logits_head_gate( + self, x: Union[torch.Tensor, Tuple[torch.Tensor, ...]], q_scale: torch.Tensor + ): + weights = self._weights_proj_bf16_in_fp32_out(x) + weights = weights * self.n_heads**-0.5 + weights = weights.unsqueeze(-1) * q_scale * self.softmax_scale + return weights + + @torch.compile(dynamic=True) + def _apply_q_scale_and_softmax_scale( + self, weights: torch.Tensor, q_scale: torch.Tensor + ): + return weights.unsqueeze(-1) * q_scale * self.softmax_scale + + def _get_q_k_bf16( + self, + q_lora: torch.Tensor, + x: torch.Tensor, + positions: torch.Tensor, + enable_dual_stream: bool, + forward_batch: ForwardBatch, + ): + if enable_dual_stream: + current_stream = torch.cuda.current_stream() + self.alt_stream.wait_stream(current_stream) + + with deep_gemm_wrapper.configure_deep_gemm_num_sms( + self.half_device_sm_count + ): + query, _ = self.wq_b(q_lora) + query = rearrange(query, "l (h d) -> l h d", d=self.head_dim) + q_rope, _ = torch.split( + query, + [self.rope_head_dim, self.head_dim - self.rope_head_dim], + dim=-1, + ) + with torch.cuda.stream(self.alt_stream): + # TODO we should also put DeepGEMM half SM here? + key, _ = self.wk(x) + key = self.k_norm(key) + + k_rope, _ = torch.split( + key, + [self.rope_head_dim, self.head_dim - self.rope_head_dim], + dim=-1, + ) + + current_stream.wait_stream(self.alt_stream) + else: + query, _ = self.wq_b(q_lora) + query = rearrange(query, "l (h d) -> l h d", d=self.head_dim) + q_rope, _ = torch.split( + query, [self.rope_head_dim, self.head_dim - self.rope_head_dim], dim=-1 + ) + key, _ = self.wk(x) + key = self.k_norm(key) + k_rope, _ = torch.split( + key, [self.rope_head_dim, self.head_dim - self.rope_head_dim], dim=-1 + ) + + q_rope, k_rope = self.rotary_emb(positions, q_rope, k_rope) + + self._update_rope_guarded(query[..., : self.rope_head_dim], q_rope) + self._update_rope_guarded(key[..., : self.rope_head_dim], k_rope) + + if enable_dual_stream: + current_stream = torch.cuda.current_stream() + self.alt_stream.wait_stream(current_stream) + query = rotate_activation(query) + + with torch.cuda.stream(self.alt_stream): + key = rotate_activation(key) + current_stream.wait_stream(self.alt_stream) + elif ( + self.alt_stream is not None + and forward_batch.attn_cp_metadata is not None + and self.dsa_enable_prefill_cp + ): + key = rotate_activation(key) + current_stream = torch.cuda.current_stream() + self.alt_stream.wait_stream(current_stream) + query = rotate_activation(query) + + with torch.cuda.stream(self.alt_stream): + key = cp_all_gather_rerange_output( + key.contiguous(), + self.cp_size, + forward_batch, + torch.cuda.current_stream(), + ) + current_stream.wait_stream(self.alt_stream) + return query, key + else: + query = rotate_activation(query) + key = rotate_activation(key) + + # allgather+rerrange + if forward_batch.attn_cp_metadata is not None and self.dsa_enable_prefill_cp: + key = cp_all_gather_rerange_output( + key.contiguous(), + self.cp_size, + forward_batch, + torch.cuda.current_stream(), + ) + return query, key + + def _get_k_bf16( + self, + x: torch.Tensor, + positions: torch.Tensor, + enable_dual_stream: bool, + ): + # Compute only key, skip query + key, _ = self.wk(x) + key = self.k_norm(key) + k_rope, _ = torch.split( + key, [self.rope_head_dim, self.head_dim - self.rope_head_dim], dim=-1 + ) + + _, k_rope = self.rotary_emb(positions, k_rope, k_rope) + self._update_rope_guarded(key[..., : self.rope_head_dim], k_rope) + key = rotate_activation(key) + + return key + + @staticmethod + def _update_rope_guarded(dst: torch.Tensor, src: torch.Tensor) -> None: + # On AMD with in-place RoPE kernels, self-aliasing can occur; + # skip write-back when src/dst tensors point to a single memory. + if src.data_ptr() == dst.data_ptr(): + return + dst.copy_(src) + + def _get_topk_paged( + self, + forward_batch: ForwardBatch, + layer_id: int, + q_fp8: torch.Tensor, + weights: torch.Tensor, + metadata: BaseIndexerMetadata, + ) -> torch.Tensor: + if TYPE_CHECKING: + assert isinstance(forward_batch.token_to_kv_pool, DSATokenToKVPool) + + page_size = forward_batch.token_to_kv_pool.page_size + # NOTE(dark): blocksize = 64 is hardcoded in deep_gemm + 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 == 1 + ), f"HIP legacy DSA path requires page_size == 1, got {page_size}" + else: + assert page_size == 64, "only support page size 64" + # NOTE(dark): this support extend/decode/decode+graph + if _is_hip and not _use_aiter_preshuffle: + block_tables = metadata.get_page_table_1() + else: + block_tables = metadata.get_page_table_64() + + max_seq_len = block_tables.shape[1] * page_size + kv_cache_fp8 = forward_batch.token_to_kv_pool.get_index_k_with_scale_buffer( + layer_id=layer_id + ) + + blocksize = page_size + if ( + forward_batch.forward_mode.is_target_verify() + or forward_batch.forward_mode.is_draft_extend(include_v2=True) + ): + seqlens_32 = metadata.get_seqlens_expanded() + else: + seqlens_32 = metadata.get_seqlens_int32() + # Reuse pre-computed schedule metadata if available (from init_forward_metadata), + # otherwise fall back to computing it here. + schedule_metadata = getattr(metadata, "paged_mqa_schedule_metadata", None) + # DeepGEMM release-0426 requires context_lens of shape [batch_size, next_n] + # to match q.shape = [batch_size, next_n, heads, head_dim]. The indexer uses + # next_n=1 with batch_size=N_total via q_fp8.unsqueeze(1) below, so mirror + # that layout here. + if seqlens_32.dim() == 2: + seqlens_32_2d = seqlens_32 + else: + seqlens_32_2d = seqlens_32.unsqueeze(-1) + if _is_cuda: + if schedule_metadata is None: + schedule_metadata = deep_gemm.get_paged_mqa_logits_metadata( + seqlens_32_2d, blocksize, self.sm_count + ) + + assert len(q_fp8.shape) == 3 + q_fp8 = q_fp8.unsqueeze(1) # the next_n dim is 1 now + assert len(kv_cache_fp8.shape) == 2 + block_kv = page_size + num_heads_kv = 1 + head_dim_with_sf = 132 + kv_cache_fp8 = kv_cache_fp8.view( + kv_cache_fp8.shape[0], block_kv, num_heads_kv, head_dim_with_sf + ) + assert len(weights.shape) == 3 + weights = weights.squeeze(2) + + # When attn_tp_size > 1 or in the MAX_LEN padding mode, padding may exist in the hidden states, + # and it is necessary to extract the actual q length. + q_offset = sum(metadata.get_dsa_extend_len_cpu()) + if _is_hip: + from aiter.ops.triton.pa_mqa_logits import deepgemm_fp8_paged_mqa_logits + + batch_size, next_n, heads, _ = q_fp8.shape + logits = torch.empty( + (batch_size * next_n, max_seq_len), + device=q_fp8.device, + dtype=torch.float32, + ) + deepgemm_fp8_paged_mqa_logits( + q_fp8, + kv_cache_fp8, + weights, + logits, + seqlens_32, + block_tables, + max_seq_len, + Preshuffle=_use_aiter_preshuffle, + KVBlockSize=block_kv, + ) + else: + logits = deep_gemm.fp8_paged_mqa_logits( + q_fp8[:q_offset], + kv_cache_fp8, + weights[:q_offset], + seqlens_32_2d, + block_tables, + schedule_metadata, + max_seq_len, + clean_logits=False, + ) + + # NOTE(dark): logits should be cleaned in topk_transform + topk_result = metadata.topk_transform(logits, self.index_topk) + # Restore possible padding exist in the hidden states. + if not _is_hip and q_offset < q_fp8.shape[0]: + pad_len = q_fp8.shape[0] - q_offset + padding = torch.full( + (pad_len, topk_result.shape[1]), + -1, + dtype=topk_result.dtype, + device=topk_result.device, + ) + topk_result = torch.cat([topk_result, padding], dim=0) + return topk_result + + def _get_mqa_logits_budget_bytes(self, device_index: int) -> int: + cached_budget = self._mqa_logits_budget_bytes.get(device_index) + if cached_budget is not None: + return cached_budget + + total_mem = torch.cuda.get_device_properties(device_index).total_memory + + total_mem_budget = int(total_mem * self._MQA_LOGITS_TOTAL_MEM_FRACTION) + mem_fraction_static = get_global_server_args().mem_fraction_static + if mem_fraction_static is None: + static_budget = total_mem_budget + else: + static_free_mem = int(total_mem * max(0.0, 1.0 - mem_fraction_static)) + static_budget = min( + int(static_free_mem * self._MQA_LOGITS_FREE_MEM_FRACTION), + total_mem_budget, + ) + static_budget = max(1, static_budget) + + # Keep the static serving-memory guard during CUDA graph capture without + # caching it. The first non-capture prefill path will cache the real + # free-memory budget below. + if get_is_capture_mode(): + return static_budget + + # Match the original free-memory guard: logits_bytes * 2 > free_mem. + # torch.cuda.mem_get_info synchronizes the host, so cache the result, + # capped by the workload-independent serving-memory headroom. + free_mem, _ = torch.cuda.mem_get_info(device_index) + budget_bytes = min( + int(free_mem * self._MQA_LOGITS_FREE_MEM_FRACTION), static_budget + ) + + budget_bytes = max(1, budget_bytes) + self._mqa_logits_budget_bytes[device_index] = budget_bytes + return budget_bytes + + def _should_chunk_mqa_logits( + self, num_q: int, num_k: int, device_index: int + ) -> Tuple[bool, int]: + """ + Detect whether we need to chunk the MQA logits computation to avoid OOM + Return: (need_chunk, logits_budget_bytes) + """ + # Quick static check for normal batches + if num_q * num_k < self._MQA_LOGITS_STATIC_SKIP_ELEMS: + return False, 0 + + logits_bytes = num_q * num_k * self._MQA_LOGITS_BYTES_PER_ELEM + logits_budget_bytes = self._get_mqa_logits_budget_bytes(device_index) + + need_chunk = logits_bytes > logits_budget_bytes + return need_chunk, logits_budget_bytes + + def _get_topk_ragged( + self, + enable_dual_stream: bool, + forward_batch: ForwardBatch, + layer_id: int, + q_fp8: torch.Tensor, + weights: torch.Tensor, + metadata: BaseIndexerMetadata, + ) -> torch.Tensor: + if TYPE_CHECKING: + assert isinstance(forward_batch.token_to_kv_pool, DSATokenToKVPool) + + assert forward_batch.forward_mode.is_extend_without_speculative() + + page_size = forward_batch.token_to_kv_pool.page_size + 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 == 1 + ), f"HIP legacy DSA path requires page_size == 1, got {page_size}" + else: + assert page_size == 64, "only support page size 64" + + assert len(weights.shape) == 3 + assert ( + forward_batch.seq_lens_cpu is not None + and forward_batch.extend_seq_lens_cpu is not None + ) + weights = weights.squeeze(-1) + + if _is_hip and not _use_aiter_preshuffle: + block_tables = metadata.get_page_table_1() + else: + block_tables = metadata.get_page_table_64() + + assert ( + forward_batch.seq_lens_cpu is not None + and forward_batch.extend_seq_lens_cpu is not None + ) + + batch_size = len(block_tables) + token_nums, _, _ = q_fp8.shape + device = q_fp8.device + device_index = device.index + assert device_index is not None, "q_fp8 must be on an indexed CUDA device" + + topk_result = torch.full( + (token_nums, self.index_topk), -1, device=device, dtype=torch.int32 + ) + if batch_size == 0: + return topk_result + + ks, ke = metadata.get_indexer_kvcache_range() + + indexer_seq_lens_cpu = metadata.get_indexer_seq_len_cpu() + seq_len_sum = torch.sum(indexer_seq_lens_cpu).item() + max_seq_len = torch.max(indexer_seq_lens_cpu).item() + k_fp8, k_scale = forward_batch.token_to_kv_pool.get_index_k_scale_buffer( + layer_id, + metadata.get_indexer_seq_len(), + block_tables, + seq_len_sum, + max_seq_len, + ) + if _is_fp8_fnuz: + k_fp8 = k_fp8.view(torch.float8_e4m3fnuz) + else: + k_fp8 = k_fp8.view(torch.float8_e4m3fn) + + k_scale = k_scale.view(torch.float32).squeeze(-1) + kv_fp8 = (k_fp8, k_scale) + + # Check if we need to chunk to avoid OOM + seq_lens_expanded = metadata.get_seqlens_expanded() + token_to_batch_idx = metadata.get_token_to_batch_idx() + q_offset = ks.shape[0] + k_offset = k_fp8.shape[0] + need_chunk, logits_budget_bytes = self._should_chunk_mqa_logits( + q_offset, k_offset, device_index + ) + + if not need_chunk: + assert q_fp8[:q_offset].shape[0] != 0 + with self._with_real_sm_count(): + if _is_hip: + from aiter.ops.triton.fp8_mqa_logits import fp8_mqa_logits + + kv, scale = kv_fp8 + logits = fp8_mqa_logits( + q_fp8[:q_offset], kv, scale, weights[:q_offset], ks, ke + ) + else: + logits = deep_gemm.fp8_mqa_logits( + q_fp8[:q_offset], + kv_fp8, + weights[:q_offset], + ks, + ke, + clean_logits=False, + ) + assert logits.shape[0] == len(seq_lens_expanded) + assert logits.shape[1] == k_offset + + raw_topk_result = metadata.topk_transform(logits, self.index_topk, ks=ks) + topk_result[:q_offset] = raw_topk_result + return topk_result + + bytes_per_row = k_offset * self._MQA_LOGITS_BYTES_PER_ELEM + max_rows = max(1, int(logits_budget_bytes // max(bytes_per_row, 1))) + max_rows = min(max_rows, q_offset) + + global_topk_offset = metadata.attn_metadata.topk_indices_offset + cu_seqlens_q_full = None + if global_topk_offset is None: + cu_seqlens_q_full = torch.ones(q_offset, dtype=torch.int32, device=device) + + assert ( + seq_lens_expanded.shape[0] == q_offset + ), f"seq_lens_expanded length mismatch: {seq_lens_expanded.shape[0]} != {q_offset}" + if global_topk_offset is not None: + assert ( + global_topk_offset.shape[0] >= q_offset + ), f"topk_indices_offset too short: {global_topk_offset.shape[0]} < {q_offset}" + + start = 0 + while start < q_offset: + end = min(start + max_rows, q_offset) + + with self._with_real_sm_count(): + if _is_hip: + from aiter.ops.triton.fp8_mqa_logits import fp8_mqa_logits + + kv, scale = kv_fp8 + logits_chunk = fp8_mqa_logits( + q_fp8[start:end], + kv, + scale, + weights[start:end], + ks[start:end], + ke[start:end], + ) + else: + logits_chunk = deep_gemm.fp8_mqa_logits( + q_fp8[start:end], + kv_fp8, + weights[start:end], + ks[start:end], + ke[start:end], + clean_logits=False, + ) + + lengths_chunk = seq_lens_expanded[start:end] + + # RAGGED: use global offset; PAGED: construct local cu_seqlens_q per chunk + if global_topk_offset is not None: + # RAGGED path + topk_offset_chunk = global_topk_offset[start:end] + cu_seqlens_q_chunk = None + batch_idx_chunk = None + else: + # PAGED path: treat each token as a length-1 sequence + topk_offset_chunk = None + cu_seqlens_q_chunk = cu_seqlens_q_full[start:end] + batch_idx_chunk = token_to_batch_idx[start:end] + + raw_topk_chunk = metadata.topk_transform( + logits_chunk, + self.index_topk, + ks=ks[start:end], + cu_seqlens_q=cu_seqlens_q_chunk, + ke_offset=lengths_chunk, + batch_idx_list=batch_idx_chunk, + topk_indices_offset_override=topk_offset_chunk, + ) + topk_result[start:end] = raw_topk_chunk + start = end + + return topk_result + + def _forward_cuda_k_only( + self, + x: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + layer_id: int, + act_quant, + enable_dual_stream: bool, + metadata: BaseIndexerMetadata, + return_indices: bool = True, + ) -> Optional[torch.Tensor]: + assert forward_batch.forward_mode.is_extend_without_speculative() + x_meta = x[0] if isinstance(x, tuple) else x + + # Fast path: only compute and store k cache, skip all q and weights ops + key = self._get_k_bf16(x, positions, enable_dual_stream) + + if not forward_batch.out_cache_loc.is_contiguous(): + forward_batch.out_cache_loc = forward_batch.out_cache_loc.contiguous() + + self._store_index_k_cache( + forward_batch=forward_batch, + layer_id=layer_id, + key=key, + act_quant=act_quant, + ) + + # MHA doesn't need topk_indices + if not return_indices: + return None + + # MLA: use dummy logits with topk kernel's fast path to generate indices + # When length <= 2048, naive_topk_cuda directly generates [0,1,...,length-1,-1,...] + seq_lens_expanded = metadata.get_seqlens_expanded() + dummy_logits = torch.zeros( + seq_lens_expanded.shape[0], + self.index_topk, + dtype=torch.float32, + device=x_meta.device, + ) + return metadata.topk_transform(dummy_logits, self.index_topk) + + def _get_topk_ragged_with_cp( + self, + forward_batch: ForwardBatch, + layer_id: int, + q_fp8: torch.Tensor, + weights: torch.Tensor, + metadata: BaseIndexerMetadata, + kv_len: int, + actual_seq_q: int, + cp_index: List[Tuple[int, int, int]] = None, + ) -> torch.Tensor: + if TYPE_CHECKING: + assert isinstance(forward_batch.token_to_kv_pool, DSATokenToKVPool) + + page_size = forward_batch.token_to_kv_pool.page_size + assert page_size == 64, "only support page size 64" + assert len(weights.shape) == 3 + weights = weights.squeeze(-1) + k_fp8_list = [] + k_scale_list = [] + ks_list = [] + ke_offset_list = [] + offset = 0 + actual_seq_q_list = [] + batch_idx_list = [] + + block_tables = metadata.get_page_table_64() + + assert ( + forward_batch.seq_lens_cpu is not None + and forward_batch.extend_seq_lens_cpu is not None + ) + if cp_index is not None: + # TODO Multi-batch support has accuracy issues + for batch_idx, start_seq_position, end_seq_position in cp_index: + pre_chunk_offset = ( + forward_batch.seq_lens_cpu[batch_idx].item() + - forward_batch.extend_seq_lens_cpu[batch_idx] + ) + start_seq_position += pre_chunk_offset + end_seq_position += pre_chunk_offset + if offset == 0 and batch_idx != 0: + offset += forward_batch.extend_seq_lens_cpu[batch_idx - 1] + k_fp8 = forward_batch.token_to_kv_pool.get_index_k_continuous( + layer_id, + end_seq_position, + block_tables[batch_idx], + ) + k_scale = forward_batch.token_to_kv_pool.get_index_k_scale_continuous( + layer_id, + end_seq_position, + block_tables[batch_idx], + ) + + extend_seq_len = end_seq_position - start_seq_position + ks = torch.full( + (extend_seq_len,), offset, dtype=torch.int32, device="cuda" + ) + k_fp8_list.append(k_fp8) + k_scale_list.append(k_scale) + ks_list.append(ks) + ke_offset = torch.arange( + start_seq_position + 1, + end_seq_position + 1, + dtype=torch.int32, + device="cuda", + ) + ke_offset_list.append(ke_offset) + actual_seq_q = torch.tensor( + [extend_seq_len], dtype=torch.int32, device="cuda" + ) + actual_seq_q_list.append(actual_seq_q) + batch_idx_list.append(batch_idx) + + k_fp8 = torch.cat(k_fp8_list, dim=0).view(torch.float8_e4m3fn) + k_scale = torch.cat(k_scale_list, dim=0).view(torch.float32).squeeze(-1) + kv_fp8 = (k_fp8, k_scale) + ks = torch.cat(ks_list, dim=0) + ke_offset = torch.cat(ke_offset_list, dim=0) + ke = ks + ke_offset + actual_seq_q = torch.cat(actual_seq_q_list, dim=0) + with self._with_real_sm_count(): + logits = deep_gemm.fp8_mqa_logits( + q_fp8, + kv_fp8, + weights, + ks, + ke, + clean_logits=False, + ) + topk_result = metadata.topk_transform( + logits, + self.index_topk, + ks=ks, + cu_seqlens_q=actual_seq_q, + ke_offset=ke_offset, + batch_idx_list=batch_idx_list, + ) + else: + kv_len = ( + forward_batch.seq_lens_cpu[0].item() + - forward_batch.extend_seq_lens_cpu[0] + + kv_len + ) + k_fp8 = forward_batch.token_to_kv_pool.get_index_k_continuous( + layer_id, + kv_len, + block_tables[0], + ) + k_scale = forward_batch.token_to_kv_pool.get_index_k_scale_continuous( + layer_id, + kv_len, + block_tables[0], + ) + + k_fp8 = k_fp8.view(torch.float8_e4m3fn) + k_scale = k_scale.view(torch.float32).squeeze(-1) + kv_fp8 = (k_fp8, k_scale) + ks = torch.full((actual_seq_q,), offset, dtype=torch.int32, device="cuda") + ke_offset = torch.arange( + (kv_len - actual_seq_q) + 1, + kv_len + 1, + dtype=torch.int32, + device="cuda", + ) + ke = ks + ke_offset + + with self._with_real_sm_count(): + logits = deep_gemm.fp8_mqa_logits( + q_fp8, + kv_fp8, + weights, + ks, + ke, + clean_logits=False, + ) + actual_seq_q = torch.tensor([actual_seq_q], dtype=torch.int32).to( + device="cuda", non_blocking=True + ) + topk_result = metadata.topk_transform( + logits, + self.index_topk, + ks=ks, + cu_seqlens_q=actual_seq_q, + ke_offset=ke_offset, + ) + + return topk_result + + def forward_indexer( + self, + q_fp8: torch.Tensor, + weights: torch.Tensor, + forward_batch: ForwardBatch, + topk: int, + layer_id: int, + ) -> Optional[torch.Tensor]: + if not _is_npu: + from sglang.srt.layers.attention.dsa.tilelang_kernel import fp8_index + + page_size = forward_batch.token_to_kv_pool.page_size + assert page_size == 64, "only support page size 64" + + assert len(weights.shape) == 3 + weights = weights.squeeze(-1) + + # logits = deep_gemm.fp8_mqa_logits(q_fp8, kv_fp8, weights, ks, ke) + k_fp8_list = [] + k_scale_list = [] + + topk_indices_list = [] + + block_tables = forward_batch.req_to_token_pool.req_to_token[ + forward_batch.req_pool_indices, : + ] + strided_indices = torch.arange( + 0, block_tables.shape[-1], page_size, device="cuda" + ) + block_tables = block_tables[:, strided_indices] // page_size + + q_len_start = 0 + + for i in range(forward_batch.batch_size): + seq_len = forward_batch.seq_lens[i].item() + q_len = ( + forward_batch.extend_seq_lens_cpu[i] + if forward_batch.forward_mode.is_extend() + else 1 + ) + q_len_end = q_len_start + q_len + + q_fp8_partial = q_fp8[q_len_start:q_len_end] + q_fp8_partial = q_fp8_partial.unsqueeze(0).contiguous() + + weights_partial = weights[q_len_start:q_len_end] + weights_partial = weights_partial.squeeze(-1).unsqueeze(0).contiguous() + + k_fp8 = forward_batch.token_to_kv_pool.get_index_k_continuous( + layer_id, + seq_len, + block_tables[i], + ) + k_scale = forward_batch.token_to_kv_pool.get_index_k_scale_continuous( + layer_id, + seq_len, + block_tables[i], + ) + + k_fp8 = k_fp8.view(torch.float8_e4m3fn).unsqueeze(0).contiguous() + k_scale = k_scale.view(torch.float32).squeeze(-1).unsqueeze(0).contiguous() + + index_score = fp8_index( + q_fp8_partial, + weights_partial, + k_fp8, + k_scale, + ) + end_pos = seq_len + topk_indices = index_score.topk(min(topk, end_pos), dim=-1)[1].squeeze(0) + + pad_len = ceil_align(topk_indices.shape[-1], 2048) - topk_indices.shape[-1] + topk_indices = torch.nn.functional.pad( + topk_indices, (0, pad_len), "constant", -1 + ) + + topk_indices_list.append(topk_indices) + + q_len_start = q_len_end + + topk_indices = torch.cat(topk_indices_list, dim=0) + return topk_indices + + def _store_index_k_cache( + self, + forward_batch: ForwardBatch, + layer_id: int, + key: torch.Tensor, + *, + act_quant=None, # fallback only + ) -> None: + """ + Store DSA indexer K cache for current step. + + Preferred: fused_store_index_k_cache(key, cache, out_cache_loc, page_size) + Fallback : act_quant(key) + token_to_kv_pool.set_index_k_scale_buffer(...) + """ + + # Fast path: JIT fused store (CUDA, page_size=64, non-fnuz) + if ( + _is_cuda + and (not _is_fp8_fnuz) + and can_use_dsa_fused_store( + key.dtype, + forward_batch.out_cache_loc.dtype, + forward_batch.token_to_kv_pool.page_size, + ) + ): + # NOTE: wrapper already normalizes shape/contiguity and asserts dtypes. + buf = forward_batch.token_to_kv_pool.get_index_k_with_scale_buffer( + layer_id=layer_id + ) + fused_store_index_k_cache( + key, + buf, + forward_batch.out_cache_loc, + forward_batch.token_to_kv_pool.page_size, + ) + return + + # Fast path: AITER fused quant + cache store + # When _use_aiter_preshuffle is True we use the new MFMA 16x16 preshuffle + # layout (page_size>=16). Otherwise we fall back to the legacy row-major + # layout with page_size=1; the same kv_cache.view works for both cases + # because page_size is 1 there. + if _use_aiter: + page_size = forward_batch.token_to_kv_pool.page_size + buf = forward_batch.token_to_kv_pool.get_index_k_with_scale_buffer( + layer_id=layer_id + ) + kv_cache = buf.view(-1, page_size, 132).view(fp8_dtype) + out_loc = forward_batch.out_cache_loc + if not out_loc.is_contiguous(): + out_loc = out_loc.contiguous() + indexer_k_quant_and_cache( + key, + kv_cache, + out_loc, + self.block_size, + self.scale_fmt, + preshuffle=_use_aiter_preshuffle, + ) + return + + # Fallback: original path + assert act_quant is not None + k_fp8, k_scale = act_quant(key, self.block_size, self.scale_fmt) + + out_loc = forward_batch.out_cache_loc + if not out_loc.is_contiguous(): + out_loc = out_loc.contiguous() + + forward_batch.token_to_kv_pool.set_index_k_scale_buffer( + layer_id=layer_id, + loc=out_loc, + index_k=k_fp8, + index_k_scale=k_scale, + ) + + def forward_xpu( + self, + x: torch.Tensor, + q_lora: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + layer_id: int, + return_indices: bool = True, + ) -> Optional[torch.Tensor]: + return self.forward_cuda( + x, q_lora, positions, forward_batch, layer_id, return_indices + ) + + def forward_cuda( + self, + x: torch.Tensor, + q_lora: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + layer_id: int, + return_indices: bool = True, + ) -> Optional[torch.Tensor]: + if _is_hip: + from sglang.srt.layers.attention.dsa.tilelang_kernel import act_quant + elif not _is_npu: + from sglang.srt.layers.attention.dsa.triton_kernel import act_quant + + if TYPE_CHECKING: + assert isinstance(forward_batch.token_to_kv_pool, DSATokenToKVPool) + + # When upstream uses fused FP8 RMSNorm+quant, activations may be passed as + # a tuple like (x_fp8, x_scale[, y]). Use `x_meta` for shape/device queries. + x_meta = x[0] if isinstance(x, tuple) else x + + metadata = forward_batch.attn_backend.get_indexer_metadata( + layer_id, forward_batch + ) + + enable_dual_stream = ( + self.alt_stream is not None + and get_is_capture_mode() + and q_lora.shape[0] > 0 + and q_lora.shape[0] <= DUAL_STREAM_TOKEN_THRESHOLD + ) + + # skip DSA if attention backend choose to skip this batch + if metadata is None: + return None + + # Determine if should skip topk based on sequence length + # We can only skip the logits computation if cuda graph is not involved + skip_logits_computation = False + if forward_batch.forward_mode.is_extend_without_speculative(): + if forward_batch.seq_lens_cpu is not None: + max_kv_len = forward_batch.seq_lens_cpu.max().item() + skip_logits_computation = max_kv_len <= self.index_topk + + # Optimization: fast path when skipping topk computation + if skip_logits_computation and (not self.dsa_enable_prefill_cp): + return maybe_capture_indexer_topk( + layer_id, + self._forward_cuda_k_only( + x, + positions, + forward_batch, + layer_id, + act_quant, + enable_dual_stream, + metadata, + return_indices, + ), + ) + + if enable_dual_stream and forward_batch.forward_mode.is_decode_or_idle(): + current_stream = torch.cuda.current_stream() + self.alt_stream.wait_stream(current_stream) + weights = self._project_and_scale_head_gates(x) + query, key = self._get_q_k_bf16( + q_lora, x, positions, enable_dual_stream, forward_batch=forward_batch + ) + q_fp8, q_scale = act_quant(query, self.block_size, self.scale_fmt) + with torch.cuda.stream(self.alt_stream): + self._store_index_k_cache( + forward_batch=forward_batch, + layer_id=layer_id, + key=key, + act_quant=act_quant, + ) + current_stream.wait_stream(self.alt_stream) + weights = self._apply_q_scale_and_softmax_scale(weights, q_scale) + else: + query, key = self._get_q_k_bf16( + q_lora, x, positions, enable_dual_stream, forward_batch=forward_batch + ) + + if enable_dual_stream: + current_stream = torch.cuda.current_stream() + self.alt_stream.wait_stream(current_stream) + + q_fp8, q_scale = act_quant(query, self.block_size, self.scale_fmt) + with torch.cuda.stream(self.alt_stream): + self._store_index_k_cache( + forward_batch=forward_batch, + layer_id=layer_id, + key=key, + act_quant=act_quant, + ) + current_stream.wait_stream(self.alt_stream) + else: + q_fp8, q_scale = act_quant(query, self.block_size, self.scale_fmt) + self._store_index_k_cache( + forward_batch=forward_batch, + layer_id=layer_id, + key=key, + act_quant=act_quant, + ) + + # aiter (ROCm gfx95): the 3-tuple (fp8, scale, bf16) from + # fused_rms_fp8_group_quant is passed directly to _get_logits_head_gate, + # which extracts the bf16 tensor via _weights_proj_bf16_in_fp32_out, + # completely skipping the FP8 dequantization path below. + if ( + _use_aiter + and _is_gfx95_supported + and isinstance(x, tuple) + and len(x) == 3 + ): + x_for_gate = x + elif isinstance(x, tuple): + assert len(x) in ( + 2, + 3, + ), "For tuple input, only (x, x_s) or (x, x_s, y) formats are accepted" + x_q, x_s = x[0], x[1] + if ( + x_s is not None + and x_q.dim() == 2 + and x_s.dim() == 2 + and x_q.shape[0] == x_s.shape[0] + ): + m, n = x_q.shape + ng = x_s.shape[1] + if ng > 0 and n % ng == 0: + group = n // ng + x_for_gate = ( + x_q.to(torch.float32) + .view(m, ng, group) + .mul_(x_s.to(torch.float32).unsqueeze(-1)) + .view(m, n) + .to(torch.bfloat16) + ) + else: + x_for_gate = x_q.to(torch.bfloat16) + else: + x_for_gate = x_q.to(torch.bfloat16) + else: + x_for_gate = x + + weights = self._get_logits_head_gate(x_for_gate, q_scale) + + if _is_cuda or _is_hip: + assert forward_batch.seq_lens_cpu is not None + if len(forward_batch.seq_lens_cpu) == 0: + # this seems b/c max-pad, no worries? + # if x.shape[0] != 0: + # print( + # "HACK: seq_lens empty but x not empty, hackily return all-invalid topk_result" + # ) + return maybe_capture_indexer_topk( + layer_id, + torch.full( + (x_meta.shape[0], self.index_topk), + -1, + dtype=torch.int, + device=x_meta.device, + ), + ) + + if ( + forward_batch.forward_mode.is_decode_or_idle() + or forward_batch.forward_mode.is_target_verify() + or forward_batch.forward_mode.is_draft_extend(include_v2=True) + ): + topk_result = self._get_topk_paged( + forward_batch, layer_id, q_fp8, weights, metadata + ) + else: + if ( + forward_batch.attn_cp_metadata is not None + and is_dsa_prefill_cp_in_seq_split() + ): + kv_len_prev = forward_batch.attn_cp_metadata.kv_len_prev + kv_len_next = forward_batch.attn_cp_metadata.kv_len_next + actual_seq_q_prev = forward_batch.attn_cp_metadata.actual_seq_q_prev + actual_seq_q_next = forward_batch.attn_cp_metadata.actual_seq_q_next + + # TODO support mutil-batch + # cp_batch_seq_index_prev = forward_batch.attn_cp_metadata["cp_batch_seq_index_prev"] + # cp_batch_seq_index_next = forward_batch.attn_cp_metadata["cp_batch_seq_index_next"] + # TODO prev, next, combined into a single call + q_fp8_prev, q_fp8_next = torch.split( + q_fp8, (q_fp8.shape[0] + 1) // 2, dim=0 + ) + weights_prev, weights_next = torch.split( + weights, (weights.shape[0] + 1) // 2, dim=0 + ) + topk_result_prev = self._get_topk_ragged_with_cp( + forward_batch, + layer_id, + q_fp8_prev, + weights_prev, + metadata, + kv_len_prev, + actual_seq_q_prev, + ) + + topk_result_next = self._get_topk_ragged_with_cp( + forward_batch, + layer_id, + q_fp8_next, + weights_next, + metadata, + kv_len_next, + actual_seq_q_next, + ) + return maybe_capture_indexer_topk( + layer_id, + torch.cat([topk_result_prev, topk_result_next], dim=0), + ) + else: + topk_result = self._get_topk_ragged( + enable_dual_stream, + forward_batch, + layer_id, + q_fp8, + weights, + metadata, + ) + else: + topk_result = self.forward_indexer( + q_fp8.contiguous(), + weights, + forward_batch, + topk=self.index_topk, + layer_id=layer_id, + ) + return maybe_capture_indexer_topk(layer_id, topk_result) + + def forward_npu( + self, + x: torch.Tensor, + q_lora: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + layer_id: int, + layer_scatter_modes=None, + dynamic_scale: torch.Tensor = None, + ) -> torch.Tensor: + if forward_batch.attn_backend.forward_metadata.seq_lens_cpu_int is None: + actual_seq_lengths_kv = forward_batch.attn_backend.forward_metadata.seq_lens + else: + actual_seq_lengths_kv = ( + forward_batch.attn_backend.forward_metadata.seq_lens_cpu_int + ) + is_prefill = ( + forward_batch.forward_mode.is_extend() + and not forward_batch.forward_mode.is_draft_extend_v2() + and not forward_batch.forward_mode.is_target_verify() + and not forward_batch.forward_mode.is_draft_extend() + ) + + bs = q_lora.shape[0] + + if self.rotary_emb.is_neox_style: + if not hasattr(forward_batch, "npu_indexer_sin_cos_cache"): + cos_sin = self.rotary_emb.cos_sin_cache[positions] + cos, sin = cos_sin.chunk(2, dim=-1) + cos = cos.repeat(1, 2).view(-1, 1, 1, self.rope_head_dim) + sin = sin.repeat(1, 2).view(-1, 1, 1, self.rope_head_dim) + forward_batch.npu_indexer_sin_cos_cache = (sin, cos) + else: + sin, cos = forward_batch.npu_indexer_sin_cos_cache + + if self.alt_stream is not None: + self.alt_stream.wait_stream(torch.npu.current_stream()) + with torch.npu.stream(self.alt_stream): + q_lora = ( + (q_lora, dynamic_scale) if dynamic_scale is not None else q_lora + ) + q = self.wq_b(q_lora)[ + 0 + ] # [bs, 1536] @ [1536, 64 * 128] = [bs, 64 * 128] + wq_b_event = self.alt_stream.record_event() + q = q.view(bs, self.n_heads, self.head_dim) # [bs, 64, 128] + q_pe, q_nope = torch.split( + q, + [self.rope_head_dim, self.head_dim - self.rope_head_dim], + dim=-1, + ) # [bs, 64, 64 + 64] + q_pe = q_pe.view(bs, self.n_heads, 1, self.rope_head_dim) + q_pe = torch_npu.npu_rotary_mul(q_pe, cos, sin).view( + bs, self.n_heads, self.rope_head_dim + ) # [bs, n, d] + q = torch.cat([q_pe, q_nope], dim=-1) + q.record_stream(self.alt_stream) + q_rope_event = self.alt_stream.record_event() + else: + q_lora = ( + (q_lora, dynamic_scale) if dynamic_scale is not None else q_lora + ) + q = self.wq_b(q_lora)[ + 0 + ] # [bs, 1536] @ [1536, 64 * 128] = [bs, 64 * 128] + q = q.view(bs, self.n_heads, self.head_dim) # [bs, 64, 128] + q_pe, q_nope = torch.split( + q, + [self.rope_head_dim, self.head_dim - self.rope_head_dim], + dim=-1, + ) # [bs, 64, 64 + 64] + q_pe = q_pe.view(bs, self.n_heads, 1, self.rope_head_dim) + q_pe = torch_npu.npu_rotary_mul(q_pe, cos, sin).view( + bs, self.n_heads, self.rope_head_dim + ) # [bs, n, d] + q = torch.cat([q_pe, q_nope], dim=-1) + + if envs.SGLANG_NPU_USE_MULTI_STREAM.get(): + indexer_weight_stream = get_indexer_weight_stream() + indexer_weight_stream.wait_stream(torch.npu.current_stream()) + with torch.npu.stream(indexer_weight_stream): + x = x.view(-1, self.hidden_size) + weights = self.weights_proj(x.float())[0].to(torch.bfloat16) + weights.record_stream(indexer_weight_stream) + weights_event = indexer_weight_stream.record_event() + else: + x = x.view(-1, self.hidden_size) + weights = self.weights_proj(x.float())[0].to(torch.bfloat16) + + k_proj = self.wk(x)[0] # [b, s, 7168] @ [7168, 128] = [b, s, 128] + k = self.k_norm(k_proj) + if ( + _use_ag_after_qlora + and layer_scatter_modes.layer_input_mode == ScatterMode.SCATTERED + and layer_scatter_modes.attn_mode == ScatterMode.TP_ATTN_FULL + ): + k = scattered_to_tp_attn_full(k, forward_batch) + k_pe, k_nope = torch.split( + k, + [self.rope_head_dim, self.head_dim - self.rope_head_dim], + dim=-1, + ) # [bs, 64 + 64] + + k_pe = k_pe.view(-1, 1, 1, self.rope_head_dim) + k_pe = torch.ops.npu.npu_rotary_mul(k_pe, cos, sin).view( + bs, 1, self.rope_head_dim + ) # [bs, 1, d] + k = torch.cat([k_pe, k_nope.unsqueeze(1)], dim=-1) # [bs, 1, 128] + + else: + if envs.SGLANG_NPU_USE_MULTI_STREAM.get(): + indexer_weight_stream = get_indexer_weight_stream() + indexer_weight_stream.wait_stream(torch.npu.current_stream()) + with torch.npu.stream(indexer_weight_stream): + x = x.view(-1, self.hidden_size) + weights = self.weights_proj(x.float())[0].to(torch.bfloat16) + weights.record_stream(indexer_weight_stream) + weights_event = indexer_weight_stream.record_event() + else: + x = x.view(-1, self.hidden_size) + weights = self.weights_proj(x.float())[0].to(torch.bfloat16) + + q_lora = (q_lora, dynamic_scale) if dynamic_scale is not None else q_lora + q = self.wq_b(q_lora)[0] # [bs, 1536] @ [1536, 64 * 128] = [bs, 64 * 128] + q = q.view(bs, self.n_heads, self.head_dim) # [bs, 64, 128] + q_pe, q_nope = torch.split( + q, + [self.rope_head_dim, self.head_dim - self.rope_head_dim], + dim=-1, + ) # [bs, 64, 64 + 64] + + k_proj = self.wk(x)[0] # [b, s, 7168] @ [7168, 128] = [b, s, 128] + k = self.k_norm(k_proj) + k_pe, k_nope = torch.split( + k, + [self.rope_head_dim, self.head_dim - self.rope_head_dim], + dim=-1, + ) # [bs, 64 + 64] + + k_pe = k_pe.unsqueeze(1) + + if layer_id == 0: + self.rotary_emb.sin_cos_cache = ( + self.rotary_emb.cos_sin_cache.index_select(0, positions) + ) + + q_pe, k_pe = self.rotary_emb(positions, q_pe, k_pe) + k_pe = k_pe.squeeze(1) + q = torch.cat([q_pe, q_nope], dim=-1) + k = torch.cat([k_pe, k_nope], dim=-1) + + if ( + is_prefill + and self.dsa_enable_prefill_cp + and forward_batch.attn_cp_metadata is not None + ): + k = cp_all_gather_rerange_output( + k.contiguous().view(-1, self.head_dim), + self.cp_size, + forward_batch, + torch.npu.current_stream(), + ) + + forward_batch.token_to_kv_pool.set_index_k_buffer( + layer_id, forward_batch.out_cache_loc, k + ) + if is_prefill: + if ( + self.dsa_enable_prefill_cp + and forward_batch.attn_cp_metadata is not None + ): + forward_batch.attn_backend.forward_metadata.actual_seq_lengths_q = ( + forward_batch.attn_cp_metadata.actual_seq_q_prev_tensor, + forward_batch.attn_cp_metadata.actual_seq_q_next_tensor, + ) + if sum(forward_batch.extend_prefix_lens_cpu) > 0: + total_kv_len_prev_tensor = ( + forward_batch.attn_cp_metadata.kv_len_prev_tensor + + forward_batch.extend_prefix_lens.squeeze() + ) + total_kv_len_next_tensor = ( + forward_batch.attn_cp_metadata.kv_len_next_tensor + + forward_batch.extend_prefix_lens.squeeze() + ) + forward_batch.attn_backend.forward_metadata.actual_seq_lengths_kv = ( + total_kv_len_prev_tensor, + total_kv_len_next_tensor, + ) + else: + forward_batch.attn_backend.forward_metadata.actual_seq_lengths_kv = ( + forward_batch.attn_cp_metadata.kv_len_prev_tensor, + forward_batch.attn_cp_metadata.kv_len_next_tensor, + ) + actual_seq_lengths_q = ( + forward_batch.attn_backend.forward_metadata.actual_seq_lengths_q + ) + actual_seq_lengths_kv = ( + forward_batch.attn_backend.forward_metadata.actual_seq_lengths_kv + ) + else: + actual_seq_lengths_kv = forward_batch.seq_lens + actual_seq_lengths_q = forward_batch.extend_seq_lens.cumsum(dim=0) + else: + if forward_batch.attn_backend.forward_metadata.actual_seq_lengths_q is None: + if ( + forward_batch.forward_mode.is_draft_extend_v2() + or forward_batch.forward_mode.is_target_verify() + or forward_batch.forward_mode.is_draft_extend() + ): + num_draft_tokens = ( + forward_batch.attn_backend.speculative_num_draft_tokens + ) + actual_seq_lengths_q = torch.arange( + num_draft_tokens, + num_draft_tokens + bs, + num_draft_tokens, + dtype=torch.int32, + device=k.device, + ) + else: + actual_seq_lengths_q = torch.tensor( + [1 + i * 1 for i in range(bs)], + dtype=torch.int32, + device=k.device, + ) + else: + actual_seq_lengths_q = ( + forward_batch.attn_backend.forward_metadata.actual_seq_lengths_q + ) + + past_key_states = forward_batch.token_to_kv_pool.get_index_k_buffer(layer_id) + + if self.rotary_emb.is_neox_style and self.alt_stream is not None: + torch.npu.current_stream().wait_event(q_rope_event) + if envs.SGLANG_NPU_USE_MULTI_STREAM.get(): + torch.npu.current_stream().wait_event(weights_event) + if ( + _use_ag_after_qlora + and layer_scatter_modes.layer_input_mode == ScatterMode.SCATTERED + and layer_scatter_modes.attn_mode == ScatterMode.TP_ATTN_FULL + ): + weights = scattered_to_tp_attn_full(weights, forward_batch) + block_table = forward_batch.attn_backend.forward_metadata.block_tables + if ( + is_prefill + and self.dsa_enable_prefill_cp + and forward_batch.attn_cp_metadata is not None + ): + block_table = block_table[: actual_seq_lengths_q[0].numel()] + topk_indices = self.do_npu_cp_balance_indexer( + q.view(-1, self.n_heads, self.head_dim), + past_key_states, + weights, + actual_seq_lengths_q, + actual_seq_lengths_kv, + block_table, + ) + return topk_indices + else: + block_table = ( + block_table[: actual_seq_lengths_q.size()[0]] + if is_prefill + else block_table + ) + + topk_indices = torch_npu.npu_lightning_indexer( + query=q.view(-1, self.n_heads, self.head_dim), + key=past_key_states, + weights=weights, + actual_seq_lengths_query=actual_seq_lengths_q.to(torch.int32), + actual_seq_lengths_key=actual_seq_lengths_kv.to(k.device).to( + torch.int32 + ), + block_table=block_table, + layout_query="TND", + layout_key="PA_BSND", + sparse_count=self.index_topk, + sparse_mode=3, + ) + return topk_indices[0] + + def do_npu_cp_balance_indexer( + self, + q, + past_key_states, + indexer_weights, + actual_seq_lengths_q, + actual_seq_lengths_kv, + block_table, + ): + q_prev, q_next = torch.split(q, (q.size(0) + 1) // 2, dim=0) + weights_prev, weights_next = None, None + if indexer_weights is not None: + weights_prev, weights_next = torch.split( + indexer_weights, (indexer_weights.size(0) + 1) // 2, dim=0 + ) + weights_prev = weights_prev.contiguous().view(-1, weights_prev.shape[-1]) + weights_next = weights_next.contiguous().view(-1, weights_next.shape[-1]) + + actual_seq_lengths_q_prev, actual_seq_lengths_q_next = actual_seq_lengths_q + actual_seq_lengths_kv_prev, actual_seq_lengths_kv_next = actual_seq_lengths_kv + + topk_indices_prev = torch_npu.npu_lightning_indexer( + query=q_prev, + key=past_key_states, + weights=weights_prev, + actual_seq_lengths_query=actual_seq_lengths_q_prev.to( + device=q.device, dtype=torch.int32 + ), + actual_seq_lengths_key=actual_seq_lengths_kv_prev.to( + device=q.device, dtype=torch.int32 + ), + block_table=block_table, + layout_query="TND", + layout_key="PA_BSND", + sparse_count=self.index_topk, + sparse_mode=3, + ) + topk_indices_next = torch_npu.npu_lightning_indexer( + query=q_next, + key=past_key_states, + weights=weights_next, + actual_seq_lengths_query=actual_seq_lengths_q_next.to( + device=q.device, dtype=torch.int32 + ), + actual_seq_lengths_key=actual_seq_lengths_kv_next.to( + device=q.device, dtype=torch.int32 + ), + block_table=block_table, + layout_query="TND", + layout_key="PA_BSND", + sparse_count=self.index_topk, + sparse_mode=3, + ) + return topk_indices_prev[0], topk_indices_next[0] + + +def scattered_to_tp_attn_full( + hidden_states: torch.Tensor, + forward_batch, +) -> torch.Tensor: + hidden_states, local_hidden_states = ( + torch.empty( + (forward_batch.input_ids.shape[0], hidden_states.shape[1]), + dtype=hidden_states.dtype, + device=hidden_states.device, + ), + hidden_states, + ) + attn_tp_all_gather_into_tensor(hidden_states, local_hidden_states.contiguous()) + return hidden_states diff --git a/python/sglang/srt/layers/attention/dsa/dsa_mtp_verification.py b/python/sglang/srt/layers/attention/dsa/dsa_mtp_verification.py new file mode 100644 index 000000000..bd1ed48f9 --- /dev/null +++ b/python/sglang/srt/layers/attention/dsa/dsa_mtp_verification.py @@ -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"], + ) diff --git a/python/sglang/srt/layers/attention/dsa/index_buf_accessor.py b/python/sglang/srt/layers/attention/dsa/index_buf_accessor.py new file mode 100644 index 000000000..95acd1658 --- /dev/null +++ b/python/sglang/srt/layers/attention/dsa/index_buf_accessor.py @@ -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) diff --git a/python/sglang/srt/layers/attention/dsa/quant_k_cache.py b/python/sglang/srt/layers/attention/dsa/quant_k_cache.py new file mode 100644 index 000000000..5454071b8 --- /dev/null +++ b/python/sglang/srt/layers/attention/dsa/quant_k_cache.py @@ -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") diff --git a/python/sglang/srt/layers/attention/dsa/tilelang_kernel.py b/python/sglang/srt/layers/attention/dsa/tilelang_kernel.py new file mode 100644 index 000000000..62509c308 --- /dev/null +++ b/python/sglang/srt/layers/attention/dsa/tilelang_kernel.py @@ -0,0 +1,2589 @@ +import functools +from functools import lru_cache +from typing import Any, Optional, Tuple + +import tilelang +import tilelang.language as T +import torch + +from sglang.srt.layers.quantization.fp8_kernel import is_fp8_fnuz +from sglang.srt.utils import is_gfx95_supported, is_hip + +tilelang.set_log_level("WARNING") + +# Workaround a tilelang bug: BaseKernelAdapter._legalize_result_idx mutates the +# `out_idx` list in place when normalising negative indices to positive ones. +# That breaks any @tilelang.jit factory that compiles two prim_funcs with +# different param counts (e.g. our unified single/dual partial kernel) — the +# second compile sees indices already-converted for the first's len(params) +# and silently builds the wrong adapter, leading to IndexError at call time. +# Patch once on import to copy the list before mutation. +from tilelang.jit.adapter.base import ( # noqa: E402 + BaseKernelAdapter as _BaseKernelAdapter, +) + +if not getattr(_BaseKernelAdapter, "_legalize_result_idx_patched", False): + _orig_legalize = _BaseKernelAdapter._legalize_result_idx + + def _legalize_result_idx_safe(self, result_idx): + if isinstance(result_idx, list): + result_idx = list(result_idx) + return _orig_legalize(self, result_idx) + + _BaseKernelAdapter._legalize_result_idx = _legalize_result_idx_safe + _BaseKernelAdapter._legalize_result_idx_patched = True + +pass_configs = { + tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True, + tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, +} +# TL_DISABLE_FAST_MATH has deprecated in v0.1.7.post1 tilelang +if hasattr(tilelang.PassConfigKey, "TL_DISABLE_FAST_MATH"): + pass_configs[tilelang.PassConfigKey.TL_DISABLE_FAST_MATH] = True +elif hasattr(tilelang.PassConfigKey, "TL_ENABLE_FAST_MATH"): + pass_configs[tilelang.PassConfigKey.TL_ENABLE_FAST_MATH] = False + +_is_hip = is_hip() +_is_gfx95_supported = is_gfx95_supported() +_is_fp8_fnuz = is_fp8_fnuz() + +BF16 = "bfloat16" +FP8 = "float8_e4m3fnuz" if _is_fp8_fnuz else "float8_e4m3fn" +FP8_DTYPE = torch.float8_e4m3fnuz if _is_fp8_fnuz else torch.float8_e4m3fn +FP32 = "float32" +INT32 = "int32" +UINT8 = "uint8" + + +def fast_log2_ceil(x): + bits_x = T.reinterpret("uint32", x) + exp_x = (bits_x >> 23) & 0xFF + man_bits = bits_x & ((1 << 23) - 1) + return T.Cast("int32", exp_x - 127 + T.if_then_else(man_bits != 0, 1, 0)) + + +def fast_pow2(x): + bits_x = (x + 127) << 23 + return T.reinterpret("float32", bits_x) + + +def fast_round_scale(amax, fp8_max_inv): + return fast_pow2(fast_log2_ceil(amax * fp8_max_inv)) + + +@lru_cache(maxsize=8) +def _pick_inner_iter(seq: int, ni: int, cu: int, block_per_cu: int) -> int: + """ + Pick the largest valid inner_iter (power-of-two divisor of ni) that keeps + enough work per CU (seq * ni / inner_iter / cu >= block_per_cu), so we avoid + under-utilization while minimizing the number of partial groups. + """ + + max_it = int(seq * ni / (cu * block_per_cu)) + it = ni + while it >= 2: + if it <= max_it and ni % it == 0: + return it + it //= 2 + return 1 + + +@tilelang.jit(pass_configs=pass_configs) +def act_quant_kernel( + N, in_dtype=BF16, out_dtype=FP8, scale_dtype=FP32, round_scale=False +): + M = T.symbolic("M") + fp8_min = -224.0 if _is_fp8_fnuz else -448.0 + fp8_max = 224.0 if _is_fp8_fnuz else 448.0 + fp8_max_inv = 1 / fp8_max + num_stages = 0 if round_scale else 2 + blk_m = 32 + group_size = 128 + + @T.prim_func + def act_quant_kernel_( + X: T.Tensor[(M, N), in_dtype], + Y: T.Tensor[(M, N), out_dtype], + S: T.Tensor[(M, T.ceildiv(N, group_size)), scale_dtype], + ): + with T.Kernel(T.ceildiv(M, blk_m), T.ceildiv(N, group_size), threads=128) as ( + pid_m, + pid_n, + ): + x_shared = T.alloc_shared((blk_m, group_size), in_dtype) + x_local = T.alloc_fragment((blk_m, group_size), in_dtype) + amax_local = T.alloc_fragment((blk_m,), scale_dtype) + s_local = T.alloc_fragment((blk_m,), scale_dtype) + y_local = T.alloc_fragment((blk_m, group_size), out_dtype) + y_shared = T.alloc_shared((blk_m, group_size), out_dtype) + + for _ in T.Pipelined(1, num_stages=num_stages): + T.copy(X[pid_m * blk_m, pid_n * group_size], x_shared) + T.copy(x_shared, x_local) + T.reduce_absmax(x_local, amax_local, dim=1) + for i in T.Parallel(blk_m): + amax_local[i] = T.max(amax_local[i], 1e-4) + if round_scale: + s_local[i] = fast_round_scale(amax_local[i], fp8_max_inv) + else: + s_local[i] = amax_local[i] * fp8_max_inv + for i, j in T.Parallel(blk_m, group_size): + y_local[i, j] = T.clamp( + x_local[i, j] / s_local[i], fp8_min, fp8_max + ) + for i in T.Parallel(blk_m): + S[pid_m * blk_m + i, pid_n] = s_local[i] + T.copy(y_local, y_shared) + T.copy(y_shared, Y[pid_m * blk_m, pid_n * group_size]) + + return act_quant_kernel_ + + +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. + + 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})" + N = x.size(-1) + if _is_fp8_fnuz: + y = torch.empty_like(x, dtype=torch.float8_e4m3fnuz) + else: + y = torch.empty_like(x, dtype=torch.float8_e4m3fn) + s = x.new_empty(*x.size()[:-1], N // block_size, dtype=torch.float32) + kernel = act_quant_kernel(N, round_scale=scale_fmt is not None) + kernel(x.view(-1, N), y.view(-1, N), s.view(-1, N // block_size)) + return y, s + + +@tilelang.jit(out_idx=[4], pass_configs=pass_configs) +def fp8_index_kernel(h: int, d: int, clear_accum=True): + b = T.symbolic("b") + m = T.symbolic("m") + n = T.symbolic("n") + + blk_n1 = 512 + blk_n2 = 128 + + @T.prim_func + def fp8_index_kernel_( + q: T.Tensor[(b, m, h, d), FP8], + q_s: T.Tensor[(b, m, h), FP32], + k: T.Tensor[(b, n, d), FP8], + k_s: T.Tensor[(b, n), FP32], + o: T.Tensor[(b, m, n), FP32], + ) -> None: + with T.Kernel(b, m, T.ceildiv(n, blk_n1)) as (i_b, i_m, i1_n): + q_smem = T.alloc_shared((h, d), FP8) + T.copy(q[i_b, i_m, 0, 0], q_smem) + + q_s_frag = T.alloc_fragment(h, FP32) + T.copy(q_s[i_b, i_m, 0], q_s_frag) + + for i2_n in T.Pipelined(blk_n1 // blk_n2, num_stages=2): + k_smem = T.alloc_shared((blk_n2, d), FP8) + T.copy(k[i_b, i1_n * blk_n1 + i2_n * blk_n2, 0], k_smem) + + k_s_frag = T.alloc_fragment(blk_n2, FP32) + T.copy(k_s[i_b, i1_n * blk_n1 + i2_n * blk_n2], k_s_frag) + + logits = T.alloc_fragment((blk_n2, h), FP32) + if not clear_accum: + T.fill(logits, 0) + T.gemm( + k_smem, + q_smem, + logits, + transpose_A=False, + transpose_B=True, + clear_accum=clear_accum, + ) + + for i_h, i3_n in T.Parallel(h, blk_n2): + logits[i3_n, i_h] = T.max(logits[i3_n, i_h], 0) * q_s_frag[i_h] + + logits_sum = T.alloc_fragment(blk_n2, FP32) + T.reduce_sum(logits, logits_sum, dim=1) + + for i3_n in T.Parallel(blk_n2): + logits_sum[i3_n] *= k_s_frag[i3_n] + + T.copy(logits_sum, o[i_b, i_m, i1_n * blk_n1 + i2_n * blk_n2]) + + return fp8_index_kernel_ + + +def fp8_index( + q: torch.Tensor, + q_s: torch.Tensor, + k: torch.Tensor, + k_s: torch.Tensor, +) -> torch.Tensor: + """ + Perform index score using FP8 precision. + + Args: + q (torch.Tensor): The Q tensor, must be contiguous. + q_s (torch.Tensor): The scaling factor for Q (float), must be contiguous. + k (torch.Tensor): The K tensor, must be contiguous. + k_s (torch.Tensor): The scaling factor for K (e8m0 here), must be contiguous. + + fp8 q @ fp8 k -> fp32 logits + relu(fp32 logits) * q_s (weights) -> fp32 logits + fp32 logits -> fp32 logits_sum + fp32 logits_sum * k_s (e8m0) -> fp32 index_score + """ + if _is_hip: + return fp8_index_kernel(q.shape[2], q.shape[3], False)(q, q_s, k, k_s) + else: + return fp8_index_kernel(q.shape[2], q.shape[3])(q, q_s, k, k_s) + + +@tilelang.jit( + out_idx=[-1], + pass_configs={ + tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, + tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True, + }, +) +def sparse_attention_fwd_kernel_v1( + num_heads, + dim, + tail_dim, + topk, + *, + kv_group=1, + sm_scale=None, + is_causal=True, + block_I=64, + num_stages=2, + threads=256, +): + assert dim == tilelang.math.next_power_of_2( + dim + ), f"haven't check padding correctness yet, dim={dim}" + assert tail_dim == tilelang.math.next_power_of_2( + tail_dim + ), f"haven't check padding correctness yet, dim={tail_dim}" + assert is_causal == True, "non-casual is not supported" + assert ( + topk % block_I == 0 + ), "otherwise will load some index=0 thus causing wrong kv to be loaded" + if sm_scale is None: + sm_scale = (1.0 / (dim + tail_dim)) ** 0.5 * 1.44269504 # log2(e) + else: + sm_scale = sm_scale * 1.44269504 # log2(e) + + batch = T.symbolic("batch") + seq_len = T.symbolic("seq_len") + seq_len_kv = T.symbolic("seq_len_kv") + + head_kv = num_heads // kv_group + q_shape = [batch, seq_len, num_heads, dim + tail_dim] + kv_shape = [batch, seq_len_kv, kv_group, dim + tail_dim] + o_shape = [batch, seq_len, num_heads, dim] + indices_shape = [batch, seq_len, kv_group, topk] + indices_dtype = "int32" + dtype = "bfloat16" + accum_dtype = "float" + + H = head_kv + padded_H = max(tilelang.math.next_power_of_2(head_kv), 16) + if padded_H != H: + assert kv_group == 1 + BI = block_I + NI = tilelang.cdiv(topk, block_I) + D = dim + D_tail = tail_dim + + if head_kv > 64: + assert head_kv % 64 == 0, "head_kv should be a multiple of 64" + REPLICATE_H = head_kv // 64 + else: + REPLICATE_H = 1 + + H_per_block = padded_H if REPLICATE_H == 1 else 64 + + @T.prim_func + def main( + Q: T.Tensor(q_shape, dtype), # type: ignore + KV: T.Tensor(kv_shape, dtype), # type: ignore + Indices: T.Tensor(indices_shape, indices_dtype), # type: ignore + Output: T.Tensor(o_shape, dtype), # type: ignore + ): + with T.Kernel(seq_len * REPLICATE_H, batch, kv_group, threads=threads) as ( + bx, + by, + bz, + ): + Q_shared = T.alloc_shared([H_per_block, D], dtype) + Q_tail_shared = T.alloc_shared([H_per_block, D_tail], dtype) + KV_shared = T.alloc_shared([BI, D], dtype) + K_tail_shared = T.alloc_shared([BI, D_tail], dtype) + O_shared = T.alloc_shared([H_per_block, D], dtype) + mask = T.alloc_fragment([BI], "bool") + + acc_o = T.alloc_fragment([H_per_block, D], accum_dtype) + acc_s = T.alloc_fragment([H_per_block, BI], accum_dtype) + S_shared = T.alloc_shared([H_per_block, BI], dtype) + sumexp = T.alloc_fragment([H_per_block], accum_dtype) + sumexp_i = T.alloc_fragment([H_per_block], accum_dtype) + alpha = T.alloc_fragment([H_per_block], accum_dtype) + m_i = T.alloc_fragment([H_per_block], accum_dtype) + m_i_prev = T.alloc_fragment([H_per_block], accum_dtype) + + T.fill(acc_o, 0) + T.fill(sumexp, 0) + T.fill(m_i, -(2**30)) # avoid -inf - inf to cause nan + + b_i, g_i = by, bz + s_i = bx if REPLICATE_H == 1 else (bx // REPLICATE_H) + q_i = s_i + max_kv_i = q_i + + H0 = g_i * padded_H + (0 if REPLICATE_H == 1 else (bx % REPLICATE_H) * 64) + H1 = H0 + H_per_block + + T.copy(Q[b_i, s_i, H0:H1, :D], Q_shared) + T.copy(Q[b_i, s_i, H0:H1, D:], Q_tail_shared) + + for i_i in T.Pipelined(NI, num_stages=num_stages): + + for bi_i in T.Parallel(BI): + mask[bi_i] = Indices[b_i, s_i, g_i, i_i * BI + bi_i] >= 0 + + for bi_i, d_i in T.Parallel(BI, D): + KV_shared[bi_i, d_i] = KV[ + b_i, Indices[b_i, s_i, g_i, i_i * BI + bi_i], g_i, d_i + ] + for bi_i, d_i in T.Parallel(BI, D_tail): + K_tail_shared[bi_i, d_i] = KV[ + b_i, Indices[b_i, s_i, g_i, i_i * BI + bi_i], g_i, D + d_i + ] + + for h_i, bi_i in T.Parallel(H_per_block, BI): + acc_s[h_i, bi_i] = T.if_then_else( + mask[bi_i], 0, -T.infinity(acc_s.dtype) + ) + T.gemm( + Q_shared, + KV_shared, + acc_s, + transpose_B=True, + policy=T.GemmWarpPolicy.FullCol, + ) + T.gemm( + Q_tail_shared, + K_tail_shared, + acc_s, + transpose_B=True, + policy=T.GemmWarpPolicy.FullCol, + ) + T.copy(m_i, m_i_prev) + T.reduce_max(acc_s, m_i, dim=1, clear=False) + for h_i in T.Parallel(H_per_block): + alpha[h_i] = T.exp2((m_i_prev[h_i] - m_i[h_i]) * sm_scale) + for h_i, bi_i in T.Parallel(H_per_block, BI): + acc_s[h_i, bi_i] = T.exp2( + acc_s[h_i, bi_i] * sm_scale - m_i[h_i] * sm_scale + ) + T.reduce_sum(acc_s, sumexp_i, dim=1) # is this a accumulate operator? + for h_i in T.Parallel(H_per_block): + sumexp[h_i] = sumexp[h_i] * alpha[h_i] + sumexp_i[h_i] + for h_i, d_i in T.Parallel(H_per_block, D): + acc_o[h_i, d_i] = acc_o[h_i, d_i] * alpha[h_i] + + T.copy(acc_s, S_shared) + T.gemm(S_shared, KV_shared, acc_o, policy=T.GemmWarpPolicy.FullCol) + + # Rescale + for h_i, d_i in T.Parallel(H_per_block, D): + acc_o[h_i, d_i] /= sumexp[h_i] + for h_i in T.Parallel(H_per_block): + sumexp[h_i] = T.log2(sumexp[h_i]) + m_i[h_i] * sm_scale + + T.copy(acc_o, O_shared) + T.copy(acc_o, Output[b_i, s_i, H0:H1, :]) + + return main + + +@tilelang.jit( + out_idx=[-1], + compile_flags=[ + "-O3", + "-Wno-deprecated-declarations", + "-U__CUDA_NO_HALF_OPERATORS__", + "-U__CUDA_NO_HALF_CONVERSIONS__", + "-U__CUDA_NO_HALF2_OPERATORS__", + "-U__CUDA_NO_BFLOAT16_CONVERSIONS__", + "--expt-relaxed-constexpr", + "--expt-extended-lambda", + "--ptxas-options=-v,--register-usage-level=10", + "-DNDEBUG", + ], +) # type: ignore +def sparse_attention_fwd_kernel_v2( + num_heads: int, + dim: int, + tail_dim: int, + topk: int, + *, + kv_group: int = 1, + sm_scale: Optional[float] = None, + block_I: int = 64, +): + assert dim == tilelang.math.next_power_of_2( + dim + ), f"haven't check padding correctness yet, dim={dim}" + assert tail_dim == tilelang.math.next_power_of_2( + tail_dim + ), f"haven't check padding correctness yet, dim={tail_dim}" + assert ( + topk % block_I == 0 + ), "otherwise will load some index=0 thus causing wrong kv to be loaded" + if sm_scale is None: + sm_scale = (1.0 / (dim + tail_dim)) ** 0.5 * 1.44269504 # log2(e) + else: + sm_scale = sm_scale * 1.44269504 # log2(e) + threads = 384 + + batch = T.symbolic("batch") + qo_len = T.symbolic("seq_len") + num_pages = T.symbolic("num_pages") + + q_shape = [batch, qo_len, num_heads, dim + tail_dim] + kv_shape = [batch, num_pages, kv_group, dim + tail_dim] + o_shape = [batch, qo_len, num_heads, dim] + indices_shape = [batch, qo_len, kv_group, topk] + + indices_dtype = "int32" + dtype = "bfloat16" + accum_dtype = "float" + + H = num_heads + padded_H = max(tilelang.math.next_power_of_2(num_heads), 16) + if padded_H != H: + assert kv_group == 1 + BI = block_I + NI = tilelang.cdiv(topk, block_I) + assert NI % 2 == 0, "NI should be a multiple of 2" + D = dim + D_tail = tail_dim + if num_heads > 64: + assert num_heads % 64 == 0, "head_kv should be a multiple of 64" + REPLICATE_H = num_heads // 64 + else: + REPLICATE_H = 1 + + H_per_block = padded_H if REPLICATE_H == 1 else 64 + + @T.prim_func + def main( + Q: T.Tensor(q_shape, dtype), # type: ignore + KV: T.Tensor(kv_shape, dtype), # type: ignore + Indices: T.Tensor(indices_shape, indices_dtype), # type: ignore + Output: T.Tensor(o_shape, dtype), # type: ignore + ): + """ + Q: [b, qo_len, H, D + D_tail] (bfloat16) + KV: [b, num_pages, kv_group, D + D_tail] (bfloat16) + Indices: [b, qo_len, kv_group, topk] (int32) + """ + + with T.Kernel(qo_len * REPLICATE_H, batch, 1, threads=threads) as (bx, by, bz): # type: ignore + Q_shared_l = T.alloc_shared([H_per_block, D // 2], dtype) + Q_shared_r = T.alloc_shared([H_per_block, D // 2], dtype) + Q_tail_shared = T.alloc_shared([H_per_block, D_tail], dtype) + KV_shared_0_l = T.alloc_shared([BI, D // 2], dtype) + KV_shared_0_r = T.alloc_shared([BI, D // 2], dtype) + KV_shared_1_l = T.alloc_shared([BI, D // 2], dtype) + KV_shared_1_r = T.alloc_shared([BI, D // 2], dtype) + K_tail_shared_0 = T.alloc_shared([BI, D_tail], dtype) + K_tail_shared_1 = T.alloc_shared([BI, D_tail], dtype) + O_shared_l = Q_shared_l + O_shared_r = Q_shared_r + is_kv_valid_0 = T.alloc_shared([BI], "bool", scope="shared") + is_kv_valid_1 = T.alloc_shared([BI], "bool", scope="shared") + + acc_o_l = T.alloc_fragment([H_per_block, D // 2], accum_dtype) + acc_o_r = T.alloc_fragment([H_per_block, D // 2], accum_dtype) + acc_s = T.alloc_fragment([H_per_block, BI], accum_dtype) + S_shared = T.alloc_shared([H_per_block, BI], dtype) + sumexp = T.alloc_fragment([H_per_block], accum_dtype) + sum_exp_shared = T.alloc_shared([H_per_block], accum_dtype) + sumexp_i = T.alloc_fragment([H_per_block], accum_dtype) + alpha_shared = T.alloc_shared([H_per_block], accum_dtype, scope="shared") + alpha_local = T.alloc_fragment([H_per_block], accum_dtype) + m_i = T.alloc_fragment([H_per_block], accum_dtype) + m_i_prev = T.alloc_fragment([H_per_block], accum_dtype) + indices_local = T.alloc_local([1], indices_dtype) + indices_tmp = T.alloc_local([1], indices_dtype) + + bar_q = T.alloc_barrier(arrive_count=384) + bar_k_0_ready = T.alloc_barrier(arrive_count=128) + bar_k_1_ready = T.alloc_barrier(arrive_count=128) + bar_k_0_free = T.alloc_barrier(arrive_count=256) + bar_k_1_free = T.alloc_barrier(arrive_count=256) + bar_sScale_and_sS_ready = T.alloc_barrier(arrive_count=256) + bar_sScale_and_sS_free = T.alloc_barrier(arrive_count=256) + + bar_0_128 = T.alloc_barrier(arrive_count=128) + bar_1_128 = T.alloc_barrier(arrive_count=128) + bar_2_128 = T.alloc_barrier(arrive_count=128) + bar_final = T.alloc_barrier(arrive_count=128) + + b_i, g_i = by, bz + s_i = bx if REPLICATE_H == 1 else bx // REPLICATE_H + + H0 = g_i * padded_H + (0 if REPLICATE_H == 1 else (bx % REPLICATE_H) * 64) + H1 = H0 + H_per_block + + tx = T.get_thread_binding() + + T.copy(Q[b_i, s_i, H0:H1, 0 : D // 2], Q_shared_l) + T.copy(Q[b_i, s_i, H0:H1, D // 2 : D], Q_shared_r) + T.copy(Q[b_i, s_i, H0:H1, D:], Q_tail_shared) + T.barrier_arrive(bar_q) + + if tx < 128: + T.set_max_nreg(240, 1) + T.fill(sumexp, 0) + T.fill(m_i, -(2**30)) # avoid -inf - inf to cause nan + T.fill(acc_o_l, 0) + T.barrier_wait(bar_q, 0) + + for i_i in T.serial(T.ceildiv(NI, 2)): + # Buffer 0 + # with sync_at(bar_0_128, 0): + T.barrier_wait(bar_k_0_ready[0], (i_i & 1)) + T.barrier_arrive(bar_0_128) + T.barrier_wait(bar_0_128, 0) + + for h_i, bi_i in T.Parallel(H_per_block, BI): + acc_s[h_i, bi_i] = T.if_then_else( + is_kv_valid_0[bi_i], 0, -T.infinity(acc_s.dtype) + ) + T.gemm( + Q_shared_l, KV_shared_0_l, acc_s, transpose_B=True, wg_wait=-1 + ) + T.gemm( + Q_shared_r, KV_shared_0_r, acc_s, transpose_B=True, wg_wait=-1 + ) + T.gemm( + Q_tail_shared, + K_tail_shared_0, + acc_s, + transpose_B=True, + wg_wait=-1, + ) + + T.wait_wgmma(0) + + if i_i != 0: + T.barrier_arrive(bar_sScale_and_sS_free) + T.barrier_wait(bar_sScale_and_sS_free, ((i_i * 2) & 1) ^ 1) + + T.copy(m_i, m_i_prev) + T.reduce_max(acc_s, m_i, dim=1, clear=False) + for h_i in T.Parallel(H_per_block): + alpha_local[h_i] = T.exp2((m_i_prev[h_i] - m_i[h_i]) * sm_scale) + for h_i, bi_i in T.Parallel(H_per_block, BI): + acc_s[h_i, bi_i] = T.exp2( + acc_s[h_i, bi_i] * sm_scale - m_i[h_i] * sm_scale + ) + T.reduce_sum( + acc_s, sumexp_i, dim=1 + ) # is this a accumulate operator? + for h_i in T.Parallel(H_per_block): + sumexp[h_i] = sumexp[h_i] * alpha_local[h_i] + sumexp_i[h_i] + for h_i, d_i in T.Parallel(H_per_block, D // 2): + acc_o_l[h_i, d_i] *= alpha_local[h_i] + T.copy(alpha_local, alpha_shared) + + T.copy(acc_s, S_shared) + T.gemm(S_shared, KV_shared_0_l, acc_o_l) + + T.barrier_arrive(bar_sScale_and_sS_ready) + T.barrier_arrive(bar_k_0_free[0]) + + # Buffer 1 + T.barrier_wait(bar_k_1_ready[0], (i_i & 1)) + T.barrier_arrive(bar_0_128) + T.barrier_wait(bar_0_128, 1) + + for h_i, bi_i in T.Parallel(H_per_block, BI): + acc_s[h_i, bi_i] = T.if_then_else( + is_kv_valid_1[bi_i], 0, -T.infinity(acc_s.dtype) + ) + T.gemm( + Q_shared_l, KV_shared_1_l, acc_s, transpose_B=True, wg_wait=-1 + ) + T.gemm( + Q_shared_r, KV_shared_1_r, acc_s, transpose_B=True, wg_wait=-1 + ) + T.gemm( + Q_tail_shared, + K_tail_shared_1, + acc_s, + transpose_B=True, + wg_wait=-1, + ) + + T.wait_wgmma(0) + + T.barrier_arrive(bar_sScale_and_sS_free) + T.barrier_wait(bar_sScale_and_sS_free, ((i_i * 2 + 1) & 1) ^ 1) + + T.copy(m_i, m_i_prev) + T.reduce_max(acc_s, m_i, dim=1, clear=False) + for h_i in T.Parallel(H_per_block): + alpha_local[h_i] = T.exp2((m_i_prev[h_i] - m_i[h_i]) * sm_scale) + for h_i, bi_i in T.Parallel(H_per_block, BI): + acc_s[h_i, bi_i] = T.exp2( + acc_s[h_i, bi_i] * sm_scale - m_i[h_i] * sm_scale + ) + T.reduce_sum( + acc_s, sumexp_i, dim=1 + ) # is this a accumulate operator? + for h_i in T.Parallel(H_per_block): + sumexp[h_i] = sumexp[h_i] * alpha_local[h_i] + sumexp_i[h_i] + for h_i, d_i in T.Parallel(H_per_block, D // 2): + acc_o_l[h_i, d_i] *= alpha_local[h_i] + T.copy(alpha_local, alpha_shared) + + T.copy(acc_s, S_shared) + T.gemm(S_shared, KV_shared_1_l, acc_o_l) + + T.barrier_arrive(bar_sScale_and_sS_ready) + T.barrier_arrive(bar_k_1_free[0]) + + # Rescale + for h_i in T.Parallel(H_per_block): + sum_exp_shared[h_i] = sumexp[h_i] + T.barrier_arrive(bar_final) + for h_i, d_i in T.Parallel(H_per_block, D // 2): + acc_o_l[h_i, d_i] /= sumexp[h_i] + for h_i in T.Parallel(H_per_block): + sumexp[h_i] = T.log2(sumexp[h_i]) + m_i[h_i] * sm_scale + T.copy(acc_o_l, O_shared_l) + T.copy(O_shared_l, Output[b_i, s_i, H0:H1, 0 : D // 2]) + elif tx >= 128 and tx < 256: + # T.set_max_nreg(168, 1) + T.fill(acc_o_r, 0) + for i_i in T.serial(T.ceildiv(NI, 2)): + # Buffer 0 + T.barrier_arrive(bar_sScale_and_sS_ready) + T.barrier_wait(bar_sScale_and_sS_ready, ((i_i * 2) & 1)) + T.barrier_arrive(bar_1_128) + T.barrier_wait(bar_1_128, 0) + for h_i, d_i in T.Parallel(H_per_block, D // 2): + acc_o_r[h_i, d_i] *= alpha_shared[h_i] + T.gemm(S_shared, KV_shared_0_r, acc_o_r) + T.barrier_arrive(bar_k_0_free[0]) + T.barrier_arrive(bar_sScale_and_sS_free) + + # Buffer 1 + T.barrier_arrive(bar_sScale_and_sS_ready) + T.barrier_wait(bar_sScale_and_sS_ready, ((i_i * 2 + 1) & 1)) + T.barrier_arrive(bar_1_128) + T.barrier_wait(bar_1_128, 1) + for h_i, d_i in T.Parallel(H_per_block, D // 2): + acc_o_r[h_i, d_i] *= alpha_shared[h_i] + T.gemm(S_shared, KV_shared_1_r, acc_o_r) + T.barrier_arrive(bar_k_1_free[0]) + if i_i != T.ceildiv(NI, 2) - 1: + T.barrier_arrive(bar_sScale_and_sS_free) + + # Rescale + T.barrier_wait(bar_final, 0) + for h_i, d_i in T.Parallel(H_per_block, D // 2): + acc_o_r[h_i, d_i] /= sum_exp_shared[h_i] + + T.copy(acc_o_r, O_shared_r) + T.copy(O_shared_r, Output[b_i, s_i, H0:H1, D // 2 : D]) + elif tx >= 256: + # producer + T.set_max_nreg(80, 0) + indices_local[0] = 0 + for i_i in T.serial(T.ceildiv(NI, 2)): + # Buffer 0 + T.barrier_wait(bar_k_0_free[0], ((i_i & 1) ^ 1)) + T.barrier_arrive(bar_2_128) + T.barrier_wait(bar_2_128, 0) + + for r in T.serial(4): + indices_tmp[0] = Indices[ + b_i, s_i, g_i, (i_i * 2) * BI + r * 16 + (tx - 256) // 8 + ] + is_kv_valid_0[r * 16 + (tx - 256) // 8] = indices_tmp[0] >= 0 + if is_kv_valid_0[r * 16 + (tx - 256) // 8]: + indices_local[0] = indices_tmp[0] + + with T.attr("default", "async_scope", 1): # type: ignore + for u in T.serial(4): + for v in T.vectorized(8): + KV_shared_0_l[ + r * 16 + (tx - 256) // 8, + 64 * u + (tx - 256) % 8 * 8 + v, + ] = KV[ + b_i, + indices_local[0], + g_i, + 64 * u + (tx - 256) % 8 * 8 + v, + ] + KV_shared_0_r[ + r * 16 + (tx - 256) // 8, + 64 * u + (tx - 256) % 8 * 8 + v, + ] = KV[ + b_i, + indices_local[0], + g_i, + D // 2 + 64 * u + (tx - 256) % 8 * 8 + v, + ] + with T.attr("default", "async_scope", 1): # type: ignore + for v in T.vectorized(8): + K_tail_shared_0[ + r * 16 + (tx - 256) // 8, (tx - 256) % 8 * 8 + v + ] = KV[ + b_i, + indices_local[0], + g_i, + D + (tx - 256) % 8 * 8 + v, + ] + + T.cp_async_barrier_noinc(bar_k_0_ready[0]) + + # Buffer 1 + T.barrier_wait(bar_k_1_free[0], ((i_i & 1) ^ 1)) + T.barrier_arrive(bar_2_128) + T.barrier_wait(bar_2_128, 1) + + for r in T.serial(4): + indices_tmp[0] = Indices[ + b_i, s_i, g_i, (i_i * 2 + 1) * BI + r * 16 + (tx - 256) // 8 + ] + is_kv_valid_1[r * 16 + (tx - 256) // 8] = indices_tmp[0] >= 0 + if is_kv_valid_1[r * 16 + (tx - 256) // 8]: + indices_local[0] = indices_tmp[0] + + with T.attr("default", "async_scope", 1): # type: ignore + for u in T.serial(4): + for v in T.vectorized(8): + KV_shared_1_l[ + r * 16 + (tx - 256) // 8, + 64 * u + (tx - 256) % 8 * 8 + v, + ] = KV[ + b_i, + indices_local[0], + g_i, + 64 * u + (tx - 256) % 8 * 8 + v, + ] + KV_shared_1_r[ + r * 16 + (tx - 256) // 8, + 64 * u + (tx - 256) % 8 * 8 + v, + ] = KV[ + b_i, + indices_local[0], + g_i, + D // 2 + 64 * u + (tx - 256) % 8 * 8 + v, + ] + with T.attr("default", "async_scope", 1): # type: ignore + for v in T.vectorized(8): + K_tail_shared_1[ + r * 16 + (tx - 256) // 8, (tx - 256) % 8 * 8 + v + ] = KV[ + b_i, + indices_local[0], + g_i, + D + (tx - 256) % 8 * 8 + v, + ] + + T.cp_async_barrier_noinc(bar_k_1_ready[0]) + + return main + + +@tilelang.jit( + out_idx=[-2, -1], + pass_configs={ + tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, + tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True, + }, +) +def sparse_mla_fwd_decode_partial( + heads, + dim, + tail_dim, + topk, + *, + kv_group=1, + sm_scale=None, + is_causal=True, + block_I=64, + inner_iter=1, + num_stages=1, + threads=256, +): + """ + grid: (seq_len * REPLICATE_H, top_k / block_I / inner_iter) + Each GPU block processes `inner_iter` consecutive KV tiles and writes one (partial_o, partial_lse) entry. + """ + + assert is_causal == True, "non-causal is not supported" + assert kv_group == 1 + assert topk % block_I == 0 + assert topk % (block_I * inner_iter) == 0, ( + f"topk ({topk}) must be divisible by block_I * inner_iter = " + f"{block_I} * {inner_iter}" + ) + + # log2(e) = 1.44269504 + if sm_scale is None: + sm_scale = (1.0 / (dim + tail_dim)) ** 0.5 * 1.44269504 + else: + sm_scale = sm_scale * 1.44269504 + + batch = 1 + seq_len = T.dynamic("seq_len") + seq_len_kv = T.dynamic("seq_len_kv") + + head_kv = heads // kv_group + padded_H = max(tilelang.math.next_power_of_2(head_kv), 16) + REPLICATE_H = (head_kv // 64) if head_kv > 64 else 1 + H_per_block = padded_H if REPLICATE_H == 1 else 64 + N_GROUPS = topk // (block_I * inner_iter) + BI = block_I + D = dim + D_tail = tail_dim + + q_shape = [batch, seq_len, heads, dim + tail_dim] + kv_shape = [batch, seq_len_kv, kv_group, dim + tail_dim] + indices_shape = [batch, seq_len, kv_group, topk] + partial_o_shape = [batch, seq_len, N_GROUPS, heads, dim] + partial_lse_shape = [batch, seq_len, N_GROUPS, heads] + indices_dtype = T.int32 + dtype = T.bfloat16 + accum_dtype = T.float32 + + _q_in_shared = inner_iter == 1 + + @T.prim_func + def main( + Q: T.Tensor(q_shape, dtype), + KV: T.Tensor(kv_shape, dtype), + Indices: T.Tensor(indices_shape, indices_dtype), + Partial_O: T.Tensor(partial_o_shape, dtype), + Partial_Lse: T.Tensor(partial_lse_shape, accum_dtype), + ): + with T.Kernel(seq_len * REPLICATE_H, N_GROUPS, threads=threads) as (bx, by): + if _q_in_shared: + Q_buf = T.alloc_shared([H_per_block, D], dtype) + Q_tail_buf = T.alloc_shared([H_per_block, D_tail], dtype) + else: + Q_buf = T.alloc_fragment([H_per_block, D], dtype) + Q_tail_buf = T.alloc_fragment([H_per_block, D_tail], dtype) + + KV_shared = T.alloc_shared([BI, D], dtype) + K_tail_shared = T.alloc_shared([BI, D_tail], dtype) + S_shared = T.alloc_shared([H_per_block, BI], dtype) + mask = T.alloc_fragment([BI], T.bool) + + acc_o = T.alloc_fragment([H_per_block, D], accum_dtype) + acc_s = T.alloc_fragment([H_per_block, BI], accum_dtype) + sumexp = T.alloc_fragment([H_per_block], accum_dtype) + sumexp_i = T.alloc_fragment([H_per_block], accum_dtype) + alpha = T.alloc_fragment([H_per_block], accum_dtype) + m_i = T.alloc_fragment([H_per_block], accum_dtype) + m_i_prev = T.alloc_fragment([H_per_block], accum_dtype) + + T.fill(acc_o, 0) + T.fill(sumexp, 0) + T.fill(m_i, -(2**30)) + + b_i, g_i = 0, 0 + s_i = bx if REPLICATE_H == 1 else (bx // REPLICATE_H) + group_i = by + H0 = 0 if REPLICATE_H == 1 else (bx % REPLICATE_H) * 64 + H1 = H0 + H_per_block + + T.copy(Q[b_i, s_i, H0:H1, :D], Q_buf) + T.copy(Q[b_i, s_i, H0:H1, D:], Q_tail_buf) + + for k_i in T.Pipelined(inner_iter, num_stages=num_stages): + topk_block_i = group_i * inner_iter + k_i + + for bi_i in T.Parallel(BI): + mask[bi_i] = Indices[b_i, s_i, g_i, topk_block_i * BI + bi_i] >= 0 + for bi_i, d_i in T.Parallel(BI, D): + idx = Indices[b_i, s_i, g_i, topk_block_i * BI + bi_i] + KV_shared[bi_i, d_i] = KV[ + b_i, T.if_then_else(idx >= 0, idx, 0), g_i, d_i + ] + for bi_i, d_i in T.Parallel(BI, D_tail): + idx = Indices[b_i, s_i, g_i, topk_block_i * BI + bi_i] + K_tail_shared[bi_i, d_i] = KV[ + b_i, T.if_then_else(idx >= 0, idx, 0), g_i, D + d_i + ] + + for h_i, bi_i in T.Parallel(H_per_block, BI): + acc_s[h_i, bi_i] = T.if_then_else( + mask[bi_i], 0, -T.infinity(acc_s.dtype) + ) + + T.gemm( + Q_buf, + KV_shared, + acc_s, + transpose_B=True, + policy=T.GemmWarpPolicy.FullCol, + ) + T.gemm( + Q_tail_buf, + K_tail_shared, + acc_s, + transpose_B=True, + policy=T.GemmWarpPolicy.FullCol, + ) + + T.copy(m_i, m_i_prev) + T.reduce_max(acc_s, m_i, dim=1, clear=False) + for h_i in T.Parallel(H_per_block): + alpha[h_i] = T.exp2((m_i_prev[h_i] - m_i[h_i]) * sm_scale) + for h_i, bi_i in T.Parallel(H_per_block, BI): + acc_s[h_i, bi_i] = T.exp2( + acc_s[h_i, bi_i] * sm_scale - m_i[h_i] * sm_scale + ) + T.reduce_sum(acc_s, sumexp_i, dim=1) + for h_i in T.Parallel(H_per_block): + sumexp[h_i] = sumexp[h_i] * alpha[h_i] + sumexp_i[h_i] + for h_i, d_i in T.Parallel(H_per_block, D): + acc_o[h_i, d_i] *= alpha[h_i] + + T.copy(acc_s, S_shared) + T.gemm(S_shared, KV_shared, acc_o, policy=T.GemmWarpPolicy.FullCol) + + # sumexp==0 (all masked), divide by 1 to get 0 and avoid nan + for h_i, d_i in T.Parallel(H_per_block, D): + acc_o[h_i, d_i] = acc_o[h_i, d_i] / T.if_then_else( + sumexp[h_i] == 0.0, 1.0, sumexp[h_i] + ) + # sumexp==0 (all masked), use large negative so combine ignores this split + for h_i in T.Parallel(H_per_block): + sumexp[h_i] = T.if_then_else( + sumexp[h_i] == 0.0, + -(2**30), + T.log2(sumexp[h_i]) + m_i[h_i] * sm_scale, + ) + + T.copy(acc_o, Partial_O[b_i, s_i, group_i, H0:H1, :]) + T.copy(sumexp, Partial_Lse[b_i, s_i, group_i, H0:H1]) + + return main + + +@tilelang.jit( + out_idx=[-1], + pass_configs={ + tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, + tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True, + }, +) +def sparse_mla_fwd_decode_combine( + heads, + dim, + topk, + head_per_block, + *, + block_I=64, + threads=256, +): + """ + grid: (seq_len * REPLICATE_H). batch=1, kv_group=1. + Each block does one tile of heads (e.g. 4 or 8 for decode). + """ + + assert heads % head_per_block == 0, f"head_per_block must divide heads" + + batch = 1 + seq_len = T.dynamic("seq_len") + + NI = topk // block_I + H_per_block = head_per_block + REPLICATE_H = heads // H_per_block + + partial_o_shape = [batch, seq_len, NI, heads, dim] + partial_lse_shape = [batch, seq_len, NI, heads] + o_shape = [batch, seq_len, heads, dim] + dtype = T.bfloat16 + accum_dtype = T.float32 + + @T.prim_func + def main( + Partial_O: T.Tensor(partial_o_shape, dtype), + Partial_Lse: T.Tensor(partial_lse_shape, accum_dtype), + Output: T.Tensor(o_shape, dtype), + ): + with T.Kernel(seq_len * REPLICATE_H, threads=threads) as (bx,): + shared_lse = T.alloc_shared([NI, H_per_block], accum_dtype) + + lse_max = T.alloc_fragment([H_per_block], accum_dtype) + lse_sum = T.alloc_fragment([H_per_block], accum_dtype) + scale = T.alloc_fragment([H_per_block, NI], accum_dtype) + acc_o = T.alloc_fragment([H_per_block, dim], accum_dtype) + + b_i = 0 + s_i = bx if REPLICATE_H == 1 else (bx // REPLICATE_H) + H0 = 0 if REPLICATE_H == 1 else (bx % REPLICATE_H) * H_per_block + H1 = H0 + H_per_block + + for k in T.serial(NI): + T.copy(Partial_Lse[b_i, s_i, k, H0:H1], shared_lse[k, :]) + + T.fill(lse_max, -(2**30)) + for k in T.serial(NI): + for h_i in T.Parallel(H_per_block): + lse_max[h_i] = T.max(lse_max[h_i], shared_lse[k, h_i]) + T.fill(lse_sum, 0) + for k in T.serial(NI): + for h_i in T.Parallel(H_per_block): + lse_sum[h_i] = lse_sum[h_i] + T.exp2( + shared_lse[k, h_i] - lse_max[h_i] + ) + for k in T.serial(NI): + for h_i in T.Parallel(H_per_block): + scale[h_i, k] = T.exp2( + shared_lse[k, h_i] - lse_max[h_i] - T.log2(lse_sum[h_i]) + ) + + T.fill(acc_o, 0) + for k in T.serial(NI): + for h_i, d_i in T.Parallel(H_per_block, dim): + acc_o[h_i, d_i] = acc_o[h_i, d_i] + scale[h_i, k] * Partial_O[ + b_i, s_i, k, H0 + h_i, d_i + ].astype(accum_dtype) + + T.copy(acc_o, Output[b_i, s_i, H0:H1, :]) + + return main + + +@tilelang.jit(out_idx=[-2, -1], pass_configs=pass_configs) +def sparse_mla_fwd_decode_partial_fp8( + num_heads: int, + d_v: int, + d_tail: int, + topk: int, + *, + sm_scale=None, + block_I=64, + inner_iter=1, + threads=256, +): + assert d_v == 512, f"only support d_v=512" + assert ( + topk % block_I == 0 + ), "otherwise will load some index=0 thus causing wrong kv to be loaded" + + # Softmax scores are in [0, 1]. We scale by fp8_max_val before FP8 cast + # to better utilize FP8 dynamic range, then apply the inverse scale after GEMM. + # This is numerically safe because softmax output is bounded by 1. + fp8_dtype = "float8_e4m3fnuz" if _is_fp8_fnuz else "float8_e4m3fn" + fp8_max_val = 240.0 if _is_fp8_fnuz else 448.0 + s_inv_scale_const = fp8_max_val + s_scale_const = 1.0 / fp8_max_val + + BI = block_I + group_size = 128 + dim_quant_fp8 = d_v + d_tail + rope_offset_fp8 = d_v + n_groups = topk // (BI * inner_iter) + + if sm_scale is None: + sm_scale = (1.0 / (d_v + d_tail)) ** 0.5 * 1.44269504 + else: + sm_scale = sm_scale * 1.44269504 + + h_per_block = 16 + # Match bf16 partial behavior: keep fixed 16-head tiles and use + # sliced T.copy on H0:H1 for tail handling. + assert ( + num_heads <= h_per_block or num_heads % h_per_block == 0 + ), "num_heads must be <=16 or divisible by 16" + head_blocks_per_seq = (num_heads + h_per_block - 1) // h_per_block + + batch = 1 + kv_group = 1 + seq_len = T.symbolic("seq_len") + num_pages = T.symbolic("num_pages") + + q_fp8_shape = [batch, seq_len, num_heads, d_v + d_tail] + kv_fp8_shape = [batch, num_pages, kv_group, dim_quant_fp8] + idx_shape = [batch, seq_len, kv_group, topk] + partial_o_shape = [batch, seq_len, n_groups, num_heads, d_v] + partial_lse_shape = [batch, seq_len, n_groups, num_heads] + + accum_dtype = T.float32 + dtype_bf16 = T.bfloat16 + + @T.prim_func + def main( + q_fp8: T.Tensor(q_fp8_shape, fp8_dtype), + kv_fp8: T.Tensor(kv_fp8_shape, fp8_dtype), + indices: T.Tensor(idx_shape, T.int32), + partial_o: T.Tensor(partial_o_shape, dtype_bf16), + partial_lse: T.Tensor(partial_lse_shape, accum_dtype), + ): + with T.Kernel(seq_len * head_blocks_per_seq, n_groups, threads=threads) as ( + bx, + by, + ): + b_i, g_i = 0, 0 + s_i = bx // head_blocks_per_seq + group_i = by + H0 = (bx % head_blocks_per_seq) * h_per_block + H1 = H0 + h_per_block + + # We intentionally split the K=512 GEMM into 4x128 tiles. + # Although this adds extra intermediate memory traffic, + # it shortens the MFMA accumulation dependency chain and improves performance. + q_tile0 = T.alloc_shared([h_per_block, group_size], fp8_dtype) + q_tile1 = T.alloc_shared([h_per_block, group_size], fp8_dtype) + q_tile2 = T.alloc_shared([h_per_block, group_size], fp8_dtype) + q_tile3 = T.alloc_shared([h_per_block, group_size], fp8_dtype) + kv_tile0 = T.alloc_shared([BI, group_size], fp8_dtype) + kv_tile1 = T.alloc_shared([BI, group_size], fp8_dtype) + kv_tile2 = T.alloc_shared([BI, group_size], fp8_dtype) + kv_tile3 = T.alloc_shared([BI, group_size], fp8_dtype) + q_tail_buf = T.alloc_shared([h_per_block, d_tail], fp8_dtype) + k_tail_shared = T.alloc_shared([BI, d_tail], fp8_dtype) + s_fp8_shared = T.alloc_shared([h_per_block, BI], fp8_dtype) + page_idx_shared = T.alloc_shared([BI], T.int32) + + mask = T.alloc_fragment([BI], T.bool) + acc_s = T.alloc_fragment([h_per_block, BI], accum_dtype) + acc_tile = T.alloc_fragment([h_per_block, BI], accum_dtype) + sv_tile = T.alloc_fragment([h_per_block, group_size], accum_dtype) + sumexp = T.alloc_fragment([h_per_block], accum_dtype) + sumexp_i = T.alloc_fragment([h_per_block], accum_dtype) + alpha = T.alloc_fragment([h_per_block], accum_dtype) + m_i = T.alloc_fragment([h_per_block], accum_dtype) + m_i_prev = T.alloc_fragment([h_per_block], accum_dtype) + inv_denom = T.alloc_fragment([h_per_block], accum_dtype) + + acc_o_tile0 = T.alloc_fragment([h_per_block, group_size], accum_dtype) + acc_o_tile1 = T.alloc_fragment([h_per_block, group_size], accum_dtype) + acc_o_tile2 = T.alloc_fragment([h_per_block, group_size], accum_dtype) + acc_o_tile3 = T.alloc_fragment([h_per_block, group_size], accum_dtype) + + T.fill(acc_o_tile0, 0) + T.fill(acc_o_tile1, 0) + T.fill(acc_o_tile2, 0) + T.fill(acc_o_tile3, 0) + T.fill(sumexp, 0) + T.fill(m_i, -(2**30)) + + T.copy(q_fp8[b_i, s_i, H0:H1, d_v:], q_tail_buf) + T.copy(q_fp8[b_i, s_i, H0:H1, 0 * group_size : 1 * group_size], q_tile0) + T.copy(q_fp8[b_i, s_i, H0:H1, 1 * group_size : 2 * group_size], q_tile1) + T.copy(q_fp8[b_i, s_i, H0:H1, 2 * group_size : 3 * group_size], q_tile2) + T.copy(q_fp8[b_i, s_i, H0:H1, 3 * group_size : 4 * group_size], q_tile3) + + for k_i in T.serial(inner_iter): + topk_block_i = group_i * inner_iter + k_i + + for bi_i in T.Parallel(BI): + idx = indices[b_i, s_i, g_i, topk_block_i * BI + bi_i] + valid = idx >= 0 + page_idx_shared[bi_i] = T.if_then_else(valid, idx, 0) + mask[bi_i] = valid + + for bi_i, j in T.Parallel(BI, group_size): + page = page_idx_shared[bi_i] + kv_tile0[bi_i, j] = kv_fp8[b_i, page, g_i, 0 * group_size + j] + kv_tile1[bi_i, j] = kv_fp8[b_i, page, g_i, 1 * group_size + j] + kv_tile2[bi_i, j] = kv_fp8[b_i, page, g_i, 2 * group_size + j] + kv_tile3[bi_i, j] = kv_fp8[b_i, page, g_i, 3 * group_size + j] + + for bi_i, j in T.Parallel(BI, d_tail): + page = page_idx_shared[bi_i] + k_tail_shared[bi_i, j] = kv_fp8[b_i, page, g_i, rope_offset_fp8 + j] + + for h_i, bi_i in T.Parallel(h_per_block, BI): + acc_s[h_i, bi_i] = T.if_then_else( + mask[bi_i], 0, -T.infinity(acc_s.dtype) + ) + + T.gemm(q_tile0, kv_tile0, acc_s, transpose_B=True, clear_accum=False) + T.gemm(q_tile1, kv_tile1, acc_tile, transpose_B=True, clear_accum=True) + for h_i, bi_i in T.Parallel(h_per_block, BI): + acc_s[h_i, bi_i] += acc_tile[h_i, bi_i] + T.gemm(q_tile2, kv_tile2, acc_tile, transpose_B=True, clear_accum=True) + for h_i, bi_i in T.Parallel(h_per_block, BI): + acc_s[h_i, bi_i] += acc_tile[h_i, bi_i] + T.gemm(q_tile3, kv_tile3, acc_tile, transpose_B=True, clear_accum=True) + for h_i, bi_i in T.Parallel(h_per_block, BI): + acc_s[h_i, bi_i] += acc_tile[h_i, bi_i] + T.gemm( + q_tail_buf, + k_tail_shared, + acc_s, + transpose_B=True, + policy=T.GemmWarpPolicy.FullCol, + ) + + T.copy(m_i, m_i_prev) + T.reduce_max(acc_s, m_i, dim=1, clear=False) + for h_i in T.Parallel(h_per_block): + alpha[h_i] = T.exp2((m_i_prev[h_i] - m_i[h_i]) * sm_scale) + for h_i, bi_i in T.Parallel(h_per_block, BI): + acc_s[h_i, bi_i] = T.exp2( + acc_s[h_i, bi_i] * sm_scale - m_i[h_i] * sm_scale + ) + T.reduce_sum(acc_s, sumexp_i, dim=1) + for h_i in T.Parallel(h_per_block): + sumexp[h_i] = sumexp[h_i] * alpha[h_i] + sumexp_i[h_i] + for h_i, j in T.Parallel(h_per_block, group_size): + acc_o_tile0[h_i, j] = acc_o_tile0[h_i, j] * alpha[h_i] + acc_o_tile1[h_i, j] = acc_o_tile1[h_i, j] * alpha[h_i] + acc_o_tile2[h_i, j] = acc_o_tile2[h_i, j] * alpha[h_i] + acc_o_tile3[h_i, j] = acc_o_tile3[h_i, j] * alpha[h_i] + + for h_i, bi_i in T.Parallel(h_per_block, BI): + s_fp8_shared[h_i, bi_i] = T.clamp( + acc_s[h_i, bi_i] * s_inv_scale_const, + -fp8_max_val, + fp8_max_val, + ) + T.gemm(s_fp8_shared, kv_tile0, sv_tile, clear_accum=True) + for h_i, j in T.Parallel(h_per_block, group_size): + acc_o_tile0[h_i, j] = ( + acc_o_tile0[h_i, j] + sv_tile[h_i, j] * s_scale_const + ) + + T.gemm(s_fp8_shared, kv_tile1, sv_tile, clear_accum=True) + for h_i, j in T.Parallel(h_per_block, group_size): + acc_o_tile1[h_i, j] = ( + acc_o_tile1[h_i, j] + sv_tile[h_i, j] * s_scale_const + ) + + T.gemm(s_fp8_shared, kv_tile2, sv_tile, clear_accum=True) + for h_i, j in T.Parallel(h_per_block, group_size): + acc_o_tile2[h_i, j] = ( + acc_o_tile2[h_i, j] + sv_tile[h_i, j] * s_scale_const + ) + + T.gemm(s_fp8_shared, kv_tile3, sv_tile, clear_accum=True) + for h_i, j in T.Parallel(h_per_block, group_size): + acc_o_tile3[h_i, j] = ( + acc_o_tile3[h_i, j] + sv_tile[h_i, j] * s_scale_const + ) + + for h_i in T.Parallel(h_per_block): + denom = T.if_then_else(sumexp[h_i] == 0.0, 1.0, sumexp[h_i]) + inv_denom[h_i] = 1.0 / denom + for h_i, j in T.Parallel(h_per_block, group_size): + acc_o_tile0[h_i, j] = acc_o_tile0[h_i, j] * inv_denom[h_i] + acc_o_tile1[h_i, j] = acc_o_tile1[h_i, j] * inv_denom[h_i] + acc_o_tile2[h_i, j] = acc_o_tile2[h_i, j] * inv_denom[h_i] + acc_o_tile3[h_i, j] = acc_o_tile3[h_i, j] * inv_denom[h_i] + + for h_i in T.Parallel(h_per_block): + sumexp[h_i] = T.if_then_else( + sumexp[h_i] == 0.0, + -(2**30), + T.log2(sumexp[h_i]) + m_i[h_i] * sm_scale, + ) + + T.copy( + acc_o_tile0, + partial_o[b_i, s_i, group_i, H0:H1, 0 * group_size : 1 * group_size], + ) + T.copy( + acc_o_tile1, + partial_o[b_i, s_i, group_i, H0:H1, 1 * group_size : 2 * group_size], + ) + T.copy( + acc_o_tile2, + partial_o[b_i, s_i, group_i, H0:H1, 2 * group_size : 3 * group_size], + ) + T.copy( + acc_o_tile3, + partial_o[b_i, s_i, group_i, H0:H1, 3 * group_size : 4 * group_size], + ) + + T.copy(sumexp, partial_lse[b_i, s_i, group_i, H0:H1]) + + return main + + +def tilelang_sparse_fwd( + q: torch.Tensor, + kv: torch.Tensor, + indices: torch.Tensor, + sm_scale: float, + d_v: int = 512, +) -> torch.Tensor: + assert q.dim() == 3 and kv.dim() == 3 and indices.dim() == 3 + num_heads = q.shape[1] + dim = q.shape[2] + tail_dim = dim - d_v + topk = indices.shape[-1] + assert topk == 2048 + + if _is_hip: + is_fp8_kv = kv.dtype in (torch.float8_e4m3fn, torch.float8_e4m3fnuz) + if is_fp8_kv: + if q.dtype != kv.dtype: + q = q.to(kv.dtype) + if _is_gfx95_supported: + block_I, threads, block_per_cu, cu = 64, 256, 2, 256 + else: + block_I, threads, block_per_cu, cu = 64, 256, 1, 304 + ni = topk // block_I + inner_iter = _pick_inner_iter(q.shape[0], ni, cu, block_per_cu) + kernel_partial = sparse_mla_fwd_decode_partial_fp8( + num_heads, + d_v, + tail_dim, + topk, + sm_scale=sm_scale, + block_I=block_I, + inner_iter=inner_iter, + threads=threads, + ) + else: + if _is_gfx95_supported: + block_I, threads, block_per_cu, cu = 64, 256, 2, 256 + else: + block_I, threads, block_per_cu, cu = 32, 128, 1, 304 + ni = topk // block_I + inner_iter = _pick_inner_iter(q.shape[0], ni, cu, block_per_cu) + kernel_partial = sparse_mla_fwd_decode_partial( + num_heads, + d_v, + tail_dim, + topk, + sm_scale=sm_scale, + block_I=block_I, + inner_iter=inner_iter, + threads=threads, + ) + partial_o_batched, partial_lse_batched = kernel_partial( + q.unsqueeze(0), kv.unsqueeze(0), indices.unsqueeze(0) + ) + n_groups = ni // inner_iter + kernel_combine = sparse_mla_fwd_decode_combine( + num_heads, + d_v, + n_groups * block_I, + head_per_block=4, + block_I=block_I, + threads=threads, + ) + out = kernel_combine(partial_o_batched, partial_lse_batched) + else: + kernel = sparse_attention_fwd_kernel_v2( + num_heads, d_v, tail_dim, topk, sm_scale=sm_scale + ) + out = kernel(q.unsqueeze(0), kv.unsqueeze(0), indices.unsqueeze(0)) # type: ignore + return out + + +@functools.cache +def fp8_paged_mqa_logits_kernel( + head_dim: int = 128, + num_heads: int = 64, + block_size: int = 64, + clear_accum: bool = True, + split_kv: int = 1, +) -> Any: + N = T.symbolic("batch_size") + L = T.symbolic("max_table_length") + S = T.symbolic("max_seq_len") + C = T.symbolic("num_blocks") + B = block_size + D = head_dim + H = num_heads + SK = int(split_kv) + BLOCK_BYTES = B * (D + 4) + SCALE_OFFSET = B * D + + assert D % 4 == 0 + assert H % 4 == 0 + assert D == 128 + assert SK >= 1 + + @tilelang.jit( + pass_configs={ + **pass_configs, + tilelang.PassConfigKey.TL_DISABLE_SAFE_MEMORY_ACCESS: True, + } + ) + def fp8_paged_mqa_logits( + q: T.Tensor[(N, H, D), FP8], + kvcache_u8: T.Tensor[(C, BLOCK_BYTES), UINT8], + weight: T.Tensor[(N, H), FP32], + seq_lens: T.Tensor[(N,), INT32], + page_table: T.Tensor[(N, L), INT32], + o: T.Tensor[(N, S), FP32], + ) -> None: + _ = N, L, S, C, D, H, B + with T.Kernel(N * SK) as bxs: + bx = bxs % N + pid_split = bxs // N + seq_len = seq_lens[bx] + np_total = T.ceildiv(seq_len, B) + stride = T.ceildiv(np_total, SK) + i_start = pid_split * stride + n_iters = T.max(0, T.min(stride, np_total - i_start)) + + q_smem = T.alloc_shared((H, D), FP8) + q_s_frag = T.alloc_fragment((H,), FP32) + T.copy(q[bx, 0, 0], q_smem) + T.copy(weight[bx, 0], q_s_frag) + + for j in T.Pipelined(n_iters, num_stages=2): + i = i_start + j + page = page_table[bx, i] + k_smem_u8 = T.alloc_shared((B * D,), UINT8) + T.copy(kvcache_u8[page, 0:SCALE_OFFSET], k_smem_u8) + k_smem = T.view(k_smem_u8, (B, D), FP8) + k_s_smem_u8 = T.alloc_shared((B * 4,), UINT8) + T.copy(kvcache_u8[page, SCALE_OFFSET:BLOCK_BYTES], k_s_smem_u8) + k_s_smem = T.view(k_s_smem_u8, (B,), FP32) + k_s_frag = T.alloc_fragment((B,), FP32) + T.copy(k_s_smem, k_s_frag) + + logits = T.alloc_fragment((B, H), FP32) + if not clear_accum: + T.fill(logits, 0.0) + T.gemm( + k_smem, + q_smem, + logits, + transpose_A=False, + transpose_B=True, + clear_accum=clear_accum, + ) + + # post processing + for h, j2 in T.Parallel(H, B): + logits[j2, h] = T.max(logits[j2, h], 0.0) * q_s_frag[h] + logits_sum = T.alloc_fragment((B,), FP32) + T.reduce_sum(logits, logits_sum, dim=1) + for j2 in T.Parallel(B): + logits_sum[j2] *= k_s_frag[j2] + T.copy(logits_sum, o[bx, i * B]) + + return fp8_paged_mqa_logits + + +def tilelang_fp8_paged_mqa_logits( + q_fp8: torch.Tensor, + kvcache_fp8: torch.Tensor, + weight: torch.Tensor, + seq_lens: torch.Tensor, + page_table: torch.Tensor, + deep_gemm_metadata: Any, + max_seq_len: int, + clean_logits: bool = True, +) -> torch.Tensor: + _ = deep_gemm_metadata + batch_size, _, num_heads, head_dim = q_fp8.shape + block_size = kvcache_fp8.shape[1] + assert head_dim == 128, "TODO" + assert block_size == 64, "TODO" + assert q_fp8.shape == (batch_size, 1, num_heads, head_dim) + assert kvcache_fp8.shape[1:] == (block_size, 1, head_dim + 4) + assert weight.shape == (batch_size, num_heads) + assert seq_lens.shape == (batch_size,) + assert page_table.shape[0] == batch_size + assert clean_logits == False + + logits = page_table.new_empty((batch_size, max_seq_len), dtype=torch.float32) + + NUM_CU = 256 + split_kv = split_kv = max(1, min(max_seq_len // block_size, NUM_CU // batch_size)) + kernel = fp8_paged_mqa_logits_kernel( + head_dim=head_dim, + num_heads=num_heads, + block_size=block_size, + clear_accum=clean_logits, + split_kv=split_kv, + ) + q_fp8 = q_fp8.view(batch_size, num_heads, head_dim) + kvcache_u8 = kvcache_fp8.view(-1, block_size * (head_dim + 4)) + kernel(q_fp8, kvcache_u8, weight, seq_lens, page_table, logits) + return logits + + +def _build_fp8_combined_view(k_cache: torch.Tensor) -> Tuple[torch.Tensor, int, int]: + """ + Reinterpret a MODEL1_FP8Sparse KV cache as a contiguous uint32 view. + Input: k_cache (num_blocks, block_size, 1, d_qk) fp8/uint8 + — per-block storage also holds scales + padding past d_qk. + Output: (num_blocks, block_pad_u32) uint32 covering the full block + stride. Same storage ashe input, no copy. + """ + k_u8 = k_cache.view(torch.uint8) if k_cache.dtype != torch.uint8 else k_cache + num_blocks = k_u8.shape[0] + block_size = k_u8.shape[1] + block_pad_u32 = k_u8.stride(0) // 4 + storage = k_u8.untyped_storage() + flat_u32 = torch.empty(0, dtype=torch.uint32, device=k_u8.device).set_( + storage, 0, (storage.nbytes() // 4,), (1,) + ) + k_combined = torch.as_strided( + flat_u32, + size=(num_blocks, block_pad_u32), + stride=(block_pad_u32, 1), + storage_offset=k_u8.storage_offset() // 4, + ) + return k_combined, num_blocks, block_size + + +_TOPK_LEN_SENTINEL_CACHE: dict = {} +_INT32_MAX = 2**30 + + +def _topk_length_sentinel(device: torch.device, batch: int) -> torch.Tensor: + """Cached `(batch,) int32 INT_MAX` tensor used when `topk_length` is None.""" + cur = _TOPK_LEN_SENTINEL_CACHE.get(device) + if cur is None or cur.numel() < batch: + cur = torch.full( + (max(batch, 256),), _INT32_MAX, dtype=torch.int32, device=device + ) + _TOPK_LEN_SENTINEL_CACHE[device] = cur + return cur[:batch] + + +@tilelang.jit( + out_idx=[-2, -1], + pass_configs={ + tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, + tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True, + tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True, + }, +) +def dpsk_v4_fp8_partial_kernel( + num_heads: int, + topk_1: int, + block_size_kv_1: int, + topk_2: int = 0, + block_size_kv_2: int = 0, + *, + dim: int = 448, + tail_dim: int = 64, + sm_scale: float = 0.0, + block_I: int = 64, + inner_iter_1: int = 1, + inner_iter_2: int = 0, + num_stages: int = 0, + threads: int = 512, +) -> Any: + """ + Read FP8 K cache directly, dequantise to BF16 in-kernel, do flash-attn + online softmax with split-K. Supports a second cache (`topk_2>0`) and + `attn_sink` is folded later by the combine kernel. + """ + log2e: float = 1.44269504 + if sm_scale <= 0.0: + sm_scale = (1.0 / (dim + tail_dim)) ** 0.5 * log2e + else: + sm_scale = sm_scale * log2e + assert dim == 448 and tail_dim == 64 + assert topk_1 % block_I == 0 + assert ( + topk_1 // block_I + ) % inner_iter_1 == 0, ( + f"NI_1={topk_1 // block_I} must be divisible by inner_iter_1={inner_iter_1}" + ) + assert block_size_kv_1 > 0 and (block_size_kv_1 & (block_size_kv_1 - 1)) == 0 + + is_dual = topk_2 > 0 + if is_dual: + assert inner_iter_2 > 0, "dual-cache call requires inner_iter_2 > 0" + assert topk_2 % block_I == 0 + assert ( + topk_2 // block_I + ) % inner_iter_2 == 0, ( + f"NI_2={topk_2 // block_I} must be divisible by inner_iter_2={inner_iter_2}" + ) + assert block_size_kv_2 > 0 and (block_size_kv_2 & (block_size_kv_2 - 1)) == 0 + + PACKED_W = dim + 2 * tail_dim + NOPE_TILE = 64 + NUM_TILES = dim // NOPE_TILE + SCALE_W = 8 + PACKED_W4 = PACKED_W // 4 + SCALE_W4 = SCALE_W // 4 + + kv_group = 1 + batch = T.symbolic("batch") + seq_len = T.symbolic("seq_len") + num_blocks_kv_1 = T.symbolic("num_blocks_kv_1") + block_pad_u32_1 = T.symbolic("block_pad_u32_1") + if is_dual: + num_blocks_kv_2 = T.symbolic("num_blocks_kv_2") + block_pad_u32_2 = T.symbolic("block_pad_u32_2") + + head_kv = num_heads // kv_group + D = dim + D_tail = tail_dim + BI = block_I + padded_H = max(tilelang.math.next_power_of_2(head_kv), 16) + if head_kv > 64: + assert head_kv % 64 == 0 + REPLICATE_H = (head_kv + 63) // 64 if head_kv > 64 else 1 + H_per_block = 64 if REPLICATE_H > 1 else padded_H + + NI_1 = topk_1 // BI + n_groups_1 = NI_1 // inner_iter_1 + NI_2 = (topk_2 // BI) if is_dual else 0 + n_groups_2 = (NI_2 // inner_iter_2) if is_dual else 0 + n_groups = n_groups_1 + n_groups_2 + + BS_KV_1 = block_size_kv_1 + NOPE_ROPE_U32_PER_BLOCK_1 = BS_KV_1 * PACKED_W4 + if is_dual: + BS_KV_2 = block_size_kv_2 + NOPE_ROPE_U32_PER_BLOCK_2 = BS_KV_2 * PACKED_W4 + + q_shape = [batch, seq_len, num_heads, D + D_tail] + k1_shape = [num_blocks_kv_1, block_pad_u32_1] + indices1_shape = [batch, seq_len, topk_1] + topk_length_shape = [batch] + partial_o_shape = [batch, seq_len, n_groups, num_heads, D + D_tail] + partial_lse_shape = [batch, seq_len, n_groups, num_heads] + if is_dual: + k2_shape = [num_blocks_kv_2, block_pad_u32_2] + indices2_shape = [batch, seq_len, topk_2] + + accum_dtype = "float" + indices_dtype = INT32 + + if is_dual: + + @T.prim_func + def main( + Q: T.Tensor(q_shape, BF16), # type: ignore + K_combined_1: T.Tensor(k1_shape, "uint32"), # type: ignore + Indices_1: T.Tensor(indices1_shape, indices_dtype), # type: ignore + Topk_length_1: T.Tensor(topk_length_shape, indices_dtype), # type: ignore + K_combined_2: T.Tensor(k2_shape, "uint32"), # type: ignore + Indices_2: T.Tensor(indices2_shape, indices_dtype), # type: ignore + Topk_length_2: T.Tensor(topk_length_shape, indices_dtype), # type: ignore + Partial_O: T.Tensor(partial_o_shape, BF16), # type: ignore + Partial_LSE: T.Tensor(partial_lse_shape, accum_dtype), # type: ignore + ) -> None: + """ + grid: (seq_len * REPLICATE_H * n_groups, batch, 1) + Each block processes `inner_iter_1` (or `inner_iter_2`) consecutive + KV tiles of one phase and writes one (partial_o, partial_lse) entry. + """ + with T.Kernel( + seq_len * REPLICATE_H * n_groups, batch, kv_group, threads=threads + ) as (bx, by, bz): + Q_shared = T.alloc_fragment([H_per_block, D], BF16) + Q_tail_shared = T.alloc_fragment([H_per_block, D_tail], BF16) + K_packed_shared = T.alloc_shared([BI, PACKED_W4], "uint32") + K_scale_shared = T.alloc_shared([BI, SCALE_W4], "uint32") + KV_shared = T.alloc_shared([BI, D], BF16) + K_tail_shared = T.alloc_shared([BI, D_tail], BF16) + S_shared = T.alloc_shared([H_per_block, BI], BF16) + page_idx_shared = T.alloc_shared([BI], INT32) + + mask = T.alloc_fragment([BI], "bool") + scale_byte_local = T.alloc_fragment([BI, NUM_TILES], "uint32") + + acc_o = T.alloc_fragment([H_per_block, D], accum_dtype) + acc_o_tail = T.alloc_fragment([H_per_block, D_tail], accum_dtype) + acc_s = T.alloc_fragment([H_per_block, BI], accum_dtype) + sumexp = T.alloc_fragment([H_per_block], accum_dtype) + sumexp_i = T.alloc_fragment([H_per_block], accum_dtype) + alpha = T.alloc_fragment([H_per_block], accum_dtype) + m_i = T.alloc_fragment([H_per_block], accum_dtype) + m_i_prev = T.alloc_fragment([H_per_block], accum_dtype) + + T.fill(acc_o, 0) + T.fill(acc_o_tail, 0) + T.fill(sumexp, 0) + T.fill(m_i, -(2**30)) + + b_i, g_i = by, bz + # bx encodes (s_i, h_replicate, group_i). + spans_per_seq = REPLICATE_H * n_groups + s_i = bx // spans_per_seq + rest = bx % spans_per_seq + group_i = rest // REPLICATE_H + h_rep = rest % REPLICATE_H + H0 = g_i * padded_H + (0 if REPLICATE_H == 1 else h_rep * 64) + H1 = H0 + H_per_block + + tk_len_1 = Topk_length_1[b_i] + tk_len_2 = Topk_length_2[b_i] + actual_n_groups_1 = T.ceildiv(tk_len_1, BI * inner_iter_1) + actual_n_groups_2 = T.ceildiv(tk_len_2, BI * inner_iter_2) + + if (group_i < n_groups_1) & (group_i < actual_n_groups_1): + # Phase 1 active: SWA cache work + Partial_O write. + T.copy(Q[b_i, s_i, H0:H1, :D], Q_shared) + T.copy(Q[b_i, s_i, H0:H1, D : D + D_tail], Q_tail_shared) + for k_i in T.Pipelined(inner_iter_1, num_stages=num_stages): + iter_i = group_i * inner_iter_1 + k_i + for bi_i in T.Parallel(BI): + pos = iter_i * BI + bi_i + idx = Indices_1[b_i, s_i, pos] + valid = (idx >= 0) & (pos < tk_len_1) + page_idx_shared[bi_i] = T.if_then_else(valid, idx, 0) + mask[bi_i] = valid + + for bi_i, w_i in T.Parallel(BI, PACKED_W4): + page = page_idx_shared[bi_i] + block_id = page // BS_KV_1 + t_in_block = page % BS_KV_1 + K_packed_shared[bi_i, w_i] = K_combined_1[ + block_id, t_in_block * PACKED_W4 + w_i + ] + + for bi_i, w_i in T.Parallel(BI, SCALE_W4): + page = page_idx_shared[bi_i] + block_id = page // BS_KV_1 + t_in_block = page % BS_KV_1 + K_scale_shared[bi_i, w_i] = K_combined_1[ + block_id, + NOPE_ROPE_U32_PER_BLOCK_1 + t_in_block * SCALE_W4 + w_i, + ] + + for bi_i, ti in T.Parallel(BI, NUM_TILES): + word_idx = ti // 4 + byte_in_word = ti % 4 + word = K_scale_shared[bi_i, word_idx] + scale_byte_local[bi_i, ti] = ( + word >> T.Cast("uint32", byte_in_word * 8) + ) & T.uint32(0xFF) + + for bi_i, d_i in T.Parallel(BI, D): + word_idx = d_i // 4 + byte_in_word = d_i % 4 + word = K_packed_shared[bi_i, word_idx] + b_u32 = ( + word >> T.Cast("uint32", byte_in_word * 8) + ) & T.uint32(0xFF) + sign_bf = (b_u32 & T.uint32(0x80)) * T.uint32(0x100) + exp_e4 = (b_u32 & T.uint32(0x78)) >> T.uint32(3) + mant_bf = (b_u32 & T.uint32(0x7)) * T.uint32(0x10) + scale_byte = scale_byte_local[bi_i, d_i // NOPE_TILE] + exp_combined = exp_e4 + scale_byte - T.uint32(7) + bf16_bits = ( + sign_bf | (exp_combined << T.uint32(7)) | mant_bf + ) + KV_shared[bi_i, d_i] = T.reinterpret( + BF16, T.Cast("uint16", bf16_bits) + ) + + for bi_i, j in T.Parallel(BI, D_tail): + abs_off = D + 2 * j + word_idx = abs_off // 4 + word_off = abs_off % 4 + word = K_packed_shared[bi_i, word_idx] + half_u32 = T.if_then_else( + word_off == 0, + word & T.uint32(0xFFFF), + (word >> T.uint32(16)) & T.uint32(0xFFFF), + ) + K_tail_shared[bi_i, j] = T.reinterpret( + BF16, T.Cast("uint16", half_u32) + ) + + for h_i, bi_i in T.Parallel(H_per_block, BI): + acc_s[h_i, bi_i] = T.if_then_else( + mask[bi_i], 0, -T.infinity(acc_s.dtype) + ) + T.gemm( + Q_shared, + KV_shared, + acc_s, + transpose_B=True, + policy=T.GemmWarpPolicy.FullRow, + ) + T.gemm( + Q_tail_shared, + K_tail_shared, + acc_s, + transpose_B=True, + policy=T.GemmWarpPolicy.FullRow, + ) + T.copy(m_i, m_i_prev) + T.reduce_max(acc_s, m_i, dim=1, clear=False) + for h_i in T.Parallel(H_per_block): + m_i[h_i] = T.max(m_i[h_i], m_i_prev[h_i]) + for h_i in T.Parallel(H_per_block): + alpha[h_i] = T.exp2((m_i_prev[h_i] - m_i[h_i]) * sm_scale) + for h_i, bi_i in T.Parallel(H_per_block, BI): + acc_s[h_i, bi_i] = T.exp2( + acc_s[h_i, bi_i] * sm_scale - m_i[h_i] * sm_scale + ) + T.reduce_sum(acc_s, sumexp_i, dim=1) + for h_i in T.Parallel(H_per_block): + sumexp[h_i] = sumexp[h_i] * alpha[h_i] + sumexp_i[h_i] + for h_i, d_i in T.Parallel(H_per_block, D): + acc_o[h_i, d_i] *= alpha[h_i] + for h_i, d_i in T.Parallel(H_per_block, D_tail): + acc_o_tail[h_i, d_i] *= alpha[h_i] + T.copy(acc_s, S_shared) + T.gemm( + S_shared, + KV_shared, + acc_o, + policy=T.GemmWarpPolicy.FullRow, + ) + T.gemm( + S_shared, + K_tail_shared, + acc_o_tail, + policy=T.GemmWarpPolicy.FullRow, + ) + # ---- finalize phase 1 (active) ---- + for h_i, d_i in T.Parallel(H_per_block, D): + acc_o[h_i, d_i] = acc_o[h_i, d_i] / T.if_then_else( + sumexp[h_i] == 0.0, 1.0, sumexp[h_i] + ) + for h_i, d_i in T.Parallel(H_per_block, D_tail): + acc_o_tail[h_i, d_i] = acc_o_tail[h_i, d_i] / T.if_then_else( + sumexp[h_i] == 0.0, 1.0, sumexp[h_i] + ) + for h_i in T.Parallel(H_per_block): + m_i[h_i] = T.if_then_else( + sumexp[h_i] == 0.0, + -(2.0**30), + T.log2(sumexp[h_i]) + m_i[h_i] * sm_scale, + ) + T.copy(acc_o, Partial_O[b_i, s_i, group_i, H0:H1, :D]) + T.copy( + acc_o_tail, + Partial_O[b_i, s_i, group_i, H0:H1, D : D + D_tail], + ) + T.copy(m_i, Partial_LSE[b_i, s_i, group_i, H0:H1]) + elif group_i < n_groups_1: + # Phase 1 skipped: m_i is still the -2^30 + T.copy(m_i, Partial_LSE[b_i, s_i, group_i, H0:H1]) + elif (group_i - n_groups_1) < actual_n_groups_2: + # Phase 2 active: c128 cache work + Partial_O write. + T.copy(Q[b_i, s_i, H0:H1, :D], Q_shared) + T.copy(Q[b_i, s_i, H0:H1, D : D + D_tail], Q_tail_shared) + for k_i in T.Pipelined(inner_iter_2, num_stages=num_stages): + iter_i = (group_i - n_groups_1) * inner_iter_2 + k_i + for bi_i in T.Parallel(BI): + pos = iter_i * BI + bi_i + idx = Indices_2[b_i, s_i, pos] + valid = (idx >= 0) & (pos < tk_len_2) + page_idx_shared[bi_i] = T.if_then_else(valid, idx, 0) + mask[bi_i] = valid + + for bi_i, w_i in T.Parallel(BI, PACKED_W4): + page = page_idx_shared[bi_i] + block_id = page // BS_KV_2 + t_in_block = page % BS_KV_2 + K_packed_shared[bi_i, w_i] = K_combined_2[ + block_id, t_in_block * PACKED_W4 + w_i + ] + + for bi_i, w_i in T.Parallel(BI, SCALE_W4): + page = page_idx_shared[bi_i] + block_id = page // BS_KV_2 + t_in_block = page % BS_KV_2 + K_scale_shared[bi_i, w_i] = K_combined_2[ + block_id, + NOPE_ROPE_U32_PER_BLOCK_2 + t_in_block * SCALE_W4 + w_i, + ] + + for bi_i, ti in T.Parallel(BI, NUM_TILES): + word_idx = ti // 4 + byte_in_word = ti % 4 + word = K_scale_shared[bi_i, word_idx] + scale_byte_local[bi_i, ti] = ( + word >> T.Cast("uint32", byte_in_word * 8) + ) & T.uint32(0xFF) + + for bi_i, d_i in T.Parallel(BI, D): + word_idx = d_i // 4 + byte_in_word = d_i % 4 + word = K_packed_shared[bi_i, word_idx] + b_u32 = ( + word >> T.Cast("uint32", byte_in_word * 8) + ) & T.uint32(0xFF) + sign_bf = (b_u32 & T.uint32(0x80)) * T.uint32(0x100) + exp_e4 = (b_u32 & T.uint32(0x78)) >> T.uint32(3) + mant_bf = (b_u32 & T.uint32(0x7)) * T.uint32(0x10) + scale_byte = scale_byte_local[bi_i, d_i // NOPE_TILE] + exp_combined = exp_e4 + scale_byte - T.uint32(7) + bf16_bits = ( + sign_bf | (exp_combined << T.uint32(7)) | mant_bf + ) + KV_shared[bi_i, d_i] = T.reinterpret( + BF16, T.Cast("uint16", bf16_bits) + ) + + for bi_i, j in T.Parallel(BI, D_tail): + abs_off = D + 2 * j + word_idx = abs_off // 4 + word_off = abs_off % 4 + word = K_packed_shared[bi_i, word_idx] + half_u32 = T.if_then_else( + word_off == 0, + word & T.uint32(0xFFFF), + (word >> T.uint32(16)) & T.uint32(0xFFFF), + ) + K_tail_shared[bi_i, j] = T.reinterpret( + BF16, T.Cast("uint16", half_u32) + ) + + for h_i, bi_i in T.Parallel(H_per_block, BI): + acc_s[h_i, bi_i] = T.if_then_else( + mask[bi_i], 0, -T.infinity(acc_s.dtype) + ) + T.gemm( + Q_shared, + KV_shared, + acc_s, + transpose_B=True, + policy=T.GemmWarpPolicy.FullRow, + ) + T.gemm( + Q_tail_shared, + K_tail_shared, + acc_s, + transpose_B=True, + policy=T.GemmWarpPolicy.FullRow, + ) + T.copy(m_i, m_i_prev) + T.reduce_max(acc_s, m_i, dim=1, clear=False) + for h_i in T.Parallel(H_per_block): + m_i[h_i] = T.max(m_i[h_i], m_i_prev[h_i]) + for h_i in T.Parallel(H_per_block): + alpha[h_i] = T.exp2((m_i_prev[h_i] - m_i[h_i]) * sm_scale) + for h_i, bi_i in T.Parallel(H_per_block, BI): + acc_s[h_i, bi_i] = T.exp2( + acc_s[h_i, bi_i] * sm_scale - m_i[h_i] * sm_scale + ) + T.reduce_sum(acc_s, sumexp_i, dim=1) + for h_i in T.Parallel(H_per_block): + sumexp[h_i] = sumexp[h_i] * alpha[h_i] + sumexp_i[h_i] + for h_i, d_i in T.Parallel(H_per_block, D): + acc_o[h_i, d_i] *= alpha[h_i] + for h_i, d_i in T.Parallel(H_per_block, D_tail): + acc_o_tail[h_i, d_i] *= alpha[h_i] + T.copy(acc_s, S_shared) + T.gemm( + S_shared, + KV_shared, + acc_o, + policy=T.GemmWarpPolicy.FullRow, + ) + T.gemm( + S_shared, + K_tail_shared, + acc_o_tail, + policy=T.GemmWarpPolicy.FullRow, + ) + # ---- finalize phase 2 (active) ---- + for h_i, d_i in T.Parallel(H_per_block, D): + acc_o[h_i, d_i] = acc_o[h_i, d_i] / T.if_then_else( + sumexp[h_i] == 0.0, 1.0, sumexp[h_i] + ) + for h_i, d_i in T.Parallel(H_per_block, D_tail): + acc_o_tail[h_i, d_i] = acc_o_tail[h_i, d_i] / T.if_then_else( + sumexp[h_i] == 0.0, 1.0, sumexp[h_i] + ) + for h_i in T.Parallel(H_per_block): + m_i[h_i] = T.if_then_else( + sumexp[h_i] == 0.0, + -(2.0**30), + T.log2(sumexp[h_i]) + m_i[h_i] * sm_scale, + ) + T.copy(acc_o, Partial_O[b_i, s_i, group_i, H0:H1, :D]) + T.copy( + acc_o_tail, + Partial_O[b_i, s_i, group_i, H0:H1, D : D + D_tail], + ) + T.copy(m_i, Partial_LSE[b_i, s_i, group_i, H0:H1]) + else: + # Phase 2 skipped: m_i is still the -2^30 + T.copy(m_i, Partial_LSE[b_i, s_i, group_i, H0:H1]) + + return main + + @T.prim_func + def main( + Q: T.Tensor(q_shape, BF16), # type: ignore + K_combined_1: T.Tensor(k1_shape, "uint32"), # type: ignore + Indices_1: T.Tensor(indices1_shape, indices_dtype), # type: ignore + Topk_length_1: T.Tensor(topk_length_shape, indices_dtype), # type: ignore + Partial_O: T.Tensor(partial_o_shape, BF16), # type: ignore + Partial_LSE: T.Tensor(partial_lse_shape, accum_dtype), # type: ignore + ) -> None: + """ + grid: (seq_len * REPLICATE_H * n_groups, batch, 1) + Each block processes `inner_iter_1` consecutive KV tiles and writes + one (partial_o, partial_lse) entry. + """ + with T.Kernel( + seq_len * REPLICATE_H * n_groups, batch, kv_group, threads=threads + ) as (bx, by, bz): + Q_shared = T.alloc_fragment([H_per_block, D], BF16) + Q_tail_shared = T.alloc_fragment([H_per_block, D_tail], BF16) + K_packed_shared = T.alloc_shared([BI, PACKED_W4], "uint32") + K_scale_shared = T.alloc_shared([BI, SCALE_W4], "uint32") + KV_shared = T.alloc_shared([BI, D], BF16) + K_tail_shared = T.alloc_shared([BI, D_tail], BF16) + S_shared = T.alloc_shared([H_per_block, BI], BF16) + page_idx_shared = T.alloc_shared([BI], INT32) + + mask = T.alloc_fragment([BI], "bool") + scale_byte_local = T.alloc_fragment([BI, NUM_TILES], "uint32") + + acc_o = T.alloc_fragment([H_per_block, D], accum_dtype) + acc_o_tail = T.alloc_fragment([H_per_block, D_tail], accum_dtype) + acc_s = T.alloc_fragment([H_per_block, BI], accum_dtype) + sumexp = T.alloc_fragment([H_per_block], accum_dtype) + sumexp_i = T.alloc_fragment([H_per_block], accum_dtype) + alpha = T.alloc_fragment([H_per_block], accum_dtype) + m_i = T.alloc_fragment([H_per_block], accum_dtype) + m_i_prev = T.alloc_fragment([H_per_block], accum_dtype) + + T.fill(acc_o, 0) + T.fill(acc_o_tail, 0) + T.fill(sumexp, 0) + T.fill(m_i, -(2**30)) + + b_i, g_i = by, bz + spans_per_seq = REPLICATE_H * n_groups + s_i = bx // spans_per_seq + rest = bx % spans_per_seq + group_i = rest // REPLICATE_H + h_rep = rest % REPLICATE_H + H0 = g_i * padded_H + (0 if REPLICATE_H == 1 else h_rep * 64) + H1 = H0 + H_per_block + + T.copy(Q[b_i, s_i, H0:H1, :D], Q_shared) + T.copy(Q[b_i, s_i, H0:H1, D : D + D_tail], Q_tail_shared) + + tk_len_1 = Topk_length_1[b_i] + + for k_i in T.Pipelined(inner_iter_1, num_stages=num_stages): + iter_i = group_i * inner_iter_1 + k_i + for bi_i in T.Parallel(BI): + pos = iter_i * BI + bi_i + idx = Indices_1[b_i, s_i, pos] + valid = (idx >= 0) & (pos < tk_len_1) + page_idx_shared[bi_i] = T.if_then_else(valid, idx, 0) + mask[bi_i] = valid + + for bi_i, w_i in T.Parallel(BI, PACKED_W4): + page = page_idx_shared[bi_i] + block_id = page // BS_KV_1 + t_in_block = page % BS_KV_1 + K_packed_shared[bi_i, w_i] = K_combined_1[ + block_id, t_in_block * PACKED_W4 + w_i + ] + + for bi_i, w_i in T.Parallel(BI, SCALE_W4): + page = page_idx_shared[bi_i] + block_id = page // BS_KV_1 + t_in_block = page % BS_KV_1 + K_scale_shared[bi_i, w_i] = K_combined_1[ + block_id, + NOPE_ROPE_U32_PER_BLOCK_1 + t_in_block * SCALE_W4 + w_i, + ] + + for bi_i, ti in T.Parallel(BI, NUM_TILES): + word_idx = ti // 4 + byte_in_word = ti % 4 + word = K_scale_shared[bi_i, word_idx] + scale_byte_local[bi_i, ti] = ( + word >> T.Cast("uint32", byte_in_word * 8) + ) & T.uint32(0xFF) + + for bi_i, d_i in T.Parallel(BI, D): + word_idx = d_i // 4 + byte_in_word = d_i % 4 + word = K_packed_shared[bi_i, word_idx] + b_u32 = (word >> T.Cast("uint32", byte_in_word * 8)) & T.uint32( + 0xFF + ) + sign_bf = (b_u32 & T.uint32(0x80)) * T.uint32(0x100) + exp_e4 = (b_u32 & T.uint32(0x78)) >> T.uint32(3) + mant_bf = (b_u32 & T.uint32(0x7)) * T.uint32(0x10) + scale_byte = scale_byte_local[bi_i, d_i // NOPE_TILE] + exp_combined = exp_e4 + scale_byte - T.uint32(7) + bf16_bits = sign_bf | (exp_combined << T.uint32(7)) | mant_bf + KV_shared[bi_i, d_i] = T.reinterpret( + BF16, T.Cast("uint16", bf16_bits) + ) + + for bi_i, j in T.Parallel(BI, D_tail): + abs_off = D + 2 * j + word_idx = abs_off // 4 + word_off = abs_off % 4 + word = K_packed_shared[bi_i, word_idx] + half_u32 = T.if_then_else( + word_off == 0, + word & T.uint32(0xFFFF), + (word >> T.uint32(16)) & T.uint32(0xFFFF), + ) + K_tail_shared[bi_i, j] = T.reinterpret( + BF16, T.Cast("uint16", half_u32) + ) + + for h_i, bi_i in T.Parallel(H_per_block, BI): + acc_s[h_i, bi_i] = T.if_then_else( + mask[bi_i], 0, -T.infinity(acc_s.dtype) + ) + T.gemm( + Q_shared, + KV_shared, + acc_s, + transpose_B=True, + policy=T.GemmWarpPolicy.FullRow, + ) + T.gemm( + Q_tail_shared, + K_tail_shared, + acc_s, + transpose_B=True, + policy=T.GemmWarpPolicy.FullRow, + ) + T.copy(m_i, m_i_prev) + T.reduce_max(acc_s, m_i, dim=1, clear=False) + for h_i in T.Parallel(H_per_block): + m_i[h_i] = T.max(m_i[h_i], m_i_prev[h_i]) + for h_i in T.Parallel(H_per_block): + alpha[h_i] = T.exp2((m_i_prev[h_i] - m_i[h_i]) * sm_scale) + for h_i, bi_i in T.Parallel(H_per_block, BI): + acc_s[h_i, bi_i] = T.exp2( + acc_s[h_i, bi_i] * sm_scale - m_i[h_i] * sm_scale + ) + T.reduce_sum(acc_s, sumexp_i, dim=1) + for h_i in T.Parallel(H_per_block): + sumexp[h_i] = sumexp[h_i] * alpha[h_i] + sumexp_i[h_i] + for h_i, d_i in T.Parallel(H_per_block, D): + acc_o[h_i, d_i] *= alpha[h_i] + for h_i, d_i in T.Parallel(H_per_block, D_tail): + acc_o_tail[h_i, d_i] *= alpha[h_i] + T.copy(acc_s, S_shared) + T.gemm(S_shared, KV_shared, acc_o, policy=T.GemmWarpPolicy.FullRow) + T.gemm( + S_shared, K_tail_shared, acc_o_tail, policy=T.GemmWarpPolicy.FullRow + ) + + for h_i, d_i in T.Parallel(H_per_block, D): + acc_o[h_i, d_i] = acc_o[h_i, d_i] / T.if_then_else( + sumexp[h_i] == 0.0, 1.0, sumexp[h_i] + ) + for h_i, d_i in T.Parallel(H_per_block, D_tail): + acc_o_tail[h_i, d_i] = acc_o_tail[h_i, d_i] / T.if_then_else( + sumexp[h_i] == 0.0, 1.0, sumexp[h_i] + ) + for h_i in T.Parallel(H_per_block): + m_i[h_i] = T.if_then_else( + sumexp[h_i] == 0.0, + -(2.0**30), + T.log2(sumexp[h_i]) + m_i[h_i] * sm_scale, + ) + T.copy(acc_o, Partial_O[b_i, s_i, group_i, H0:H1, :D]) + T.copy(acc_o_tail, Partial_O[b_i, s_i, group_i, H0:H1, D : D + D_tail]) + T.copy(m_i, Partial_LSE[b_i, s_i, group_i, H0:H1]) + + return main + + +@tilelang.jit( + out_idx=[-2, -1], + pass_configs={ + tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, + tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True, + tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True, + }, +) +def dpsk_v4_combine_kernel( + num_heads: int, + n_groups_1: int, + n_groups_2: int = 0, + *, + block_I: int = 64, + inner_iter_1: int = 1, + inner_iter_2: int = 1, + dim: int = 448, + tail_dim: int = 64, + head_per_block: int = 16, + threads: int = 256, + use_attn_sink: bool = False, +) -> Any: + """ + Combine `n_groups` flash-attention partials into the final output. + + Inputs: + Partial_O : (batch, seq_len, n_groups, num_heads, dim+tail_dim) bf16 + Partial_LSE : (batch, seq_len, n_groups, num_heads) fp32, log2 form + Topk_length_1: (batch,) int32, actual phase-1 length + Topk_length_2: (batch,) int32, actual phase-2 length (dual only) + Attn_sink : (num_heads,) fp32 + Outputs: + Output : (batch, seq_len, num_heads, dim+tail_dim) bf16 + LSE : (batch, seq_len, num_heads) fp32, natural log + + Each grid block handles `head_per_block` heads of one (batch, seq) row. + """ + log2e: float = 1.44269504 + ln2: float = 0.69314718 + assert num_heads % head_per_block == 0 + + is_dual = n_groups_2 > 0 + n_groups = n_groups_1 + n_groups_2 + + H_per_block = head_per_block + HEAD_BLOCKS = num_heads // H_per_block + DT = dim + tail_dim + + batch = T.symbolic("batch") + seq_len = T.symbolic("seq_len") + + accum_dtype = "float" + + if is_dual: + + @T.prim_func + def main( + Partial_O: T.Tensor( + [batch, seq_len, n_groups, num_heads, DT], BF16 + ), # type: ignore + Partial_LSE: T.Tensor( + [batch, seq_len, n_groups, num_heads], accum_dtype + ), # type: ignore + Topk_length_1: T.Tensor([batch], INT32), # type: ignore + Topk_length_2: T.Tensor([batch], INT32), # type: ignore + Attn_sink: T.Tensor([num_heads], FP32), # type: ignore + Output: T.Tensor([batch, seq_len, num_heads, DT], BF16), # type: ignore + LSE: T.Tensor([batch, seq_len, num_heads], accum_dtype), # type: ignore + ) -> None: + with T.Kernel(seq_len * HEAD_BLOCKS, batch, threads=threads) as ( + bx, + by, + ): + shared_lse = T.alloc_shared([n_groups, H_per_block], accum_dtype) + lse_max = T.alloc_fragment([H_per_block], accum_dtype) + lse_sum = T.alloc_fragment([H_per_block], accum_dtype) + scale = T.alloc_fragment([H_per_block, n_groups], accum_dtype) + acc_o = T.alloc_fragment([H_per_block, DT], accum_dtype) + attn_sink_frag = T.alloc_fragment([H_per_block], accum_dtype) + o_scale_frag = T.alloc_fragment([H_per_block], accum_dtype) + final_lse = T.alloc_fragment([H_per_block], accum_dtype) + + b_i = by + s_i = bx // HEAD_BLOCKS + head_block = bx % HEAD_BLOCKS + H0 = head_block * H_per_block + H1 = H0 + H_per_block + + # Clamp to the captured-shape upper bounds so callers passing + # the INT32_MAX sentinel (= "all valid") still iterate exactly + # n_groups groups, not 33M. + actual_n_groups_1 = T.min( + T.ceildiv(Topk_length_1[b_i], block_I * inner_iter_1), + n_groups_1, + ) + actual_n_groups_2 = T.min( + T.ceildiv(Topk_length_2[b_i], block_I * inner_iter_2), + n_groups - n_groups_1, + ) + actual_n_groups = actual_n_groups_1 + actual_n_groups_2 + + # Pass 1: load only active groups' LSE into compact slots. + for k_c in T.serial(actual_n_groups): + k = T.if_then_else( + k_c < actual_n_groups_1, + k_c, + n_groups_1 + (k_c - actual_n_groups_1), + ) + T.copy(Partial_LSE[b_i, s_i, k, H0:H1], shared_lse[k_c, :]) + + T.fill(lse_max, -(2**30)) + for k_c in T.serial(actual_n_groups): + for h_i in T.Parallel(H_per_block): + lse_max[h_i] = T.max(lse_max[h_i], shared_lse[k_c, h_i]) + T.fill(lse_sum, 0) + for k_c in T.serial(actual_n_groups): + for h_i in T.Parallel(H_per_block): + lse_sum[h_i] = lse_sum[h_i] + T.exp2( + shared_lse[k_c, h_i] - lse_max[h_i] + ) + for k_c in T.serial(actual_n_groups): + for h_i in T.Parallel(H_per_block): + scale[h_i, k_c] = T.exp2( + shared_lse[k_c, h_i] - lse_max[h_i] - T.log2(lse_sum[h_i]) + ) + + T.fill(acc_o, 0) + for k_c in T.serial(actual_n_groups): + k = T.if_then_else( + k_c < actual_n_groups_1, + k_c, + n_groups_1 + (k_c - actual_n_groups_1), + ) + for h_i, d_i in T.Parallel(H_per_block, DT): + acc_o[h_i, d_i] = acc_o[h_i, d_i] + scale[h_i, k_c] * Partial_O[ + b_i, s_i, k, H0 + h_i, d_i + ].astype(accum_dtype) + + for h_i in T.Parallel(H_per_block): + empty = lse_max[h_i] <= -(2**29) + final_lse[h_i] = T.if_then_else( + empty, + T.infinity(accum_dtype), + (lse_max[h_i] + T.log2(lse_sum[h_i])) * ln2, + ) + + if use_attn_sink: + for h_i in T.Parallel(H_per_block): + attn_sink_frag[h_i] = Attn_sink[H0 + h_i] + for h_i in T.Parallel(H_per_block): + empty = lse_max[h_i] <= -(2**29) + o_scale_frag[h_i] = T.if_then_else( + empty, + 0.0, + 1.0 + / ( + 1.0 + + T.exp2((attn_sink_frag[h_i] - final_lse[h_i]) * log2e) + ), + ) + for h_i, d_i in T.Parallel(H_per_block, DT): + acc_o[h_i, d_i] = acc_o[h_i, d_i] * o_scale_frag[h_i] + + T.copy(acc_o, Output[b_i, s_i, H0:H1, :]) + T.copy(final_lse, LSE[b_i, s_i, H0:H1]) + + return main + + @T.prim_func + def main( + Partial_O: T.Tensor( + [batch, seq_len, n_groups, num_heads, DT], BF16 + ), # type: ignore + Partial_LSE: T.Tensor( + [batch, seq_len, n_groups, num_heads], accum_dtype + ), # type: ignore + Attn_sink: T.Tensor([num_heads], FP32), # type: ignore + Output: T.Tensor([batch, seq_len, num_heads, DT], BF16), # type: ignore + LSE: T.Tensor([batch, seq_len, num_heads], accum_dtype), # type: ignore + ) -> None: + with T.Kernel(seq_len * HEAD_BLOCKS, batch, threads=threads) as (bx, by): + shared_lse = T.alloc_shared([n_groups, H_per_block], accum_dtype) + + lse_max = T.alloc_fragment([H_per_block], accum_dtype) + lse_sum = T.alloc_fragment([H_per_block], accum_dtype) + scale = T.alloc_fragment([H_per_block, n_groups], accum_dtype) + acc_o = T.alloc_fragment([H_per_block, DT], accum_dtype) + attn_sink_frag = T.alloc_fragment([H_per_block], accum_dtype) + o_scale_frag = T.alloc_fragment([H_per_block], accum_dtype) + final_lse = T.alloc_fragment([H_per_block], accum_dtype) + + b_i = by + s_i = bx // HEAD_BLOCKS + head_block = bx % HEAD_BLOCKS + H0 = head_block * H_per_block + H1 = H0 + H_per_block + + for k in T.serial(n_groups): + T.copy(Partial_LSE[b_i, s_i, k, H0:H1], shared_lse[k, :]) + + T.fill(lse_max, -(2**30)) + for k in T.serial(n_groups): + for h_i in T.Parallel(H_per_block): + lse_max[h_i] = T.max(lse_max[h_i], shared_lse[k, h_i]) + T.fill(lse_sum, 0) + for k in T.serial(n_groups): + for h_i in T.Parallel(H_per_block): + lse_sum[h_i] = lse_sum[h_i] + T.exp2( + shared_lse[k, h_i] - lse_max[h_i] + ) + for k in T.serial(n_groups): + for h_i in T.Parallel(H_per_block): + scale[h_i, k] = T.exp2( + shared_lse[k, h_i] - lse_max[h_i] - T.log2(lse_sum[h_i]) + ) + + T.fill(acc_o, 0) + for k in T.serial(n_groups): + for h_i, d_i in T.Parallel(H_per_block, DT): + acc_o[h_i, d_i] = acc_o[h_i, d_i] + scale[h_i, k] * Partial_O[ + b_i, s_i, k, H0 + h_i, d_i + ].astype(accum_dtype) + + for h_i in T.Parallel(H_per_block): + empty = lse_max[h_i] <= -(2**29) + final_lse[h_i] = T.if_then_else( + empty, + T.infinity(accum_dtype), + (lse_max[h_i] + T.log2(lse_sum[h_i])) * ln2, + ) + + if use_attn_sink: + for h_i in T.Parallel(H_per_block): + attn_sink_frag[h_i] = Attn_sink[H0 + h_i] + for h_i in T.Parallel(H_per_block): + empty = lse_max[h_i] <= -(2**29) + o_scale_frag[h_i] = T.if_then_else( + empty, + 0.0, + 1.0 + / ( + 1.0 + T.exp2((attn_sink_frag[h_i] - final_lse[h_i]) * log2e) + ), + ) + for h_i, d_i in T.Parallel(H_per_block, DT): + acc_o[h_i, d_i] = acc_o[h_i, d_i] * o_scale_frag[h_i] + + T.copy(acc_o, Output[b_i, s_i, H0:H1, :]) + T.copy(final_lse, LSE[b_i, s_i, H0:H1]) + + return main + + +""" +2-stage attention kernel (partial + combine) over an FP8 KV cache, +with optional second cache (`extra_k_cache`). +""" + + +def dpsk_v4_fp8_attention_fwd( + q: torch.Tensor, + k_cache: torch.Tensor, + block_table: Optional[torch.Tensor], + cache_seqlens: Optional[torch.Tensor], + head_dim_v: int, + tile_scheduler_metadata: Any, + num_splits: None = None, + softmax_scale: Optional[float] = None, + causal: bool = False, + is_fp8_kvcache: bool = False, + indices: Optional[torch.Tensor] = None, + attn_sink: Optional[torch.Tensor] = None, + extra_k_cache: Optional[torch.Tensor] = None, + extra_indices_in_kvcache: Optional[torch.Tensor] = None, + topk_length: Optional[torch.Tensor] = None, + extra_topk_length: Optional[torch.Tensor] = None, +) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Follows the original `flash_mla.flash_mla_with_kvcache` signature. + """ + if _is_gfx95_supported: + block_I, threads, num_stages, block_per_cu, cu = 64, 512, 0, 2, 256 + else: + block_I, threads, num_stages, block_per_cu, cu = 32, 128, 1, 1, 304 + + batch, seq_len, num_heads, _ = q.shape + # Partial grid is (seq_len * REPLICATE_H * n_groups, batch, kv_group); the + # heuristic in _pick_inner_iter assumes `total_blocks = seq * ni / inner_iter`, + # so `seq` must include REPLICATE_H or n_groups doubles for medium batches. + replicate_h = max((num_heads + 63) // 64, 1) + seq = batch * seq_len * replicate_h + + k1, _, bs_kv_1 = _build_fp8_combined_view(k_cache) + topk_1 = indices.shape[-1] + ni_1 = topk_1 // block_I + tk_len_1 = ( + topk_length + if topk_length is not None + else _topk_length_sentinel(q.device, batch) + ) + if attn_sink is None: + attn_sink = torch.full( + (num_heads,), float("-inf"), dtype=torch.float32, device=q.device + ) + + has_extra = extra_k_cache is not None + if not has_extra: + inner_iter_1 = _pick_inner_iter(seq, ni_1, cu, block_per_cu) + inner_iter_2 = 1 + n_groups_1 = ni_1 // inner_iter_1 + n_groups_2 = 0 + partial = dpsk_v4_fp8_partial_kernel( + num_heads, + topk_1, + bs_kv_1, + sm_scale=softmax_scale, + block_I=block_I, + inner_iter_1=inner_iter_1, + num_stages=num_stages, + threads=threads, + ) + partial_o, partial_lse = partial(q, k1, indices, tk_len_1) + else: + k2, _, bs_kv_2 = _build_fp8_combined_view(extra_k_cache) + topk_2 = extra_indices_in_kvcache.shape[-1] + ni_2 = topk_2 // block_I + # Each phase picks its own optimal split-K independently — kernel + # body uses two T.Pipelined loops with separate compile-time iter + # counts, no shared-divisor constraint. + inner_iter_1 = _pick_inner_iter(seq, ni_1, cu, block_per_cu) + inner_iter_2 = _pick_inner_iter(seq, ni_2, cu, block_per_cu) + n_groups_1 = ni_1 // inner_iter_1 + n_groups_2 = ni_2 // inner_iter_2 + tk_len_2 = ( + extra_topk_length + if extra_topk_length is not None + else _topk_length_sentinel(q.device, batch) + ) + partial = dpsk_v4_fp8_partial_kernel( + num_heads, + topk_1, + bs_kv_1, + topk_2, + bs_kv_2, + sm_scale=softmax_scale, + block_I=block_I, + inner_iter_1=inner_iter_1, + inner_iter_2=inner_iter_2, + num_stages=num_stages, + threads=threads, + ) + partial_o, partial_lse = partial( + q, + k1, + indices, + tk_len_1, + k2, + extra_indices_in_kvcache, + tk_len_2, + ) + + combine = dpsk_v4_combine_kernel( + num_heads, + n_groups_1, + n_groups_2, + block_I=block_I, + inner_iter_1=inner_iter_1, + inner_iter_2=inner_iter_2, + head_per_block=4, + threads=256, + use_attn_sink=True, + ) + if has_extra: + return combine(partial_o, partial_lse, tk_len_1, tk_len_2, attn_sink) + return combine(partial_o, partial_lse, attn_sink) diff --git a/python/sglang/srt/layers/attention/dsa/transform_index.py b/python/sglang/srt/layers/attention/dsa/transform_index.py new file mode 100644 index 000000000..10b1068f5 --- /dev/null +++ b/python/sglang/srt/layers/attention/dsa/transform_index.py @@ -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") diff --git a/python/sglang/srt/layers/attention/dsa/triton_kernel.py b/python/sglang/srt/layers/attention/dsa/triton_kernel.py new file mode 100644 index 000000000..0d2969804 --- /dev/null +++ b/python/sglang/srt/layers/attention/dsa/triton_kernel.py @@ -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, + ) diff --git a/python/sglang/srt/layers/attention/dsa/utils.py b/python/sglang/srt/layers/attention/dsa/utils.py new file mode 100644 index 000000000..b49631c97 --- /dev/null +++ b/python/sglang/srt/layers/attention/dsa/utils.py @@ -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 diff --git a/python/sglang/srt/layers/attention/dsa_backend.py b/python/sglang/srt/layers/attention/dsa_backend.py new file mode 100644 index 000000000..a2c062b91 --- /dev/null +++ b/python/sglang/srt/layers/attention/dsa_backend.py @@ -0,0 +1,2528 @@ +from __future__ import annotations + +from dataclasses import dataclass +from enum import IntEnum, auto +from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Tuple, TypeAlias + +import torch + +from sglang.srt.configs.model_config import get_dsa_index_topk, is_deepseek_dsa +from sglang.srt.environ import envs +from sglang.srt.layers.attention.base_attn_backend import AttentionBackend +from sglang.srt.layers.attention.dsa.dequant_k_cache import dequantize_k_cache_paged +from sglang.srt.layers.attention.dsa.dsa_backend_mtp_precompute import ( + DeepseekSparseAttnBackendMTPPrecomputeMixin, + PrecomputedMetadata, + compute_cu_seqlens, +) +from sglang.srt.layers.attention.dsa.dsa_indexer import BaseIndexerMetadata +from sglang.srt.layers.attention.dsa.quant_k_cache import quantize_k_cache +from sglang.srt.layers.attention.dsa.transform_index import ( + transform_index_page_table_decode, + transform_index_page_table_prefill, +) +from sglang.srt.layers.attention.dsa.utils import ( + can_dsa_prefill_cp_round_robin_split, + compute_dsa_seqlens, + dsa_cp_round_robin_split_data, + dsa_cp_round_robin_split_q_seqs, + is_dsa_enable_prefill_cp, + pad_dsa_cache_seqlens, +) +from sglang.srt.layers.attention.utils import ( + concat_mla_absorb_q_general, + mla_quantize_and_rope_for_fp8, + seqlens_expand_triton, +) +from sglang.srt.layers.dp_attention import get_attention_tp_size +from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode +from sglang.srt.utils import is_cuda, is_hip + +if TYPE_CHECKING: + from sglang.srt.layers.radix_attention import RadixAttention + from sglang.srt.model_executor.model_runner import ModelRunner + from sglang.srt.speculative.spec_info import SpecInput + + +_is_hip = is_hip() + +if _is_hip: + from sglang.srt.layers.attention.dsa.triton_kernel import get_valid_kv_indices + + try: + from aiter import ( # noqa: F401 + flash_attn_varlen_func, + mha_batch_prefill_func, + paged_attention_ragged, + ) + from aiter.mla import mla_decode_fwd, mla_prefill_fwd # noqa: F401 + except ImportError: + print( + "aiter is AMD specific kernel library. Please make sure aiter is installed on your AMD device." + ) +else: + from sglang.jit_kernel.flash_attention import ( + flash_attn_varlen_func, + flash_attn_with_kvcache, + ) + + +def _to_2d_context_lens(seqlens_32: torch.Tensor, batch_size: int) -> torch.Tensor: + # Always normalize to (N_total, 1) layout, to avoid deadlock at deep_gemm.fp8_paged_mqa_logits + if seqlens_32.dim() == 2: + if seqlens_32.size(1) == 1: + return seqlens_32 + # Fall through and re-flatten if the caller already gave us a (bs, next_n) + # view — we want (N_total, 1) regardless. + seqlens_32 = seqlens_32.reshape(-1) + return seqlens_32.contiguous().view(-1, 1) + + +# Reuse this workspace buffer across all DSA backend instances +global_workspace_buffer = None + +# Control whether to use fused metadata copy kernel for cuda graph replay (default: enabled) +# Set SGLANG_USE_FUSED_METADATA_COPY=0 or false to disable +_USE_FUSED_METADATA_COPY = envs.SGLANG_USE_FUSED_METADATA_COPY.get() and not _is_hip + + +@dataclass(frozen=True) +class DSAFlashMLAMetadata: + """Metadata only needed by FlashMLA""" + + flashmla_metadata: torch.Tensor + num_splits: torch.Tensor + + def slice(self, sli): + return DSAFlashMLAMetadata( + flashmla_metadata=self.flashmla_metadata, + num_splits=self.num_splits[sli], + ) + + def copy_(self, other: "DSAFlashMLAMetadata"): + self.flashmla_metadata.copy_(other.flashmla_metadata) + self.num_splits.copy_(other.num_splits) + + +@dataclass(frozen=True) +class DSAMetadata: + page_size: int + + # Sequence lengths for the forward batch + cache_seqlens_int32: torch.Tensor + # Maximum sequence length for query + max_seq_len_q: int + # Maximum sequence length for key + max_seq_len_k: int + # Cumulative sequence lengths for query + cu_seqlens_q: torch.Tensor + # Cumulative sequence lengths for key + cu_seqlens_k: torch.Tensor + # Page table, the index of KV Cache Tables/Blocks + # this table is always with page_size = 1 + page_table_1: torch.Tensor + + # NOTE(dark): This will property be used in: + # 1. dense decode/prefill, we use paged flash attention, need real_page_table + # 2. sparse decode/prefill, indexer need real_page_table to compute the score + real_page_table: torch.Tensor + + # DSA metadata (dsa prefill are expanded) + dsa_cache_seqlens_int32: torch.Tensor # this seqlens is clipped to `topk` + dsa_cu_seqlens_q: torch.Tensor # must be arange(0, len(dsa_cu_seqlens_k)) + dsa_cu_seqlens_k: torch.Tensor # cumsum of `dsa_cache_seqlens_int32` + dsa_extend_seq_lens_list: List[int] + dsa_seqlens_expanded: torch.Tensor # expanded, unclipped `seqlens` + dsa_max_seqlen_q: Literal[1] = 1 # always 1 for decode, variable for extend + + flashmla_metadata: Optional[DSAFlashMLAMetadata] = None + # DeepGEMM schedule metadata for paged MQA logits (decode/target_verify/draft_extend only). + # Precomputed once per forward batch and reused across layers. + paged_mqa_schedule_metadata: Optional[torch.Tensor] = None + # The sum of sequence lengths for key, prefill only + seq_lens_sum: Optional[int] = None + # The flattened 1D page table with shape (seq_lens_sum,), prefill only + # this table is always with page_size = 1 + page_table_1_flattened: Optional[torch.Tensor] = None + # The offset of topk indices in ragged kv, prefill only + # shape: (seq_lens_sum,) + topk_indices_offset: Optional[torch.Tensor] = None + + # k_start and k_end in kv cache for each token. + indexer_k_start_end: Optional[Tuple[torch.Tensor, torch.Tensor]] = None + # seq lens for each batch. + indexer_seq_lens_cpu: Optional[torch.Tensor] = None + # seq lens for each batch. + indexer_seq_lens: Optional[torch.Tensor] = None + # batch index for each token. + token_to_batch_idx: Optional[torch.Tensor] = None + + +class TopkTransformMethod(IntEnum): + # Transform topk indices to indices to the page table (page_size = 1) + PAGED = auto() + # Transform topk indices to indices to ragged kv (non-paged) + RAGGED = auto() + + +@torch.compile +def _compiled_cat(tensors: list[torch.Tensor], dim: int = -1) -> torch.Tensor: + return torch.cat(tensors, dim=dim) + + +def _cat(tensors: list[torch.Tensor], dim: int = -1) -> torch.Tensor: + """ + Concatenate two tensors along the last dimension. + Use this function to concatenate q_nope and q_rope or k_nope and k_rope. + """ + assert len(tensors) == 2 + + qk_nope, qk_rope = tensors + assert qk_nope.ndim == 3 and qk_rope.ndim == 3 + + torch._dynamo.mark_dynamic(qk_nope, 0) + torch._dynamo.mark_dynamic(qk_rope, 0) + + return _compiled_cat([qk_nope, qk_rope], dim=dim) + + +@dataclass(frozen=True) +class DSAIndexerMetadata(BaseIndexerMetadata): + attn_metadata: DSAMetadata + topk_transform_method: TopkTransformMethod + paged_mqa_schedule_metadata: Optional[torch.Tensor] = None + force_unfused_topk: bool = False + + def get_seqlens_int32(self) -> torch.Tensor: + return self.attn_metadata.cache_seqlens_int32 + + def get_page_table_64(self) -> torch.Tensor: + return self.attn_metadata.real_page_table + + def get_page_table_1(self) -> torch.Tensor: + return self.attn_metadata.page_table_1 + + def get_seqlens_expanded(self) -> torch.Tensor: + return self.attn_metadata.dsa_seqlens_expanded + + def get_cu_seqlens_k(self) -> torch.Tensor: + return self.attn_metadata.cu_seqlens_k + + def get_indexer_kvcache_range(self) -> Tuple[torch.Tensor, torch.Tensor]: + return self.attn_metadata.indexer_k_start_end + + def get_indexer_seq_len(self) -> torch.Tensor: + return self.attn_metadata.indexer_seq_lens + + def get_indexer_seq_len_cpu(self) -> torch.Tensor: + return self.attn_metadata.indexer_seq_lens_cpu + + def get_dsa_extend_len_cpu(self) -> List[int]: + return self.attn_metadata.dsa_extend_seq_lens_list + + def get_token_to_batch_idx(self) -> torch.Tensor: + return self.attn_metadata.token_to_batch_idx + + def topk_transform( + self, + logits: torch.Tensor, + topk: int, + ks: Optional[torch.Tensor] = None, + cu_seqlens_q: torch.Tensor = None, + ke_offset: torch.Tensor = None, + batch_idx_list: List[int] = None, + topk_indices_offset_override: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + from sgl_kernel import ( + fast_topk_transform_fused, + fast_topk_transform_ragged_fused, + fast_topk_v2, + ) + + if topk_indices_offset_override is not None: + cu_topk_indices_offset = topk_indices_offset_override + cu_seqlens_q_topk = None + elif cu_seqlens_q is not None: + cu_seqlens_q = cu_seqlens_q.to(torch.int32) + cu_seqlens_q_topk = compute_cu_seqlens(cu_seqlens_q) + cu_topk_indices_offset = torch.repeat_interleave( + cu_seqlens_q_topk[:-1], + cu_seqlens_q, + ) + else: + cu_seqlens_q_topk = self.attn_metadata.cu_seqlens_q + cu_topk_indices_offset = self.attn_metadata.topk_indices_offset + if ke_offset is not None: + seq_lens_topk = ke_offset + else: + seq_lens_topk = self.get_seqlens_expanded() + if batch_idx_list is not None: + page_table_size_1 = self.attn_metadata.page_table_1[batch_idx_list] + else: + page_table_size_1 = self.attn_metadata.page_table_1 + + if not envs.SGLANG_DSA_FUSE_TOPK.get() or self.force_unfused_topk: + return fast_topk_v2(logits, seq_lens_topk, topk, row_starts=ks) + elif self.topk_transform_method == TopkTransformMethod.PAGED: + # NOTE(dark): if fused, we return a transformed page table directly + return fast_topk_transform_fused( + score=logits, + lengths=seq_lens_topk, + page_table_size_1=page_table_size_1, + cu_seqlens_q=cu_seqlens_q_topk, + topk=topk, + row_starts=ks, + ) + elif self.topk_transform_method == TopkTransformMethod.RAGGED: + if cu_topk_indices_offset is None: + raise RuntimeError( + "RAGGED topk_transform requires topk_indices_offset; " + "expected extend-without-speculative metadata." + ) + return fast_topk_transform_ragged_fused( + score=logits, + lengths=seq_lens_topk, + topk_indices_offset=cu_topk_indices_offset, + topk=topk, + row_starts=ks, + ) + else: + assert False, f"Unsupported {self.topk_transform_method = }" + + +_DSA_IMPL_T: TypeAlias = Literal[ + "flashmla_sparse", "flashmla_kv", "fa3", "tilelang", "trtllm" +] + + +class DeepseekSparseAttnBackend( + DeepseekSparseAttnBackendMTPPrecomputeMixin, AttentionBackend +): + def __init__( + self, + model_runner: ModelRunner, + skip_prefill: bool = False, + speculative_step_id=0, + topk=0, + speculative_num_steps=0, + ): + super().__init__() + self.forward_metadata: DSAMetadata + self.device = model_runner.device + assert isinstance(model_runner.page_size, int) + self.real_page_size = model_runner.page_size + self.num_splits = ( + 1 if model_runner.server_args.enable_deterministic_inference else 0 + ) + self.use_dsa = is_deepseek_dsa(model_runner.model_config.hf_config) + assert self.use_dsa, "DSA backend only supports DeepSeek DSA" + self.dsa_kv_cache_store_fp8 = ( + model_runner.token_to_kv_pool.dsa_kv_cache_store_fp8 + ) + self.dsa_index_topk = get_dsa_index_topk(model_runner.model_config.hf_config) + self.max_context_len = model_runner.model_config.context_len + self.num_q_heads = ( + model_runner.model_config.num_attention_heads // get_attention_tp_size() + ) + self.kv_cache_dim = model_runner.token_to_kv_pool.kv_cache_dim + self.qk_nope_head_dim = model_runner.model_config.qk_nope_head_dim + self.kv_lora_rank = model_runner.model_config.kv_lora_rank + self.qk_rope_head_dim = model_runner.model_config.qk_rope_head_dim + + assert model_runner.req_to_token_pool is not None + self.req_to_token = model_runner.req_to_token_pool.req_to_token + + self.use_mha: bool = False + self.dsa_prefill_impl: _DSA_IMPL_T = ( + model_runner.server_args.dsa_prefill_backend + ) + self.dsa_decode_impl: _DSA_IMPL_T = model_runner.server_args.dsa_decode_backend + if self.num_q_heads <= 64: + self.flashmla_kv_num_q_heads = 64 + elif self.num_q_heads <= 128: + self.flashmla_kv_num_q_heads = 128 + else: + # Keep original head count if it exceeds current padded variants. + self.flashmla_kv_num_q_heads = self.num_q_heads + self.enable_auto_select_prefill_impl = self.dsa_prefill_impl == "flashmla_auto" + + self._arange_buf = torch.arange(16384, device=self.device, dtype=torch.int32) + + if _is_hip: + max_bs = model_runner.req_to_token_pool.size + + self.kv_indptr = torch.zeros( + (max_bs + 1,), dtype=torch.int32, device=model_runner.device + ) + + self.kv_indices = torch.zeros( + max_bs * self.dsa_index_topk, + dtype=torch.int32, + device=self.device, + ) + # Aiter mla_decode_fwd supports num_heads multiples of 16 in range [16, 128]. + # For models with fewer heads per GPU (e.g. GLM-5 64 heads / TP8 = 8), need to pad the heads to 16. + self.need_pad_heads = self.num_q_heads < 16 + self.head_repeat_factor = ( + 16 // self.num_q_heads if self.num_q_heads < 16 else 1 + ) + + # Speculative decoding + self.topk = model_runner.server_args.speculative_eagle_topk or 0 + self.speculative_num_steps = speculative_num_steps + self.speculative_num_draft_tokens = ( + model_runner.server_args.speculative_num_draft_tokens + ) + self.speculative_step_id = speculative_step_id + + self.device_capability = torch.cuda.get_device_capability() + self.device_sm_major = self.device_capability[0] + self.kv_cache_dtype = model_runner.kv_cache_dtype + + # Allocate global workspace buffer for TRT-LLM kernels (ragged attention on SM100/B200, or trtllm decode) + if self.device_sm_major >= 10 or self.dsa_decode_impl == "trtllm": + global global_workspace_buffer + if global_workspace_buffer is None: + global_workspace_buffer = torch.empty( + envs.SGLANG_FLASHINFER_WORKSPACE_SIZE.get(), + dtype=torch.uint8, + device=model_runner.device, + ) + self.workspace_buffer = global_workspace_buffer + else: + self.workspace_buffer = None + + def get_device_int32_arange(self, l: int) -> torch.Tensor: + if l > len(self._arange_buf): + next_pow_of_2 = 1 << (l - 1).bit_length() + self._arange_buf = torch.arange( + next_pow_of_2, device=self.device, dtype=torch.int32 + ) + return self._arange_buf[:l] + + def _transform_table_1_to_real(self, page_table: torch.Tensor) -> torch.Tensor: + page_size = self.real_page_size + if page_size == 1: + return page_table + max_seqlen_k = page_table.shape[1] + strided_indices = torch.arange( + 0, max_seqlen_k, page_size, device=page_table.device, dtype=torch.int32 + ) + return page_table[:, strided_indices] // page_size + + def init_forward_metadata(self, forward_batch: ForwardBatch): + """Init the metadata for a forward pass.""" + batch_size = forward_batch.batch_size + device = forward_batch.seq_lens.device + + if forward_batch.forward_mode.is_target_verify(): + draft_token_num = self.speculative_num_draft_tokens + else: + draft_token_num = 0 + + cache_seqlens_int32 = (forward_batch.seq_lens + draft_token_num).to(torch.int32) + cu_seqlens_k = compute_cu_seqlens(cache_seqlens_int32) + assert forward_batch.seq_lens_cpu is not None + max_seqlen_k = int(forward_batch.seq_lens_cpu.max().item() + draft_token_num) + # [b, max_seqlen_k] + page_table = forward_batch.req_to_token_pool.req_to_token[ + forward_batch.req_pool_indices, :max_seqlen_k + ] + + page_table_1_flattened = None + topk_indices_offset = None + + # Centralized dispatch: decide all strategies for this batch + self.set_dsa_prefill_impl(forward_batch) + dsa_impl_for_batch = ( + self.dsa_decode_impl + if ( + forward_batch.forward_mode.is_decode_or_idle() + or forward_batch.forward_mode.is_target_verify() + or forward_batch.forward_mode.is_draft_extend(include_v2=True) + ) + else self.dsa_prefill_impl + ) + use_flashmla_kv = (not self.use_mha) and dsa_impl_for_batch == "flashmla_kv" + topk_transform_method = self.get_topk_transform_method( + forward_batch.forward_mode + ) + # Batch indices selected when cp enabled: After splitting multiple sequences, + # a certain cp rank may not have some of these sequences. + # We use bs_idx_cpu to mark which sequences are finally selected by the current cp rank, + # a default value of None indicates that all sequences are selected. + bs_idx_cpu = None + # seq_len_cpu of selected sequences + indexer_seq_lens_cpu = forward_batch.seq_lens_cpu + indexer_seq_lens = forward_batch.seq_lens + + if forward_batch.forward_mode.is_decode_or_idle(): + extend_seq_lens_cpu = [1] * batch_size + max_seqlen_q = 1 + cu_seqlens_q = self.get_device_int32_arange(batch_size + 1) + seqlens_expanded = cache_seqlens_int32 + elif forward_batch.forward_mode.is_target_verify(): + max_seqlen_q = 1 + cu_seqlens_q = torch.arange( + 0, + batch_size * self.speculative_num_draft_tokens + 1, + 1, + dtype=torch.int32, + device=device, + ) + extend_seq_lens_cpu = [self.speculative_num_draft_tokens] * batch_size + forward_batch.extend_seq_lens_cpu = extend_seq_lens_cpu + + seqlens_expanded = seqlens_expand_triton( + torch.tensor(extend_seq_lens_cpu, dtype=torch.int32, device=device), + cache_seqlens_int32, + self.speculative_num_draft_tokens * batch_size, + self.speculative_num_draft_tokens, + ) + page_table = torch.repeat_interleave( + page_table, repeats=self.speculative_num_draft_tokens, dim=0 + ) + elif forward_batch.forward_mode.is_draft_extend(include_v2=True): + assert ( + forward_batch.extend_seq_lens_cpu is not None + and forward_batch.extend_seq_lens is not None + and forward_batch.extend_prefix_lens_cpu is not None + ), "All of them must not be None" + + extend_seq_lens_cpu = forward_batch.extend_seq_lens_cpu + assert forward_batch.extend_seq_lens is not None + + max_seqlen_q = 1 + cu_seqlens_q = torch.arange( + 0, + forward_batch.extend_num_tokens + 1, + 1, + dtype=torch.int32, + device=device, + ) + + seqlens_expanded = seqlens_expand_triton( + forward_batch.extend_seq_lens, + cache_seqlens_int32, + sum(extend_seq_lens_cpu), + self.speculative_num_draft_tokens, + ) + if forward_batch.forward_mode.is_draft_extend_v2(): + # DRAFT_EXTEND_V2: V2 worker pre-fills draft KV cache with ALL speculated + # tokens upfront. All requests extend by the same fixed + # (speculative_num_draft_tokens). Use scalar to avoid GPU sync. + page_table = torch.repeat_interleave( + page_table, repeats=self.speculative_num_draft_tokens, dim=0 + ) + else: + # DRAFT_EXTEND (v1): V1 worker extends by (num_correct_drafts + 1) per request + # after verification. Lengths vary per request based on how many tokens + # were accepted. + page_table = torch.repeat_interleave( + page_table, repeats=forward_batch.extend_seq_lens, dim=0 + ) + elif forward_batch.forward_mode.is_extend(): + assert ( + forward_batch.extend_seq_lens_cpu is not None + and forward_batch.extend_seq_lens is not None + and forward_batch.extend_prefix_lens_cpu is not None + ), "All of them must not be None" + extend_seq_lens_cpu = forward_batch.extend_seq_lens_cpu + assert forward_batch.extend_seq_lens is not None + extend_seq_lens = forward_batch.extend_seq_lens + + seqlens_expanded = torch.cat( + [ + torch.arange( + kv_len - qo_len + 1, + kv_len + 1, + dtype=torch.int32, + device=device, + ) + for qo_len, kv_len in zip( + forward_batch.extend_seq_lens_cpu, + forward_batch.seq_lens_cpu.tolist(), + strict=True, + ) + ] + ) + + if can_dsa_prefill_cp_round_robin_split(forward_batch): + seqlens_expanded = dsa_cp_round_robin_split_data(seqlens_expanded) + extend_seq_lens_cpu, extend_seq_lens, bs_idx_cpu, bs_idx = ( + dsa_cp_round_robin_split_q_seqs( + extend_seq_lens_cpu, extend_seq_lens + ) + ) + indexer_seq_lens_cpu = indexer_seq_lens_cpu[bs_idx_cpu] + indexer_seq_lens = indexer_seq_lens[bs_idx] + cache_seqlens_int32 = cache_seqlens_int32[bs_idx] + cu_seqlens_k = compute_cu_seqlens(cache_seqlens_int32) + max_seqlen_k = ( + int(indexer_seq_lens_cpu.max().item() + draft_token_num) + if len(indexer_seq_lens_cpu) != 0 + else 0 + ) + page_table = page_table[bs_idx, :max_seqlen_k] + + if ( + any(forward_batch.extend_prefix_lens_cpu) + or forward_batch.forward_mode == ForwardMode.DRAFT_EXTEND + or bs_idx_cpu is not None + ): + max_seqlen_q = ( + max(extend_seq_lens_cpu) if len(extend_seq_lens_cpu) != 0 else 1 + ) + cu_seqlens_q = compute_cu_seqlens(extend_seq_lens.to(torch.int32)) + else: + max_seqlen_q = max_seqlen_k + cu_seqlens_q = cu_seqlens_k + + # Check if MHA FP8 dequantization is needed + mha_dequantize_needed = ( + self.use_mha + and forward_batch.token_to_kv_pool.dtype == torch.float8_e4m3fn + ) + forward_batch.using_mha_one_shot_fp8_dequant = mha_dequantize_needed + + # page_table_1_flattened is only used when prefix sharing is enabled: + has_prefix_sharing = any(forward_batch.extend_prefix_lens_cpu) + if has_prefix_sharing and ( + topk_transform_method == TopkTransformMethod.RAGGED + or mha_dequantize_needed + ): + page_table_1_flattened = torch.cat( + [ + page_table[i, :kv_len] + for i, kv_len in enumerate( + indexer_seq_lens_cpu.tolist(), + ) + ] + ) + assert page_table_1_flattened.shape[0] == sum( + indexer_seq_lens_cpu + ), f"{page_table_1_flattened.shape[0] = } must be the same as {sum(indexer_seq_lens_cpu) = }" + + # Validate indices when logical tokens exceed physical capacity + # This is likely to be triggered by PP with high kv reuse & parallelism + kv_cache_capacity = ( + forward_batch.token_to_kv_pool.size + + forward_batch.token_to_kv_pool.page_size + ) + if forward_batch.seq_lens_sum > kv_cache_capacity: + max_idx = page_table_1_flattened.max().item() + assert max_idx < kv_cache_capacity, ( + f"Invalid page table index: max={max_idx}, " + f"kv_cache_capacity={kv_cache_capacity}" + ) + + if topk_transform_method == TopkTransformMethod.RAGGED: + topk_indices_offset = torch.repeat_interleave( + cu_seqlens_k[:-1], + extend_seq_lens, + ) + else: + assert False, f"Unsupported {forward_batch.forward_mode = }" + + indexer_k_start_end, token_to_batch_idx = self._cal_indexer_k_start_end( + forward_batch, bs_idx_cpu + ) + # 1D, expanded seqlens (1D means cheap to compute, so always compute it) + dsa_cache_seqlens_int32 = compute_dsa_seqlens( + original_seq_lens=seqlens_expanded, + dsa_index_topk=self.dsa_index_topk, + ) + dsa_cache_seqlens_int32 = pad_dsa_cache_seqlens( + forward_batch, dsa_cache_seqlens_int32 + ) + dsa_cu_seqlens_k = compute_cu_seqlens(dsa_cache_seqlens_int32) + dsa_cu_seqlens_q = self.get_device_int32_arange(len(dsa_cu_seqlens_k)) + + paged_mqa_schedule_metadata = None + # DeepGEMM paged MQA logits path needs a schedule metadata tensor. + # Compute it once per forward batch and reuse it across layers. + if is_cuda() and ( + forward_batch.forward_mode.is_decode_or_idle() + or forward_batch.forward_mode.is_target_verify() + or forward_batch.forward_mode.is_draft_extend(include_v2=True) + ): + try: + import deep_gemm + + # NOTE: DeepGEMM paged path uses block_size=64. + seqlens_32 = ( + seqlens_expanded + if ( + forward_batch.forward_mode.is_target_verify() + or forward_batch.forward_mode.is_draft_extend(include_v2=True) + ) + else cache_seqlens_int32 + ) + seqlens_32_2d = _to_2d_context_lens( + seqlens_32, forward_batch.batch_size + ) + paged_mqa_schedule_metadata = deep_gemm.get_paged_mqa_logits_metadata( + seqlens_32_2d, 64, deep_gemm.get_num_sms() + ) + except (ImportError, ModuleNotFoundError): + paged_mqa_schedule_metadata = None + + metadata = DSAMetadata( + page_size=self.real_page_size, + cache_seqlens_int32=cache_seqlens_int32, + max_seq_len_q=max_seqlen_q, + max_seq_len_k=max_seqlen_k, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + seq_lens_sum=forward_batch.seq_lens_sum, + page_table_1=page_table, + page_table_1_flattened=page_table_1_flattened, + flashmla_metadata=( + self._compute_flashmla_metadata( + cache_seqlens=dsa_cache_seqlens_int32, + seq_len_q=1, + ) + if use_flashmla_kv + else None + ), + paged_mqa_schedule_metadata=paged_mqa_schedule_metadata, + dsa_cache_seqlens_int32=dsa_cache_seqlens_int32, + dsa_cu_seqlens_q=dsa_cu_seqlens_q, + dsa_cu_seqlens_k=dsa_cu_seqlens_k, + dsa_seqlens_expanded=seqlens_expanded, + dsa_extend_seq_lens_list=extend_seq_lens_cpu, + real_page_table=self._transform_table_1_to_real(page_table), + dsa_max_seqlen_q=1, + topk_indices_offset=topk_indices_offset, + indexer_k_start_end=indexer_k_start_end, + indexer_seq_lens_cpu=indexer_seq_lens_cpu, + indexer_seq_lens=indexer_seq_lens, + token_to_batch_idx=token_to_batch_idx, + ) + self.forward_metadata = metadata + + def _cal_indexer_k_start_end( + self, + forward_batch: ForwardBatch, + bs_idx: Optional[List[int]] = None, + ): + if not forward_batch.forward_mode.is_extend_without_speculative(): + return None, None + if forward_batch.batch_size == 0 or (bs_idx is not None and len(bs_idx) == 0): + empty_t = torch.empty(0, dtype=torch.int32, device=self.device) + return (empty_t, empty_t), empty_t + + # Suppose there are two requests, with extend_seq_len = [3, 2] + # and seq_lens = [10, 4] + # The logits matrix looks like this, with * representing the valid logits + # and - representing the invalid logits: + # + # ********--|---- + # *********-|---- + # **********|---- + # ----------|***- + # ----------|**** + # + # ks = [0, 0, 0, 10, 10] + # ke = [8, 9, 10, 13, 14] + ks_list = [] + ke_list = [] + token_to_batch_idx = [] + + q_offset = 0 + k_offset = 0 + + assert ( + forward_batch.seq_lens_cpu is not None + and forward_batch.extend_seq_lens_cpu is not None + ) + for i in range(forward_batch.batch_size): + seq_len = forward_batch.seq_lens_cpu[i].item() + assert isinstance(seq_len, int) + extend_seq_len = forward_batch.extend_seq_lens_cpu[i] + ks = torch.full( + (extend_seq_len,), k_offset, dtype=torch.int32, device=self.device + ) + kv_len = seq_len + if forward_batch.forward_mode.is_target_verify(): + kv_len += self.speculative_num_draft_tokens + seq_lens_expanded = torch.arange( + kv_len - extend_seq_len + 1, + kv_len + 1, + dtype=torch.int32, + device=self.device, + ) + ke = ks + seq_lens_expanded + ks_list.append(ks) + ke_list.append(ke) + + # bi: The index within the selected batch bs_idx. Entries that were not selected are ignored. + bi = bs_idx.index(i) if (bs_idx is not None and i in bs_idx) else i + tb = torch.full( + (extend_seq_len,), bi, dtype=torch.int32, device=self.device + ) + token_to_batch_idx.append(tb) + + if bs_idx is None or i in bs_idx: # skip batch not included in bs_idx + q_offset += extend_seq_len + k_offset += seq_len + + ks = torch.cat(ks_list, dim=0) + ke = torch.cat(ke_list, dim=0) + token_to_batch_idx = torch.cat(token_to_batch_idx, dim=0) + if bs_idx is not None: + assert can_dsa_prefill_cp_round_robin_split(forward_batch) + ks = dsa_cp_round_robin_split_data(ks) + ke = dsa_cp_round_robin_split_data(ke) + token_to_batch_idx = dsa_cp_round_robin_split_data(token_to_batch_idx) + return (ks, ke), token_to_batch_idx + + def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int): + """Initialize CUDA graph state for the attention backend. + + Args: + max_bs (int): Maximum batch size to support in CUDA graphs + + This creates fixed-size tensors that will be reused during CUDA graph replay + to avoid memory allocations. + """ + self.decode_cuda_graph_metadata: Dict = { + "cache_seqlens": torch.ones( + max_num_tokens, dtype=torch.int32, device=self.device + ), + "cu_seqlens_q": torch.arange( + 0, max_bs + 1, dtype=torch.int32, device=self.device + ), + "cu_seqlens_k": torch.zeros( + max_bs + 1, dtype=torch.int32, device=self.device + ), + # fake page_table for sparse_prefill + # Add extra columns for speculative draft tokens to avoid + # overflow during target_verify when max_seqlen_k = seq_len + num_draft_tokens + "page_table": torch.zeros( + max_num_tokens, + self.max_context_len + (self.speculative_num_draft_tokens or 0), + dtype=torch.int32, + device=self.device, + ), + "flashmla_metadata": ( + self._compute_flashmla_metadata( + cache_seqlens=torch.ones( + max_num_tokens, dtype=torch.int32, device=self.device + ), + seq_len_q=1, + ) + if self.dsa_decode_impl == "flashmla_kv" + else None + ), + } + + def init_forward_metadata_capture_cuda_graph( + self, + bs: int, + num_tokens: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + encoder_lens: Optional[torch.Tensor], + forward_mode: ForwardMode, + spec_info: Optional[SpecInput], + ): + self.set_dsa_prefill_impl(forward_batch=None) + + """Initialize forward metadata for capturing CUDA graph.""" + if forward_mode.is_decode_or_idle(): + # Normal Decode + # Get sequence information + cache_seqlens_int32 = seq_lens.to(torch.int32) + cu_seqlens_k = compute_cu_seqlens(cache_seqlens_int32) + + # Use max context length for seq_len_k + page_table_1 = self.decode_cuda_graph_metadata["page_table"][:bs, :] + max_seqlen_q = 1 + max_seqlen_k = page_table_1.shape[1] + + # Precompute page table + # Precompute cumulative sequence lengths + + # NOTE(dark): this is always arange, since we are decoding + cu_seqlens_q = self.decode_cuda_graph_metadata["cu_seqlens_q"][: bs + 1] + dsa_cache_seqlens_int32 = compute_dsa_seqlens( + cache_seqlens_int32, dsa_index_topk=self.dsa_index_topk + ) + + seqlens_expanded = cache_seqlens_int32 + dsa_extend_seq_lens_list = [1] * num_tokens + if self.dsa_decode_impl == "flashmla_kv": + flashmla_metadata = self.decode_cuda_graph_metadata[ + "flashmla_metadata" + ].slice(slice(0, num_tokens + 1)) + flashmla_metadata.copy_( + self._compute_flashmla_metadata( + cache_seqlens=dsa_cache_seqlens_int32, + seq_len_q=1, + ) + ) + else: + flashmla_metadata = None + elif forward_mode.is_target_verify() or forward_mode.is_draft_extend( + include_v2=True + ): + cache_seqlens_int32 = (seq_lens + self.speculative_num_draft_tokens).to( + torch.int32 + ) + cu_seqlens_k = compute_cu_seqlens(cache_seqlens_int32) + max_seqlen_q = 1 + page_table_1 = self.decode_cuda_graph_metadata["page_table"][ + : bs * self.speculative_num_draft_tokens, : + ] + max_seqlen_k = page_table_1.shape[1] + + cu_seqlens_q = torch.arange( + 0, + bs * self.speculative_num_draft_tokens + 1, + 1, + dtype=torch.int32, + device=self.device, + ) + + 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.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, + ) + ] + ) + dsa_cache_seqlens_int32 = compute_dsa_seqlens( + seqlens_expanded, dsa_index_topk=self.dsa_index_topk + ) + dsa_extend_seq_lens_list = [1] * bs * self.speculative_num_draft_tokens + + if self.dsa_decode_impl == "flashmla_kv": + flashmla_metadata = self.decode_cuda_graph_metadata[ + "flashmla_metadata" + ].slice(slice(0, bs * self.speculative_num_draft_tokens + 1)) + + flashmla_metadata.copy_( + self._compute_flashmla_metadata( + cache_seqlens=dsa_cache_seqlens_int32, + seq_len_q=1, + ) + ) + else: + flashmla_metadata = None + + dsa_cu_seqlens_k = compute_cu_seqlens(dsa_cache_seqlens_int32) + dsa_cu_seqlens_q = self.get_device_int32_arange(len(dsa_cu_seqlens_k)) + real_page_table = self._transform_table_1_to_real(page_table_1) + + paged_mqa_schedule_metadata = None + if is_cuda() and ( + forward_mode.is_decode_or_idle() + or forward_mode.is_target_verify() + or forward_mode.is_draft_extend(include_v2=True) + ): + try: + import deep_gemm + + seqlens_32 = ( + seqlens_expanded + if ( + forward_mode.is_target_verify() + or forward_mode.is_draft_extend(include_v2=True) + ) + else cache_seqlens_int32 + ) + seqlens_32_2d = _to_2d_context_lens(seqlens_32, bs) + paged_mqa_schedule_metadata = deep_gemm.get_paged_mqa_logits_metadata( + seqlens_32_2d, 64, deep_gemm.get_num_sms() + ) + except (ImportError, ModuleNotFoundError): + paged_mqa_schedule_metadata = None + + metadata = DSAMetadata( + page_size=self.real_page_size, + cache_seqlens_int32=cache_seqlens_int32, + max_seq_len_q=max_seqlen_q, + max_seq_len_k=max_seqlen_k, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + page_table_1=page_table_1, + flashmla_metadata=flashmla_metadata, + paged_mqa_schedule_metadata=paged_mqa_schedule_metadata, + dsa_cache_seqlens_int32=dsa_cache_seqlens_int32, + dsa_cu_seqlens_q=dsa_cu_seqlens_q, + dsa_cu_seqlens_k=dsa_cu_seqlens_k, + dsa_seqlens_expanded=seqlens_expanded, + real_page_table=real_page_table, + dsa_extend_seq_lens_list=dsa_extend_seq_lens_list, + ) + self.decode_cuda_graph_metadata[bs] = metadata + self.forward_metadata = metadata + + def init_forward_metadata_replay_cuda_graph( + self, + bs: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + seq_lens_sum: int, + encoder_lens: Optional[torch.Tensor], + forward_mode: ForwardMode, + spec_info: Optional[SpecInput], + seq_lens_cpu: Optional[torch.Tensor], + out_cache_loc: Optional[torch.Tensor] = None, + actual_forward_mode: Optional[ForwardMode] = None, + ): + """Initialize forward metadata for replaying CUDA graph.""" + assert seq_lens_cpu is not None + + self.set_dsa_prefill_impl(forward_batch=None) + + seq_lens = seq_lens[:bs] + seq_lens_cpu = seq_lens_cpu[:bs] + req_pool_indices = req_pool_indices[:bs] + + # Normal Decode + metadata: DSAMetadata = self.decode_cuda_graph_metadata[bs] + if forward_mode.is_decode_or_idle(): + # Normal Decode + max_len = int(seq_lens_cpu.max().item()) + + cache_seqlens = seq_lens.to(torch.int32) + metadata.cache_seqlens_int32.copy_(cache_seqlens) + metadata.cu_seqlens_k[1:].copy_( + torch.cumsum(cache_seqlens, dim=0, dtype=torch.int32) + ) + page_indices = self.req_to_token[req_pool_indices, :max_len] + metadata.page_table_1[:, :max_len].copy_(page_indices) + dsa_cache_seqlens = compute_dsa_seqlens( + cache_seqlens, dsa_index_topk=self.dsa_index_topk + ) + metadata.dsa_cache_seqlens_int32.copy_(dsa_cache_seqlens) + seqlens_expanded = cache_seqlens + elif forward_mode.is_target_verify(): + max_seqlen_k = int( + seq_lens_cpu.max().item() + self.speculative_num_draft_tokens + ) + + cache_seqlens = (seq_lens + self.speculative_num_draft_tokens).to( + torch.int32 + ) + metadata.cache_seqlens_int32.copy_(cache_seqlens) + metadata.cu_seqlens_k[1:].copy_( + torch.cumsum(cache_seqlens, dim=0, dtype=torch.int32) + ) + 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 + ) + metadata.page_table_1[:, :max_seqlen_k].copy_(page_indices) + extend_seq_lens_cpu = [self.speculative_num_draft_tokens] * bs + + seqlens_expanded = seqlens_expand_triton( + torch.tensor( + extend_seq_lens_cpu, dtype=torch.int32, device=self.device + ), + cache_seqlens, + self.speculative_num_draft_tokens * bs, + self.speculative_num_draft_tokens, + ) + metadata.dsa_seqlens_expanded.copy_(seqlens_expanded) + dsa_cache_seqlens = compute_dsa_seqlens( + seqlens_expanded, self.dsa_index_topk + ) + metadata.dsa_cache_seqlens_int32.copy_(dsa_cache_seqlens) + elif forward_mode.is_draft_extend(include_v2=True): + max_seqlen_k = int(seq_lens_cpu.max().item()) + cache_seqlens = seq_lens.to(torch.int32) + metadata.cache_seqlens_int32.copy_(cache_seqlens) + metadata.cu_seqlens_k[1:].copy_( + torch.cumsum(cache_seqlens, dim=0, dtype=torch.int32) + ) + + extend_seq_lens = spec_info.num_accept_tokens[:bs] + extend_seq_lens_cpu = extend_seq_lens.tolist() + + 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 + ) + metadata.page_table_1[: page_indices.shape[0], :max_seqlen_k].copy_( + page_indices + ) + + seqlens_expanded = seqlens_expand_triton( + extend_seq_lens, + cache_seqlens, + sum(extend_seq_lens_cpu), + self.speculative_num_draft_tokens, + ) + metadata.dsa_seqlens_expanded[: seqlens_expanded.shape[0]].copy_( + seqlens_expanded + ) + dsa_cache_seqlens = compute_dsa_seqlens( + seqlens_expanded, self.dsa_index_topk + ) + metadata.dsa_cache_seqlens_int32[: seqlens_expanded.shape[0]].copy_( + dsa_cache_seqlens + ) + + # Update DeepGEMM paged MQA schedule metadata outside the captured graph. + if is_cuda() and ( + forward_mode.is_decode_or_idle() + or forward_mode.is_target_verify() + or forward_mode.is_draft_extend(include_v2=True) + ): + try: + import deep_gemm + + seqlens_32 = ( + seqlens_expanded + if ( + forward_mode.is_target_verify() + or forward_mode.is_draft_extend(include_v2=True) + ) + else metadata.cache_seqlens_int32 + ) + seqlens_32_2d = _to_2d_context_lens(seqlens_32, bs) + new_schedule = deep_gemm.get_paged_mqa_logits_metadata( + seqlens_32_2d, 64, deep_gemm.get_num_sms() + ) + if metadata.paged_mqa_schedule_metadata is None: + object.__setattr__( + metadata, "paged_mqa_schedule_metadata", new_schedule + ) + else: + metadata.paged_mqa_schedule_metadata.copy_(new_schedule) + except (ImportError, ModuleNotFoundError): + object.__setattr__(metadata, "paged_mqa_schedule_metadata", None) + seqlens_expanded_size = seqlens_expanded.shape[0] + assert ( + metadata.dsa_cache_seqlens_int32 is not None + and metadata.dsa_cu_seqlens_k is not None + and self.dsa_index_topk is not None + ) + + metadata.dsa_cu_seqlens_k[1 : 1 + seqlens_expanded_size].copy_( + torch.cumsum(dsa_cache_seqlens, dim=0, dtype=torch.int32) + ) + # NOTE(dark): (dsa-) cu_seqlens_q is always arange, no need to copy + + assert self.real_page_size == metadata.page_size + if self.real_page_size > 1: + real_table = self._transform_table_1_to_real(page_indices) + new_rows = real_table.shape[0] + new_cols = real_table.shape[1] + metadata.real_page_table[:new_rows, :new_cols].copy_(real_table) + else: + assert metadata.real_page_table is metadata.page_table_1 + + if self.dsa_decode_impl == "flashmla_kv": + flashmla_metadata = metadata.flashmla_metadata.slice( + slice(0, seqlens_expanded_size + 1) + ) + flashmla_metadata.copy_( + self._compute_flashmla_metadata( + cache_seqlens=dsa_cache_seqlens, + seq_len_q=1, + ) + ) + + self.forward_metadata = metadata + + def init_forward_metadata_replay_cuda_graph_from_precomputed( + self, + bs: int, + precomputed: PrecomputedMetadata, + forward_mode: ForwardMode, + ): + """Fast path: copy precomputed metadata to this backend's metadata. + + This function only performs copy operations, no computation. + + Args: + bs: Batch size + precomputed: Precomputed metadata to copy from + forward_mode: Forward mode + """ + self.set_dsa_prefill_impl(forward_batch=None) + + metadata = self.decode_cuda_graph_metadata[bs] + + # Track whether fused kernel succeeded + fused_kernel_succeeded = False + + # Use fused CUDA kernel for all copy operations + if _USE_FUSED_METADATA_COPY: + try: + from sglang.jit_kernel.fused_metadata_copy import ( + fused_metadata_copy_cuda, + ) + + # Map forward_mode to integer enum + if forward_mode.is_decode_or_idle(): + mode_int = 0 # DECODE + elif forward_mode.is_target_verify(): + mode_int = 1 # TARGET_VERIFY + elif forward_mode.is_draft_extend(): + mode_int = 2 # DRAFT_EXTEND + else: + raise ValueError(f"Unsupported forward_mode: {forward_mode}") + + # Prepare FlashMLA tensors if needed + flashmla_num_splits_src = None + flashmla_num_splits_dst = None + flashmla_metadata_src = None + flashmla_metadata_dst = None + if precomputed.flashmla_metadata is not None: + flashmla_num_splits_src = precomputed.flashmla_metadata.num_splits + flashmla_num_splits_dst = metadata.flashmla_metadata.num_splits + flashmla_metadata_src = ( + precomputed.flashmla_metadata.flashmla_metadata + ) + flashmla_metadata_dst = metadata.flashmla_metadata.flashmla_metadata + + # Call fused kernel + fused_metadata_copy_cuda( + # Source tensors + precomputed.cache_seqlens, + precomputed.cu_seqlens_k, + precomputed.page_indices, + precomputed.dsa_cache_seqlens, + precomputed.seqlens_expanded, + precomputed.dsa_cu_seqlens_k, + precomputed.real_page_table, + flashmla_num_splits_src, + flashmla_metadata_src, + # Destination tensors + metadata.cache_seqlens_int32, + metadata.cu_seqlens_k, + metadata.page_table_1, + metadata.dsa_cache_seqlens_int32, + metadata.dsa_seqlens_expanded, + metadata.dsa_cu_seqlens_k, + ( + metadata.real_page_table + if precomputed.real_page_table is not None + else None + ), + flashmla_num_splits_dst, + flashmla_metadata_dst, + # Parameters + mode_int, + bs, + precomputed.max_len, + precomputed.max_seqlen_k, + precomputed.seqlens_expanded_size, + ) + + # Successfully used fused kernel + fused_kernel_succeeded = True + + except ImportError: + print( + "Warning: Fused metadata copy kernel not available, falling back to individual copies." + ) + except Exception as e: + print( + f"Warning: Fused metadata copy kernel failed with error: {e}, falling back to individual copies." + ) + + # Fallback to individual copy operations if fused kernel disabled or failed + if not fused_kernel_succeeded: + # Copy basic seqlens + metadata.cache_seqlens_int32.copy_(precomputed.cache_seqlens) + metadata.cu_seqlens_k[1:].copy_(precomputed.cu_seqlens_k[1:]) + + # Mode-specific copy logic + if forward_mode.is_decode_or_idle(): + # Decode mode + metadata.page_table_1[:, : precomputed.max_len].copy_( + precomputed.page_indices + ) + metadata.dsa_cache_seqlens_int32.copy_(precomputed.dsa_cache_seqlens) + # seqlens_expanded is same as cache_seqlens (already copied) + + elif forward_mode.is_target_verify(): + # Target verify mode + metadata.page_table_1[:, : precomputed.max_seqlen_k].copy_( + precomputed.page_indices + ) + metadata.dsa_seqlens_expanded.copy_(precomputed.seqlens_expanded) + metadata.dsa_cache_seqlens_int32.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 + metadata.page_table_1[:rows, :cols].copy_(precomputed.page_indices) + + size = precomputed.seqlens_expanded_size + metadata.dsa_seqlens_expanded[:size].copy_(precomputed.seqlens_expanded) + metadata.dsa_cache_seqlens_int32[:size].copy_( + precomputed.dsa_cache_seqlens + ) + + # Copy DSA cu_seqlens + size = precomputed.seqlens_expanded_size + metadata.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 + metadata.real_page_table[:rows, :cols].copy_( + precomputed.real_page_table + ) + + # Copy FlashMLA metadata in fallback path + if precomputed.flashmla_metadata is not None: + size = precomputed.seqlens_expanded_size + flashmla_metadata = metadata.flashmla_metadata.slice(slice(0, size + 1)) + flashmla_metadata.copy_(precomputed.flashmla_metadata) + + # Refresh DeepGEMM paged MQA schedule metadata for the actual seqlens of + # this replay (the captured graph holds stale data otherwise, which can + # deadlock the kernel when the runtime work decomposition diverges from + # the captured one). + if is_cuda(): + try: + import deep_gemm + + if forward_mode.is_decode_or_idle(): + seqlens_32 = metadata.cache_seqlens_int32 + else: + seqlens_32 = metadata.dsa_seqlens_expanded[ + : precomputed.seqlens_expanded_size + ] + seqlens_32_2d = _to_2d_context_lens(seqlens_32, bs) + new_schedule = deep_gemm.get_paged_mqa_logits_metadata( + seqlens_32_2d, 64, deep_gemm.get_num_sms() + ) + if metadata.paged_mqa_schedule_metadata is None: + object.__setattr__( + metadata, "paged_mqa_schedule_metadata", new_schedule + ) + else: + metadata.paged_mqa_schedule_metadata.copy_(new_schedule) + except (ImportError, ModuleNotFoundError): + pass + + self.forward_metadata = metadata + + def forward_extend( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer: RadixAttention, + forward_batch: ForwardBatch, + save_kv_cache=True, + # For multi-head latent attention + q_rope: Optional[torch.Tensor] = None, + k_rope: Optional[torch.Tensor] = None, + topk_indices: Optional[torch.Tensor] = None, + cos_sin_cache: Optional[torch.Tensor] = None, + is_neox: Optional[bool] = False, + llama_4_scaling: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + + causal = not layer.is_cross_attention + metadata = self.forward_metadata + assert causal, "DSA is causal only" + + dsa_impl = ( + self.dsa_decode_impl + if ( + forward_batch.forward_mode.is_target_verify() + or forward_batch.forward_mode.is_draft_extend(include_v2=True) + ) + else self.dsa_prefill_impl + ) + + if dsa_impl == "trtllm" and not self.use_mha: + return self._forward_trtllm( + q, + k, + v, + layer, + forward_batch, + metadata.dsa_cache_seqlens_int32, + save_kv_cache, + q_rope, + k_rope, + topk_indices, + cos_sin_cache, + is_neox, + llama_4_scaling, + is_prefill=True, + ) + + if k is not None: + assert v is not None + if save_kv_cache: + cache_loc = ( + forward_batch.out_cache_loc + if not layer.is_cross_attention + else forward_batch.encoder_out_cache_loc + ) + forward_batch.token_to_kv_pool.set_mla_kv_buffer( # type: ignore + layer, + cache_loc, + k, + k_rope, + ) + + # Use MHA kernel if in MHA_ONE_SHOT mode + if self.use_mha: + assert k is not None and v is not None + assert q_rope is None, "MHA_ONE_SHOT path should not pass q_rope" + assert ( + layer.tp_k_head_num == layer.tp_q_head_num > 1 + ), "MHA_ONE_SHOT requires dense multi-head config" + return self._forward_standard_mha( + q=q, + k=k, + v=v, + layer=layer, + forward_batch=forward_batch, + metadata=metadata, + ) + + # Do absorbed multi-latent attention (MLA path) + assert q_rope is not None + kv_cache = forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id) + + if q_rope is not None: + q_nope = q.view(-1, layer.tp_q_head_num, layer.v_head_dim) + q_rope = q_rope.view( + -1, layer.tp_q_head_num, layer.head_dim - layer.v_head_dim + ) + else: + q_all = q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim) + q_nope = q_all[:, :, : layer.v_head_dim] + q_rope = q_all[:, :, layer.v_head_dim :] + + # Align topk_indices with q dimensions + # This handles cases where q is padded (TP + partial DP attention) + if topk_indices is not None: + topk_indices = self._pad_topk_indices(topk_indices, q_nope.shape[0]) + + # NOTE(dark): here, we use page size = 1 + topk_transform_method = self.get_topk_transform_method( + forward_batch.forward_mode + ) + if envs.SGLANG_DSA_FUSE_TOPK.get(): + page_table_1 = topk_indices + else: + if topk_transform_method == TopkTransformMethod.RAGGED: + topk_indices_offset = metadata.topk_indices_offset + assert topk_indices_offset is not None + mask = topk_indices != -1 + topk_indices_offset = ( + topk_indices_offset.unsqueeze(1) + if topk_indices_offset.ndim == 1 + else topk_indices_offset + ) + topk_indices = torch.where( + mask, topk_indices + topk_indices_offset, topk_indices + ) + elif topk_transform_method == TopkTransformMethod.PAGED: + assert metadata.dsa_extend_seq_lens_list is not None + page_table_1 = transform_index_page_table_prefill( + page_table=metadata.page_table_1, + topk_indices=topk_indices, + extend_lens_cpu=metadata.dsa_extend_seq_lens_list, + page_size=1, + ) + + # todo hisparse: to cover more backends + if forward_batch.hisparse_coordinator is not None: + page_table_1 = ( + forward_batch.token_to_kv_pool.translate_loc_to_hisparse_device( + page_table_1 + ) + ) + + if dsa_impl == "tilelang": + if q_rope is not None: + q_all = concat_mla_absorb_q_general(q_nope, q_rope) + return self._forward_tilelang( + q_all=q_all, + kv_cache=kv_cache, + page_table_1=page_table_1, + sm_scale=layer.scaling, + v_head_dim=layer.v_head_dim, + ) + elif dsa_impl == "flashmla_sparse": + if q_rope is not None: + q_all = concat_mla_absorb_q_general(q_nope, q_rope) + + if topk_transform_method == TopkTransformMethod.RAGGED: + if any(forward_batch.extend_prefix_lens_cpu): + page_table_1_flattened = ( + self.forward_metadata.page_table_1_flattened + ) + assert page_table_1_flattened is not None + kv_cache = dequantize_k_cache_paged( + kv_cache, page_table_1_flattened + ) + else: + kv_cache = _cat([k, k_rope], dim=-1) + page_table_1 = topk_indices + + return self._forward_flashmla_sparse( + q_all=q_all, + kv_cache=kv_cache, + page_table_1=page_table_1, + sm_scale=layer.scaling, + v_head_dim=layer.v_head_dim, + ) + elif dsa_impl == "flashmla_kv": + if q_rope is not None: + q_all = concat_mla_absorb_q_general(q_nope, q_rope) + return self._forward_flashmla_kv( + q_all=q_all, + kv_cache=kv_cache, + sm_scale=layer.scaling, + v_head_dim=layer.v_head_dim, + # TODO optimize args + layer=layer, + metadata=metadata, + page_table_1=page_table_1, + ) + elif dsa_impl == "fa3": + return self._forward_fa3( + q_rope=q_rope, + kv_cache=kv_cache, + v_head_dim=layer.v_head_dim, + q_nope=q_nope, + page_table=page_table_1, + cache_seqlens=metadata.dsa_cache_seqlens_int32, + cu_seqlens_q=metadata.dsa_cu_seqlens_q, + cu_seqlens_k=metadata.dsa_cu_seqlens_k, + max_seqlen_q=metadata.dsa_max_seqlen_q, + sm_scale=layer.scaling, + logit_cap=layer.logit_cap, + page_size=1, + ) + elif dsa_impl == "aiter": + if q_rope is not None: + q_all = torch.cat([q_nope, q_rope], dim=-1) + return self._forward_aiter_extend( + q_all=q_all, + kv_cache=kv_cache, + page_table_1=page_table_1, + layer=layer, + ) + else: + raise ValueError( + f"Unsupported {dsa_impl = } for forward_extend. Consider using an other attention backend." + ) + + def forward_decode( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer: RadixAttention, + forward_batch: ForwardBatch, + save_kv_cache=True, + # For multi-head latent attention + q_rope: Optional[torch.Tensor] = None, + k_rope: Optional[torch.Tensor] = None, + topk_indices: Optional[torch.Tensor] = None, + cos_sin_cache: Optional[torch.Tensor] = None, + is_neox: Optional[bool] = False, + llama_4_scaling: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + + causal = not layer.is_cross_attention + metadata = self.forward_metadata + assert causal, "DSA is causal only" + + if self.dsa_decode_impl == "trtllm": + return self._forward_trtllm( + q, + k, + v, + layer, + forward_batch, + metadata.cache_seqlens_int32, + save_kv_cache, + q_rope, + k_rope, + topk_indices, + cos_sin_cache, + is_neox, + llama_4_scaling, + ) + + if k is not None: + assert v is not None + if save_kv_cache: + cache_loc = ( + forward_batch.out_cache_loc + if not layer.is_cross_attention + else forward_batch.encoder_out_cache_loc + ) + forward_batch.token_to_kv_pool.set_mla_kv_buffer( # type: ignore + layer, + cache_loc, + k, + k_rope, + ) + + # Do absorbed multi-latent attention + kv_cache = forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id) + if q_rope is not None: + q_nope = q.view(-1, layer.tp_q_head_num, layer.v_head_dim) + q_rope = q_rope.view( + -1, layer.tp_q_head_num, layer.head_dim - layer.v_head_dim + ) + # Caller passed split q_nope / q_rope; we'll need to concat below if + # the chosen impl wants q_all. + q_all = None + else: + # Caller passed already-concatenated q (q_all = q). Reuse it directly + # via a zero-copy view; the impl-specific blocks below will skip the + # otherwise redundant concat_mla_absorb_q_general call. + q_all = q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim) + q_nope = q_all[:, :, : layer.v_head_dim] + q_rope = q_all[:, :, layer.v_head_dim :] + + # Align topk_indices with q dimensions + if topk_indices is not None: + topk_indices = self._pad_topk_indices(topk_indices, q_nope.shape[0]) + + if forward_batch.hisparse_coordinator is not None: + page_table_1 = forward_batch.hisparse_coordinator.swap_in_selected_pages( + forward_batch.req_pool_indices, + forward_batch.seq_lens, + topk_indices, + layer.layer_id, + ) + elif envs.SGLANG_DSA_FUSE_TOPK.get(): + page_table_1 = topk_indices + else: + page_table_1 = transform_index_page_table_decode( + page_table=metadata.page_table_1, + topk_indices=topk_indices, + page_size=1, + ) + + if self.dsa_decode_impl == "flashmla_sparse": + if q_rope is not None: + q_all = concat_mla_absorb_q_general(q_nope, q_rope) + return self._forward_flashmla_sparse( + q_all=q_all, + kv_cache=kv_cache, + page_table_1=page_table_1, + sm_scale=layer.scaling, + v_head_dim=layer.v_head_dim, + ) + elif self.dsa_decode_impl == "flashmla_kv": + if q_rope is not None: + q_all = concat_mla_absorb_q_general(q_nope, q_rope) + return self._forward_flashmla_kv( + q_all=q_all, + kv_cache=kv_cache, + sm_scale=layer.scaling, + v_head_dim=layer.v_head_dim, + # TODO optimize args + layer=layer, + metadata=metadata, + page_table_1=page_table_1, + ) + elif self.dsa_decode_impl == "tilelang": + # Cat-skip (HIP-only): when caller passes q_rope=None on HIP, q_all + # has already been set to a zero-copy view of q in the else branch + # above and we can reuse it directly. The `not _is_hip` clause keeps + # CUDA / MUSA paths byte-identical to pre-patch by always re-cat. + if q_all is None or not _is_hip: + q_all = concat_mla_absorb_q_general(q_nope, q_rope) + return self._forward_tilelang( + q_all=q_all, + kv_cache=kv_cache, + page_table_1=page_table_1, + sm_scale=layer.scaling, + v_head_dim=layer.v_head_dim, + ) + elif self.dsa_decode_impl == "fa3": + return self._forward_fa3( + q_rope=q_rope, + kv_cache=kv_cache, + v_head_dim=layer.v_head_dim, + q_nope=q_nope, + page_table=page_table_1, + cache_seqlens=metadata.dsa_cache_seqlens_int32, + cu_seqlens_q=metadata.dsa_cu_seqlens_q, + cu_seqlens_k=metadata.dsa_cu_seqlens_k, + max_seqlen_q=metadata.dsa_max_seqlen_q, + sm_scale=layer.scaling, + logit_cap=layer.logit_cap, + page_size=1, + ) + elif self.dsa_decode_impl == "aiter": + if q_all is None or not _is_hip: + q_all = torch.cat([q_nope, q_rope], dim=-1) + return self._forward_aiter( + q_all=q_all, + kv_cache=kv_cache, + page_table_1=page_table_1, + layer=layer, + metadata=metadata, + bs=forward_batch.batch_size, + ) + + else: + assert False, f"Unsupported {self.dsa_decode_impl = }" + + def _forward_fa3( + self, + q_rope: torch.Tensor, + kv_cache: torch.Tensor, + v_head_dim: int, + q_nope: torch.Tensor, + page_table: torch.Tensor, + cache_seqlens: torch.Tensor, + cu_seqlens_q: torch.Tensor, + cu_seqlens_k: torch.Tensor, + max_seqlen_q: int, + sm_scale: float, + logit_cap: float, + page_size: int, + ) -> torch.Tensor: + k_rope_cache = kv_cache[:, :, v_head_dim:] + c_kv_cache = kv_cache[:, :, :v_head_dim] + qk_rope_dim = k_rope_cache.shape[-1] + k_rope_cache = k_rope_cache.view(-1, page_size, 1, qk_rope_dim) + c_kv_cache = c_kv_cache.view(-1, page_size, 1, v_head_dim) + o = flash_attn_with_kvcache( + q=q_rope, + k_cache=k_rope_cache, + v_cache=c_kv_cache, + qv=q_nope, + page_table=page_table, + cache_seqlens=cache_seqlens, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k_new=cu_seqlens_k, + max_seqlen_q=max_seqlen_q, + softmax_scale=sm_scale, + causal=True, + softcap=logit_cap, + return_softmax_lse=False, + num_splits=self.num_splits, + ) + return o # type: ignore + + def _forward_flashmla_sparse( + self, + q_all: torch.Tensor, + kv_cache: torch.Tensor, + v_head_dim: int, + page_table_1: torch.Tensor, + sm_scale: float, + ) -> torch.Tensor: + from sgl_kernel.flash_mla import flash_mla_sparse_fwd + + # FlashMLA sparse kernel requires num_heads to be a multiple of 64 (Hopper) or 128 (Blackwell) + # When using TP, num_heads might be smaller (e.g., 256//8=32) + num_tokens, num_heads, head_dim = q_all.shape + + # Determine required padding based on GPU architecture (use cached value) + required_padding = 128 if self.device_sm_major >= 10 else 64 + + need_padding = num_heads % required_padding != 0 + + if need_padding: + assert required_padding % num_heads == 0, ( + f"num_heads {num_heads} cannot be padded to {required_padding}. " + f"TP size may be too large for this model." + ) + + # Pad q to required size + q_padded = q_all.new_zeros((num_tokens, required_padding, head_dim)) + q_padded[:, :num_heads, :] = q_all + q_input = q_padded + else: + q_input = q_all + + # indices shape must be (s_q, h_kv=1, topk), keep h_kv=1 unchanged + indices_input = page_table_1.unsqueeze(1) + + o, _, _ = flash_mla_sparse_fwd( + q=q_input, + kv=kv_cache, + indices=indices_input, + sm_scale=sm_scale, + d_v=v_head_dim, + ) + + # Trim output back to original num_heads if we padded + if need_padding: + o = o[:, :num_heads, :] + + return o + + def _forward_flashmla_kv( + self, + q_all: torch.Tensor, + kv_cache: torch.Tensor, + v_head_dim: int, + sm_scale: float, + layer, + metadata: DSAMetadata, + page_table_1, + ) -> torch.Tensor: + from sgl_kernel.flash_mla import flash_mla_with_kvcache + + cache_seqlens = metadata.dsa_cache_seqlens_int32 + assert metadata.flashmla_metadata is not None + + # TODO the 2nd dim is seq_len_q, need to be >1 when MTP + q_all = q_all.view(-1, 1, layer.tp_q_head_num, layer.head_dim) + num_q_heads = q_all.shape[2] + target_q_heads = self.flashmla_kv_num_q_heads + if target_q_heads != num_q_heads: + # Pad q heads to match FlashMLA decode supported head-count variants. + q_input = q_all.new_zeros( + q_all.shape[0], q_all.shape[1], target_q_heads, q_all.shape[3] + ) + q_input[:, :, :num_q_heads, :] = q_all + else: + q_input = q_all + + kv_cache = kv_cache.view(-1, self.real_page_size, 1, self.kv_cache_dim) + assert self.real_page_size == 64, "only page size 64 is supported" + + if not self.dsa_kv_cache_store_fp8: + # inefficiently quantize the whole cache + kv_cache = quantize_k_cache(kv_cache) + + indices = page_table_1.unsqueeze(1) + assert ( + indices.shape[-1] == self.dsa_index_topk + ) # requirement of FlashMLA decode kernel + + o, _ = flash_mla_with_kvcache( + q=q_input, + k_cache=kv_cache, + cache_seqlens=cache_seqlens, + head_dim_v=v_head_dim, + tile_scheduler_metadata=metadata.flashmla_metadata.flashmla_metadata, + num_splits=metadata.flashmla_metadata.num_splits, + softmax_scale=sm_scale, + indices=indices, + # doc says it is not used, but if pass in None then error + block_table=torch.empty( + (q_all.shape[0], 0), dtype=torch.int32, device=q_all.device + ), + is_fp8_kvcache=True, + ) + + if target_q_heads != num_q_heads: + o = o[:, :, :num_q_heads, :] + + return o + + def _forward_standard_mha( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer: RadixAttention, + forward_batch: ForwardBatch, + metadata: DSAMetadata, + ) -> torch.Tensor: + """Standard MHA using FlashAttention varlen for MHA_ONE_SHOT mode.""" + q = q.view(-1, layer.tp_q_head_num, layer.head_dim) + k = k.view(-1, layer.tp_k_head_num, layer.head_dim) + v = v.view(-1, layer.tp_v_head_num, layer.v_head_dim) + + # MHA_ONE_SHOT: k/v include all tokens (prefix + current) + cu_seqlens_q = metadata.cu_seqlens_q + cu_seqlens_k = metadata.cu_seqlens_k + max_seqlen_k = metadata.max_seq_len_k + causal = True + + # Verify batch sizes match (length of cu_seqlens should be batch_size + 1) + assert len(cu_seqlens_q) == len(cu_seqlens_k), ( + f"batch_size mismatch: cu_seqlens_q has {len(cu_seqlens_q)-1} requests, " + f"cu_seqlens_k has {len(cu_seqlens_k)-1} requests" + ) + + # Use TRTLLm ragged attention for SM100 (Blackwell/B200) to avoid FA4 accuracy issues + if self.device_sm_major >= 10: + import flashinfer + + seq_lens = metadata.cache_seqlens_int32 + return flashinfer.prefill.trtllm_ragged_attention_deepseek( + query=q, + key=k, + value=v, + workspace_buffer=self.workspace_buffer, + seq_lens=seq_lens, + max_q_len=metadata.max_seq_len_q, + max_kv_len=max_seqlen_k, + bmm1_scale=layer.scaling, + bmm2_scale=1.0, + o_sf_scale=1.0, + batch_size=forward_batch.batch_size, + window_left=-1, + cum_seq_lens_q=cu_seqlens_q, + cum_seq_lens_kv=cu_seqlens_k, + enable_pdl=False, + is_causal=causal, + return_lse=False, + skip_softmax_threshold_scale_factor=envs.SGLANG_SKIP_SOFTMAX_PREFILL_THRESHOLD_SCALE_FACTOR.get(), + ) + + # Use FA3 for SM90 (Hopper/H200) + return flash_attn_varlen_func( + q=q, + k=k, + v=v, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=metadata.max_seq_len_q, + max_seqlen_k=max_seqlen_k, + softmax_scale=layer.scaling, + causal=causal, + ) + + def _forward_tilelang( + self, + q_all: torch.Tensor, + kv_cache: torch.Tensor, + v_head_dim: int, + page_table_1: torch.Tensor, + sm_scale: float, + ) -> torch.Tensor: + from sglang.srt.layers.attention.dsa.tilelang_kernel import tilelang_sparse_fwd + + return tilelang_sparse_fwd( + q=q_all, + kv=kv_cache, + indices=page_table_1.unsqueeze(1), + sm_scale=sm_scale, + d_v=v_head_dim, + ) + + def _forward_aiter( + self, + q_all: torch.Tensor, + kv_cache: torch.Tensor, + page_table_1: torch.Tensor, + layer: RadixAttention, + metadata: DSAMetadata, + bs: int, + ) -> torch.Tensor: + q = q_all.reshape(-1, layer.tp_q_head_num * layer.head_dim) + + if layer.head_dim != layer.v_head_dim: + o = q.new_empty((q.shape[0], layer.tp_q_head_num * layer.v_head_dim)) + else: + o = torch.empty_like(q) + + if self.need_pad_heads: + q_kernel = q.view( + -1, layer.tp_q_head_num, layer.head_dim + ).repeat_interleave(self.head_repeat_factor, dim=1) + o_kernel = q.new_empty( + ( + q.shape[0], + layer.tp_q_head_num * self.head_repeat_factor, + layer.v_head_dim, + ) + ) + else: + q_kernel = q.view(-1, layer.tp_q_head_num, layer.head_dim) + o_kernel = o.view(-1, layer.tp_q_head_num, layer.v_head_dim) + + kv_indptr = self.kv_indptr + + non_minus1_mask = page_table_1 != -1 + non_minus1_counts = non_minus1_mask.sum(dim=1) + kv_indptr[1 : bs + 1] = torch.cumsum(non_minus1_counts, dim=0) + + kv_indices = self.kv_indices + get_valid_kv_indices(page_table_1, kv_indptr, kv_indices, bs) + + mla_decode_fwd( + q_kernel, + kv_cache.view(-1, 1, 1, layer.head_dim), + o_kernel, + metadata.cu_seqlens_q, + kv_indptr, + kv_indices, + metadata.cu_seqlens_q, + metadata.max_seq_len_q, + sm_scale=layer.scaling, + logit_cap=layer.logit_cap, + ) + + if self.need_pad_heads: + o = o_kernel[:, :: self.head_repeat_factor, :] + + return o + + def _forward_aiter_extend( + self, + q_all: torch.Tensor, + kv_cache: torch.Tensor, + page_table_1: torch.Tensor, + layer: RadixAttention, + ) -> torch.Tensor: + num_tokens = q_all.shape[0] + q = q_all.reshape(-1, layer.tp_q_head_num * layer.head_dim) + + if layer.head_dim != layer.v_head_dim: + o = q.new_empty((num_tokens, layer.tp_q_head_num * layer.v_head_dim)) + else: + o = torch.empty_like(q) + + if self.need_pad_heads: + q_kernel = q.view( + -1, layer.tp_q_head_num, layer.head_dim + ).repeat_interleave(self.head_repeat_factor, dim=1) + o_kernel = q.new_empty( + ( + num_tokens, + layer.tp_q_head_num * self.head_repeat_factor, + layer.v_head_dim, + ) + ) + else: + q_kernel = q.view(-1, layer.tp_q_head_num, layer.head_dim) + o_kernel = o.view(-1, layer.tp_q_head_num, layer.v_head_dim) + + non_minus1_mask = page_table_1 != -1 + non_minus1_counts = non_minus1_mask.sum(dim=1) + + kv_indptr = torch.zeros(num_tokens + 1, dtype=torch.int32, device=self.device) + kv_indptr[1:] = torch.cumsum(non_minus1_counts, dim=0) + + # Allocate kv_indices with upper-bound size (num_tokens * topk) + topk = page_table_1.shape[1] + kv_indices = torch.zeros( + num_tokens * topk, dtype=torch.int32, device=self.device + ) + + # Use get_valid_kv_indices kernel to extract valid indices + get_valid_kv_indices(page_table_1, kv_indptr, kv_indices, num_tokens) + + # Build cu_seqlens_q for extend: each token is treated as seq_len_q=1 + cu_seqlens_q = torch.arange( + 0, num_tokens + 1, dtype=torch.int32, device=self.device + ) + # TODO support more forward_mode + mla_decode_fwd( + q_kernel, + kv_cache.view(-1, 1, 1, layer.head_dim), + o_kernel, + cu_seqlens_q, + kv_indptr, + kv_indices, + cu_seqlens_q, + 1, # max_seq_len_q = 1 for per-token attention + sm_scale=layer.scaling, + logit_cap=layer.logit_cap, + ) + + if self.need_pad_heads: + o = o_kernel[:, :: self.head_repeat_factor, :] + + return o + + def _forward_trtllm( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer: RadixAttention, + forward_batch: ForwardBatch, + seq_lens: torch.Tensor, + save_kv_cache=True, + # For multi-head latent attention + q_rope: Optional[torch.Tensor] = None, + k_rope: Optional[torch.Tensor] = None, + topk_indices: Optional[torch.Tensor] = None, + cos_sin_cache: Optional[torch.Tensor] = None, + is_neox: Optional[bool] = False, + llama_4_scaling: Optional[torch.Tensor] = None, + is_prefill: bool = False, + ) -> torch.Tensor: + """Forward using TRT-LLM sparse MLA kernel.""" + import flashinfer.decode + + metadata = self.forward_metadata + + merge_query = q_rope is not None + if self.kv_cache_dtype == torch.float8_e4m3fn: + # For FP8 path, we quantize the query and rope parts and merge them into a single tensor + # Note: rope application in deepseek_v2.py:forward_absorb_prepare is skipped for FP8 decode path of this trtllm_mla backend + assert q_rope is not None, "For FP8 path q_rope should not be None." + assert k_rope is not None, "For FP8 path k_rope should not be None." + assert ( + cos_sin_cache is not None + ), "For FP8 path cos_sin_cache should not be None." + + q, k, k_rope = mla_quantize_and_rope_for_fp8( + q, + q_rope, + k.squeeze(1), + k_rope.squeeze(1), + forward_batch.positions, + cos_sin_cache, + is_neox, + self.kv_lora_rank, + self.qk_rope_head_dim, + ) + merge_query = False + + # Save KV cache if requested + if save_kv_cache: + assert ( + k is not None and k_rope is not None + ), "For populating trtllm_mla kv cache, both k_nope and k_rope should be not None." + cache_loc = ( + forward_batch.out_cache_loc + if not layer.is_cross_attention + else forward_batch.encoder_out_cache_loc + ) + forward_batch.token_to_kv_pool.set_mla_kv_buffer( + layer, cache_loc, k, k_rope + ) + + k_cache = forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id) + kv_cache = k_cache.view(-1, self.real_page_size, self.kv_cache_dim).unsqueeze(1) + + if merge_query: + q_nope = q.view(-1, layer.tp_q_head_num, layer.v_head_dim) + q_rope_reshaped = q_rope.view( + -1, layer.tp_q_head_num, layer.head_dim - layer.v_head_dim + ) + q_all = concat_mla_absorb_q_general(q_nope, q_rope_reshaped) + else: + q_all = q.view(-1, layer.tp_q_head_num, layer.head_dim) + + # Align topk_indices with q dimensions + if topk_indices is not None: + topk_indices = self._pad_topk_indices(topk_indices, q.shape[0]) + + if envs.SGLANG_DSA_FUSE_TOPK.get(): + page_table_1 = topk_indices + elif is_prefill: + page_table_1 = transform_index_page_table_prefill( + page_table=metadata.page_table_1, + topk_indices=topk_indices, + extend_lens_cpu=metadata.dsa_extend_seq_lens_list, + page_size=1, + ) + else: + page_table_1 = transform_index_page_table_decode( + page_table=metadata.page_table_1, + topk_indices=topk_indices, + page_size=1, + ) + + q_scale = 1.0 + k_scale = ( + layer.k_scale_float + if getattr(layer, "k_scale_float", None) is not None + else 1.0 + ) + bmm1_scale = q_scale * k_scale * layer.scaling + + batch_size = page_table_1.shape[0] + _, num_heads, head_dim = q_all.shape + + q = q_all.view(batch_size, 1, num_heads, head_dim) + kv = kv_cache.view(-1, 1, self.real_page_size, self.kv_cache_dim) + block_tables = page_table_1.unsqueeze(1) + seq_lens = metadata.cache_seqlens_int32 if seq_lens is None else seq_lens + + out = flashinfer.decode.trtllm_batch_decode_with_kv_cache_mla( + query=q, + kv_cache=kv, + workspace_buffer=self.workspace_buffer, + qk_nope_head_dim=self.qk_nope_head_dim, + kv_lora_rank=self.kv_lora_rank, + qk_rope_head_dim=self.qk_rope_head_dim, + block_tables=block_tables, + seq_lens=seq_lens, + max_seq_len=metadata.max_seq_len_k, + sparse_mla_top_k=self.dsa_index_topk, + bmm1_scale=bmm1_scale, + backend="trtllm-gen", + skip_softmax_threshold_scale_factor=envs.SGLANG_SKIP_SOFTMAX_DECODE_THRESHOLD_SCALE_FACTOR.get(), + ) + # Output: [batch, q_len=1, heads, v_dim] -> [batch, heads, v_dim] + return out.squeeze(1) + + def _pad_topk_indices( + self, topk_indices: torch.Tensor, num_tokens: int + ) -> torch.Tensor: + current_tokens = topk_indices.shape[0] + if current_tokens == num_tokens: + return topk_indices + + assert current_tokens <= num_tokens, ( + f"topk_indices rows ({current_tokens}) > num_tokens ({num_tokens}); " + "this indicates a mismatch between indexer output and q layout." + ) + + pad_size = num_tokens - current_tokens + padding = torch.full( + (pad_size, topk_indices.shape[1]), + -1, + dtype=topk_indices.dtype, + device=topk_indices.device, + ) + return torch.cat([topk_indices, padding], dim=0) + + def get_cuda_graph_seq_len_fill_value(self): + """Get the fill value for sequence length in CUDA graph.""" + return 1 + + def set_dsa_prefill_impl(self, forward_batch: Optional[ForwardBatch] = None): + """ + Decide all attention prefill dispatch strategies for this batch. + """ + from sglang.srt.utils import get_device_sm, is_blackwell + + # Decide MHA vs MLA + if forward_batch and forward_batch.forward_mode.is_extend_without_speculative(): + # Check if sequence meets criteria for MHA_ONE_SHOT + assert forward_batch.seq_lens_cpu is not None + max_kv_len = forward_batch.seq_lens_cpu.max().item() + sum_seq_lens = sum(forward_batch.seq_lens_cpu) + device_sm = get_device_sm() + + # Requirements: H200/B200, short sequences, supported dtype, fits in chunk + self.use_mha = ( + ( + device_sm == 90 or (device_sm >= 100 and device_sm < 110) + ) # SM90/SM100 only + and max_kv_len + <= envs.SGLANG_DSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD.get() # Short enough for MHA + and forward_batch.token_to_kv_pool.dtype + in [torch.bfloat16, torch.float8_e4m3fn] + and sum_seq_lens + <= forward_batch.get_max_chunk_capacity() # Fits in chunk + and (not is_dsa_enable_prefill_cp()) # CP not enabled + and (forward_batch.hisparse_coordinator is None) + ) + else: + self.use_mha = False # Decode/verify always use MLA + + # Set MLA implementation only if not using MHA + if not self.use_mha and self.enable_auto_select_prefill_impl: + if self.dsa_kv_cache_store_fp8: + if ( + is_blackwell() + and forward_batch is not None + and forward_batch.forward_mode == ForwardMode.EXTEND + ): + total_kv_tokens = forward_batch.seq_lens_sum + total_q_tokens = forward_batch.extend_num_tokens + # Heuristic based on benchmarking flashmla_kv vs flashmla_sparse + dequantize_k_cache_paged + if total_kv_tokens < total_q_tokens * 512: + self.dsa_prefill_impl = "flashmla_sparse" + return + self.dsa_prefill_impl = "flashmla_kv" + else: + # bf16 kv cache + self.dsa_prefill_impl = "flashmla_sparse" + + def get_topk_transform_method( + self, forward_mode: Optional[ForwardMode] = None + ) -> TopkTransformMethod: + """ + SGLANG_DSA_FUSE_TOPK controls whether to fuse the topk transform into the topk kernel. + This method is used to select the topk transform method which can be fused or unfused. + """ + if ( + # disable for MTP + self.dsa_kv_cache_store_fp8 + and self.dsa_prefill_impl == "flashmla_sparse" + and forward_mode == ForwardMode.EXTEND + ): + topk_transform_method = TopkTransformMethod.RAGGED + else: + topk_transform_method = TopkTransformMethod.PAGED + return topk_transform_method + + def get_indexer_metadata( + self, layer_id: int, forward_batch: ForwardBatch + ) -> DSAIndexerMetadata: + force_unfused = ( + forward_batch.hisparse_coordinator is not None + and forward_batch.forward_mode.is_decode_or_idle() + ) + return DSAIndexerMetadata( + attn_metadata=self.forward_metadata, + topk_transform_method=self.get_topk_transform_method( + forward_batch.forward_mode + ), + paged_mqa_schedule_metadata=self.forward_metadata.paged_mqa_schedule_metadata, + force_unfused_topk=force_unfused, + ) + + def _compute_flashmla_metadata(self, cache_seqlens: torch.Tensor, seq_len_q: int): + from sgl_kernel.flash_mla import get_mla_metadata + + num_heads_q = self.flashmla_kv_num_q_heads + + flashmla_metadata, num_splits = get_mla_metadata( + cache_seqlens=cache_seqlens, + # TODO doc says `num_q_tokens_per_q_seq * num_heads_q // num_heads_k` + # but the name looks like need seq_len_q? + num_q_tokens_per_head_k=seq_len_q * num_heads_q // 1, + num_heads_k=1, + num_heads_q=num_heads_q, + is_fp8_kvcache=True, + topk=self.dsa_index_topk, + ) + + return DSAFlashMLAMetadata( + flashmla_metadata=flashmla_metadata, + num_splits=num_splits, + ) + + +class DeepseekSparseAttnMultiStepBackend: + + def __init__( + self, model_runner: ModelRunner, topk: int, speculative_num_steps: int + ): + self.model_runner = model_runner + self.topk = topk + self.speculative_num_steps = speculative_num_steps + self.attn_backends = [] + for i in range(self.speculative_num_steps - 1): + self.attn_backends.append( + DeepseekSparseAttnBackend( + model_runner, + speculative_step_id=i, + topk=self.topk, + speculative_num_steps=self.speculative_num_steps, + ) + ) + + def init_forward_metadata(self, forward_batch: ForwardBatch): + for i in range(self.speculative_num_steps - 1): + self.attn_backends[i].init_forward_metadata(forward_batch) + + def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int): + for i in range(self.speculative_num_steps - 1): + self.attn_backends[i].init_cuda_graph_state(max_bs, max_num_tokens) + + def init_forward_metadata_capture_cuda_graph(self, forward_batch: ForwardBatch): + for i in range(self.speculative_num_steps - 1): + self.attn_backends[i].init_forward_metadata_capture_cuda_graph( + forward_batch.batch_size, + forward_batch.batch_size * self.topk, + forward_batch.req_pool_indices, + forward_batch.seq_lens, + encoder_lens=None, + forward_mode=ForwardMode.DECODE, + spec_info=forward_batch.spec_info, + ) + + def init_forward_metadata_replay_cuda_graph( + self, forward_batch: ForwardBatch, bs: int + ): + if envs.SGLANG_DSA_ENABLE_MTP_PRECOMPUTE_METADATA.get(): + # Precompute metadata once (shared across all backends) + precomputed = self.attn_backends[0]._precompute_replay_metadata( + bs=bs, + req_pool_indices=forward_batch.req_pool_indices, + seq_lens=forward_batch.seq_lens, + seq_lens_cpu=forward_batch.seq_lens_cpu, + forward_mode=ForwardMode.DECODE, + spec_info=forward_batch.spec_info, + ) + + # Use multi-backend fused copy when we have 3 or more backends + # This is 3x faster than calling the single-backend copy 3 times + if self.speculative_num_steps > 3: + try: + from sglang.jit_kernel.fused_metadata_copy import ( + fused_metadata_copy_multi_cuda, + ) + + metadata0 = self.attn_backends[0].decode_cuda_graph_metadata[bs] + metadata1 = self.attn_backends[1].decode_cuda_graph_metadata[bs] + metadata2 = self.attn_backends[2].decode_cuda_graph_metadata[bs] + + # Set dsa_prefill_impl for first 3 backends (required by the method) + for i in range(3): + self.attn_backends[i].set_dsa_prefill_impl(forward_batch=None) + + # Prepare FlashMLA tensors if needed + flashmla_num_splits_src = None + flashmla_metadata_src = None + flashmla_num_splits_dst0 = None + flashmla_num_splits_dst1 = None + flashmla_num_splits_dst2 = None + flashmla_metadata_dst0 = None + flashmla_metadata_dst1 = None + flashmla_metadata_dst2 = None + + if precomputed.flashmla_metadata is not None: + flashmla_num_splits_src = ( + precomputed.flashmla_metadata.num_splits + ) + flashmla_metadata_src = ( + precomputed.flashmla_metadata.flashmla_metadata + ) + flashmla_num_splits_dst0 = ( + metadata0.flashmla_metadata.num_splits + ) + flashmla_num_splits_dst1 = ( + metadata1.flashmla_metadata.num_splits + ) + flashmla_num_splits_dst2 = ( + metadata2.flashmla_metadata.num_splits + ) + flashmla_metadata_dst0 = ( + metadata0.flashmla_metadata.flashmla_metadata + ) + flashmla_metadata_dst1 = ( + metadata1.flashmla_metadata.flashmla_metadata + ) + flashmla_metadata_dst2 = ( + metadata2.flashmla_metadata.flashmla_metadata + ) + + # Call the multi-backend fused kernel for first 3 backends + fused_metadata_copy_multi_cuda( + # Source tensors + precomputed.cache_seqlens, + precomputed.cu_seqlens_k, + precomputed.page_indices, + precomputed.dsa_cache_seqlens, + precomputed.dsa_cu_seqlens_k, + precomputed.real_page_table, + flashmla_num_splits_src, + flashmla_metadata_src, + # Destination tensors for backend 0 + metadata0.cache_seqlens_int32, + metadata0.cu_seqlens_k, + metadata0.page_table_1, + metadata0.dsa_cache_seqlens_int32, + metadata0.dsa_cu_seqlens_k, + ( + metadata0.real_page_table + if precomputed.real_page_table is not None + else None + ), + flashmla_num_splits_dst0, + flashmla_metadata_dst0, + # Destination tensors for backend 1 + metadata1.cache_seqlens_int32, + metadata1.cu_seqlens_k, + metadata1.page_table_1, + metadata1.dsa_cache_seqlens_int32, + metadata1.dsa_cu_seqlens_k, + ( + metadata1.real_page_table + if precomputed.real_page_table is not None + else None + ), + flashmla_num_splits_dst1, + flashmla_metadata_dst1, + # Destination tensors for backend 2 + metadata2.cache_seqlens_int32, + metadata2.cu_seqlens_k, + metadata2.page_table_1, + metadata2.dsa_cache_seqlens_int32, + metadata2.dsa_cu_seqlens_k, + ( + metadata2.real_page_table + if precomputed.real_page_table is not None + else None + ), + flashmla_num_splits_dst2, + flashmla_metadata_dst2, + # Parameters + bs, + precomputed.max_len, + precomputed.seqlens_expanded_size, + ) + + # Copy remaining backends one by one (if > 3 backends) + for i in range(3, self.speculative_num_steps - 1): + self.attn_backends[ + i + ].init_forward_metadata_replay_cuda_graph_from_precomputed( + bs=bs, + precomputed=precomputed, + forward_mode=ForwardMode.DECODE, + ) + except (ImportError, Exception) as e: + # Fallback to loop if multi-backend kernel not available or fails + if isinstance(e, ImportError): + print( + "Warning: Multi-backend fused metadata copy kernel not available, falling back to loop." + ) + else: + print( + f"Warning: Multi-backend fused metadata copy kernel failed with error: {e}, falling back to loop." + ) + for i in range(self.speculative_num_steps - 1): + self.attn_backends[ + i + ].init_forward_metadata_replay_cuda_graph_from_precomputed( + bs=bs, + precomputed=precomputed, + forward_mode=ForwardMode.DECODE, + ) + else: + # Less than 3 backends: copy to each backend individually + for i in range(self.speculative_num_steps - 1): + self.attn_backends[ + i + ].init_forward_metadata_replay_cuda_graph_from_precomputed( + bs=bs, + precomputed=precomputed, + forward_mode=ForwardMode.DECODE, + ) + else: + # Fallback: compute metadata separately for each backend + for i in range(self.speculative_num_steps - 1): + self.attn_backends[i].init_forward_metadata_replay_cuda_graph( + bs=bs, + req_pool_indices=forward_batch.req_pool_indices, + seq_lens=forward_batch.seq_lens, + seq_lens_sum=forward_batch.seq_lens_sum, + encoder_lens=None, + forward_mode=ForwardMode.DECODE, + spec_info=forward_batch.spec_info, + seq_lens_cpu=forward_batch.seq_lens_cpu, + out_cache_loc=None, + ) + + +# Backward-compat aliases (deprecated: use DSA class names) +DeepseekSparseAttnBackend = DeepseekSparseAttnBackend +DeepseekSparseAttnMultiStepBackend = DeepseekSparseAttnMultiStepBackend +DSAMetadata = DSAMetadata +DSAFlashMLAMetadata = DSAFlashMLAMetadata +DSAIndexerMetadata = DSAIndexerMetadata diff --git a/python/sglang/srt/layers/attention/dsv4/compress_hip.py b/python/sglang/srt/layers/attention/dsv4/compress_hip.py index b796de3a5..1c69f7e46 100644 --- a/python/sglang/srt/layers/attention/dsv4/compress_hip.py +++ b/python/sglang/srt/layers/attention/dsv4/compress_hip.py @@ -10,8 +10,8 @@ import triton import triton.language as tl 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.nsa.nsa_indexer import rotate_activation from sglang.srt.layers.deepseek_v4_rope import ( apply_rotary_emb_triton, fused_norm_rope_inplace_triton, diff --git a/python/sglang/srt/layers/attention/dsv4/compressor.py b/python/sglang/srt/layers/attention/dsv4/compressor.py index 56174f0e5..428b5af14 100644 --- a/python/sglang/srt/layers/attention/dsv4/compressor.py +++ b/python/sglang/srt/layers/attention/dsv4/compressor.py @@ -15,11 +15,11 @@ from sglang.jit_kernel.deepseek_v4 import ( ) from sglang.srt.configs.deepseek_v4 import DeepSeekV4Config 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 ( 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.layernorm import RMSNorm from sglang.srt.layers.linear import ReplicatedLinear @@ -71,7 +71,7 @@ class CompressorBackendMixin: compress_ratio: int, is_paged: bool = False, ) -> 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 ( 4, @@ -358,7 +358,7 @@ class Compressor(nn.Module): kv_score = linear_bf16_fp32(x, self.wkv_gate.weight) # 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, get_attention_cp_size(), diff --git a/python/sglang/srt/layers/attention/dsv4/indexer.py b/python/sglang/srt/layers/attention/dsv4/indexer.py index a6a5b9af8..facf5ef3d 100644 --- a/python/sglang/srt/layers/attention/dsv4/indexer.py +++ b/python/sglang/srt/layers/attention/dsv4/indexer.py @@ -368,7 +368,7 @@ class C4IndexerBackendMixin: assert len(weights.shape) == 3 weights = weights.squeeze(2) 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, ) elif envs.SGLANG_FP8_PAGED_MQA_LOGITS_TORCH.get(): diff --git a/python/sglang/srt/layers/attention/hip_flash_mla.py b/python/sglang/srt/layers/attention/hip_flash_mla.py index 8f1fbd117..c22d4f38f 100644 --- a/python/sglang/srt/layers/attention/hip_flash_mla.py +++ b/python/sglang/srt/layers/attention/hip_flash_mla.py @@ -12,7 +12,7 @@ def flash_mla_with_kvcache_entrypoint(backend: str, **kwargs): if is_hip(): 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, ) diff --git a/python/sglang/srt/layers/attention/hybrid_attn_backend.py b/python/sglang/srt/layers/attention/hybrid_attn_backend.py index 57e10daa6..69e80149e 100644 --- a/python/sglang/srt/layers/attention/hybrid_attn_backend.py +++ b/python/sglang/srt/layers/attention/hybrid_attn_backend.py @@ -3,7 +3,7 @@ from typing import Optional import torch 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.model_executor.forward_batch_info import ForwardBatch, ForwardMode from sglang.srt.model_executor.model_runner import ModelRunner diff --git a/python/sglang/srt/layers/attention/nsa/__init__.py b/python/sglang/srt/layers/attention/nsa/__init__.py new file mode 100644 index 000000000..29c4139eb --- /dev/null +++ b/python/sglang/srt/layers/attention/nsa/__init__.py @@ -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 diff --git a/python/sglang/srt/layers/attention/nsa/dequant_k_cache.py b/python/sglang/srt/layers/attention/nsa/dequant_k_cache.py index 5c69d135b..43c080a6c 100644 --- a/python/sglang/srt/layers/attention/nsa/dequant_k_cache.py +++ b/python/sglang/srt/layers/attention/nsa/dequant_k_cache.py @@ -1,289 +1,10 @@ -import torch -import triton -import triton.language as tl +# [Deprecated] Re-export shim for backward compatibility. Use dsa.dequant_k_cache instead. +import warnings - -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 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") +warnings.warn( + "sglang.srt.layers.attention.nsa.dequant_k_cache is deprecated; " + "use sglang.srt.layers.attention.dsa.dequant_k_cache instead.", + DeprecationWarning, + stacklevel=2, +) +from sglang.srt.layers.attention.dsa.dequant_k_cache import * # noqa: F401, F403 diff --git a/python/sglang/srt/layers/attention/nsa/index_buf_accessor.py b/python/sglang/srt/layers/attention/nsa/index_buf_accessor.py index db1e80ee3..7ee993b4f 100644 --- a/python/sglang/srt/layers/attention/nsa/index_buf_accessor.py +++ b/python/sglang/srt/layers/attention/nsa/index_buf_accessor.py @@ -1,814 +1,10 @@ -from typing import TYPE_CHECKING - -import torch -import triton -import triton.language as tl - -from sglang.srt.layers.attention.nsa.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 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): - 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) +# [Deprecated] Re-export shim for backward compatibility. Use dsa.index_buf_accessor instead. +import warnings + +warnings.warn( + "sglang.srt.layers.attention.nsa.index_buf_accessor is deprecated; " + "use sglang.srt.layers.attention.dsa.index_buf_accessor instead.", + DeprecationWarning, + stacklevel=2, +) +from sglang.srt.layers.attention.dsa.index_buf_accessor import * # noqa: F401, F403 diff --git a/python/sglang/srt/layers/attention/nsa/nsa_backend_mtp_precompute.py b/python/sglang/srt/layers/attention/nsa/nsa_backend_mtp_precompute.py index 1e61c8416..ec8be8914 100644 --- a/python/sglang/srt/layers/attention/nsa/nsa_backend_mtp_precompute.py +++ b/python/sglang/srt/layers/attention/nsa/nsa_backend_mtp_precompute.py @@ -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 -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.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, - ) +warnings.warn( + "sglang.srt.layers.attention.nsa.nsa_backend_mtp_precompute is deprecated; " + "use sglang.srt.layers.attention.dsa.dsa_backend_mtp_precompute instead.", + DeprecationWarning, + stacklevel=2, +) +from sglang.srt.layers.attention.dsa.dsa_backend_mtp_precompute import * # noqa: F401, F403 diff --git a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py index 2902b668a..6d7fbba31 100644 --- a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py +++ b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py @@ -1,1746 +1,10 @@ -from __future__ import annotations +# [Deprecated] Re-export shim for backward compatibility. Use dsa.dsa_indexer instead. +import warnings -import contextlib -import logging -from abc import ABC, abstractmethod -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union - -import torch -from einops import rearrange - -from sglang.jit_kernel.fused_store_index_cache import ( - can_use_nsa_fused_store, - fused_store_index_k_cache, +warnings.warn( + "sglang.srt.layers.attention.nsa.nsa_indexer is deprecated; " + "use sglang.srt.layers.attention.dsa.dsa_indexer instead.", + DeprecationWarning, + stacklevel=2, ) -from sglang.srt.environ import envs -from sglang.srt.layers.attention.nsa.utils import ( - aiter_can_use_preshuffle_paged_mqa, - is_nsa_enable_prefill_cp, - is_nsa_prefill_cp_in_seq_split, -) -from sglang.srt.layers.dp_attention import attn_tp_all_gather_into_tensor -from sglang.srt.layers.layernorm import LayerNorm -from sglang.srt.layers.quantization.fp8_kernel import fp8_dtype, is_fp8_fnuz -from sglang.srt.layers.utils import MultiPlatformOp -from sglang.srt.state_capturer.indexer_topk import ( - maybe_capture_indexer_topk, -) -from sglang.srt.utils import ( - add_prefix, - ceil_align, - get_bool_env_var, - is_cuda, - is_gfx95_supported, - is_hip, - is_npu, -) - -logger = logging.getLogger(__name__) - -global _use_multi_stream -_is_cuda = is_cuda() -_is_hip = is_hip() -_is_npu = is_npu() -_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip -_is_fp8_fnuz = is_fp8_fnuz() -_is_gfx95_supported = is_gfx95_supported() -# Whether the aiter preshuffle paged-MQA path (page_size=64 + Preshuffle=True + -# KVBlockSize=64) can be used. Falls back to the legacy page_size=1 / KVBlockSize=1 -# path when the gluon kernel is unavailable (Triton<3.5 and no AOT bundle). -_use_aiter_preshuffle = aiter_can_use_preshuffle_paged_mqa() -if _use_aiter and not _use_aiter_preshuffle: - logger.warning( - "ROCm NSA indexer: aiter preshuffle paged-MQA path is unavailable " - "(needs Triton>=3.5.0 or AITER_ENABLE_AOT_GLUON_PA_MQA_LOGITS=1); " - "falling back to legacy page_size=1 / KVBlockSize=1 path." - ) -if _is_cuda: - try: - import deep_gemm - except ImportError as e: - deep_gemm = e - -if _use_aiter: - from aiter.ops.cache import indexer_k_quant_and_cache - -if is_npu(): - import torch_npu - from sglang.srt.hardware_backend.npu.utils import get_indexer_weight_stream - -from sglang.srt.distributed import ( - get_attn_context_model_parallel_rank, - get_attn_context_model_parallel_world_size, -) -from sglang.srt.distributed.parallel_state import get_pp_group -from sglang.srt.layers import deep_gemm_wrapper -from sglang.srt.layers.communicator import ScatterMode -from sglang.srt.layers.linear import ReplicatedLinear -from sglang.srt.layers.quantization.base_config import QuantizationConfig -from sglang.srt.layers.rotary_embedding import get_rope_wrapper -from sglang.srt.layers.utils.cp_utils import cp_all_gather_rerange_output -from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode -from sglang.srt.model_executor.forward_batch_info import ForwardBatch -from sglang.srt.server_args import get_global_server_args - -_use_ag_after_qlora = envs.SGLANG_USE_AG_AFTER_QLORA.get() -if TYPE_CHECKING: - from sglang.srt.mem_cache.memory_pool import NSATokenToKVPool - - -DUAL_STREAM_TOKEN_THRESHOLD = 1024 if _is_cuda else 0 - - -class BaseIndexerMetadata(ABC): - @abstractmethod - def get_seqlens_int32(self) -> torch.Tensor: - """ - Return: (batch_size,) int32 tensor - """ - - @abstractmethod - def get_page_table_64(self) -> torch.Tensor: - """ - Return: (batch_size, num_blocks) int32, page table. - The page size of the table is 64. - """ - - @abstractmethod - def get_page_table_1(self) -> torch.Tensor: - """ - Return: (batch_size, num_blocks) int32, page table. - The page size of the table is 1. - """ - - @abstractmethod - def get_seqlens_expanded(self) -> torch.Tensor: - """ - Return: (sum_extend_seq_len,) int32 tensor - """ - - def get_indexer_kvcache_range(self) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Return: (tokens, ), (tokens, ) int32, k_start and k_end in kv cache(token,xxx) for each token. - """ - - def get_indexer_seq_len_cpu(self) -> torch.Tensor: - """ - Return: seq lens for each batch. - """ - - def get_indexer_seq_len(self) -> torch.Tensor: - """ - Return: seq lens for each batch. - """ - - def get_nsa_extend_len_cpu(self) -> List[int]: - """ - Return: extend seq lens for each batch. - """ - - def get_token_to_batch_idx(self) -> torch.Tensor: - """ - Return: batch idx for each token. - """ - - @abstractmethod - def topk_transform( - self, - logits: torch.Tensor, - topk: int, - ) -> torch.Tensor: - """ - Perform topk selection on the logits and possibly transform the result. - - NOTE that attention backend may override this function to do some - transformation, which means the result of this topk_transform may not - be the topk indices of the input logits. - - Return: Anything, since it will be passed to the attention backend - for further processing on sparse attention computation. - Don't assume it is the topk indices of the input logits. - """ - - -def rotate_activation(x: torch.Tensor) -> torch.Tensor: - # from sgl_kernel import hadamard_transform - if _is_hip: - from fast_hadamard_transform import hadamard_transform - else: - from sglang.jit_kernel.hadamard import hadamard_transform - - hidden_size = x.size(-1) - assert ( - hidden_size & (hidden_size - 1) - ) == 0, "Hidden size must be a power of 2 for Hadamard transform." - return hadamard_transform(x, scale=hidden_size**-0.5) - - -class Indexer(MultiPlatformOp): - _MQA_LOGITS_BYTES_PER_ELEM = 4 - _MQA_LOGITS_STATIC_SKIP_ELEMS = 8_000_000 - _MQA_LOGITS_FREE_MEM_FRACTION = 0.5 - _MQA_LOGITS_TOTAL_MEM_FRACTION = 0.3 - _mqa_logits_budget_bytes: Dict[int, int] = {} - - def __init__( - self, - hidden_size: int, - index_n_heads: int, - index_head_dim: int, - rope_head_dim: int, - index_topk: int, - q_lora_rank: int, - max_position_embeddings: int, - rope_theta: float, - layer_id: int, - scale_fmt: Optional[str], - block_size: int = 128, - rope_scaling: Optional[Dict[str, Any]] = None, - is_neox_style: bool = True, - prefix: str = "", - quant_config: Optional[QuantizationConfig] = None, - alt_stream: Optional[torch.cuda.Stream] = None, - ): - super().__init__() - self.hidden_size = hidden_size - self.n_heads = index_n_heads - self.head_dim = index_head_dim - self.rope_head_dim = rope_head_dim - self.index_topk = index_topk - self.q_lora_rank = q_lora_rank - self.layer_id = layer_id - self.alt_stream = alt_stream - self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp() - if self.nsa_enable_prefill_cp: - self.cp_size = get_attn_context_model_parallel_world_size() - self.cp_rank = get_attn_context_model_parallel_rank() - else: - self.cp_size = None - self.cp_rank = None - if _is_cuda: - self.sm_count = deep_gemm.get_num_sms() - self.half_device_sm_count = ceil_align(self.sm_count // 2, 8) - pp_size = get_global_server_args().pp_size - self.logits_with_pp_recv = pp_size > 1 and not get_pp_group().is_last_rank - else: - self.logits_with_pp_recv = False - - self.wq_b = ReplicatedLinear( - self.q_lora_rank, - self.n_heads * self.head_dim, - bias=False, - quant_config=quant_config, - prefix=add_prefix("wq_b", prefix), - ) - - self.wk = ReplicatedLinear( - self.hidden_size, - self.head_dim, - bias=False, - quant_config=quant_config, - prefix=add_prefix("wk", prefix), - ) - self.weights_proj = ReplicatedLinear( - self.hidden_size, - self.n_heads, - bias=False, - params_dtype=torch.bfloat16, - prefix=add_prefix("weights_proj", prefix), - ) - self.k_norm = LayerNorm( - self.head_dim, dtype=torch.bfloat16 if _use_aiter else torch.float32 - ) - self.rotary_emb = get_rope_wrapper( - rope_head_dim, - rotary_dim=rope_head_dim, - max_position=max_position_embeddings, - base=rope_theta, # type: ignore - rope_scaling=rope_scaling, - is_neox_style=is_neox_style, - device=get_global_server_args().device, - ) - self.block_size = block_size - self.scale_fmt = scale_fmt - self.softmax_scale = self.head_dim**-0.5 - - @contextlib.contextmanager - def _with_real_sm_count(self): - # When pipeline parallelism is enabled, each PP rank initiates a recv operation after the _pp_launch_batch - # request to receive the PP proxy tensor or output from the previous stage, occupying one SM resource. - # Model execution runs in parallel with the recv operation, so the SMs available to the indexer must be reduced - # by 1. Currently, the last rank starts the send result + recv request only after waiting for execution results. - if self.logits_with_pp_recv: - pp_recv_sm_count = 1 - with deep_gemm_wrapper.configure_deep_gemm_num_sms( - self.sm_count - pp_recv_sm_count - ): - yield - else: - yield - - def _weights_proj_bf16_in_fp32_out( - self, x: Union[torch.Tensor, Tuple[torch.Tensor, ...]] - ) -> torch.Tensor: - # aiter (ROCm gfx95): extract the passthrough bf16 tensor from the - # 3-tuple (fp8, scale, bf16) produced by fused_rms_fp8_group_quant, - # avoiding an expensive FP8-to-bf16 dequantization. - if _use_aiter and _is_gfx95_supported and isinstance(x, tuple) and len(x) == 3: - x = x[2] - if _is_cuda: - return torch.mm(x, self.weights_proj.weight.t(), out_dtype=torch.float32) - - weights, _ = self.weights_proj(x) - if _is_hip: - # Return bf16; multiplying with q_scale promotes back to fp32. - return weights - return weights.float() - - @torch.compile(dynamic=True) - def _project_and_scale_head_gates( - self, x: Union[torch.Tensor, Tuple[torch.Tensor, ...]] - ): - weights = self._weights_proj_bf16_in_fp32_out(x) - weights = weights * self.n_heads**-0.5 - return weights - - @torch.compile(dynamic=True) - def _get_logits_head_gate( - self, x: Union[torch.Tensor, Tuple[torch.Tensor, ...]], q_scale: torch.Tensor - ): - weights = self._weights_proj_bf16_in_fp32_out(x) - weights = weights * self.n_heads**-0.5 - weights = weights.unsqueeze(-1) * q_scale * self.softmax_scale - return weights - - @torch.compile(dynamic=True) - def _apply_q_scale_and_softmax_scale( - self, weights: torch.Tensor, q_scale: torch.Tensor - ): - return weights.unsqueeze(-1) * q_scale * self.softmax_scale - - def _get_q_k_bf16( - self, - q_lora: torch.Tensor, - x: torch.Tensor, - positions: torch.Tensor, - enable_dual_stream: bool, - forward_batch: ForwardBatch, - ): - if enable_dual_stream: - current_stream = torch.cuda.current_stream() - self.alt_stream.wait_stream(current_stream) - - with deep_gemm_wrapper.configure_deep_gemm_num_sms( - self.half_device_sm_count - ): - query, _ = self.wq_b(q_lora) - query = rearrange(query, "l (h d) -> l h d", d=self.head_dim) - q_rope, _ = torch.split( - query, - [self.rope_head_dim, self.head_dim - self.rope_head_dim], - dim=-1, - ) - with torch.cuda.stream(self.alt_stream): - # TODO we should also put DeepGEMM half SM here? - key, _ = self.wk(x) - key = self.k_norm(key) - - k_rope, _ = torch.split( - key, - [self.rope_head_dim, self.head_dim - self.rope_head_dim], - dim=-1, - ) - - current_stream.wait_stream(self.alt_stream) - else: - query, _ = self.wq_b(q_lora) - query = rearrange(query, "l (h d) -> l h d", d=self.head_dim) - q_rope, _ = torch.split( - query, [self.rope_head_dim, self.head_dim - self.rope_head_dim], dim=-1 - ) - key, _ = self.wk(x) - key = self.k_norm(key) - k_rope, _ = torch.split( - key, [self.rope_head_dim, self.head_dim - self.rope_head_dim], dim=-1 - ) - - q_rope, k_rope = self.rotary_emb(positions, q_rope, k_rope) - - self._update_rope_guarded(query[..., : self.rope_head_dim], q_rope) - self._update_rope_guarded(key[..., : self.rope_head_dim], k_rope) - - if enable_dual_stream: - current_stream = torch.cuda.current_stream() - self.alt_stream.wait_stream(current_stream) - query = rotate_activation(query) - - with torch.cuda.stream(self.alt_stream): - key = rotate_activation(key) - current_stream.wait_stream(self.alt_stream) - elif ( - self.alt_stream is not None - and forward_batch.attn_cp_metadata is not None - and self.nsa_enable_prefill_cp - ): - key = rotate_activation(key) - current_stream = torch.cuda.current_stream() - self.alt_stream.wait_stream(current_stream) - query = rotate_activation(query) - - with torch.cuda.stream(self.alt_stream): - key = cp_all_gather_rerange_output( - key.contiguous(), - self.cp_size, - forward_batch, - torch.cuda.current_stream(), - ) - current_stream.wait_stream(self.alt_stream) - return query, key - else: - query = rotate_activation(query) - key = rotate_activation(key) - - # allgather+rerrange - if forward_batch.attn_cp_metadata is not None and self.nsa_enable_prefill_cp: - key = cp_all_gather_rerange_output( - key.contiguous(), - self.cp_size, - forward_batch, - torch.cuda.current_stream(), - ) - return query, key - - def _get_k_bf16( - self, - x: torch.Tensor, - positions: torch.Tensor, - enable_dual_stream: bool, - ): - # Compute only key, skip query - key, _ = self.wk(x) - key = self.k_norm(key) - k_rope, _ = torch.split( - key, [self.rope_head_dim, self.head_dim - self.rope_head_dim], dim=-1 - ) - - _, k_rope = self.rotary_emb(positions, k_rope, k_rope) - self._update_rope_guarded(key[..., : self.rope_head_dim], k_rope) - key = rotate_activation(key) - - return key - - @staticmethod - def _update_rope_guarded(dst: torch.Tensor, src: torch.Tensor) -> None: - # On AMD with in-place RoPE kernels, self-aliasing can occur; - # skip write-back when src/dst tensors point to a single memory. - if src.data_ptr() == dst.data_ptr(): - return - dst.copy_(src) - - def _get_topk_paged( - self, - forward_batch: ForwardBatch, - layer_id: int, - q_fp8: torch.Tensor, - weights: torch.Tensor, - metadata: BaseIndexerMetadata, - ) -> torch.Tensor: - if TYPE_CHECKING: - assert isinstance(forward_batch.token_to_kv_pool, NSATokenToKVPool) - - page_size = forward_batch.token_to_kv_pool.page_size - # NOTE(dark): blocksize = 64 is hardcoded in deep_gemm - 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 == 1 - ), f"HIP legacy NSA path requires page_size == 1, got {page_size}" - else: - assert page_size == 64, "only support page size 64" - # NOTE(dark): this support extend/decode/decode+graph - if _is_hip and not _use_aiter_preshuffle: - block_tables = metadata.get_page_table_1() - else: - block_tables = metadata.get_page_table_64() - - max_seq_len = block_tables.shape[1] * page_size - kv_cache_fp8 = forward_batch.token_to_kv_pool.get_index_k_with_scale_buffer( - layer_id=layer_id - ) - - blocksize = page_size - if ( - forward_batch.forward_mode.is_target_verify() - or forward_batch.forward_mode.is_draft_extend(include_v2=True) - ): - seqlens_32 = metadata.get_seqlens_expanded() - else: - seqlens_32 = metadata.get_seqlens_int32() - # Reuse pre-computed schedule metadata if available (from init_forward_metadata), - # otherwise fall back to computing it here. - schedule_metadata = getattr(metadata, "paged_mqa_schedule_metadata", None) - # DeepGEMM release-0426 requires context_lens of shape [batch_size, next_n] - # to match q.shape = [batch_size, next_n, heads, head_dim]. The indexer uses - # next_n=1 with batch_size=N_total via q_fp8.unsqueeze(1) below, so mirror - # that layout here. - if seqlens_32.dim() == 2: - seqlens_32_2d = seqlens_32 - else: - seqlens_32_2d = seqlens_32.unsqueeze(-1) - if _is_cuda: - if schedule_metadata is None: - schedule_metadata = deep_gemm.get_paged_mqa_logits_metadata( - seqlens_32_2d, blocksize, self.sm_count - ) - - assert len(q_fp8.shape) == 3 - q_fp8 = q_fp8.unsqueeze(1) # the next_n dim is 1 now - assert len(kv_cache_fp8.shape) == 2 - block_kv = page_size - num_heads_kv = 1 - head_dim_with_sf = 132 - kv_cache_fp8 = kv_cache_fp8.view( - kv_cache_fp8.shape[0], block_kv, num_heads_kv, head_dim_with_sf - ) - assert len(weights.shape) == 3 - weights = weights.squeeze(2) - - # When attn_tp_size > 1 or in the MAX_LEN padding mode, padding may exist in the hidden states, - # and it is necessary to extract the actual q length. - q_offset = sum(metadata.get_nsa_extend_len_cpu()) - if _is_hip: - from aiter.ops.triton.pa_mqa_logits import deepgemm_fp8_paged_mqa_logits - - batch_size, next_n, heads, _ = q_fp8.shape - logits = torch.empty( - (batch_size * next_n, max_seq_len), - device=q_fp8.device, - dtype=torch.float32, - ) - deepgemm_fp8_paged_mqa_logits( - q_fp8, - kv_cache_fp8, - weights, - logits, - seqlens_32, - block_tables, - max_seq_len, - Preshuffle=_use_aiter_preshuffle, - KVBlockSize=block_kv, - ) - else: - logits = deep_gemm.fp8_paged_mqa_logits( - q_fp8[:q_offset], - kv_cache_fp8, - weights[:q_offset], - seqlens_32_2d, - block_tables, - schedule_metadata, - max_seq_len, - clean_logits=False, - ) - - # NOTE(dark): logits should be cleaned in topk_transform - topk_result = metadata.topk_transform(logits, self.index_topk) - # Restore possible padding exist in the hidden states. - if not _is_hip and q_offset < q_fp8.shape[0]: - pad_len = q_fp8.shape[0] - q_offset - padding = torch.full( - (pad_len, topk_result.shape[1]), - -1, - dtype=topk_result.dtype, - device=topk_result.device, - ) - topk_result = torch.cat([topk_result, padding], dim=0) - return topk_result - - def _get_mqa_logits_budget_bytes(self, device_index: int) -> int: - cached_budget = self._mqa_logits_budget_bytes.get(device_index) - if cached_budget is not None: - return cached_budget - - total_mem = torch.cuda.get_device_properties(device_index).total_memory - - total_mem_budget = int(total_mem * self._MQA_LOGITS_TOTAL_MEM_FRACTION) - mem_fraction_static = get_global_server_args().mem_fraction_static - if mem_fraction_static is None: - static_budget = total_mem_budget - else: - static_free_mem = int(total_mem * max(0.0, 1.0 - mem_fraction_static)) - static_budget = min( - int(static_free_mem * self._MQA_LOGITS_FREE_MEM_FRACTION), - total_mem_budget, - ) - static_budget = max(1, static_budget) - - # Keep the static serving-memory guard during CUDA graph capture without - # caching it. The first non-capture prefill path will cache the real - # free-memory budget below. - if get_is_capture_mode(): - return static_budget - - # Match the original free-memory guard: logits_bytes * 2 > free_mem. - # torch.cuda.mem_get_info synchronizes the host, so cache the result, - # capped by the workload-independent serving-memory headroom. - free_mem, _ = torch.cuda.mem_get_info(device_index) - budget_bytes = min( - int(free_mem * self._MQA_LOGITS_FREE_MEM_FRACTION), static_budget - ) - - budget_bytes = max(1, budget_bytes) - self._mqa_logits_budget_bytes[device_index] = budget_bytes - return budget_bytes - - def _should_chunk_mqa_logits( - self, num_q: int, num_k: int, device_index: int - ) -> Tuple[bool, int]: - """ - Detect whether we need to chunk the MQA logits computation to avoid OOM - Return: (need_chunk, logits_budget_bytes) - """ - # Quick static check for normal batches - if num_q * num_k < self._MQA_LOGITS_STATIC_SKIP_ELEMS: - return False, 0 - - logits_bytes = num_q * num_k * self._MQA_LOGITS_BYTES_PER_ELEM - logits_budget_bytes = self._get_mqa_logits_budget_bytes(device_index) - - need_chunk = logits_bytes > logits_budget_bytes - return need_chunk, logits_budget_bytes - - def _get_topk_ragged( - self, - enable_dual_stream: bool, - forward_batch: ForwardBatch, - layer_id: int, - q_fp8: torch.Tensor, - weights: torch.Tensor, - metadata: BaseIndexerMetadata, - ) -> torch.Tensor: - if TYPE_CHECKING: - assert isinstance(forward_batch.token_to_kv_pool, NSATokenToKVPool) - - assert forward_batch.forward_mode.is_extend_without_speculative() - - page_size = forward_batch.token_to_kv_pool.page_size - 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 == 1 - ), f"HIP legacy NSA path requires page_size == 1, got {page_size}" - else: - assert page_size == 64, "only support page size 64" - - assert len(weights.shape) == 3 - assert ( - forward_batch.seq_lens_cpu is not None - and forward_batch.extend_seq_lens_cpu is not None - ) - weights = weights.squeeze(-1) - - if _is_hip and not _use_aiter_preshuffle: - block_tables = metadata.get_page_table_1() - else: - block_tables = metadata.get_page_table_64() - - assert ( - forward_batch.seq_lens_cpu is not None - and forward_batch.extend_seq_lens_cpu is not None - ) - - batch_size = len(block_tables) - token_nums, _, _ = q_fp8.shape - device = q_fp8.device - device_index = device.index - assert device_index is not None, "q_fp8 must be on an indexed CUDA device" - - topk_result = torch.full( - (token_nums, self.index_topk), -1, device=device, dtype=torch.int32 - ) - if batch_size == 0: - return topk_result - - ks, ke = metadata.get_indexer_kvcache_range() - - indexer_seq_lens_cpu = metadata.get_indexer_seq_len_cpu() - seq_len_sum = torch.sum(indexer_seq_lens_cpu).item() - max_seq_len = torch.max(indexer_seq_lens_cpu).item() - k_fp8, k_scale = forward_batch.token_to_kv_pool.get_index_k_scale_buffer( - layer_id, - metadata.get_indexer_seq_len(), - block_tables, - seq_len_sum, - max_seq_len, - ) - if _is_fp8_fnuz: - k_fp8 = k_fp8.view(torch.float8_e4m3fnuz) - else: - k_fp8 = k_fp8.view(torch.float8_e4m3fn) - - k_scale = k_scale.view(torch.float32).squeeze(-1) - kv_fp8 = (k_fp8, k_scale) - - # Check if we need to chunk to avoid OOM - seq_lens_expanded = metadata.get_seqlens_expanded() - token_to_batch_idx = metadata.get_token_to_batch_idx() - q_offset = ks.shape[0] - k_offset = k_fp8.shape[0] - need_chunk, logits_budget_bytes = self._should_chunk_mqa_logits( - q_offset, k_offset, device_index - ) - - if not need_chunk: - assert q_fp8[:q_offset].shape[0] != 0 - with self._with_real_sm_count(): - if _is_hip: - from aiter.ops.triton.fp8_mqa_logits import fp8_mqa_logits - - kv, scale = kv_fp8 - logits = fp8_mqa_logits( - q_fp8[:q_offset], kv, scale, weights[:q_offset], ks, ke - ) - else: - logits = deep_gemm.fp8_mqa_logits( - q_fp8[:q_offset], - kv_fp8, - weights[:q_offset], - ks, - ke, - clean_logits=False, - ) - assert logits.shape[0] == len(seq_lens_expanded) - assert logits.shape[1] == k_offset - - raw_topk_result = metadata.topk_transform(logits, self.index_topk, ks=ks) - topk_result[:q_offset] = raw_topk_result - return topk_result - - bytes_per_row = k_offset * self._MQA_LOGITS_BYTES_PER_ELEM - max_rows = max(1, int(logits_budget_bytes // max(bytes_per_row, 1))) - max_rows = min(max_rows, q_offset) - - global_topk_offset = metadata.attn_metadata.topk_indices_offset - cu_seqlens_q_full = None - if global_topk_offset is None: - cu_seqlens_q_full = torch.ones(q_offset, dtype=torch.int32, device=device) - - assert ( - seq_lens_expanded.shape[0] == q_offset - ), f"seq_lens_expanded length mismatch: {seq_lens_expanded.shape[0]} != {q_offset}" - if global_topk_offset is not None: - assert ( - global_topk_offset.shape[0] >= q_offset - ), f"topk_indices_offset too short: {global_topk_offset.shape[0]} < {q_offset}" - - start = 0 - while start < q_offset: - end = min(start + max_rows, q_offset) - - with self._with_real_sm_count(): - if _is_hip: - from aiter.ops.triton.fp8_mqa_logits import fp8_mqa_logits - - kv, scale = kv_fp8 - logits_chunk = fp8_mqa_logits( - q_fp8[start:end], - kv, - scale, - weights[start:end], - ks[start:end], - ke[start:end], - ) - else: - logits_chunk = deep_gemm.fp8_mqa_logits( - q_fp8[start:end], - kv_fp8, - weights[start:end], - ks[start:end], - ke[start:end], - clean_logits=False, - ) - - lengths_chunk = seq_lens_expanded[start:end] - - # RAGGED: use global offset; PAGED: construct local cu_seqlens_q per chunk - if global_topk_offset is not None: - # RAGGED path - topk_offset_chunk = global_topk_offset[start:end] - cu_seqlens_q_chunk = None - batch_idx_chunk = None - else: - # PAGED path: treat each token as a length-1 sequence - topk_offset_chunk = None - cu_seqlens_q_chunk = cu_seqlens_q_full[start:end] - batch_idx_chunk = token_to_batch_idx[start:end] - - raw_topk_chunk = metadata.topk_transform( - logits_chunk, - self.index_topk, - ks=ks[start:end], - cu_seqlens_q=cu_seqlens_q_chunk, - ke_offset=lengths_chunk, - batch_idx_list=batch_idx_chunk, - topk_indices_offset_override=topk_offset_chunk, - ) - topk_result[start:end] = raw_topk_chunk - start = end - - return topk_result - - def _forward_cuda_k_only( - self, - x: torch.Tensor, - positions: torch.Tensor, - forward_batch: ForwardBatch, - layer_id: int, - act_quant, - enable_dual_stream: bool, - metadata: BaseIndexerMetadata, - return_indices: bool = True, - ) -> Optional[torch.Tensor]: - assert forward_batch.forward_mode.is_extend_without_speculative() - x_meta = x[0] if isinstance(x, tuple) else x - - # Fast path: only compute and store k cache, skip all q and weights ops - key = self._get_k_bf16(x, positions, enable_dual_stream) - - if not forward_batch.out_cache_loc.is_contiguous(): - forward_batch.out_cache_loc = forward_batch.out_cache_loc.contiguous() - - self._store_index_k_cache( - forward_batch=forward_batch, - layer_id=layer_id, - key=key, - act_quant=act_quant, - ) - - # MHA doesn't need topk_indices - if not return_indices: - return None - - # MLA: use dummy logits with topk kernel's fast path to generate indices - # When length <= 2048, naive_topk_cuda directly generates [0,1,...,length-1,-1,...] - seq_lens_expanded = metadata.get_seqlens_expanded() - dummy_logits = torch.zeros( - seq_lens_expanded.shape[0], - self.index_topk, - dtype=torch.float32, - device=x_meta.device, - ) - return metadata.topk_transform(dummy_logits, self.index_topk) - - def _get_topk_ragged_with_cp( - self, - forward_batch: ForwardBatch, - layer_id: int, - q_fp8: torch.Tensor, - weights: torch.Tensor, - metadata: BaseIndexerMetadata, - kv_len: int, - actual_seq_q: int, - cp_index: List[Tuple[int, int, int]] = None, - ) -> torch.Tensor: - if TYPE_CHECKING: - assert isinstance(forward_batch.token_to_kv_pool, NSATokenToKVPool) - - page_size = forward_batch.token_to_kv_pool.page_size - assert page_size == 64, "only support page size 64" - assert len(weights.shape) == 3 - weights = weights.squeeze(-1) - k_fp8_list = [] - k_scale_list = [] - ks_list = [] - ke_offset_list = [] - offset = 0 - actual_seq_q_list = [] - batch_idx_list = [] - - block_tables = metadata.get_page_table_64() - - assert ( - forward_batch.seq_lens_cpu is not None - and forward_batch.extend_seq_lens_cpu is not None - ) - if cp_index is not None: - # TODO Multi-batch support has accuracy issues - for batch_idx, start_seq_position, end_seq_position in cp_index: - pre_chunk_offset = ( - forward_batch.seq_lens_cpu[batch_idx].item() - - forward_batch.extend_seq_lens_cpu[batch_idx] - ) - start_seq_position += pre_chunk_offset - end_seq_position += pre_chunk_offset - if offset == 0 and batch_idx != 0: - offset += forward_batch.extend_seq_lens_cpu[batch_idx - 1] - k_fp8 = forward_batch.token_to_kv_pool.get_index_k_continuous( - layer_id, - end_seq_position, - block_tables[batch_idx], - ) - k_scale = forward_batch.token_to_kv_pool.get_index_k_scale_continuous( - layer_id, - end_seq_position, - block_tables[batch_idx], - ) - - extend_seq_len = end_seq_position - start_seq_position - ks = torch.full( - (extend_seq_len,), offset, dtype=torch.int32, device="cuda" - ) - k_fp8_list.append(k_fp8) - k_scale_list.append(k_scale) - ks_list.append(ks) - ke_offset = torch.arange( - start_seq_position + 1, - end_seq_position + 1, - dtype=torch.int32, - device="cuda", - ) - ke_offset_list.append(ke_offset) - actual_seq_q = torch.tensor( - [extend_seq_len], dtype=torch.int32, device="cuda" - ) - actual_seq_q_list.append(actual_seq_q) - batch_idx_list.append(batch_idx) - - k_fp8 = torch.cat(k_fp8_list, dim=0).view(torch.float8_e4m3fn) - k_scale = torch.cat(k_scale_list, dim=0).view(torch.float32).squeeze(-1) - kv_fp8 = (k_fp8, k_scale) - ks = torch.cat(ks_list, dim=0) - ke_offset = torch.cat(ke_offset_list, dim=0) - ke = ks + ke_offset - actual_seq_q = torch.cat(actual_seq_q_list, dim=0) - with self._with_real_sm_count(): - logits = deep_gemm.fp8_mqa_logits( - q_fp8, - kv_fp8, - weights, - ks, - ke, - clean_logits=False, - ) - topk_result = metadata.topk_transform( - logits, - self.index_topk, - ks=ks, - cu_seqlens_q=actual_seq_q, - ke_offset=ke_offset, - batch_idx_list=batch_idx_list, - ) - else: - kv_len = ( - forward_batch.seq_lens_cpu[0].item() - - forward_batch.extend_seq_lens_cpu[0] - + kv_len - ) - k_fp8 = forward_batch.token_to_kv_pool.get_index_k_continuous( - layer_id, - kv_len, - block_tables[0], - ) - k_scale = forward_batch.token_to_kv_pool.get_index_k_scale_continuous( - layer_id, - kv_len, - block_tables[0], - ) - - k_fp8 = k_fp8.view(torch.float8_e4m3fn) - k_scale = k_scale.view(torch.float32).squeeze(-1) - kv_fp8 = (k_fp8, k_scale) - ks = torch.full((actual_seq_q,), offset, dtype=torch.int32, device="cuda") - ke_offset = torch.arange( - (kv_len - actual_seq_q) + 1, - kv_len + 1, - dtype=torch.int32, - device="cuda", - ) - ke = ks + ke_offset - - with self._with_real_sm_count(): - logits = deep_gemm.fp8_mqa_logits( - q_fp8, - kv_fp8, - weights, - ks, - ke, - clean_logits=False, - ) - actual_seq_q = torch.tensor([actual_seq_q], dtype=torch.int32).to( - device="cuda", non_blocking=True - ) - topk_result = metadata.topk_transform( - logits, - self.index_topk, - ks=ks, - cu_seqlens_q=actual_seq_q, - ke_offset=ke_offset, - ) - - return topk_result - - def forward_indexer( - self, - q_fp8: torch.Tensor, - weights: torch.Tensor, - forward_batch: ForwardBatch, - topk: int, - layer_id: int, - ) -> Optional[torch.Tensor]: - if not _is_npu: - from sglang.srt.layers.attention.nsa.tilelang_kernel import fp8_index - - page_size = forward_batch.token_to_kv_pool.page_size - assert page_size == 64, "only support page size 64" - - assert len(weights.shape) == 3 - weights = weights.squeeze(-1) - - # logits = deep_gemm.fp8_mqa_logits(q_fp8, kv_fp8, weights, ks, ke) - k_fp8_list = [] - k_scale_list = [] - - topk_indices_list = [] - - block_tables = forward_batch.req_to_token_pool.req_to_token[ - forward_batch.req_pool_indices, : - ] - strided_indices = torch.arange( - 0, block_tables.shape[-1], page_size, device="cuda" - ) - block_tables = block_tables[:, strided_indices] // page_size - - q_len_start = 0 - - for i in range(forward_batch.batch_size): - seq_len = forward_batch.seq_lens[i].item() - q_len = ( - forward_batch.extend_seq_lens_cpu[i] - if forward_batch.forward_mode.is_extend() - else 1 - ) - q_len_end = q_len_start + q_len - - q_fp8_partial = q_fp8[q_len_start:q_len_end] - q_fp8_partial = q_fp8_partial.unsqueeze(0).contiguous() - - weights_partial = weights[q_len_start:q_len_end] - weights_partial = weights_partial.squeeze(-1).unsqueeze(0).contiguous() - - k_fp8 = forward_batch.token_to_kv_pool.get_index_k_continuous( - layer_id, - seq_len, - block_tables[i], - ) - k_scale = forward_batch.token_to_kv_pool.get_index_k_scale_continuous( - layer_id, - seq_len, - block_tables[i], - ) - - k_fp8 = k_fp8.view(torch.float8_e4m3fn).unsqueeze(0).contiguous() - k_scale = k_scale.view(torch.float32).squeeze(-1).unsqueeze(0).contiguous() - - index_score = fp8_index( - q_fp8_partial, - weights_partial, - k_fp8, - k_scale, - ) - end_pos = seq_len - topk_indices = index_score.topk(min(topk, end_pos), dim=-1)[1].squeeze(0) - - pad_len = ceil_align(topk_indices.shape[-1], 2048) - topk_indices.shape[-1] - topk_indices = torch.nn.functional.pad( - topk_indices, (0, pad_len), "constant", -1 - ) - - topk_indices_list.append(topk_indices) - - q_len_start = q_len_end - - topk_indices = torch.cat(topk_indices_list, dim=0) - return topk_indices - - def _store_index_k_cache( - self, - forward_batch: ForwardBatch, - layer_id: int, - key: torch.Tensor, - *, - act_quant=None, # fallback only - ) -> None: - """ - Store NSA indexer K cache for current step. - - Preferred: fused_store_index_k_cache(key, cache, out_cache_loc, page_size) - Fallback : act_quant(key) + token_to_kv_pool.set_index_k_scale_buffer(...) - """ - - # Fast path: JIT fused store (CUDA, page_size=64, non-fnuz) - if ( - _is_cuda - and (not _is_fp8_fnuz) - and can_use_nsa_fused_store( - key.dtype, - forward_batch.out_cache_loc.dtype, - forward_batch.token_to_kv_pool.page_size, - ) - ): - # NOTE: wrapper already normalizes shape/contiguity and asserts dtypes. - buf = forward_batch.token_to_kv_pool.get_index_k_with_scale_buffer( - layer_id=layer_id - ) - fused_store_index_k_cache( - key, - buf, - forward_batch.out_cache_loc, - forward_batch.token_to_kv_pool.page_size, - ) - return - - # Fast path: AITER fused quant + cache store - # When _use_aiter_preshuffle is True we use the new MFMA 16x16 preshuffle - # layout (page_size>=16). Otherwise we fall back to the legacy row-major - # layout with page_size=1; the same kv_cache.view works for both cases - # because page_size is 1 there. - if _use_aiter: - page_size = forward_batch.token_to_kv_pool.page_size - buf = forward_batch.token_to_kv_pool.get_index_k_with_scale_buffer( - layer_id=layer_id - ) - kv_cache = buf.view(-1, page_size, 132).view(fp8_dtype) - out_loc = forward_batch.out_cache_loc - if not out_loc.is_contiguous(): - out_loc = out_loc.contiguous() - indexer_k_quant_and_cache( - key, - kv_cache, - out_loc, - self.block_size, - self.scale_fmt, - preshuffle=_use_aiter_preshuffle, - ) - return - - # Fallback: original path - assert act_quant is not None - k_fp8, k_scale = act_quant(key, self.block_size, self.scale_fmt) - - out_loc = forward_batch.out_cache_loc - if not out_loc.is_contiguous(): - out_loc = out_loc.contiguous() - - forward_batch.token_to_kv_pool.set_index_k_scale_buffer( - layer_id=layer_id, - loc=out_loc, - index_k=k_fp8, - index_k_scale=k_scale, - ) - - def forward_xpu( - self, - x: torch.Tensor, - q_lora: torch.Tensor, - positions: torch.Tensor, - forward_batch: ForwardBatch, - layer_id: int, - return_indices: bool = True, - ) -> Optional[torch.Tensor]: - return self.forward_cuda( - x, q_lora, positions, forward_batch, layer_id, return_indices - ) - - def forward_cuda( - self, - x: torch.Tensor, - q_lora: torch.Tensor, - positions: torch.Tensor, - forward_batch: ForwardBatch, - layer_id: int, - return_indices: bool = True, - ) -> Optional[torch.Tensor]: - if _is_hip: - from sglang.srt.layers.attention.nsa.tilelang_kernel import act_quant - elif not _is_npu: - from sglang.srt.layers.attention.nsa.triton_kernel import act_quant - - if TYPE_CHECKING: - assert isinstance(forward_batch.token_to_kv_pool, NSATokenToKVPool) - - # When upstream uses fused FP8 RMSNorm+quant, activations may be passed as - # a tuple like (x_fp8, x_scale[, y]). Use `x_meta` for shape/device queries. - x_meta = x[0] if isinstance(x, tuple) else x - - metadata = forward_batch.attn_backend.get_indexer_metadata( - layer_id, forward_batch - ) - - enable_dual_stream = ( - self.alt_stream is not None - and get_is_capture_mode() - and q_lora.shape[0] > 0 - and q_lora.shape[0] <= DUAL_STREAM_TOKEN_THRESHOLD - ) - - # skip NSA if attention backend choose to skip this batch - if metadata is None: - return None - - # Determine if should skip topk based on sequence length - # We can only skip the logits computation if cuda graph is not involved - skip_logits_computation = False - if forward_batch.forward_mode.is_extend_without_speculative(): - if forward_batch.seq_lens_cpu is not None: - max_kv_len = forward_batch.seq_lens_cpu.max().item() - skip_logits_computation = max_kv_len <= self.index_topk - - # Optimization: fast path when skipping topk computation - if skip_logits_computation and (not self.nsa_enable_prefill_cp): - return maybe_capture_indexer_topk( - layer_id, - self._forward_cuda_k_only( - x, - positions, - forward_batch, - layer_id, - act_quant, - enable_dual_stream, - metadata, - return_indices, - ), - ) - - if enable_dual_stream and forward_batch.forward_mode.is_decode_or_idle(): - current_stream = torch.cuda.current_stream() - self.alt_stream.wait_stream(current_stream) - weights = self._project_and_scale_head_gates(x) - query, key = self._get_q_k_bf16( - q_lora, x, positions, enable_dual_stream, forward_batch=forward_batch - ) - q_fp8, q_scale = act_quant(query, self.block_size, self.scale_fmt) - with torch.cuda.stream(self.alt_stream): - self._store_index_k_cache( - forward_batch=forward_batch, - layer_id=layer_id, - key=key, - act_quant=act_quant, - ) - current_stream.wait_stream(self.alt_stream) - weights = self._apply_q_scale_and_softmax_scale(weights, q_scale) - else: - query, key = self._get_q_k_bf16( - q_lora, x, positions, enable_dual_stream, forward_batch=forward_batch - ) - - if enable_dual_stream: - current_stream = torch.cuda.current_stream() - self.alt_stream.wait_stream(current_stream) - - q_fp8, q_scale = act_quant(query, self.block_size, self.scale_fmt) - with torch.cuda.stream(self.alt_stream): - self._store_index_k_cache( - forward_batch=forward_batch, - layer_id=layer_id, - key=key, - act_quant=act_quant, - ) - current_stream.wait_stream(self.alt_stream) - else: - q_fp8, q_scale = act_quant(query, self.block_size, self.scale_fmt) - self._store_index_k_cache( - forward_batch=forward_batch, - layer_id=layer_id, - key=key, - act_quant=act_quant, - ) - - # aiter (ROCm gfx95): the 3-tuple (fp8, scale, bf16) from - # fused_rms_fp8_group_quant is passed directly to _get_logits_head_gate, - # which extracts the bf16 tensor via _weights_proj_bf16_in_fp32_out, - # completely skipping the FP8 dequantization path below. - if ( - _use_aiter - and _is_gfx95_supported - and isinstance(x, tuple) - and len(x) == 3 - ): - x_for_gate = x - elif isinstance(x, tuple): - assert len(x) in ( - 2, - 3, - ), "For tuple input, only (x, x_s) or (x, x_s, y) formats are accepted" - x_q, x_s = x[0], x[1] - if ( - x_s is not None - and x_q.dim() == 2 - and x_s.dim() == 2 - and x_q.shape[0] == x_s.shape[0] - ): - m, n = x_q.shape - ng = x_s.shape[1] - if ng > 0 and n % ng == 0: - group = n // ng - x_for_gate = ( - x_q.to(torch.float32) - .view(m, ng, group) - .mul_(x_s.to(torch.float32).unsqueeze(-1)) - .view(m, n) - .to(torch.bfloat16) - ) - else: - x_for_gate = x_q.to(torch.bfloat16) - else: - x_for_gate = x_q.to(torch.bfloat16) - else: - x_for_gate = x - - weights = self._get_logits_head_gate(x_for_gate, q_scale) - - if _is_cuda or _is_hip: - assert forward_batch.seq_lens_cpu is not None - if len(forward_batch.seq_lens_cpu) == 0: - # this seems b/c max-pad, no worries? - # if x.shape[0] != 0: - # print( - # "HACK: seq_lens empty but x not empty, hackily return all-invalid topk_result" - # ) - return maybe_capture_indexer_topk( - layer_id, - torch.full( - (x_meta.shape[0], self.index_topk), - -1, - dtype=torch.int, - device=x_meta.device, - ), - ) - - if ( - forward_batch.forward_mode.is_decode_or_idle() - or forward_batch.forward_mode.is_target_verify() - or forward_batch.forward_mode.is_draft_extend(include_v2=True) - ): - topk_result = self._get_topk_paged( - forward_batch, layer_id, q_fp8, weights, metadata - ) - else: - if ( - forward_batch.attn_cp_metadata is not None - and is_nsa_prefill_cp_in_seq_split() - ): - kv_len_prev = forward_batch.attn_cp_metadata.kv_len_prev - kv_len_next = forward_batch.attn_cp_metadata.kv_len_next - actual_seq_q_prev = forward_batch.attn_cp_metadata.actual_seq_q_prev - actual_seq_q_next = forward_batch.attn_cp_metadata.actual_seq_q_next - - # TODO support mutil-batch - # cp_batch_seq_index_prev = forward_batch.attn_cp_metadata["cp_batch_seq_index_prev"] - # cp_batch_seq_index_next = forward_batch.attn_cp_metadata["cp_batch_seq_index_next"] - # TODO prev, next, combined into a single call - q_fp8_prev, q_fp8_next = torch.split( - q_fp8, (q_fp8.shape[0] + 1) // 2, dim=0 - ) - weights_prev, weights_next = torch.split( - weights, (weights.shape[0] + 1) // 2, dim=0 - ) - topk_result_prev = self._get_topk_ragged_with_cp( - forward_batch, - layer_id, - q_fp8_prev, - weights_prev, - metadata, - kv_len_prev, - actual_seq_q_prev, - ) - - topk_result_next = self._get_topk_ragged_with_cp( - forward_batch, - layer_id, - q_fp8_next, - weights_next, - metadata, - kv_len_next, - actual_seq_q_next, - ) - return maybe_capture_indexer_topk( - layer_id, - torch.cat([topk_result_prev, topk_result_next], dim=0), - ) - else: - topk_result = self._get_topk_ragged( - enable_dual_stream, - forward_batch, - layer_id, - q_fp8, - weights, - metadata, - ) - else: - topk_result = self.forward_indexer( - q_fp8.contiguous(), - weights, - forward_batch, - topk=self.index_topk, - layer_id=layer_id, - ) - return maybe_capture_indexer_topk(layer_id, topk_result) - - def forward_npu( - self, - x: torch.Tensor, - q_lora: torch.Tensor, - positions: torch.Tensor, - forward_batch: ForwardBatch, - layer_id: int, - layer_scatter_modes=None, - dynamic_scale: torch.Tensor = None, - ) -> torch.Tensor: - if forward_batch.attn_backend.forward_metadata.seq_lens_cpu_int is None: - actual_seq_lengths_kv = forward_batch.attn_backend.forward_metadata.seq_lens - else: - actual_seq_lengths_kv = ( - forward_batch.attn_backend.forward_metadata.seq_lens_cpu_int - ) - is_prefill = ( - forward_batch.forward_mode.is_extend() - and not forward_batch.forward_mode.is_draft_extend_v2() - and not forward_batch.forward_mode.is_target_verify() - and not forward_batch.forward_mode.is_draft_extend() - ) - - bs = q_lora.shape[0] - - if self.rotary_emb.is_neox_style: - if not hasattr(forward_batch, "npu_indexer_sin_cos_cache"): - cos_sin = self.rotary_emb.cos_sin_cache[positions] - cos, sin = cos_sin.chunk(2, dim=-1) - cos = cos.repeat(1, 2).view(-1, 1, 1, self.rope_head_dim) - sin = sin.repeat(1, 2).view(-1, 1, 1, self.rope_head_dim) - forward_batch.npu_indexer_sin_cos_cache = (sin, cos) - else: - sin, cos = forward_batch.npu_indexer_sin_cos_cache - - if self.alt_stream is not None: - self.alt_stream.wait_stream(torch.npu.current_stream()) - with torch.npu.stream(self.alt_stream): - q_lora = ( - (q_lora, dynamic_scale) if dynamic_scale is not None else q_lora - ) - q = self.wq_b(q_lora)[ - 0 - ] # [bs, 1536] @ [1536, 64 * 128] = [bs, 64 * 128] - wq_b_event = self.alt_stream.record_event() - q = q.view(bs, self.n_heads, self.head_dim) # [bs, 64, 128] - q_pe, q_nope = torch.split( - q, - [self.rope_head_dim, self.head_dim - self.rope_head_dim], - dim=-1, - ) # [bs, 64, 64 + 64] - q_pe = q_pe.view(bs, self.n_heads, 1, self.rope_head_dim) - q_pe = torch_npu.npu_rotary_mul(q_pe, cos, sin).view( - bs, self.n_heads, self.rope_head_dim - ) # [bs, n, d] - q = torch.cat([q_pe, q_nope], dim=-1) - q.record_stream(self.alt_stream) - q_rope_event = self.alt_stream.record_event() - else: - q_lora = ( - (q_lora, dynamic_scale) if dynamic_scale is not None else q_lora - ) - q = self.wq_b(q_lora)[ - 0 - ] # [bs, 1536] @ [1536, 64 * 128] = [bs, 64 * 128] - q = q.view(bs, self.n_heads, self.head_dim) # [bs, 64, 128] - q_pe, q_nope = torch.split( - q, - [self.rope_head_dim, self.head_dim - self.rope_head_dim], - dim=-1, - ) # [bs, 64, 64 + 64] - q_pe = q_pe.view(bs, self.n_heads, 1, self.rope_head_dim) - q_pe = torch_npu.npu_rotary_mul(q_pe, cos, sin).view( - bs, self.n_heads, self.rope_head_dim - ) # [bs, n, d] - q = torch.cat([q_pe, q_nope], dim=-1) - - if envs.SGLANG_NPU_USE_MULTI_STREAM.get(): - indexer_weight_stream = get_indexer_weight_stream() - indexer_weight_stream.wait_stream(torch.npu.current_stream()) - with torch.npu.stream(indexer_weight_stream): - x = x.view(-1, self.hidden_size) - weights = self.weights_proj(x.float())[0].to(torch.bfloat16) - weights.record_stream(indexer_weight_stream) - weights_event = indexer_weight_stream.record_event() - else: - x = x.view(-1, self.hidden_size) - weights = self.weights_proj(x.float())[0].to(torch.bfloat16) - - k_proj = self.wk(x)[0] # [b, s, 7168] @ [7168, 128] = [b, s, 128] - k = self.k_norm(k_proj) - if ( - _use_ag_after_qlora - and layer_scatter_modes.layer_input_mode == ScatterMode.SCATTERED - and layer_scatter_modes.attn_mode == ScatterMode.TP_ATTN_FULL - ): - k = scattered_to_tp_attn_full(k, forward_batch) - k_pe, k_nope = torch.split( - k, - [self.rope_head_dim, self.head_dim - self.rope_head_dim], - dim=-1, - ) # [bs, 64 + 64] - - k_pe = k_pe.view(-1, 1, 1, self.rope_head_dim) - k_pe = torch.ops.npu.npu_rotary_mul(k_pe, cos, sin).view( - bs, 1, self.rope_head_dim - ) # [bs, 1, d] - k = torch.cat([k_pe, k_nope.unsqueeze(1)], dim=-1) # [bs, 1, 128] - - else: - if envs.SGLANG_NPU_USE_MULTI_STREAM.get(): - indexer_weight_stream = get_indexer_weight_stream() - indexer_weight_stream.wait_stream(torch.npu.current_stream()) - with torch.npu.stream(indexer_weight_stream): - x = x.view(-1, self.hidden_size) - weights = self.weights_proj(x.float())[0].to(torch.bfloat16) - weights.record_stream(indexer_weight_stream) - weights_event = indexer_weight_stream.record_event() - else: - x = x.view(-1, self.hidden_size) - weights = self.weights_proj(x.float())[0].to(torch.bfloat16) - - q_lora = (q_lora, dynamic_scale) if dynamic_scale is not None else q_lora - q = self.wq_b(q_lora)[0] # [bs, 1536] @ [1536, 64 * 128] = [bs, 64 * 128] - q = q.view(bs, self.n_heads, self.head_dim) # [bs, 64, 128] - q_pe, q_nope = torch.split( - q, - [self.rope_head_dim, self.head_dim - self.rope_head_dim], - dim=-1, - ) # [bs, 64, 64 + 64] - - k_proj = self.wk(x)[0] # [b, s, 7168] @ [7168, 128] = [b, s, 128] - k = self.k_norm(k_proj) - k_pe, k_nope = torch.split( - k, - [self.rope_head_dim, self.head_dim - self.rope_head_dim], - dim=-1, - ) # [bs, 64 + 64] - - k_pe = k_pe.unsqueeze(1) - - if layer_id == 0: - self.rotary_emb.sin_cos_cache = ( - self.rotary_emb.cos_sin_cache.index_select(0, positions) - ) - - q_pe, k_pe = self.rotary_emb(positions, q_pe, k_pe) - k_pe = k_pe.squeeze(1) - q = torch.cat([q_pe, q_nope], dim=-1) - k = torch.cat([k_pe, k_nope], dim=-1) - - if ( - is_prefill - and self.nsa_enable_prefill_cp - and forward_batch.attn_cp_metadata is not None - ): - k = cp_all_gather_rerange_output( - k.contiguous().view(-1, self.head_dim), - self.cp_size, - forward_batch, - torch.npu.current_stream(), - ) - - forward_batch.token_to_kv_pool.set_index_k_buffer( - layer_id, forward_batch.out_cache_loc, k - ) - if is_prefill: - if ( - self.nsa_enable_prefill_cp - and forward_batch.attn_cp_metadata is not None - ): - forward_batch.attn_backend.forward_metadata.actual_seq_lengths_q = ( - forward_batch.attn_cp_metadata.actual_seq_q_prev_tensor, - forward_batch.attn_cp_metadata.actual_seq_q_next_tensor, - ) - if sum(forward_batch.extend_prefix_lens_cpu) > 0: - total_kv_len_prev_tensor = ( - forward_batch.attn_cp_metadata.kv_len_prev_tensor - + forward_batch.extend_prefix_lens.squeeze() - ) - total_kv_len_next_tensor = ( - forward_batch.attn_cp_metadata.kv_len_next_tensor - + forward_batch.extend_prefix_lens.squeeze() - ) - forward_batch.attn_backend.forward_metadata.actual_seq_lengths_kv = ( - total_kv_len_prev_tensor, - total_kv_len_next_tensor, - ) - else: - forward_batch.attn_backend.forward_metadata.actual_seq_lengths_kv = ( - forward_batch.attn_cp_metadata.kv_len_prev_tensor, - forward_batch.attn_cp_metadata.kv_len_next_tensor, - ) - actual_seq_lengths_q = ( - forward_batch.attn_backend.forward_metadata.actual_seq_lengths_q - ) - actual_seq_lengths_kv = ( - forward_batch.attn_backend.forward_metadata.actual_seq_lengths_kv - ) - else: - actual_seq_lengths_kv = forward_batch.seq_lens - actual_seq_lengths_q = forward_batch.extend_seq_lens.cumsum(dim=0) - else: - if forward_batch.attn_backend.forward_metadata.actual_seq_lengths_q is None: - if ( - forward_batch.forward_mode.is_draft_extend_v2() - or forward_batch.forward_mode.is_target_verify() - or forward_batch.forward_mode.is_draft_extend() - ): - num_draft_tokens = ( - forward_batch.attn_backend.speculative_num_draft_tokens - ) - actual_seq_lengths_q = torch.arange( - num_draft_tokens, - num_draft_tokens + bs, - num_draft_tokens, - dtype=torch.int32, - device=k.device, - ) - else: - actual_seq_lengths_q = torch.tensor( - [1 + i * 1 for i in range(bs)], - dtype=torch.int32, - device=k.device, - ) - else: - actual_seq_lengths_q = ( - forward_batch.attn_backend.forward_metadata.actual_seq_lengths_q - ) - - past_key_states = forward_batch.token_to_kv_pool.get_index_k_buffer(layer_id) - - if self.rotary_emb.is_neox_style and self.alt_stream is not None: - torch.npu.current_stream().wait_event(q_rope_event) - if envs.SGLANG_NPU_USE_MULTI_STREAM.get(): - torch.npu.current_stream().wait_event(weights_event) - if ( - _use_ag_after_qlora - and layer_scatter_modes.layer_input_mode == ScatterMode.SCATTERED - and layer_scatter_modes.attn_mode == ScatterMode.TP_ATTN_FULL - ): - weights = scattered_to_tp_attn_full(weights, forward_batch) - block_table = forward_batch.attn_backend.forward_metadata.block_tables - if ( - is_prefill - and self.nsa_enable_prefill_cp - and forward_batch.attn_cp_metadata is not None - ): - block_table = block_table[: actual_seq_lengths_q[0].numel()] - topk_indices = self.do_npu_cp_balance_indexer( - q.view(-1, self.n_heads, self.head_dim), - past_key_states, - weights, - actual_seq_lengths_q, - actual_seq_lengths_kv, - block_table, - ) - return topk_indices - else: - block_table = ( - block_table[: actual_seq_lengths_q.size()[0]] - if is_prefill - else block_table - ) - - topk_indices = torch_npu.npu_lightning_indexer( - query=q.view(-1, self.n_heads, self.head_dim), - key=past_key_states, - weights=weights, - actual_seq_lengths_query=actual_seq_lengths_q.to(torch.int32), - actual_seq_lengths_key=actual_seq_lengths_kv.to(k.device).to( - torch.int32 - ), - block_table=block_table, - layout_query="TND", - layout_key="PA_BSND", - sparse_count=self.index_topk, - sparse_mode=3, - ) - return topk_indices[0] - - def do_npu_cp_balance_indexer( - self, - q, - past_key_states, - indexer_weights, - actual_seq_lengths_q, - actual_seq_lengths_kv, - block_table, - ): - q_prev, q_next = torch.split(q, (q.size(0) + 1) // 2, dim=0) - weights_prev, weights_next = None, None - if indexer_weights is not None: - weights_prev, weights_next = torch.split( - indexer_weights, (indexer_weights.size(0) + 1) // 2, dim=0 - ) - weights_prev = weights_prev.contiguous().view(-1, weights_prev.shape[-1]) - weights_next = weights_next.contiguous().view(-1, weights_next.shape[-1]) - - actual_seq_lengths_q_prev, actual_seq_lengths_q_next = actual_seq_lengths_q - actual_seq_lengths_kv_prev, actual_seq_lengths_kv_next = actual_seq_lengths_kv - - topk_indices_prev = torch_npu.npu_lightning_indexer( - query=q_prev, - key=past_key_states, - weights=weights_prev, - actual_seq_lengths_query=actual_seq_lengths_q_prev.to( - device=q.device, dtype=torch.int32 - ), - actual_seq_lengths_key=actual_seq_lengths_kv_prev.to( - device=q.device, dtype=torch.int32 - ), - block_table=block_table, - layout_query="TND", - layout_key="PA_BSND", - sparse_count=self.index_topk, - sparse_mode=3, - ) - topk_indices_next = torch_npu.npu_lightning_indexer( - query=q_next, - key=past_key_states, - weights=weights_next, - actual_seq_lengths_query=actual_seq_lengths_q_next.to( - device=q.device, dtype=torch.int32 - ), - actual_seq_lengths_key=actual_seq_lengths_kv_next.to( - device=q.device, dtype=torch.int32 - ), - block_table=block_table, - layout_query="TND", - layout_key="PA_BSND", - sparse_count=self.index_topk, - sparse_mode=3, - ) - return topk_indices_prev[0], topk_indices_next[0] - - -def scattered_to_tp_attn_full( - hidden_states: torch.Tensor, - forward_batch, -) -> torch.Tensor: - hidden_states, local_hidden_states = ( - torch.empty( - (forward_batch.input_ids.shape[0], hidden_states.shape[1]), - dtype=hidden_states.dtype, - device=hidden_states.device, - ), - hidden_states, - ) - attn_tp_all_gather_into_tensor(hidden_states, local_hidden_states.contiguous()) - return hidden_states +from sglang.srt.layers.attention.dsa.dsa_indexer import * # noqa: F401, F403 diff --git a/python/sglang/srt/layers/attention/nsa/nsa_mtp_verification.py b/python/sglang/srt/layers/attention/nsa/nsa_mtp_verification.py index b957d4ba8..ee111ce43 100644 --- a/python/sglang/srt/layers/attention/nsa/nsa_mtp_verification.py +++ b/python/sglang/srt/layers/attention/nsa/nsa_mtp_verification.py @@ -1,407 +1,10 @@ -""" -Verification utilities for NSA backend fused metadata copy operations. +# [Deprecated] Re-export shim for backward compatibility. Use dsa.dsa_mtp_verification instead. +import warnings -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 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"], - ) +warnings.warn( + "sglang.srt.layers.attention.nsa.nsa_mtp_verification is deprecated; " + "use sglang.srt.layers.attention.dsa.dsa_mtp_verification instead.", + DeprecationWarning, + stacklevel=2, +) +from sglang.srt.layers.attention.dsa.dsa_mtp_verification import * # noqa: F401, F403 diff --git a/python/sglang/srt/layers/attention/nsa/quant_k_cache.py b/python/sglang/srt/layers/attention/nsa/quant_k_cache.py index 5454071b8..65bca9e39 100644 --- a/python/sglang/srt/layers/attention/nsa/quant_k_cache.py +++ b/python/sglang/srt/layers/attention/nsa/quant_k_cache.py @@ -1,449 +1,10 @@ -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") +# [Deprecated] Re-export shim for backward compatibility. Use dsa.quant_k_cache instead. +import warnings + +warnings.warn( + "sglang.srt.layers.attention.nsa.quant_k_cache is deprecated; " + "use sglang.srt.layers.attention.dsa.quant_k_cache instead.", + DeprecationWarning, + stacklevel=2, +) +from sglang.srt.layers.attention.dsa.quant_k_cache import * # noqa: F401, F403 diff --git a/python/sglang/srt/layers/attention/nsa/tilelang_kernel.py b/python/sglang/srt/layers/attention/nsa/tilelang_kernel.py index 62509c308..7e653f7ec 100644 --- a/python/sglang/srt/layers/attention/nsa/tilelang_kernel.py +++ b/python/sglang/srt/layers/attention/nsa/tilelang_kernel.py @@ -1,2589 +1,10 @@ -import functools -from functools import lru_cache -from typing import Any, Optional, Tuple +# [Deprecated] Re-export shim for backward compatibility. Use dsa.tilelang_kernel instead. +import warnings -import tilelang -import tilelang.language as T -import torch - -from sglang.srt.layers.quantization.fp8_kernel import is_fp8_fnuz -from sglang.srt.utils import is_gfx95_supported, is_hip - -tilelang.set_log_level("WARNING") - -# Workaround a tilelang bug: BaseKernelAdapter._legalize_result_idx mutates the -# `out_idx` list in place when normalising negative indices to positive ones. -# That breaks any @tilelang.jit factory that compiles two prim_funcs with -# different param counts (e.g. our unified single/dual partial kernel) — the -# second compile sees indices already-converted for the first's len(params) -# and silently builds the wrong adapter, leading to IndexError at call time. -# Patch once on import to copy the list before mutation. -from tilelang.jit.adapter.base import ( # noqa: E402 - BaseKernelAdapter as _BaseKernelAdapter, +warnings.warn( + "sglang.srt.layers.attention.nsa.tilelang_kernel is deprecated; " + "use sglang.srt.layers.attention.dsa.tilelang_kernel instead.", + DeprecationWarning, + stacklevel=2, ) - -if not getattr(_BaseKernelAdapter, "_legalize_result_idx_patched", False): - _orig_legalize = _BaseKernelAdapter._legalize_result_idx - - def _legalize_result_idx_safe(self, result_idx): - if isinstance(result_idx, list): - result_idx = list(result_idx) - return _orig_legalize(self, result_idx) - - _BaseKernelAdapter._legalize_result_idx = _legalize_result_idx_safe - _BaseKernelAdapter._legalize_result_idx_patched = True - -pass_configs = { - tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True, - tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, -} -# TL_DISABLE_FAST_MATH has deprecated in v0.1.7.post1 tilelang -if hasattr(tilelang.PassConfigKey, "TL_DISABLE_FAST_MATH"): - pass_configs[tilelang.PassConfigKey.TL_DISABLE_FAST_MATH] = True -elif hasattr(tilelang.PassConfigKey, "TL_ENABLE_FAST_MATH"): - pass_configs[tilelang.PassConfigKey.TL_ENABLE_FAST_MATH] = False - -_is_hip = is_hip() -_is_gfx95_supported = is_gfx95_supported() -_is_fp8_fnuz = is_fp8_fnuz() - -BF16 = "bfloat16" -FP8 = "float8_e4m3fnuz" if _is_fp8_fnuz else "float8_e4m3fn" -FP8_DTYPE = torch.float8_e4m3fnuz if _is_fp8_fnuz else torch.float8_e4m3fn -FP32 = "float32" -INT32 = "int32" -UINT8 = "uint8" - - -def fast_log2_ceil(x): - bits_x = T.reinterpret("uint32", x) - exp_x = (bits_x >> 23) & 0xFF - man_bits = bits_x & ((1 << 23) - 1) - return T.Cast("int32", exp_x - 127 + T.if_then_else(man_bits != 0, 1, 0)) - - -def fast_pow2(x): - bits_x = (x + 127) << 23 - return T.reinterpret("float32", bits_x) - - -def fast_round_scale(amax, fp8_max_inv): - return fast_pow2(fast_log2_ceil(amax * fp8_max_inv)) - - -@lru_cache(maxsize=8) -def _pick_inner_iter(seq: int, ni: int, cu: int, block_per_cu: int) -> int: - """ - Pick the largest valid inner_iter (power-of-two divisor of ni) that keeps - enough work per CU (seq * ni / inner_iter / cu >= block_per_cu), so we avoid - under-utilization while minimizing the number of partial groups. - """ - - max_it = int(seq * ni / (cu * block_per_cu)) - it = ni - while it >= 2: - if it <= max_it and ni % it == 0: - return it - it //= 2 - return 1 - - -@tilelang.jit(pass_configs=pass_configs) -def act_quant_kernel( - N, in_dtype=BF16, out_dtype=FP8, scale_dtype=FP32, round_scale=False -): - M = T.symbolic("M") - fp8_min = -224.0 if _is_fp8_fnuz else -448.0 - fp8_max = 224.0 if _is_fp8_fnuz else 448.0 - fp8_max_inv = 1 / fp8_max - num_stages = 0 if round_scale else 2 - blk_m = 32 - group_size = 128 - - @T.prim_func - def act_quant_kernel_( - X: T.Tensor[(M, N), in_dtype], - Y: T.Tensor[(M, N), out_dtype], - S: T.Tensor[(M, T.ceildiv(N, group_size)), scale_dtype], - ): - with T.Kernel(T.ceildiv(M, blk_m), T.ceildiv(N, group_size), threads=128) as ( - pid_m, - pid_n, - ): - x_shared = T.alloc_shared((blk_m, group_size), in_dtype) - x_local = T.alloc_fragment((blk_m, group_size), in_dtype) - amax_local = T.alloc_fragment((blk_m,), scale_dtype) - s_local = T.alloc_fragment((blk_m,), scale_dtype) - y_local = T.alloc_fragment((blk_m, group_size), out_dtype) - y_shared = T.alloc_shared((blk_m, group_size), out_dtype) - - for _ in T.Pipelined(1, num_stages=num_stages): - T.copy(X[pid_m * blk_m, pid_n * group_size], x_shared) - T.copy(x_shared, x_local) - T.reduce_absmax(x_local, amax_local, dim=1) - for i in T.Parallel(blk_m): - amax_local[i] = T.max(amax_local[i], 1e-4) - if round_scale: - s_local[i] = fast_round_scale(amax_local[i], fp8_max_inv) - else: - s_local[i] = amax_local[i] * fp8_max_inv - for i, j in T.Parallel(blk_m, group_size): - y_local[i, j] = T.clamp( - x_local[i, j] / s_local[i], fp8_min, fp8_max - ) - for i in T.Parallel(blk_m): - S[pid_m * blk_m + i, pid_n] = s_local[i] - T.copy(y_local, y_shared) - T.copy(y_shared, Y[pid_m * blk_m, pid_n * group_size]) - - return act_quant_kernel_ - - -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. - - 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})" - N = x.size(-1) - if _is_fp8_fnuz: - y = torch.empty_like(x, dtype=torch.float8_e4m3fnuz) - else: - y = torch.empty_like(x, dtype=torch.float8_e4m3fn) - s = x.new_empty(*x.size()[:-1], N // block_size, dtype=torch.float32) - kernel = act_quant_kernel(N, round_scale=scale_fmt is not None) - kernel(x.view(-1, N), y.view(-1, N), s.view(-1, N // block_size)) - return y, s - - -@tilelang.jit(out_idx=[4], pass_configs=pass_configs) -def fp8_index_kernel(h: int, d: int, clear_accum=True): - b = T.symbolic("b") - m = T.symbolic("m") - n = T.symbolic("n") - - blk_n1 = 512 - blk_n2 = 128 - - @T.prim_func - def fp8_index_kernel_( - q: T.Tensor[(b, m, h, d), FP8], - q_s: T.Tensor[(b, m, h), FP32], - k: T.Tensor[(b, n, d), FP8], - k_s: T.Tensor[(b, n), FP32], - o: T.Tensor[(b, m, n), FP32], - ) -> None: - with T.Kernel(b, m, T.ceildiv(n, blk_n1)) as (i_b, i_m, i1_n): - q_smem = T.alloc_shared((h, d), FP8) - T.copy(q[i_b, i_m, 0, 0], q_smem) - - q_s_frag = T.alloc_fragment(h, FP32) - T.copy(q_s[i_b, i_m, 0], q_s_frag) - - for i2_n in T.Pipelined(blk_n1 // blk_n2, num_stages=2): - k_smem = T.alloc_shared((blk_n2, d), FP8) - T.copy(k[i_b, i1_n * blk_n1 + i2_n * blk_n2, 0], k_smem) - - k_s_frag = T.alloc_fragment(blk_n2, FP32) - T.copy(k_s[i_b, i1_n * blk_n1 + i2_n * blk_n2], k_s_frag) - - logits = T.alloc_fragment((blk_n2, h), FP32) - if not clear_accum: - T.fill(logits, 0) - T.gemm( - k_smem, - q_smem, - logits, - transpose_A=False, - transpose_B=True, - clear_accum=clear_accum, - ) - - for i_h, i3_n in T.Parallel(h, blk_n2): - logits[i3_n, i_h] = T.max(logits[i3_n, i_h], 0) * q_s_frag[i_h] - - logits_sum = T.alloc_fragment(blk_n2, FP32) - T.reduce_sum(logits, logits_sum, dim=1) - - for i3_n in T.Parallel(blk_n2): - logits_sum[i3_n] *= k_s_frag[i3_n] - - T.copy(logits_sum, o[i_b, i_m, i1_n * blk_n1 + i2_n * blk_n2]) - - return fp8_index_kernel_ - - -def fp8_index( - q: torch.Tensor, - q_s: torch.Tensor, - k: torch.Tensor, - k_s: torch.Tensor, -) -> torch.Tensor: - """ - Perform index score using FP8 precision. - - Args: - q (torch.Tensor): The Q tensor, must be contiguous. - q_s (torch.Tensor): The scaling factor for Q (float), must be contiguous. - k (torch.Tensor): The K tensor, must be contiguous. - k_s (torch.Tensor): The scaling factor for K (e8m0 here), must be contiguous. - - fp8 q @ fp8 k -> fp32 logits - relu(fp32 logits) * q_s (weights) -> fp32 logits - fp32 logits -> fp32 logits_sum - fp32 logits_sum * k_s (e8m0) -> fp32 index_score - """ - if _is_hip: - return fp8_index_kernel(q.shape[2], q.shape[3], False)(q, q_s, k, k_s) - else: - return fp8_index_kernel(q.shape[2], q.shape[3])(q, q_s, k, k_s) - - -@tilelang.jit( - out_idx=[-1], - pass_configs={ - tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, - tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True, - }, -) -def sparse_attention_fwd_kernel_v1( - num_heads, - dim, - tail_dim, - topk, - *, - kv_group=1, - sm_scale=None, - is_causal=True, - block_I=64, - num_stages=2, - threads=256, -): - assert dim == tilelang.math.next_power_of_2( - dim - ), f"haven't check padding correctness yet, dim={dim}" - assert tail_dim == tilelang.math.next_power_of_2( - tail_dim - ), f"haven't check padding correctness yet, dim={tail_dim}" - assert is_causal == True, "non-casual is not supported" - assert ( - topk % block_I == 0 - ), "otherwise will load some index=0 thus causing wrong kv to be loaded" - if sm_scale is None: - sm_scale = (1.0 / (dim + tail_dim)) ** 0.5 * 1.44269504 # log2(e) - else: - sm_scale = sm_scale * 1.44269504 # log2(e) - - batch = T.symbolic("batch") - seq_len = T.symbolic("seq_len") - seq_len_kv = T.symbolic("seq_len_kv") - - head_kv = num_heads // kv_group - q_shape = [batch, seq_len, num_heads, dim + tail_dim] - kv_shape = [batch, seq_len_kv, kv_group, dim + tail_dim] - o_shape = [batch, seq_len, num_heads, dim] - indices_shape = [batch, seq_len, kv_group, topk] - indices_dtype = "int32" - dtype = "bfloat16" - accum_dtype = "float" - - H = head_kv - padded_H = max(tilelang.math.next_power_of_2(head_kv), 16) - if padded_H != H: - assert kv_group == 1 - BI = block_I - NI = tilelang.cdiv(topk, block_I) - D = dim - D_tail = tail_dim - - if head_kv > 64: - assert head_kv % 64 == 0, "head_kv should be a multiple of 64" - REPLICATE_H = head_kv // 64 - else: - REPLICATE_H = 1 - - H_per_block = padded_H if REPLICATE_H == 1 else 64 - - @T.prim_func - def main( - Q: T.Tensor(q_shape, dtype), # type: ignore - KV: T.Tensor(kv_shape, dtype), # type: ignore - Indices: T.Tensor(indices_shape, indices_dtype), # type: ignore - Output: T.Tensor(o_shape, dtype), # type: ignore - ): - with T.Kernel(seq_len * REPLICATE_H, batch, kv_group, threads=threads) as ( - bx, - by, - bz, - ): - Q_shared = T.alloc_shared([H_per_block, D], dtype) - Q_tail_shared = T.alloc_shared([H_per_block, D_tail], dtype) - KV_shared = T.alloc_shared([BI, D], dtype) - K_tail_shared = T.alloc_shared([BI, D_tail], dtype) - O_shared = T.alloc_shared([H_per_block, D], dtype) - mask = T.alloc_fragment([BI], "bool") - - acc_o = T.alloc_fragment([H_per_block, D], accum_dtype) - acc_s = T.alloc_fragment([H_per_block, BI], accum_dtype) - S_shared = T.alloc_shared([H_per_block, BI], dtype) - sumexp = T.alloc_fragment([H_per_block], accum_dtype) - sumexp_i = T.alloc_fragment([H_per_block], accum_dtype) - alpha = T.alloc_fragment([H_per_block], accum_dtype) - m_i = T.alloc_fragment([H_per_block], accum_dtype) - m_i_prev = T.alloc_fragment([H_per_block], accum_dtype) - - T.fill(acc_o, 0) - T.fill(sumexp, 0) - T.fill(m_i, -(2**30)) # avoid -inf - inf to cause nan - - b_i, g_i = by, bz - s_i = bx if REPLICATE_H == 1 else (bx // REPLICATE_H) - q_i = s_i - max_kv_i = q_i - - H0 = g_i * padded_H + (0 if REPLICATE_H == 1 else (bx % REPLICATE_H) * 64) - H1 = H0 + H_per_block - - T.copy(Q[b_i, s_i, H0:H1, :D], Q_shared) - T.copy(Q[b_i, s_i, H0:H1, D:], Q_tail_shared) - - for i_i in T.Pipelined(NI, num_stages=num_stages): - - for bi_i in T.Parallel(BI): - mask[bi_i] = Indices[b_i, s_i, g_i, i_i * BI + bi_i] >= 0 - - for bi_i, d_i in T.Parallel(BI, D): - KV_shared[bi_i, d_i] = KV[ - b_i, Indices[b_i, s_i, g_i, i_i * BI + bi_i], g_i, d_i - ] - for bi_i, d_i in T.Parallel(BI, D_tail): - K_tail_shared[bi_i, d_i] = KV[ - b_i, Indices[b_i, s_i, g_i, i_i * BI + bi_i], g_i, D + d_i - ] - - for h_i, bi_i in T.Parallel(H_per_block, BI): - acc_s[h_i, bi_i] = T.if_then_else( - mask[bi_i], 0, -T.infinity(acc_s.dtype) - ) - T.gemm( - Q_shared, - KV_shared, - acc_s, - transpose_B=True, - policy=T.GemmWarpPolicy.FullCol, - ) - T.gemm( - Q_tail_shared, - K_tail_shared, - acc_s, - transpose_B=True, - policy=T.GemmWarpPolicy.FullCol, - ) - T.copy(m_i, m_i_prev) - T.reduce_max(acc_s, m_i, dim=1, clear=False) - for h_i in T.Parallel(H_per_block): - alpha[h_i] = T.exp2((m_i_prev[h_i] - m_i[h_i]) * sm_scale) - for h_i, bi_i in T.Parallel(H_per_block, BI): - acc_s[h_i, bi_i] = T.exp2( - acc_s[h_i, bi_i] * sm_scale - m_i[h_i] * sm_scale - ) - T.reduce_sum(acc_s, sumexp_i, dim=1) # is this a accumulate operator? - for h_i in T.Parallel(H_per_block): - sumexp[h_i] = sumexp[h_i] * alpha[h_i] + sumexp_i[h_i] - for h_i, d_i in T.Parallel(H_per_block, D): - acc_o[h_i, d_i] = acc_o[h_i, d_i] * alpha[h_i] - - T.copy(acc_s, S_shared) - T.gemm(S_shared, KV_shared, acc_o, policy=T.GemmWarpPolicy.FullCol) - - # Rescale - for h_i, d_i in T.Parallel(H_per_block, D): - acc_o[h_i, d_i] /= sumexp[h_i] - for h_i in T.Parallel(H_per_block): - sumexp[h_i] = T.log2(sumexp[h_i]) + m_i[h_i] * sm_scale - - T.copy(acc_o, O_shared) - T.copy(acc_o, Output[b_i, s_i, H0:H1, :]) - - return main - - -@tilelang.jit( - out_idx=[-1], - compile_flags=[ - "-O3", - "-Wno-deprecated-declarations", - "-U__CUDA_NO_HALF_OPERATORS__", - "-U__CUDA_NO_HALF_CONVERSIONS__", - "-U__CUDA_NO_HALF2_OPERATORS__", - "-U__CUDA_NO_BFLOAT16_CONVERSIONS__", - "--expt-relaxed-constexpr", - "--expt-extended-lambda", - "--ptxas-options=-v,--register-usage-level=10", - "-DNDEBUG", - ], -) # type: ignore -def sparse_attention_fwd_kernel_v2( - num_heads: int, - dim: int, - tail_dim: int, - topk: int, - *, - kv_group: int = 1, - sm_scale: Optional[float] = None, - block_I: int = 64, -): - assert dim == tilelang.math.next_power_of_2( - dim - ), f"haven't check padding correctness yet, dim={dim}" - assert tail_dim == tilelang.math.next_power_of_2( - tail_dim - ), f"haven't check padding correctness yet, dim={tail_dim}" - assert ( - topk % block_I == 0 - ), "otherwise will load some index=0 thus causing wrong kv to be loaded" - if sm_scale is None: - sm_scale = (1.0 / (dim + tail_dim)) ** 0.5 * 1.44269504 # log2(e) - else: - sm_scale = sm_scale * 1.44269504 # log2(e) - threads = 384 - - batch = T.symbolic("batch") - qo_len = T.symbolic("seq_len") - num_pages = T.symbolic("num_pages") - - q_shape = [batch, qo_len, num_heads, dim + tail_dim] - kv_shape = [batch, num_pages, kv_group, dim + tail_dim] - o_shape = [batch, qo_len, num_heads, dim] - indices_shape = [batch, qo_len, kv_group, topk] - - indices_dtype = "int32" - dtype = "bfloat16" - accum_dtype = "float" - - H = num_heads - padded_H = max(tilelang.math.next_power_of_2(num_heads), 16) - if padded_H != H: - assert kv_group == 1 - BI = block_I - NI = tilelang.cdiv(topk, block_I) - assert NI % 2 == 0, "NI should be a multiple of 2" - D = dim - D_tail = tail_dim - if num_heads > 64: - assert num_heads % 64 == 0, "head_kv should be a multiple of 64" - REPLICATE_H = num_heads // 64 - else: - REPLICATE_H = 1 - - H_per_block = padded_H if REPLICATE_H == 1 else 64 - - @T.prim_func - def main( - Q: T.Tensor(q_shape, dtype), # type: ignore - KV: T.Tensor(kv_shape, dtype), # type: ignore - Indices: T.Tensor(indices_shape, indices_dtype), # type: ignore - Output: T.Tensor(o_shape, dtype), # type: ignore - ): - """ - Q: [b, qo_len, H, D + D_tail] (bfloat16) - KV: [b, num_pages, kv_group, D + D_tail] (bfloat16) - Indices: [b, qo_len, kv_group, topk] (int32) - """ - - with T.Kernel(qo_len * REPLICATE_H, batch, 1, threads=threads) as (bx, by, bz): # type: ignore - Q_shared_l = T.alloc_shared([H_per_block, D // 2], dtype) - Q_shared_r = T.alloc_shared([H_per_block, D // 2], dtype) - Q_tail_shared = T.alloc_shared([H_per_block, D_tail], dtype) - KV_shared_0_l = T.alloc_shared([BI, D // 2], dtype) - KV_shared_0_r = T.alloc_shared([BI, D // 2], dtype) - KV_shared_1_l = T.alloc_shared([BI, D // 2], dtype) - KV_shared_1_r = T.alloc_shared([BI, D // 2], dtype) - K_tail_shared_0 = T.alloc_shared([BI, D_tail], dtype) - K_tail_shared_1 = T.alloc_shared([BI, D_tail], dtype) - O_shared_l = Q_shared_l - O_shared_r = Q_shared_r - is_kv_valid_0 = T.alloc_shared([BI], "bool", scope="shared") - is_kv_valid_1 = T.alloc_shared([BI], "bool", scope="shared") - - acc_o_l = T.alloc_fragment([H_per_block, D // 2], accum_dtype) - acc_o_r = T.alloc_fragment([H_per_block, D // 2], accum_dtype) - acc_s = T.alloc_fragment([H_per_block, BI], accum_dtype) - S_shared = T.alloc_shared([H_per_block, BI], dtype) - sumexp = T.alloc_fragment([H_per_block], accum_dtype) - sum_exp_shared = T.alloc_shared([H_per_block], accum_dtype) - sumexp_i = T.alloc_fragment([H_per_block], accum_dtype) - alpha_shared = T.alloc_shared([H_per_block], accum_dtype, scope="shared") - alpha_local = T.alloc_fragment([H_per_block], accum_dtype) - m_i = T.alloc_fragment([H_per_block], accum_dtype) - m_i_prev = T.alloc_fragment([H_per_block], accum_dtype) - indices_local = T.alloc_local([1], indices_dtype) - indices_tmp = T.alloc_local([1], indices_dtype) - - bar_q = T.alloc_barrier(arrive_count=384) - bar_k_0_ready = T.alloc_barrier(arrive_count=128) - bar_k_1_ready = T.alloc_barrier(arrive_count=128) - bar_k_0_free = T.alloc_barrier(arrive_count=256) - bar_k_1_free = T.alloc_barrier(arrive_count=256) - bar_sScale_and_sS_ready = T.alloc_barrier(arrive_count=256) - bar_sScale_and_sS_free = T.alloc_barrier(arrive_count=256) - - bar_0_128 = T.alloc_barrier(arrive_count=128) - bar_1_128 = T.alloc_barrier(arrive_count=128) - bar_2_128 = T.alloc_barrier(arrive_count=128) - bar_final = T.alloc_barrier(arrive_count=128) - - b_i, g_i = by, bz - s_i = bx if REPLICATE_H == 1 else bx // REPLICATE_H - - H0 = g_i * padded_H + (0 if REPLICATE_H == 1 else (bx % REPLICATE_H) * 64) - H1 = H0 + H_per_block - - tx = T.get_thread_binding() - - T.copy(Q[b_i, s_i, H0:H1, 0 : D // 2], Q_shared_l) - T.copy(Q[b_i, s_i, H0:H1, D // 2 : D], Q_shared_r) - T.copy(Q[b_i, s_i, H0:H1, D:], Q_tail_shared) - T.barrier_arrive(bar_q) - - if tx < 128: - T.set_max_nreg(240, 1) - T.fill(sumexp, 0) - T.fill(m_i, -(2**30)) # avoid -inf - inf to cause nan - T.fill(acc_o_l, 0) - T.barrier_wait(bar_q, 0) - - for i_i in T.serial(T.ceildiv(NI, 2)): - # Buffer 0 - # with sync_at(bar_0_128, 0): - T.barrier_wait(bar_k_0_ready[0], (i_i & 1)) - T.barrier_arrive(bar_0_128) - T.barrier_wait(bar_0_128, 0) - - for h_i, bi_i in T.Parallel(H_per_block, BI): - acc_s[h_i, bi_i] = T.if_then_else( - is_kv_valid_0[bi_i], 0, -T.infinity(acc_s.dtype) - ) - T.gemm( - Q_shared_l, KV_shared_0_l, acc_s, transpose_B=True, wg_wait=-1 - ) - T.gemm( - Q_shared_r, KV_shared_0_r, acc_s, transpose_B=True, wg_wait=-1 - ) - T.gemm( - Q_tail_shared, - K_tail_shared_0, - acc_s, - transpose_B=True, - wg_wait=-1, - ) - - T.wait_wgmma(0) - - if i_i != 0: - T.barrier_arrive(bar_sScale_and_sS_free) - T.barrier_wait(bar_sScale_and_sS_free, ((i_i * 2) & 1) ^ 1) - - T.copy(m_i, m_i_prev) - T.reduce_max(acc_s, m_i, dim=1, clear=False) - for h_i in T.Parallel(H_per_block): - alpha_local[h_i] = T.exp2((m_i_prev[h_i] - m_i[h_i]) * sm_scale) - for h_i, bi_i in T.Parallel(H_per_block, BI): - acc_s[h_i, bi_i] = T.exp2( - acc_s[h_i, bi_i] * sm_scale - m_i[h_i] * sm_scale - ) - T.reduce_sum( - acc_s, sumexp_i, dim=1 - ) # is this a accumulate operator? - for h_i in T.Parallel(H_per_block): - sumexp[h_i] = sumexp[h_i] * alpha_local[h_i] + sumexp_i[h_i] - for h_i, d_i in T.Parallel(H_per_block, D // 2): - acc_o_l[h_i, d_i] *= alpha_local[h_i] - T.copy(alpha_local, alpha_shared) - - T.copy(acc_s, S_shared) - T.gemm(S_shared, KV_shared_0_l, acc_o_l) - - T.barrier_arrive(bar_sScale_and_sS_ready) - T.barrier_arrive(bar_k_0_free[0]) - - # Buffer 1 - T.barrier_wait(bar_k_1_ready[0], (i_i & 1)) - T.barrier_arrive(bar_0_128) - T.barrier_wait(bar_0_128, 1) - - for h_i, bi_i in T.Parallel(H_per_block, BI): - acc_s[h_i, bi_i] = T.if_then_else( - is_kv_valid_1[bi_i], 0, -T.infinity(acc_s.dtype) - ) - T.gemm( - Q_shared_l, KV_shared_1_l, acc_s, transpose_B=True, wg_wait=-1 - ) - T.gemm( - Q_shared_r, KV_shared_1_r, acc_s, transpose_B=True, wg_wait=-1 - ) - T.gemm( - Q_tail_shared, - K_tail_shared_1, - acc_s, - transpose_B=True, - wg_wait=-1, - ) - - T.wait_wgmma(0) - - T.barrier_arrive(bar_sScale_and_sS_free) - T.barrier_wait(bar_sScale_and_sS_free, ((i_i * 2 + 1) & 1) ^ 1) - - T.copy(m_i, m_i_prev) - T.reduce_max(acc_s, m_i, dim=1, clear=False) - for h_i in T.Parallel(H_per_block): - alpha_local[h_i] = T.exp2((m_i_prev[h_i] - m_i[h_i]) * sm_scale) - for h_i, bi_i in T.Parallel(H_per_block, BI): - acc_s[h_i, bi_i] = T.exp2( - acc_s[h_i, bi_i] * sm_scale - m_i[h_i] * sm_scale - ) - T.reduce_sum( - acc_s, sumexp_i, dim=1 - ) # is this a accumulate operator? - for h_i in T.Parallel(H_per_block): - sumexp[h_i] = sumexp[h_i] * alpha_local[h_i] + sumexp_i[h_i] - for h_i, d_i in T.Parallel(H_per_block, D // 2): - acc_o_l[h_i, d_i] *= alpha_local[h_i] - T.copy(alpha_local, alpha_shared) - - T.copy(acc_s, S_shared) - T.gemm(S_shared, KV_shared_1_l, acc_o_l) - - T.barrier_arrive(bar_sScale_and_sS_ready) - T.barrier_arrive(bar_k_1_free[0]) - - # Rescale - for h_i in T.Parallel(H_per_block): - sum_exp_shared[h_i] = sumexp[h_i] - T.barrier_arrive(bar_final) - for h_i, d_i in T.Parallel(H_per_block, D // 2): - acc_o_l[h_i, d_i] /= sumexp[h_i] - for h_i in T.Parallel(H_per_block): - sumexp[h_i] = T.log2(sumexp[h_i]) + m_i[h_i] * sm_scale - T.copy(acc_o_l, O_shared_l) - T.copy(O_shared_l, Output[b_i, s_i, H0:H1, 0 : D // 2]) - elif tx >= 128 and tx < 256: - # T.set_max_nreg(168, 1) - T.fill(acc_o_r, 0) - for i_i in T.serial(T.ceildiv(NI, 2)): - # Buffer 0 - T.barrier_arrive(bar_sScale_and_sS_ready) - T.barrier_wait(bar_sScale_and_sS_ready, ((i_i * 2) & 1)) - T.barrier_arrive(bar_1_128) - T.barrier_wait(bar_1_128, 0) - for h_i, d_i in T.Parallel(H_per_block, D // 2): - acc_o_r[h_i, d_i] *= alpha_shared[h_i] - T.gemm(S_shared, KV_shared_0_r, acc_o_r) - T.barrier_arrive(bar_k_0_free[0]) - T.barrier_arrive(bar_sScale_and_sS_free) - - # Buffer 1 - T.barrier_arrive(bar_sScale_and_sS_ready) - T.barrier_wait(bar_sScale_and_sS_ready, ((i_i * 2 + 1) & 1)) - T.barrier_arrive(bar_1_128) - T.barrier_wait(bar_1_128, 1) - for h_i, d_i in T.Parallel(H_per_block, D // 2): - acc_o_r[h_i, d_i] *= alpha_shared[h_i] - T.gemm(S_shared, KV_shared_1_r, acc_o_r) - T.barrier_arrive(bar_k_1_free[0]) - if i_i != T.ceildiv(NI, 2) - 1: - T.barrier_arrive(bar_sScale_and_sS_free) - - # Rescale - T.barrier_wait(bar_final, 0) - for h_i, d_i in T.Parallel(H_per_block, D // 2): - acc_o_r[h_i, d_i] /= sum_exp_shared[h_i] - - T.copy(acc_o_r, O_shared_r) - T.copy(O_shared_r, Output[b_i, s_i, H0:H1, D // 2 : D]) - elif tx >= 256: - # producer - T.set_max_nreg(80, 0) - indices_local[0] = 0 - for i_i in T.serial(T.ceildiv(NI, 2)): - # Buffer 0 - T.barrier_wait(bar_k_0_free[0], ((i_i & 1) ^ 1)) - T.barrier_arrive(bar_2_128) - T.barrier_wait(bar_2_128, 0) - - for r in T.serial(4): - indices_tmp[0] = Indices[ - b_i, s_i, g_i, (i_i * 2) * BI + r * 16 + (tx - 256) // 8 - ] - is_kv_valid_0[r * 16 + (tx - 256) // 8] = indices_tmp[0] >= 0 - if is_kv_valid_0[r * 16 + (tx - 256) // 8]: - indices_local[0] = indices_tmp[0] - - with T.attr("default", "async_scope", 1): # type: ignore - for u in T.serial(4): - for v in T.vectorized(8): - KV_shared_0_l[ - r * 16 + (tx - 256) // 8, - 64 * u + (tx - 256) % 8 * 8 + v, - ] = KV[ - b_i, - indices_local[0], - g_i, - 64 * u + (tx - 256) % 8 * 8 + v, - ] - KV_shared_0_r[ - r * 16 + (tx - 256) // 8, - 64 * u + (tx - 256) % 8 * 8 + v, - ] = KV[ - b_i, - indices_local[0], - g_i, - D // 2 + 64 * u + (tx - 256) % 8 * 8 + v, - ] - with T.attr("default", "async_scope", 1): # type: ignore - for v in T.vectorized(8): - K_tail_shared_0[ - r * 16 + (tx - 256) // 8, (tx - 256) % 8 * 8 + v - ] = KV[ - b_i, - indices_local[0], - g_i, - D + (tx - 256) % 8 * 8 + v, - ] - - T.cp_async_barrier_noinc(bar_k_0_ready[0]) - - # Buffer 1 - T.barrier_wait(bar_k_1_free[0], ((i_i & 1) ^ 1)) - T.barrier_arrive(bar_2_128) - T.barrier_wait(bar_2_128, 1) - - for r in T.serial(4): - indices_tmp[0] = Indices[ - b_i, s_i, g_i, (i_i * 2 + 1) * BI + r * 16 + (tx - 256) // 8 - ] - is_kv_valid_1[r * 16 + (tx - 256) // 8] = indices_tmp[0] >= 0 - if is_kv_valid_1[r * 16 + (tx - 256) // 8]: - indices_local[0] = indices_tmp[0] - - with T.attr("default", "async_scope", 1): # type: ignore - for u in T.serial(4): - for v in T.vectorized(8): - KV_shared_1_l[ - r * 16 + (tx - 256) // 8, - 64 * u + (tx - 256) % 8 * 8 + v, - ] = KV[ - b_i, - indices_local[0], - g_i, - 64 * u + (tx - 256) % 8 * 8 + v, - ] - KV_shared_1_r[ - r * 16 + (tx - 256) // 8, - 64 * u + (tx - 256) % 8 * 8 + v, - ] = KV[ - b_i, - indices_local[0], - g_i, - D // 2 + 64 * u + (tx - 256) % 8 * 8 + v, - ] - with T.attr("default", "async_scope", 1): # type: ignore - for v in T.vectorized(8): - K_tail_shared_1[ - r * 16 + (tx - 256) // 8, (tx - 256) % 8 * 8 + v - ] = KV[ - b_i, - indices_local[0], - g_i, - D + (tx - 256) % 8 * 8 + v, - ] - - T.cp_async_barrier_noinc(bar_k_1_ready[0]) - - return main - - -@tilelang.jit( - out_idx=[-2, -1], - pass_configs={ - tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, - tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True, - }, -) -def sparse_mla_fwd_decode_partial( - heads, - dim, - tail_dim, - topk, - *, - kv_group=1, - sm_scale=None, - is_causal=True, - block_I=64, - inner_iter=1, - num_stages=1, - threads=256, -): - """ - grid: (seq_len * REPLICATE_H, top_k / block_I / inner_iter) - Each GPU block processes `inner_iter` consecutive KV tiles and writes one (partial_o, partial_lse) entry. - """ - - assert is_causal == True, "non-causal is not supported" - assert kv_group == 1 - assert topk % block_I == 0 - assert topk % (block_I * inner_iter) == 0, ( - f"topk ({topk}) must be divisible by block_I * inner_iter = " - f"{block_I} * {inner_iter}" - ) - - # log2(e) = 1.44269504 - if sm_scale is None: - sm_scale = (1.0 / (dim + tail_dim)) ** 0.5 * 1.44269504 - else: - sm_scale = sm_scale * 1.44269504 - - batch = 1 - seq_len = T.dynamic("seq_len") - seq_len_kv = T.dynamic("seq_len_kv") - - head_kv = heads // kv_group - padded_H = max(tilelang.math.next_power_of_2(head_kv), 16) - REPLICATE_H = (head_kv // 64) if head_kv > 64 else 1 - H_per_block = padded_H if REPLICATE_H == 1 else 64 - N_GROUPS = topk // (block_I * inner_iter) - BI = block_I - D = dim - D_tail = tail_dim - - q_shape = [batch, seq_len, heads, dim + tail_dim] - kv_shape = [batch, seq_len_kv, kv_group, dim + tail_dim] - indices_shape = [batch, seq_len, kv_group, topk] - partial_o_shape = [batch, seq_len, N_GROUPS, heads, dim] - partial_lse_shape = [batch, seq_len, N_GROUPS, heads] - indices_dtype = T.int32 - dtype = T.bfloat16 - accum_dtype = T.float32 - - _q_in_shared = inner_iter == 1 - - @T.prim_func - def main( - Q: T.Tensor(q_shape, dtype), - KV: T.Tensor(kv_shape, dtype), - Indices: T.Tensor(indices_shape, indices_dtype), - Partial_O: T.Tensor(partial_o_shape, dtype), - Partial_Lse: T.Tensor(partial_lse_shape, accum_dtype), - ): - with T.Kernel(seq_len * REPLICATE_H, N_GROUPS, threads=threads) as (bx, by): - if _q_in_shared: - Q_buf = T.alloc_shared([H_per_block, D], dtype) - Q_tail_buf = T.alloc_shared([H_per_block, D_tail], dtype) - else: - Q_buf = T.alloc_fragment([H_per_block, D], dtype) - Q_tail_buf = T.alloc_fragment([H_per_block, D_tail], dtype) - - KV_shared = T.alloc_shared([BI, D], dtype) - K_tail_shared = T.alloc_shared([BI, D_tail], dtype) - S_shared = T.alloc_shared([H_per_block, BI], dtype) - mask = T.alloc_fragment([BI], T.bool) - - acc_o = T.alloc_fragment([H_per_block, D], accum_dtype) - acc_s = T.alloc_fragment([H_per_block, BI], accum_dtype) - sumexp = T.alloc_fragment([H_per_block], accum_dtype) - sumexp_i = T.alloc_fragment([H_per_block], accum_dtype) - alpha = T.alloc_fragment([H_per_block], accum_dtype) - m_i = T.alloc_fragment([H_per_block], accum_dtype) - m_i_prev = T.alloc_fragment([H_per_block], accum_dtype) - - T.fill(acc_o, 0) - T.fill(sumexp, 0) - T.fill(m_i, -(2**30)) - - b_i, g_i = 0, 0 - s_i = bx if REPLICATE_H == 1 else (bx // REPLICATE_H) - group_i = by - H0 = 0 if REPLICATE_H == 1 else (bx % REPLICATE_H) * 64 - H1 = H0 + H_per_block - - T.copy(Q[b_i, s_i, H0:H1, :D], Q_buf) - T.copy(Q[b_i, s_i, H0:H1, D:], Q_tail_buf) - - for k_i in T.Pipelined(inner_iter, num_stages=num_stages): - topk_block_i = group_i * inner_iter + k_i - - for bi_i in T.Parallel(BI): - mask[bi_i] = Indices[b_i, s_i, g_i, topk_block_i * BI + bi_i] >= 0 - for bi_i, d_i in T.Parallel(BI, D): - idx = Indices[b_i, s_i, g_i, topk_block_i * BI + bi_i] - KV_shared[bi_i, d_i] = KV[ - b_i, T.if_then_else(idx >= 0, idx, 0), g_i, d_i - ] - for bi_i, d_i in T.Parallel(BI, D_tail): - idx = Indices[b_i, s_i, g_i, topk_block_i * BI + bi_i] - K_tail_shared[bi_i, d_i] = KV[ - b_i, T.if_then_else(idx >= 0, idx, 0), g_i, D + d_i - ] - - for h_i, bi_i in T.Parallel(H_per_block, BI): - acc_s[h_i, bi_i] = T.if_then_else( - mask[bi_i], 0, -T.infinity(acc_s.dtype) - ) - - T.gemm( - Q_buf, - KV_shared, - acc_s, - transpose_B=True, - policy=T.GemmWarpPolicy.FullCol, - ) - T.gemm( - Q_tail_buf, - K_tail_shared, - acc_s, - transpose_B=True, - policy=T.GemmWarpPolicy.FullCol, - ) - - T.copy(m_i, m_i_prev) - T.reduce_max(acc_s, m_i, dim=1, clear=False) - for h_i in T.Parallel(H_per_block): - alpha[h_i] = T.exp2((m_i_prev[h_i] - m_i[h_i]) * sm_scale) - for h_i, bi_i in T.Parallel(H_per_block, BI): - acc_s[h_i, bi_i] = T.exp2( - acc_s[h_i, bi_i] * sm_scale - m_i[h_i] * sm_scale - ) - T.reduce_sum(acc_s, sumexp_i, dim=1) - for h_i in T.Parallel(H_per_block): - sumexp[h_i] = sumexp[h_i] * alpha[h_i] + sumexp_i[h_i] - for h_i, d_i in T.Parallel(H_per_block, D): - acc_o[h_i, d_i] *= alpha[h_i] - - T.copy(acc_s, S_shared) - T.gemm(S_shared, KV_shared, acc_o, policy=T.GemmWarpPolicy.FullCol) - - # sumexp==0 (all masked), divide by 1 to get 0 and avoid nan - for h_i, d_i in T.Parallel(H_per_block, D): - acc_o[h_i, d_i] = acc_o[h_i, d_i] / T.if_then_else( - sumexp[h_i] == 0.0, 1.0, sumexp[h_i] - ) - # sumexp==0 (all masked), use large negative so combine ignores this split - for h_i in T.Parallel(H_per_block): - sumexp[h_i] = T.if_then_else( - sumexp[h_i] == 0.0, - -(2**30), - T.log2(sumexp[h_i]) + m_i[h_i] * sm_scale, - ) - - T.copy(acc_o, Partial_O[b_i, s_i, group_i, H0:H1, :]) - T.copy(sumexp, Partial_Lse[b_i, s_i, group_i, H0:H1]) - - return main - - -@tilelang.jit( - out_idx=[-1], - pass_configs={ - tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, - tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True, - }, -) -def sparse_mla_fwd_decode_combine( - heads, - dim, - topk, - head_per_block, - *, - block_I=64, - threads=256, -): - """ - grid: (seq_len * REPLICATE_H). batch=1, kv_group=1. - Each block does one tile of heads (e.g. 4 or 8 for decode). - """ - - assert heads % head_per_block == 0, f"head_per_block must divide heads" - - batch = 1 - seq_len = T.dynamic("seq_len") - - NI = topk // block_I - H_per_block = head_per_block - REPLICATE_H = heads // H_per_block - - partial_o_shape = [batch, seq_len, NI, heads, dim] - partial_lse_shape = [batch, seq_len, NI, heads] - o_shape = [batch, seq_len, heads, dim] - dtype = T.bfloat16 - accum_dtype = T.float32 - - @T.prim_func - def main( - Partial_O: T.Tensor(partial_o_shape, dtype), - Partial_Lse: T.Tensor(partial_lse_shape, accum_dtype), - Output: T.Tensor(o_shape, dtype), - ): - with T.Kernel(seq_len * REPLICATE_H, threads=threads) as (bx,): - shared_lse = T.alloc_shared([NI, H_per_block], accum_dtype) - - lse_max = T.alloc_fragment([H_per_block], accum_dtype) - lse_sum = T.alloc_fragment([H_per_block], accum_dtype) - scale = T.alloc_fragment([H_per_block, NI], accum_dtype) - acc_o = T.alloc_fragment([H_per_block, dim], accum_dtype) - - b_i = 0 - s_i = bx if REPLICATE_H == 1 else (bx // REPLICATE_H) - H0 = 0 if REPLICATE_H == 1 else (bx % REPLICATE_H) * H_per_block - H1 = H0 + H_per_block - - for k in T.serial(NI): - T.copy(Partial_Lse[b_i, s_i, k, H0:H1], shared_lse[k, :]) - - T.fill(lse_max, -(2**30)) - for k in T.serial(NI): - for h_i in T.Parallel(H_per_block): - lse_max[h_i] = T.max(lse_max[h_i], shared_lse[k, h_i]) - T.fill(lse_sum, 0) - for k in T.serial(NI): - for h_i in T.Parallel(H_per_block): - lse_sum[h_i] = lse_sum[h_i] + T.exp2( - shared_lse[k, h_i] - lse_max[h_i] - ) - for k in T.serial(NI): - for h_i in T.Parallel(H_per_block): - scale[h_i, k] = T.exp2( - shared_lse[k, h_i] - lse_max[h_i] - T.log2(lse_sum[h_i]) - ) - - T.fill(acc_o, 0) - for k in T.serial(NI): - for h_i, d_i in T.Parallel(H_per_block, dim): - acc_o[h_i, d_i] = acc_o[h_i, d_i] + scale[h_i, k] * Partial_O[ - b_i, s_i, k, H0 + h_i, d_i - ].astype(accum_dtype) - - T.copy(acc_o, Output[b_i, s_i, H0:H1, :]) - - return main - - -@tilelang.jit(out_idx=[-2, -1], pass_configs=pass_configs) -def sparse_mla_fwd_decode_partial_fp8( - num_heads: int, - d_v: int, - d_tail: int, - topk: int, - *, - sm_scale=None, - block_I=64, - inner_iter=1, - threads=256, -): - assert d_v == 512, f"only support d_v=512" - assert ( - topk % block_I == 0 - ), "otherwise will load some index=0 thus causing wrong kv to be loaded" - - # Softmax scores are in [0, 1]. We scale by fp8_max_val before FP8 cast - # to better utilize FP8 dynamic range, then apply the inverse scale after GEMM. - # This is numerically safe because softmax output is bounded by 1. - fp8_dtype = "float8_e4m3fnuz" if _is_fp8_fnuz else "float8_e4m3fn" - fp8_max_val = 240.0 if _is_fp8_fnuz else 448.0 - s_inv_scale_const = fp8_max_val - s_scale_const = 1.0 / fp8_max_val - - BI = block_I - group_size = 128 - dim_quant_fp8 = d_v + d_tail - rope_offset_fp8 = d_v - n_groups = topk // (BI * inner_iter) - - if sm_scale is None: - sm_scale = (1.0 / (d_v + d_tail)) ** 0.5 * 1.44269504 - else: - sm_scale = sm_scale * 1.44269504 - - h_per_block = 16 - # Match bf16 partial behavior: keep fixed 16-head tiles and use - # sliced T.copy on H0:H1 for tail handling. - assert ( - num_heads <= h_per_block or num_heads % h_per_block == 0 - ), "num_heads must be <=16 or divisible by 16" - head_blocks_per_seq = (num_heads + h_per_block - 1) // h_per_block - - batch = 1 - kv_group = 1 - seq_len = T.symbolic("seq_len") - num_pages = T.symbolic("num_pages") - - q_fp8_shape = [batch, seq_len, num_heads, d_v + d_tail] - kv_fp8_shape = [batch, num_pages, kv_group, dim_quant_fp8] - idx_shape = [batch, seq_len, kv_group, topk] - partial_o_shape = [batch, seq_len, n_groups, num_heads, d_v] - partial_lse_shape = [batch, seq_len, n_groups, num_heads] - - accum_dtype = T.float32 - dtype_bf16 = T.bfloat16 - - @T.prim_func - def main( - q_fp8: T.Tensor(q_fp8_shape, fp8_dtype), - kv_fp8: T.Tensor(kv_fp8_shape, fp8_dtype), - indices: T.Tensor(idx_shape, T.int32), - partial_o: T.Tensor(partial_o_shape, dtype_bf16), - partial_lse: T.Tensor(partial_lse_shape, accum_dtype), - ): - with T.Kernel(seq_len * head_blocks_per_seq, n_groups, threads=threads) as ( - bx, - by, - ): - b_i, g_i = 0, 0 - s_i = bx // head_blocks_per_seq - group_i = by - H0 = (bx % head_blocks_per_seq) * h_per_block - H1 = H0 + h_per_block - - # We intentionally split the K=512 GEMM into 4x128 tiles. - # Although this adds extra intermediate memory traffic, - # it shortens the MFMA accumulation dependency chain and improves performance. - q_tile0 = T.alloc_shared([h_per_block, group_size], fp8_dtype) - q_tile1 = T.alloc_shared([h_per_block, group_size], fp8_dtype) - q_tile2 = T.alloc_shared([h_per_block, group_size], fp8_dtype) - q_tile3 = T.alloc_shared([h_per_block, group_size], fp8_dtype) - kv_tile0 = T.alloc_shared([BI, group_size], fp8_dtype) - kv_tile1 = T.alloc_shared([BI, group_size], fp8_dtype) - kv_tile2 = T.alloc_shared([BI, group_size], fp8_dtype) - kv_tile3 = T.alloc_shared([BI, group_size], fp8_dtype) - q_tail_buf = T.alloc_shared([h_per_block, d_tail], fp8_dtype) - k_tail_shared = T.alloc_shared([BI, d_tail], fp8_dtype) - s_fp8_shared = T.alloc_shared([h_per_block, BI], fp8_dtype) - page_idx_shared = T.alloc_shared([BI], T.int32) - - mask = T.alloc_fragment([BI], T.bool) - acc_s = T.alloc_fragment([h_per_block, BI], accum_dtype) - acc_tile = T.alloc_fragment([h_per_block, BI], accum_dtype) - sv_tile = T.alloc_fragment([h_per_block, group_size], accum_dtype) - sumexp = T.alloc_fragment([h_per_block], accum_dtype) - sumexp_i = T.alloc_fragment([h_per_block], accum_dtype) - alpha = T.alloc_fragment([h_per_block], accum_dtype) - m_i = T.alloc_fragment([h_per_block], accum_dtype) - m_i_prev = T.alloc_fragment([h_per_block], accum_dtype) - inv_denom = T.alloc_fragment([h_per_block], accum_dtype) - - acc_o_tile0 = T.alloc_fragment([h_per_block, group_size], accum_dtype) - acc_o_tile1 = T.alloc_fragment([h_per_block, group_size], accum_dtype) - acc_o_tile2 = T.alloc_fragment([h_per_block, group_size], accum_dtype) - acc_o_tile3 = T.alloc_fragment([h_per_block, group_size], accum_dtype) - - T.fill(acc_o_tile0, 0) - T.fill(acc_o_tile1, 0) - T.fill(acc_o_tile2, 0) - T.fill(acc_o_tile3, 0) - T.fill(sumexp, 0) - T.fill(m_i, -(2**30)) - - T.copy(q_fp8[b_i, s_i, H0:H1, d_v:], q_tail_buf) - T.copy(q_fp8[b_i, s_i, H0:H1, 0 * group_size : 1 * group_size], q_tile0) - T.copy(q_fp8[b_i, s_i, H0:H1, 1 * group_size : 2 * group_size], q_tile1) - T.copy(q_fp8[b_i, s_i, H0:H1, 2 * group_size : 3 * group_size], q_tile2) - T.copy(q_fp8[b_i, s_i, H0:H1, 3 * group_size : 4 * group_size], q_tile3) - - for k_i in T.serial(inner_iter): - topk_block_i = group_i * inner_iter + k_i - - for bi_i in T.Parallel(BI): - idx = indices[b_i, s_i, g_i, topk_block_i * BI + bi_i] - valid = idx >= 0 - page_idx_shared[bi_i] = T.if_then_else(valid, idx, 0) - mask[bi_i] = valid - - for bi_i, j in T.Parallel(BI, group_size): - page = page_idx_shared[bi_i] - kv_tile0[bi_i, j] = kv_fp8[b_i, page, g_i, 0 * group_size + j] - kv_tile1[bi_i, j] = kv_fp8[b_i, page, g_i, 1 * group_size + j] - kv_tile2[bi_i, j] = kv_fp8[b_i, page, g_i, 2 * group_size + j] - kv_tile3[bi_i, j] = kv_fp8[b_i, page, g_i, 3 * group_size + j] - - for bi_i, j in T.Parallel(BI, d_tail): - page = page_idx_shared[bi_i] - k_tail_shared[bi_i, j] = kv_fp8[b_i, page, g_i, rope_offset_fp8 + j] - - for h_i, bi_i in T.Parallel(h_per_block, BI): - acc_s[h_i, bi_i] = T.if_then_else( - mask[bi_i], 0, -T.infinity(acc_s.dtype) - ) - - T.gemm(q_tile0, kv_tile0, acc_s, transpose_B=True, clear_accum=False) - T.gemm(q_tile1, kv_tile1, acc_tile, transpose_B=True, clear_accum=True) - for h_i, bi_i in T.Parallel(h_per_block, BI): - acc_s[h_i, bi_i] += acc_tile[h_i, bi_i] - T.gemm(q_tile2, kv_tile2, acc_tile, transpose_B=True, clear_accum=True) - for h_i, bi_i in T.Parallel(h_per_block, BI): - acc_s[h_i, bi_i] += acc_tile[h_i, bi_i] - T.gemm(q_tile3, kv_tile3, acc_tile, transpose_B=True, clear_accum=True) - for h_i, bi_i in T.Parallel(h_per_block, BI): - acc_s[h_i, bi_i] += acc_tile[h_i, bi_i] - T.gemm( - q_tail_buf, - k_tail_shared, - acc_s, - transpose_B=True, - policy=T.GemmWarpPolicy.FullCol, - ) - - T.copy(m_i, m_i_prev) - T.reduce_max(acc_s, m_i, dim=1, clear=False) - for h_i in T.Parallel(h_per_block): - alpha[h_i] = T.exp2((m_i_prev[h_i] - m_i[h_i]) * sm_scale) - for h_i, bi_i in T.Parallel(h_per_block, BI): - acc_s[h_i, bi_i] = T.exp2( - acc_s[h_i, bi_i] * sm_scale - m_i[h_i] * sm_scale - ) - T.reduce_sum(acc_s, sumexp_i, dim=1) - for h_i in T.Parallel(h_per_block): - sumexp[h_i] = sumexp[h_i] * alpha[h_i] + sumexp_i[h_i] - for h_i, j in T.Parallel(h_per_block, group_size): - acc_o_tile0[h_i, j] = acc_o_tile0[h_i, j] * alpha[h_i] - acc_o_tile1[h_i, j] = acc_o_tile1[h_i, j] * alpha[h_i] - acc_o_tile2[h_i, j] = acc_o_tile2[h_i, j] * alpha[h_i] - acc_o_tile3[h_i, j] = acc_o_tile3[h_i, j] * alpha[h_i] - - for h_i, bi_i in T.Parallel(h_per_block, BI): - s_fp8_shared[h_i, bi_i] = T.clamp( - acc_s[h_i, bi_i] * s_inv_scale_const, - -fp8_max_val, - fp8_max_val, - ) - T.gemm(s_fp8_shared, kv_tile0, sv_tile, clear_accum=True) - for h_i, j in T.Parallel(h_per_block, group_size): - acc_o_tile0[h_i, j] = ( - acc_o_tile0[h_i, j] + sv_tile[h_i, j] * s_scale_const - ) - - T.gemm(s_fp8_shared, kv_tile1, sv_tile, clear_accum=True) - for h_i, j in T.Parallel(h_per_block, group_size): - acc_o_tile1[h_i, j] = ( - acc_o_tile1[h_i, j] + sv_tile[h_i, j] * s_scale_const - ) - - T.gemm(s_fp8_shared, kv_tile2, sv_tile, clear_accum=True) - for h_i, j in T.Parallel(h_per_block, group_size): - acc_o_tile2[h_i, j] = ( - acc_o_tile2[h_i, j] + sv_tile[h_i, j] * s_scale_const - ) - - T.gemm(s_fp8_shared, kv_tile3, sv_tile, clear_accum=True) - for h_i, j in T.Parallel(h_per_block, group_size): - acc_o_tile3[h_i, j] = ( - acc_o_tile3[h_i, j] + sv_tile[h_i, j] * s_scale_const - ) - - for h_i in T.Parallel(h_per_block): - denom = T.if_then_else(sumexp[h_i] == 0.0, 1.0, sumexp[h_i]) - inv_denom[h_i] = 1.0 / denom - for h_i, j in T.Parallel(h_per_block, group_size): - acc_o_tile0[h_i, j] = acc_o_tile0[h_i, j] * inv_denom[h_i] - acc_o_tile1[h_i, j] = acc_o_tile1[h_i, j] * inv_denom[h_i] - acc_o_tile2[h_i, j] = acc_o_tile2[h_i, j] * inv_denom[h_i] - acc_o_tile3[h_i, j] = acc_o_tile3[h_i, j] * inv_denom[h_i] - - for h_i in T.Parallel(h_per_block): - sumexp[h_i] = T.if_then_else( - sumexp[h_i] == 0.0, - -(2**30), - T.log2(sumexp[h_i]) + m_i[h_i] * sm_scale, - ) - - T.copy( - acc_o_tile0, - partial_o[b_i, s_i, group_i, H0:H1, 0 * group_size : 1 * group_size], - ) - T.copy( - acc_o_tile1, - partial_o[b_i, s_i, group_i, H0:H1, 1 * group_size : 2 * group_size], - ) - T.copy( - acc_o_tile2, - partial_o[b_i, s_i, group_i, H0:H1, 2 * group_size : 3 * group_size], - ) - T.copy( - acc_o_tile3, - partial_o[b_i, s_i, group_i, H0:H1, 3 * group_size : 4 * group_size], - ) - - T.copy(sumexp, partial_lse[b_i, s_i, group_i, H0:H1]) - - return main - - -def tilelang_sparse_fwd( - q: torch.Tensor, - kv: torch.Tensor, - indices: torch.Tensor, - sm_scale: float, - d_v: int = 512, -) -> torch.Tensor: - assert q.dim() == 3 and kv.dim() == 3 and indices.dim() == 3 - num_heads = q.shape[1] - dim = q.shape[2] - tail_dim = dim - d_v - topk = indices.shape[-1] - assert topk == 2048 - - if _is_hip: - is_fp8_kv = kv.dtype in (torch.float8_e4m3fn, torch.float8_e4m3fnuz) - if is_fp8_kv: - if q.dtype != kv.dtype: - q = q.to(kv.dtype) - if _is_gfx95_supported: - block_I, threads, block_per_cu, cu = 64, 256, 2, 256 - else: - block_I, threads, block_per_cu, cu = 64, 256, 1, 304 - ni = topk // block_I - inner_iter = _pick_inner_iter(q.shape[0], ni, cu, block_per_cu) - kernel_partial = sparse_mla_fwd_decode_partial_fp8( - num_heads, - d_v, - tail_dim, - topk, - sm_scale=sm_scale, - block_I=block_I, - inner_iter=inner_iter, - threads=threads, - ) - else: - if _is_gfx95_supported: - block_I, threads, block_per_cu, cu = 64, 256, 2, 256 - else: - block_I, threads, block_per_cu, cu = 32, 128, 1, 304 - ni = topk // block_I - inner_iter = _pick_inner_iter(q.shape[0], ni, cu, block_per_cu) - kernel_partial = sparse_mla_fwd_decode_partial( - num_heads, - d_v, - tail_dim, - topk, - sm_scale=sm_scale, - block_I=block_I, - inner_iter=inner_iter, - threads=threads, - ) - partial_o_batched, partial_lse_batched = kernel_partial( - q.unsqueeze(0), kv.unsqueeze(0), indices.unsqueeze(0) - ) - n_groups = ni // inner_iter - kernel_combine = sparse_mla_fwd_decode_combine( - num_heads, - d_v, - n_groups * block_I, - head_per_block=4, - block_I=block_I, - threads=threads, - ) - out = kernel_combine(partial_o_batched, partial_lse_batched) - else: - kernel = sparse_attention_fwd_kernel_v2( - num_heads, d_v, tail_dim, topk, sm_scale=sm_scale - ) - out = kernel(q.unsqueeze(0), kv.unsqueeze(0), indices.unsqueeze(0)) # type: ignore - return out - - -@functools.cache -def fp8_paged_mqa_logits_kernel( - head_dim: int = 128, - num_heads: int = 64, - block_size: int = 64, - clear_accum: bool = True, - split_kv: int = 1, -) -> Any: - N = T.symbolic("batch_size") - L = T.symbolic("max_table_length") - S = T.symbolic("max_seq_len") - C = T.symbolic("num_blocks") - B = block_size - D = head_dim - H = num_heads - SK = int(split_kv) - BLOCK_BYTES = B * (D + 4) - SCALE_OFFSET = B * D - - assert D % 4 == 0 - assert H % 4 == 0 - assert D == 128 - assert SK >= 1 - - @tilelang.jit( - pass_configs={ - **pass_configs, - tilelang.PassConfigKey.TL_DISABLE_SAFE_MEMORY_ACCESS: True, - } - ) - def fp8_paged_mqa_logits( - q: T.Tensor[(N, H, D), FP8], - kvcache_u8: T.Tensor[(C, BLOCK_BYTES), UINT8], - weight: T.Tensor[(N, H), FP32], - seq_lens: T.Tensor[(N,), INT32], - page_table: T.Tensor[(N, L), INT32], - o: T.Tensor[(N, S), FP32], - ) -> None: - _ = N, L, S, C, D, H, B - with T.Kernel(N * SK) as bxs: - bx = bxs % N - pid_split = bxs // N - seq_len = seq_lens[bx] - np_total = T.ceildiv(seq_len, B) - stride = T.ceildiv(np_total, SK) - i_start = pid_split * stride - n_iters = T.max(0, T.min(stride, np_total - i_start)) - - q_smem = T.alloc_shared((H, D), FP8) - q_s_frag = T.alloc_fragment((H,), FP32) - T.copy(q[bx, 0, 0], q_smem) - T.copy(weight[bx, 0], q_s_frag) - - for j in T.Pipelined(n_iters, num_stages=2): - i = i_start + j - page = page_table[bx, i] - k_smem_u8 = T.alloc_shared((B * D,), UINT8) - T.copy(kvcache_u8[page, 0:SCALE_OFFSET], k_smem_u8) - k_smem = T.view(k_smem_u8, (B, D), FP8) - k_s_smem_u8 = T.alloc_shared((B * 4,), UINT8) - T.copy(kvcache_u8[page, SCALE_OFFSET:BLOCK_BYTES], k_s_smem_u8) - k_s_smem = T.view(k_s_smem_u8, (B,), FP32) - k_s_frag = T.alloc_fragment((B,), FP32) - T.copy(k_s_smem, k_s_frag) - - logits = T.alloc_fragment((B, H), FP32) - if not clear_accum: - T.fill(logits, 0.0) - T.gemm( - k_smem, - q_smem, - logits, - transpose_A=False, - transpose_B=True, - clear_accum=clear_accum, - ) - - # post processing - for h, j2 in T.Parallel(H, B): - logits[j2, h] = T.max(logits[j2, h], 0.0) * q_s_frag[h] - logits_sum = T.alloc_fragment((B,), FP32) - T.reduce_sum(logits, logits_sum, dim=1) - for j2 in T.Parallel(B): - logits_sum[j2] *= k_s_frag[j2] - T.copy(logits_sum, o[bx, i * B]) - - return fp8_paged_mqa_logits - - -def tilelang_fp8_paged_mqa_logits( - q_fp8: torch.Tensor, - kvcache_fp8: torch.Tensor, - weight: torch.Tensor, - seq_lens: torch.Tensor, - page_table: torch.Tensor, - deep_gemm_metadata: Any, - max_seq_len: int, - clean_logits: bool = True, -) -> torch.Tensor: - _ = deep_gemm_metadata - batch_size, _, num_heads, head_dim = q_fp8.shape - block_size = kvcache_fp8.shape[1] - assert head_dim == 128, "TODO" - assert block_size == 64, "TODO" - assert q_fp8.shape == (batch_size, 1, num_heads, head_dim) - assert kvcache_fp8.shape[1:] == (block_size, 1, head_dim + 4) - assert weight.shape == (batch_size, num_heads) - assert seq_lens.shape == (batch_size,) - assert page_table.shape[0] == batch_size - assert clean_logits == False - - logits = page_table.new_empty((batch_size, max_seq_len), dtype=torch.float32) - - NUM_CU = 256 - split_kv = split_kv = max(1, min(max_seq_len // block_size, NUM_CU // batch_size)) - kernel = fp8_paged_mqa_logits_kernel( - head_dim=head_dim, - num_heads=num_heads, - block_size=block_size, - clear_accum=clean_logits, - split_kv=split_kv, - ) - q_fp8 = q_fp8.view(batch_size, num_heads, head_dim) - kvcache_u8 = kvcache_fp8.view(-1, block_size * (head_dim + 4)) - kernel(q_fp8, kvcache_u8, weight, seq_lens, page_table, logits) - return logits - - -def _build_fp8_combined_view(k_cache: torch.Tensor) -> Tuple[torch.Tensor, int, int]: - """ - Reinterpret a MODEL1_FP8Sparse KV cache as a contiguous uint32 view. - Input: k_cache (num_blocks, block_size, 1, d_qk) fp8/uint8 - — per-block storage also holds scales + padding past d_qk. - Output: (num_blocks, block_pad_u32) uint32 covering the full block - stride. Same storage ashe input, no copy. - """ - k_u8 = k_cache.view(torch.uint8) if k_cache.dtype != torch.uint8 else k_cache - num_blocks = k_u8.shape[0] - block_size = k_u8.shape[1] - block_pad_u32 = k_u8.stride(0) // 4 - storage = k_u8.untyped_storage() - flat_u32 = torch.empty(0, dtype=torch.uint32, device=k_u8.device).set_( - storage, 0, (storage.nbytes() // 4,), (1,) - ) - k_combined = torch.as_strided( - flat_u32, - size=(num_blocks, block_pad_u32), - stride=(block_pad_u32, 1), - storage_offset=k_u8.storage_offset() // 4, - ) - return k_combined, num_blocks, block_size - - -_TOPK_LEN_SENTINEL_CACHE: dict = {} -_INT32_MAX = 2**30 - - -def _topk_length_sentinel(device: torch.device, batch: int) -> torch.Tensor: - """Cached `(batch,) int32 INT_MAX` tensor used when `topk_length` is None.""" - cur = _TOPK_LEN_SENTINEL_CACHE.get(device) - if cur is None or cur.numel() < batch: - cur = torch.full( - (max(batch, 256),), _INT32_MAX, dtype=torch.int32, device=device - ) - _TOPK_LEN_SENTINEL_CACHE[device] = cur - return cur[:batch] - - -@tilelang.jit( - out_idx=[-2, -1], - pass_configs={ - tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, - tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True, - tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True, - }, -) -def dpsk_v4_fp8_partial_kernel( - num_heads: int, - topk_1: int, - block_size_kv_1: int, - topk_2: int = 0, - block_size_kv_2: int = 0, - *, - dim: int = 448, - tail_dim: int = 64, - sm_scale: float = 0.0, - block_I: int = 64, - inner_iter_1: int = 1, - inner_iter_2: int = 0, - num_stages: int = 0, - threads: int = 512, -) -> Any: - """ - Read FP8 K cache directly, dequantise to BF16 in-kernel, do flash-attn - online softmax with split-K. Supports a second cache (`topk_2>0`) and - `attn_sink` is folded later by the combine kernel. - """ - log2e: float = 1.44269504 - if sm_scale <= 0.0: - sm_scale = (1.0 / (dim + tail_dim)) ** 0.5 * log2e - else: - sm_scale = sm_scale * log2e - assert dim == 448 and tail_dim == 64 - assert topk_1 % block_I == 0 - assert ( - topk_1 // block_I - ) % inner_iter_1 == 0, ( - f"NI_1={topk_1 // block_I} must be divisible by inner_iter_1={inner_iter_1}" - ) - assert block_size_kv_1 > 0 and (block_size_kv_1 & (block_size_kv_1 - 1)) == 0 - - is_dual = topk_2 > 0 - if is_dual: - assert inner_iter_2 > 0, "dual-cache call requires inner_iter_2 > 0" - assert topk_2 % block_I == 0 - assert ( - topk_2 // block_I - ) % inner_iter_2 == 0, ( - f"NI_2={topk_2 // block_I} must be divisible by inner_iter_2={inner_iter_2}" - ) - assert block_size_kv_2 > 0 and (block_size_kv_2 & (block_size_kv_2 - 1)) == 0 - - PACKED_W = dim + 2 * tail_dim - NOPE_TILE = 64 - NUM_TILES = dim // NOPE_TILE - SCALE_W = 8 - PACKED_W4 = PACKED_W // 4 - SCALE_W4 = SCALE_W // 4 - - kv_group = 1 - batch = T.symbolic("batch") - seq_len = T.symbolic("seq_len") - num_blocks_kv_1 = T.symbolic("num_blocks_kv_1") - block_pad_u32_1 = T.symbolic("block_pad_u32_1") - if is_dual: - num_blocks_kv_2 = T.symbolic("num_blocks_kv_2") - block_pad_u32_2 = T.symbolic("block_pad_u32_2") - - head_kv = num_heads // kv_group - D = dim - D_tail = tail_dim - BI = block_I - padded_H = max(tilelang.math.next_power_of_2(head_kv), 16) - if head_kv > 64: - assert head_kv % 64 == 0 - REPLICATE_H = (head_kv + 63) // 64 if head_kv > 64 else 1 - H_per_block = 64 if REPLICATE_H > 1 else padded_H - - NI_1 = topk_1 // BI - n_groups_1 = NI_1 // inner_iter_1 - NI_2 = (topk_2 // BI) if is_dual else 0 - n_groups_2 = (NI_2 // inner_iter_2) if is_dual else 0 - n_groups = n_groups_1 + n_groups_2 - - BS_KV_1 = block_size_kv_1 - NOPE_ROPE_U32_PER_BLOCK_1 = BS_KV_1 * PACKED_W4 - if is_dual: - BS_KV_2 = block_size_kv_2 - NOPE_ROPE_U32_PER_BLOCK_2 = BS_KV_2 * PACKED_W4 - - q_shape = [batch, seq_len, num_heads, D + D_tail] - k1_shape = [num_blocks_kv_1, block_pad_u32_1] - indices1_shape = [batch, seq_len, topk_1] - topk_length_shape = [batch] - partial_o_shape = [batch, seq_len, n_groups, num_heads, D + D_tail] - partial_lse_shape = [batch, seq_len, n_groups, num_heads] - if is_dual: - k2_shape = [num_blocks_kv_2, block_pad_u32_2] - indices2_shape = [batch, seq_len, topk_2] - - accum_dtype = "float" - indices_dtype = INT32 - - if is_dual: - - @T.prim_func - def main( - Q: T.Tensor(q_shape, BF16), # type: ignore - K_combined_1: T.Tensor(k1_shape, "uint32"), # type: ignore - Indices_1: T.Tensor(indices1_shape, indices_dtype), # type: ignore - Topk_length_1: T.Tensor(topk_length_shape, indices_dtype), # type: ignore - K_combined_2: T.Tensor(k2_shape, "uint32"), # type: ignore - Indices_2: T.Tensor(indices2_shape, indices_dtype), # type: ignore - Topk_length_2: T.Tensor(topk_length_shape, indices_dtype), # type: ignore - Partial_O: T.Tensor(partial_o_shape, BF16), # type: ignore - Partial_LSE: T.Tensor(partial_lse_shape, accum_dtype), # type: ignore - ) -> None: - """ - grid: (seq_len * REPLICATE_H * n_groups, batch, 1) - Each block processes `inner_iter_1` (or `inner_iter_2`) consecutive - KV tiles of one phase and writes one (partial_o, partial_lse) entry. - """ - with T.Kernel( - seq_len * REPLICATE_H * n_groups, batch, kv_group, threads=threads - ) as (bx, by, bz): - Q_shared = T.alloc_fragment([H_per_block, D], BF16) - Q_tail_shared = T.alloc_fragment([H_per_block, D_tail], BF16) - K_packed_shared = T.alloc_shared([BI, PACKED_W4], "uint32") - K_scale_shared = T.alloc_shared([BI, SCALE_W4], "uint32") - KV_shared = T.alloc_shared([BI, D], BF16) - K_tail_shared = T.alloc_shared([BI, D_tail], BF16) - S_shared = T.alloc_shared([H_per_block, BI], BF16) - page_idx_shared = T.alloc_shared([BI], INT32) - - mask = T.alloc_fragment([BI], "bool") - scale_byte_local = T.alloc_fragment([BI, NUM_TILES], "uint32") - - acc_o = T.alloc_fragment([H_per_block, D], accum_dtype) - acc_o_tail = T.alloc_fragment([H_per_block, D_tail], accum_dtype) - acc_s = T.alloc_fragment([H_per_block, BI], accum_dtype) - sumexp = T.alloc_fragment([H_per_block], accum_dtype) - sumexp_i = T.alloc_fragment([H_per_block], accum_dtype) - alpha = T.alloc_fragment([H_per_block], accum_dtype) - m_i = T.alloc_fragment([H_per_block], accum_dtype) - m_i_prev = T.alloc_fragment([H_per_block], accum_dtype) - - T.fill(acc_o, 0) - T.fill(acc_o_tail, 0) - T.fill(sumexp, 0) - T.fill(m_i, -(2**30)) - - b_i, g_i = by, bz - # bx encodes (s_i, h_replicate, group_i). - spans_per_seq = REPLICATE_H * n_groups - s_i = bx // spans_per_seq - rest = bx % spans_per_seq - group_i = rest // REPLICATE_H - h_rep = rest % REPLICATE_H - H0 = g_i * padded_H + (0 if REPLICATE_H == 1 else h_rep * 64) - H1 = H0 + H_per_block - - tk_len_1 = Topk_length_1[b_i] - tk_len_2 = Topk_length_2[b_i] - actual_n_groups_1 = T.ceildiv(tk_len_1, BI * inner_iter_1) - actual_n_groups_2 = T.ceildiv(tk_len_2, BI * inner_iter_2) - - if (group_i < n_groups_1) & (group_i < actual_n_groups_1): - # Phase 1 active: SWA cache work + Partial_O write. - T.copy(Q[b_i, s_i, H0:H1, :D], Q_shared) - T.copy(Q[b_i, s_i, H0:H1, D : D + D_tail], Q_tail_shared) - for k_i in T.Pipelined(inner_iter_1, num_stages=num_stages): - iter_i = group_i * inner_iter_1 + k_i - for bi_i in T.Parallel(BI): - pos = iter_i * BI + bi_i - idx = Indices_1[b_i, s_i, pos] - valid = (idx >= 0) & (pos < tk_len_1) - page_idx_shared[bi_i] = T.if_then_else(valid, idx, 0) - mask[bi_i] = valid - - for bi_i, w_i in T.Parallel(BI, PACKED_W4): - page = page_idx_shared[bi_i] - block_id = page // BS_KV_1 - t_in_block = page % BS_KV_1 - K_packed_shared[bi_i, w_i] = K_combined_1[ - block_id, t_in_block * PACKED_W4 + w_i - ] - - for bi_i, w_i in T.Parallel(BI, SCALE_W4): - page = page_idx_shared[bi_i] - block_id = page // BS_KV_1 - t_in_block = page % BS_KV_1 - K_scale_shared[bi_i, w_i] = K_combined_1[ - block_id, - NOPE_ROPE_U32_PER_BLOCK_1 + t_in_block * SCALE_W4 + w_i, - ] - - for bi_i, ti in T.Parallel(BI, NUM_TILES): - word_idx = ti // 4 - byte_in_word = ti % 4 - word = K_scale_shared[bi_i, word_idx] - scale_byte_local[bi_i, ti] = ( - word >> T.Cast("uint32", byte_in_word * 8) - ) & T.uint32(0xFF) - - for bi_i, d_i in T.Parallel(BI, D): - word_idx = d_i // 4 - byte_in_word = d_i % 4 - word = K_packed_shared[bi_i, word_idx] - b_u32 = ( - word >> T.Cast("uint32", byte_in_word * 8) - ) & T.uint32(0xFF) - sign_bf = (b_u32 & T.uint32(0x80)) * T.uint32(0x100) - exp_e4 = (b_u32 & T.uint32(0x78)) >> T.uint32(3) - mant_bf = (b_u32 & T.uint32(0x7)) * T.uint32(0x10) - scale_byte = scale_byte_local[bi_i, d_i // NOPE_TILE] - exp_combined = exp_e4 + scale_byte - T.uint32(7) - bf16_bits = ( - sign_bf | (exp_combined << T.uint32(7)) | mant_bf - ) - KV_shared[bi_i, d_i] = T.reinterpret( - BF16, T.Cast("uint16", bf16_bits) - ) - - for bi_i, j in T.Parallel(BI, D_tail): - abs_off = D + 2 * j - word_idx = abs_off // 4 - word_off = abs_off % 4 - word = K_packed_shared[bi_i, word_idx] - half_u32 = T.if_then_else( - word_off == 0, - word & T.uint32(0xFFFF), - (word >> T.uint32(16)) & T.uint32(0xFFFF), - ) - K_tail_shared[bi_i, j] = T.reinterpret( - BF16, T.Cast("uint16", half_u32) - ) - - for h_i, bi_i in T.Parallel(H_per_block, BI): - acc_s[h_i, bi_i] = T.if_then_else( - mask[bi_i], 0, -T.infinity(acc_s.dtype) - ) - T.gemm( - Q_shared, - KV_shared, - acc_s, - transpose_B=True, - policy=T.GemmWarpPolicy.FullRow, - ) - T.gemm( - Q_tail_shared, - K_tail_shared, - acc_s, - transpose_B=True, - policy=T.GemmWarpPolicy.FullRow, - ) - T.copy(m_i, m_i_prev) - T.reduce_max(acc_s, m_i, dim=1, clear=False) - for h_i in T.Parallel(H_per_block): - m_i[h_i] = T.max(m_i[h_i], m_i_prev[h_i]) - for h_i in T.Parallel(H_per_block): - alpha[h_i] = T.exp2((m_i_prev[h_i] - m_i[h_i]) * sm_scale) - for h_i, bi_i in T.Parallel(H_per_block, BI): - acc_s[h_i, bi_i] = T.exp2( - acc_s[h_i, bi_i] * sm_scale - m_i[h_i] * sm_scale - ) - T.reduce_sum(acc_s, sumexp_i, dim=1) - for h_i in T.Parallel(H_per_block): - sumexp[h_i] = sumexp[h_i] * alpha[h_i] + sumexp_i[h_i] - for h_i, d_i in T.Parallel(H_per_block, D): - acc_o[h_i, d_i] *= alpha[h_i] - for h_i, d_i in T.Parallel(H_per_block, D_tail): - acc_o_tail[h_i, d_i] *= alpha[h_i] - T.copy(acc_s, S_shared) - T.gemm( - S_shared, - KV_shared, - acc_o, - policy=T.GemmWarpPolicy.FullRow, - ) - T.gemm( - S_shared, - K_tail_shared, - acc_o_tail, - policy=T.GemmWarpPolicy.FullRow, - ) - # ---- finalize phase 1 (active) ---- - for h_i, d_i in T.Parallel(H_per_block, D): - acc_o[h_i, d_i] = acc_o[h_i, d_i] / T.if_then_else( - sumexp[h_i] == 0.0, 1.0, sumexp[h_i] - ) - for h_i, d_i in T.Parallel(H_per_block, D_tail): - acc_o_tail[h_i, d_i] = acc_o_tail[h_i, d_i] / T.if_then_else( - sumexp[h_i] == 0.0, 1.0, sumexp[h_i] - ) - for h_i in T.Parallel(H_per_block): - m_i[h_i] = T.if_then_else( - sumexp[h_i] == 0.0, - -(2.0**30), - T.log2(sumexp[h_i]) + m_i[h_i] * sm_scale, - ) - T.copy(acc_o, Partial_O[b_i, s_i, group_i, H0:H1, :D]) - T.copy( - acc_o_tail, - Partial_O[b_i, s_i, group_i, H0:H1, D : D + D_tail], - ) - T.copy(m_i, Partial_LSE[b_i, s_i, group_i, H0:H1]) - elif group_i < n_groups_1: - # Phase 1 skipped: m_i is still the -2^30 - T.copy(m_i, Partial_LSE[b_i, s_i, group_i, H0:H1]) - elif (group_i - n_groups_1) < actual_n_groups_2: - # Phase 2 active: c128 cache work + Partial_O write. - T.copy(Q[b_i, s_i, H0:H1, :D], Q_shared) - T.copy(Q[b_i, s_i, H0:H1, D : D + D_tail], Q_tail_shared) - for k_i in T.Pipelined(inner_iter_2, num_stages=num_stages): - iter_i = (group_i - n_groups_1) * inner_iter_2 + k_i - for bi_i in T.Parallel(BI): - pos = iter_i * BI + bi_i - idx = Indices_2[b_i, s_i, pos] - valid = (idx >= 0) & (pos < tk_len_2) - page_idx_shared[bi_i] = T.if_then_else(valid, idx, 0) - mask[bi_i] = valid - - for bi_i, w_i in T.Parallel(BI, PACKED_W4): - page = page_idx_shared[bi_i] - block_id = page // BS_KV_2 - t_in_block = page % BS_KV_2 - K_packed_shared[bi_i, w_i] = K_combined_2[ - block_id, t_in_block * PACKED_W4 + w_i - ] - - for bi_i, w_i in T.Parallel(BI, SCALE_W4): - page = page_idx_shared[bi_i] - block_id = page // BS_KV_2 - t_in_block = page % BS_KV_2 - K_scale_shared[bi_i, w_i] = K_combined_2[ - block_id, - NOPE_ROPE_U32_PER_BLOCK_2 + t_in_block * SCALE_W4 + w_i, - ] - - for bi_i, ti in T.Parallel(BI, NUM_TILES): - word_idx = ti // 4 - byte_in_word = ti % 4 - word = K_scale_shared[bi_i, word_idx] - scale_byte_local[bi_i, ti] = ( - word >> T.Cast("uint32", byte_in_word * 8) - ) & T.uint32(0xFF) - - for bi_i, d_i in T.Parallel(BI, D): - word_idx = d_i // 4 - byte_in_word = d_i % 4 - word = K_packed_shared[bi_i, word_idx] - b_u32 = ( - word >> T.Cast("uint32", byte_in_word * 8) - ) & T.uint32(0xFF) - sign_bf = (b_u32 & T.uint32(0x80)) * T.uint32(0x100) - exp_e4 = (b_u32 & T.uint32(0x78)) >> T.uint32(3) - mant_bf = (b_u32 & T.uint32(0x7)) * T.uint32(0x10) - scale_byte = scale_byte_local[bi_i, d_i // NOPE_TILE] - exp_combined = exp_e4 + scale_byte - T.uint32(7) - bf16_bits = ( - sign_bf | (exp_combined << T.uint32(7)) | mant_bf - ) - KV_shared[bi_i, d_i] = T.reinterpret( - BF16, T.Cast("uint16", bf16_bits) - ) - - for bi_i, j in T.Parallel(BI, D_tail): - abs_off = D + 2 * j - word_idx = abs_off // 4 - word_off = abs_off % 4 - word = K_packed_shared[bi_i, word_idx] - half_u32 = T.if_then_else( - word_off == 0, - word & T.uint32(0xFFFF), - (word >> T.uint32(16)) & T.uint32(0xFFFF), - ) - K_tail_shared[bi_i, j] = T.reinterpret( - BF16, T.Cast("uint16", half_u32) - ) - - for h_i, bi_i in T.Parallel(H_per_block, BI): - acc_s[h_i, bi_i] = T.if_then_else( - mask[bi_i], 0, -T.infinity(acc_s.dtype) - ) - T.gemm( - Q_shared, - KV_shared, - acc_s, - transpose_B=True, - policy=T.GemmWarpPolicy.FullRow, - ) - T.gemm( - Q_tail_shared, - K_tail_shared, - acc_s, - transpose_B=True, - policy=T.GemmWarpPolicy.FullRow, - ) - T.copy(m_i, m_i_prev) - T.reduce_max(acc_s, m_i, dim=1, clear=False) - for h_i in T.Parallel(H_per_block): - m_i[h_i] = T.max(m_i[h_i], m_i_prev[h_i]) - for h_i in T.Parallel(H_per_block): - alpha[h_i] = T.exp2((m_i_prev[h_i] - m_i[h_i]) * sm_scale) - for h_i, bi_i in T.Parallel(H_per_block, BI): - acc_s[h_i, bi_i] = T.exp2( - acc_s[h_i, bi_i] * sm_scale - m_i[h_i] * sm_scale - ) - T.reduce_sum(acc_s, sumexp_i, dim=1) - for h_i in T.Parallel(H_per_block): - sumexp[h_i] = sumexp[h_i] * alpha[h_i] + sumexp_i[h_i] - for h_i, d_i in T.Parallel(H_per_block, D): - acc_o[h_i, d_i] *= alpha[h_i] - for h_i, d_i in T.Parallel(H_per_block, D_tail): - acc_o_tail[h_i, d_i] *= alpha[h_i] - T.copy(acc_s, S_shared) - T.gemm( - S_shared, - KV_shared, - acc_o, - policy=T.GemmWarpPolicy.FullRow, - ) - T.gemm( - S_shared, - K_tail_shared, - acc_o_tail, - policy=T.GemmWarpPolicy.FullRow, - ) - # ---- finalize phase 2 (active) ---- - for h_i, d_i in T.Parallel(H_per_block, D): - acc_o[h_i, d_i] = acc_o[h_i, d_i] / T.if_then_else( - sumexp[h_i] == 0.0, 1.0, sumexp[h_i] - ) - for h_i, d_i in T.Parallel(H_per_block, D_tail): - acc_o_tail[h_i, d_i] = acc_o_tail[h_i, d_i] / T.if_then_else( - sumexp[h_i] == 0.0, 1.0, sumexp[h_i] - ) - for h_i in T.Parallel(H_per_block): - m_i[h_i] = T.if_then_else( - sumexp[h_i] == 0.0, - -(2.0**30), - T.log2(sumexp[h_i]) + m_i[h_i] * sm_scale, - ) - T.copy(acc_o, Partial_O[b_i, s_i, group_i, H0:H1, :D]) - T.copy( - acc_o_tail, - Partial_O[b_i, s_i, group_i, H0:H1, D : D + D_tail], - ) - T.copy(m_i, Partial_LSE[b_i, s_i, group_i, H0:H1]) - else: - # Phase 2 skipped: m_i is still the -2^30 - T.copy(m_i, Partial_LSE[b_i, s_i, group_i, H0:H1]) - - return main - - @T.prim_func - def main( - Q: T.Tensor(q_shape, BF16), # type: ignore - K_combined_1: T.Tensor(k1_shape, "uint32"), # type: ignore - Indices_1: T.Tensor(indices1_shape, indices_dtype), # type: ignore - Topk_length_1: T.Tensor(topk_length_shape, indices_dtype), # type: ignore - Partial_O: T.Tensor(partial_o_shape, BF16), # type: ignore - Partial_LSE: T.Tensor(partial_lse_shape, accum_dtype), # type: ignore - ) -> None: - """ - grid: (seq_len * REPLICATE_H * n_groups, batch, 1) - Each block processes `inner_iter_1` consecutive KV tiles and writes - one (partial_o, partial_lse) entry. - """ - with T.Kernel( - seq_len * REPLICATE_H * n_groups, batch, kv_group, threads=threads - ) as (bx, by, bz): - Q_shared = T.alloc_fragment([H_per_block, D], BF16) - Q_tail_shared = T.alloc_fragment([H_per_block, D_tail], BF16) - K_packed_shared = T.alloc_shared([BI, PACKED_W4], "uint32") - K_scale_shared = T.alloc_shared([BI, SCALE_W4], "uint32") - KV_shared = T.alloc_shared([BI, D], BF16) - K_tail_shared = T.alloc_shared([BI, D_tail], BF16) - S_shared = T.alloc_shared([H_per_block, BI], BF16) - page_idx_shared = T.alloc_shared([BI], INT32) - - mask = T.alloc_fragment([BI], "bool") - scale_byte_local = T.alloc_fragment([BI, NUM_TILES], "uint32") - - acc_o = T.alloc_fragment([H_per_block, D], accum_dtype) - acc_o_tail = T.alloc_fragment([H_per_block, D_tail], accum_dtype) - acc_s = T.alloc_fragment([H_per_block, BI], accum_dtype) - sumexp = T.alloc_fragment([H_per_block], accum_dtype) - sumexp_i = T.alloc_fragment([H_per_block], accum_dtype) - alpha = T.alloc_fragment([H_per_block], accum_dtype) - m_i = T.alloc_fragment([H_per_block], accum_dtype) - m_i_prev = T.alloc_fragment([H_per_block], accum_dtype) - - T.fill(acc_o, 0) - T.fill(acc_o_tail, 0) - T.fill(sumexp, 0) - T.fill(m_i, -(2**30)) - - b_i, g_i = by, bz - spans_per_seq = REPLICATE_H * n_groups - s_i = bx // spans_per_seq - rest = bx % spans_per_seq - group_i = rest // REPLICATE_H - h_rep = rest % REPLICATE_H - H0 = g_i * padded_H + (0 if REPLICATE_H == 1 else h_rep * 64) - H1 = H0 + H_per_block - - T.copy(Q[b_i, s_i, H0:H1, :D], Q_shared) - T.copy(Q[b_i, s_i, H0:H1, D : D + D_tail], Q_tail_shared) - - tk_len_1 = Topk_length_1[b_i] - - for k_i in T.Pipelined(inner_iter_1, num_stages=num_stages): - iter_i = group_i * inner_iter_1 + k_i - for bi_i in T.Parallel(BI): - pos = iter_i * BI + bi_i - idx = Indices_1[b_i, s_i, pos] - valid = (idx >= 0) & (pos < tk_len_1) - page_idx_shared[bi_i] = T.if_then_else(valid, idx, 0) - mask[bi_i] = valid - - for bi_i, w_i in T.Parallel(BI, PACKED_W4): - page = page_idx_shared[bi_i] - block_id = page // BS_KV_1 - t_in_block = page % BS_KV_1 - K_packed_shared[bi_i, w_i] = K_combined_1[ - block_id, t_in_block * PACKED_W4 + w_i - ] - - for bi_i, w_i in T.Parallel(BI, SCALE_W4): - page = page_idx_shared[bi_i] - block_id = page // BS_KV_1 - t_in_block = page % BS_KV_1 - K_scale_shared[bi_i, w_i] = K_combined_1[ - block_id, - NOPE_ROPE_U32_PER_BLOCK_1 + t_in_block * SCALE_W4 + w_i, - ] - - for bi_i, ti in T.Parallel(BI, NUM_TILES): - word_idx = ti // 4 - byte_in_word = ti % 4 - word = K_scale_shared[bi_i, word_idx] - scale_byte_local[bi_i, ti] = ( - word >> T.Cast("uint32", byte_in_word * 8) - ) & T.uint32(0xFF) - - for bi_i, d_i in T.Parallel(BI, D): - word_idx = d_i // 4 - byte_in_word = d_i % 4 - word = K_packed_shared[bi_i, word_idx] - b_u32 = (word >> T.Cast("uint32", byte_in_word * 8)) & T.uint32( - 0xFF - ) - sign_bf = (b_u32 & T.uint32(0x80)) * T.uint32(0x100) - exp_e4 = (b_u32 & T.uint32(0x78)) >> T.uint32(3) - mant_bf = (b_u32 & T.uint32(0x7)) * T.uint32(0x10) - scale_byte = scale_byte_local[bi_i, d_i // NOPE_TILE] - exp_combined = exp_e4 + scale_byte - T.uint32(7) - bf16_bits = sign_bf | (exp_combined << T.uint32(7)) | mant_bf - KV_shared[bi_i, d_i] = T.reinterpret( - BF16, T.Cast("uint16", bf16_bits) - ) - - for bi_i, j in T.Parallel(BI, D_tail): - abs_off = D + 2 * j - word_idx = abs_off // 4 - word_off = abs_off % 4 - word = K_packed_shared[bi_i, word_idx] - half_u32 = T.if_then_else( - word_off == 0, - word & T.uint32(0xFFFF), - (word >> T.uint32(16)) & T.uint32(0xFFFF), - ) - K_tail_shared[bi_i, j] = T.reinterpret( - BF16, T.Cast("uint16", half_u32) - ) - - for h_i, bi_i in T.Parallel(H_per_block, BI): - acc_s[h_i, bi_i] = T.if_then_else( - mask[bi_i], 0, -T.infinity(acc_s.dtype) - ) - T.gemm( - Q_shared, - KV_shared, - acc_s, - transpose_B=True, - policy=T.GemmWarpPolicy.FullRow, - ) - T.gemm( - Q_tail_shared, - K_tail_shared, - acc_s, - transpose_B=True, - policy=T.GemmWarpPolicy.FullRow, - ) - T.copy(m_i, m_i_prev) - T.reduce_max(acc_s, m_i, dim=1, clear=False) - for h_i in T.Parallel(H_per_block): - m_i[h_i] = T.max(m_i[h_i], m_i_prev[h_i]) - for h_i in T.Parallel(H_per_block): - alpha[h_i] = T.exp2((m_i_prev[h_i] - m_i[h_i]) * sm_scale) - for h_i, bi_i in T.Parallel(H_per_block, BI): - acc_s[h_i, bi_i] = T.exp2( - acc_s[h_i, bi_i] * sm_scale - m_i[h_i] * sm_scale - ) - T.reduce_sum(acc_s, sumexp_i, dim=1) - for h_i in T.Parallel(H_per_block): - sumexp[h_i] = sumexp[h_i] * alpha[h_i] + sumexp_i[h_i] - for h_i, d_i in T.Parallel(H_per_block, D): - acc_o[h_i, d_i] *= alpha[h_i] - for h_i, d_i in T.Parallel(H_per_block, D_tail): - acc_o_tail[h_i, d_i] *= alpha[h_i] - T.copy(acc_s, S_shared) - T.gemm(S_shared, KV_shared, acc_o, policy=T.GemmWarpPolicy.FullRow) - T.gemm( - S_shared, K_tail_shared, acc_o_tail, policy=T.GemmWarpPolicy.FullRow - ) - - for h_i, d_i in T.Parallel(H_per_block, D): - acc_o[h_i, d_i] = acc_o[h_i, d_i] / T.if_then_else( - sumexp[h_i] == 0.0, 1.0, sumexp[h_i] - ) - for h_i, d_i in T.Parallel(H_per_block, D_tail): - acc_o_tail[h_i, d_i] = acc_o_tail[h_i, d_i] / T.if_then_else( - sumexp[h_i] == 0.0, 1.0, sumexp[h_i] - ) - for h_i in T.Parallel(H_per_block): - m_i[h_i] = T.if_then_else( - sumexp[h_i] == 0.0, - -(2.0**30), - T.log2(sumexp[h_i]) + m_i[h_i] * sm_scale, - ) - T.copy(acc_o, Partial_O[b_i, s_i, group_i, H0:H1, :D]) - T.copy(acc_o_tail, Partial_O[b_i, s_i, group_i, H0:H1, D : D + D_tail]) - T.copy(m_i, Partial_LSE[b_i, s_i, group_i, H0:H1]) - - return main - - -@tilelang.jit( - out_idx=[-2, -1], - pass_configs={ - tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, - tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True, - tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True, - }, -) -def dpsk_v4_combine_kernel( - num_heads: int, - n_groups_1: int, - n_groups_2: int = 0, - *, - block_I: int = 64, - inner_iter_1: int = 1, - inner_iter_2: int = 1, - dim: int = 448, - tail_dim: int = 64, - head_per_block: int = 16, - threads: int = 256, - use_attn_sink: bool = False, -) -> Any: - """ - Combine `n_groups` flash-attention partials into the final output. - - Inputs: - Partial_O : (batch, seq_len, n_groups, num_heads, dim+tail_dim) bf16 - Partial_LSE : (batch, seq_len, n_groups, num_heads) fp32, log2 form - Topk_length_1: (batch,) int32, actual phase-1 length - Topk_length_2: (batch,) int32, actual phase-2 length (dual only) - Attn_sink : (num_heads,) fp32 - Outputs: - Output : (batch, seq_len, num_heads, dim+tail_dim) bf16 - LSE : (batch, seq_len, num_heads) fp32, natural log - - Each grid block handles `head_per_block` heads of one (batch, seq) row. - """ - log2e: float = 1.44269504 - ln2: float = 0.69314718 - assert num_heads % head_per_block == 0 - - is_dual = n_groups_2 > 0 - n_groups = n_groups_1 + n_groups_2 - - H_per_block = head_per_block - HEAD_BLOCKS = num_heads // H_per_block - DT = dim + tail_dim - - batch = T.symbolic("batch") - seq_len = T.symbolic("seq_len") - - accum_dtype = "float" - - if is_dual: - - @T.prim_func - def main( - Partial_O: T.Tensor( - [batch, seq_len, n_groups, num_heads, DT], BF16 - ), # type: ignore - Partial_LSE: T.Tensor( - [batch, seq_len, n_groups, num_heads], accum_dtype - ), # type: ignore - Topk_length_1: T.Tensor([batch], INT32), # type: ignore - Topk_length_2: T.Tensor([batch], INT32), # type: ignore - Attn_sink: T.Tensor([num_heads], FP32), # type: ignore - Output: T.Tensor([batch, seq_len, num_heads, DT], BF16), # type: ignore - LSE: T.Tensor([batch, seq_len, num_heads], accum_dtype), # type: ignore - ) -> None: - with T.Kernel(seq_len * HEAD_BLOCKS, batch, threads=threads) as ( - bx, - by, - ): - shared_lse = T.alloc_shared([n_groups, H_per_block], accum_dtype) - lse_max = T.alloc_fragment([H_per_block], accum_dtype) - lse_sum = T.alloc_fragment([H_per_block], accum_dtype) - scale = T.alloc_fragment([H_per_block, n_groups], accum_dtype) - acc_o = T.alloc_fragment([H_per_block, DT], accum_dtype) - attn_sink_frag = T.alloc_fragment([H_per_block], accum_dtype) - o_scale_frag = T.alloc_fragment([H_per_block], accum_dtype) - final_lse = T.alloc_fragment([H_per_block], accum_dtype) - - b_i = by - s_i = bx // HEAD_BLOCKS - head_block = bx % HEAD_BLOCKS - H0 = head_block * H_per_block - H1 = H0 + H_per_block - - # Clamp to the captured-shape upper bounds so callers passing - # the INT32_MAX sentinel (= "all valid") still iterate exactly - # n_groups groups, not 33M. - actual_n_groups_1 = T.min( - T.ceildiv(Topk_length_1[b_i], block_I * inner_iter_1), - n_groups_1, - ) - actual_n_groups_2 = T.min( - T.ceildiv(Topk_length_2[b_i], block_I * inner_iter_2), - n_groups - n_groups_1, - ) - actual_n_groups = actual_n_groups_1 + actual_n_groups_2 - - # Pass 1: load only active groups' LSE into compact slots. - for k_c in T.serial(actual_n_groups): - k = T.if_then_else( - k_c < actual_n_groups_1, - k_c, - n_groups_1 + (k_c - actual_n_groups_1), - ) - T.copy(Partial_LSE[b_i, s_i, k, H0:H1], shared_lse[k_c, :]) - - T.fill(lse_max, -(2**30)) - for k_c in T.serial(actual_n_groups): - for h_i in T.Parallel(H_per_block): - lse_max[h_i] = T.max(lse_max[h_i], shared_lse[k_c, h_i]) - T.fill(lse_sum, 0) - for k_c in T.serial(actual_n_groups): - for h_i in T.Parallel(H_per_block): - lse_sum[h_i] = lse_sum[h_i] + T.exp2( - shared_lse[k_c, h_i] - lse_max[h_i] - ) - for k_c in T.serial(actual_n_groups): - for h_i in T.Parallel(H_per_block): - scale[h_i, k_c] = T.exp2( - shared_lse[k_c, h_i] - lse_max[h_i] - T.log2(lse_sum[h_i]) - ) - - T.fill(acc_o, 0) - for k_c in T.serial(actual_n_groups): - k = T.if_then_else( - k_c < actual_n_groups_1, - k_c, - n_groups_1 + (k_c - actual_n_groups_1), - ) - for h_i, d_i in T.Parallel(H_per_block, DT): - acc_o[h_i, d_i] = acc_o[h_i, d_i] + scale[h_i, k_c] * Partial_O[ - b_i, s_i, k, H0 + h_i, d_i - ].astype(accum_dtype) - - for h_i in T.Parallel(H_per_block): - empty = lse_max[h_i] <= -(2**29) - final_lse[h_i] = T.if_then_else( - empty, - T.infinity(accum_dtype), - (lse_max[h_i] + T.log2(lse_sum[h_i])) * ln2, - ) - - if use_attn_sink: - for h_i in T.Parallel(H_per_block): - attn_sink_frag[h_i] = Attn_sink[H0 + h_i] - for h_i in T.Parallel(H_per_block): - empty = lse_max[h_i] <= -(2**29) - o_scale_frag[h_i] = T.if_then_else( - empty, - 0.0, - 1.0 - / ( - 1.0 - + T.exp2((attn_sink_frag[h_i] - final_lse[h_i]) * log2e) - ), - ) - for h_i, d_i in T.Parallel(H_per_block, DT): - acc_o[h_i, d_i] = acc_o[h_i, d_i] * o_scale_frag[h_i] - - T.copy(acc_o, Output[b_i, s_i, H0:H1, :]) - T.copy(final_lse, LSE[b_i, s_i, H0:H1]) - - return main - - @T.prim_func - def main( - Partial_O: T.Tensor( - [batch, seq_len, n_groups, num_heads, DT], BF16 - ), # type: ignore - Partial_LSE: T.Tensor( - [batch, seq_len, n_groups, num_heads], accum_dtype - ), # type: ignore - Attn_sink: T.Tensor([num_heads], FP32), # type: ignore - Output: T.Tensor([batch, seq_len, num_heads, DT], BF16), # type: ignore - LSE: T.Tensor([batch, seq_len, num_heads], accum_dtype), # type: ignore - ) -> None: - with T.Kernel(seq_len * HEAD_BLOCKS, batch, threads=threads) as (bx, by): - shared_lse = T.alloc_shared([n_groups, H_per_block], accum_dtype) - - lse_max = T.alloc_fragment([H_per_block], accum_dtype) - lse_sum = T.alloc_fragment([H_per_block], accum_dtype) - scale = T.alloc_fragment([H_per_block, n_groups], accum_dtype) - acc_o = T.alloc_fragment([H_per_block, DT], accum_dtype) - attn_sink_frag = T.alloc_fragment([H_per_block], accum_dtype) - o_scale_frag = T.alloc_fragment([H_per_block], accum_dtype) - final_lse = T.alloc_fragment([H_per_block], accum_dtype) - - b_i = by - s_i = bx // HEAD_BLOCKS - head_block = bx % HEAD_BLOCKS - H0 = head_block * H_per_block - H1 = H0 + H_per_block - - for k in T.serial(n_groups): - T.copy(Partial_LSE[b_i, s_i, k, H0:H1], shared_lse[k, :]) - - T.fill(lse_max, -(2**30)) - for k in T.serial(n_groups): - for h_i in T.Parallel(H_per_block): - lse_max[h_i] = T.max(lse_max[h_i], shared_lse[k, h_i]) - T.fill(lse_sum, 0) - for k in T.serial(n_groups): - for h_i in T.Parallel(H_per_block): - lse_sum[h_i] = lse_sum[h_i] + T.exp2( - shared_lse[k, h_i] - lse_max[h_i] - ) - for k in T.serial(n_groups): - for h_i in T.Parallel(H_per_block): - scale[h_i, k] = T.exp2( - shared_lse[k, h_i] - lse_max[h_i] - T.log2(lse_sum[h_i]) - ) - - T.fill(acc_o, 0) - for k in T.serial(n_groups): - for h_i, d_i in T.Parallel(H_per_block, DT): - acc_o[h_i, d_i] = acc_o[h_i, d_i] + scale[h_i, k] * Partial_O[ - b_i, s_i, k, H0 + h_i, d_i - ].astype(accum_dtype) - - for h_i in T.Parallel(H_per_block): - empty = lse_max[h_i] <= -(2**29) - final_lse[h_i] = T.if_then_else( - empty, - T.infinity(accum_dtype), - (lse_max[h_i] + T.log2(lse_sum[h_i])) * ln2, - ) - - if use_attn_sink: - for h_i in T.Parallel(H_per_block): - attn_sink_frag[h_i] = Attn_sink[H0 + h_i] - for h_i in T.Parallel(H_per_block): - empty = lse_max[h_i] <= -(2**29) - o_scale_frag[h_i] = T.if_then_else( - empty, - 0.0, - 1.0 - / ( - 1.0 + T.exp2((attn_sink_frag[h_i] - final_lse[h_i]) * log2e) - ), - ) - for h_i, d_i in T.Parallel(H_per_block, DT): - acc_o[h_i, d_i] = acc_o[h_i, d_i] * o_scale_frag[h_i] - - T.copy(acc_o, Output[b_i, s_i, H0:H1, :]) - T.copy(final_lse, LSE[b_i, s_i, H0:H1]) - - return main - - -""" -2-stage attention kernel (partial + combine) over an FP8 KV cache, -with optional second cache (`extra_k_cache`). -""" - - -def dpsk_v4_fp8_attention_fwd( - q: torch.Tensor, - k_cache: torch.Tensor, - block_table: Optional[torch.Tensor], - cache_seqlens: Optional[torch.Tensor], - head_dim_v: int, - tile_scheduler_metadata: Any, - num_splits: None = None, - softmax_scale: Optional[float] = None, - causal: bool = False, - is_fp8_kvcache: bool = False, - indices: Optional[torch.Tensor] = None, - attn_sink: Optional[torch.Tensor] = None, - extra_k_cache: Optional[torch.Tensor] = None, - extra_indices_in_kvcache: Optional[torch.Tensor] = None, - topk_length: Optional[torch.Tensor] = None, - extra_topk_length: Optional[torch.Tensor] = None, -) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Follows the original `flash_mla.flash_mla_with_kvcache` signature. - """ - if _is_gfx95_supported: - block_I, threads, num_stages, block_per_cu, cu = 64, 512, 0, 2, 256 - else: - block_I, threads, num_stages, block_per_cu, cu = 32, 128, 1, 1, 304 - - batch, seq_len, num_heads, _ = q.shape - # Partial grid is (seq_len * REPLICATE_H * n_groups, batch, kv_group); the - # heuristic in _pick_inner_iter assumes `total_blocks = seq * ni / inner_iter`, - # so `seq` must include REPLICATE_H or n_groups doubles for medium batches. - replicate_h = max((num_heads + 63) // 64, 1) - seq = batch * seq_len * replicate_h - - k1, _, bs_kv_1 = _build_fp8_combined_view(k_cache) - topk_1 = indices.shape[-1] - ni_1 = topk_1 // block_I - tk_len_1 = ( - topk_length - if topk_length is not None - else _topk_length_sentinel(q.device, batch) - ) - if attn_sink is None: - attn_sink = torch.full( - (num_heads,), float("-inf"), dtype=torch.float32, device=q.device - ) - - has_extra = extra_k_cache is not None - if not has_extra: - inner_iter_1 = _pick_inner_iter(seq, ni_1, cu, block_per_cu) - inner_iter_2 = 1 - n_groups_1 = ni_1 // inner_iter_1 - n_groups_2 = 0 - partial = dpsk_v4_fp8_partial_kernel( - num_heads, - topk_1, - bs_kv_1, - sm_scale=softmax_scale, - block_I=block_I, - inner_iter_1=inner_iter_1, - num_stages=num_stages, - threads=threads, - ) - partial_o, partial_lse = partial(q, k1, indices, tk_len_1) - else: - k2, _, bs_kv_2 = _build_fp8_combined_view(extra_k_cache) - topk_2 = extra_indices_in_kvcache.shape[-1] - ni_2 = topk_2 // block_I - # Each phase picks its own optimal split-K independently — kernel - # body uses two T.Pipelined loops with separate compile-time iter - # counts, no shared-divisor constraint. - inner_iter_1 = _pick_inner_iter(seq, ni_1, cu, block_per_cu) - inner_iter_2 = _pick_inner_iter(seq, ni_2, cu, block_per_cu) - n_groups_1 = ni_1 // inner_iter_1 - n_groups_2 = ni_2 // inner_iter_2 - tk_len_2 = ( - extra_topk_length - if extra_topk_length is not None - else _topk_length_sentinel(q.device, batch) - ) - partial = dpsk_v4_fp8_partial_kernel( - num_heads, - topk_1, - bs_kv_1, - topk_2, - bs_kv_2, - sm_scale=softmax_scale, - block_I=block_I, - inner_iter_1=inner_iter_1, - inner_iter_2=inner_iter_2, - num_stages=num_stages, - threads=threads, - ) - partial_o, partial_lse = partial( - q, - k1, - indices, - tk_len_1, - k2, - extra_indices_in_kvcache, - tk_len_2, - ) - - combine = dpsk_v4_combine_kernel( - num_heads, - n_groups_1, - n_groups_2, - block_I=block_I, - inner_iter_1=inner_iter_1, - inner_iter_2=inner_iter_2, - head_per_block=4, - threads=256, - use_attn_sink=True, - ) - if has_extra: - return combine(partial_o, partial_lse, tk_len_1, tk_len_2, attn_sink) - return combine(partial_o, partial_lse, attn_sink) +from sglang.srt.layers.attention.dsa.tilelang_kernel import * # noqa: F401, F403 diff --git a/python/sglang/srt/layers/attention/nsa/transform_index.py b/python/sglang/srt/layers/attention/nsa/transform_index.py index 10b1068f5..67ccaa772 100644 --- a/python/sglang/srt/layers/attention/nsa/transform_index.py +++ b/python/sglang/srt/layers/attention/nsa/transform_index.py @@ -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 -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") +warnings.warn( + "sglang.srt.layers.attention.nsa.transform_index is deprecated; " + "use sglang.srt.layers.attention.dsa.transform_index instead.", + DeprecationWarning, + stacklevel=2, +) +from sglang.srt.layers.attention.dsa.transform_index import * # noqa: F401, F403 diff --git a/python/sglang/srt/layers/attention/nsa/triton_kernel.py b/python/sglang/srt/layers/attention/nsa/triton_kernel.py index 0d2969804..efa6c354f 100644 --- a/python/sglang/srt/layers/attention/nsa/triton_kernel.py +++ b/python/sglang/srt/layers/attention/nsa/triton_kernel.py @@ -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 -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, - ) +warnings.warn( + "sglang.srt.layers.attention.nsa.triton_kernel is deprecated; " + "use sglang.srt.layers.attention.dsa.triton_kernel instead.", + DeprecationWarning, + stacklevel=2, +) +from sglang.srt.layers.attention.dsa.triton_kernel import * # noqa: F401, F403 diff --git a/python/sglang/srt/layers/attention/nsa/utils.py b/python/sglang/srt/layers/attention/nsa/utils.py index b49a3e261..328241841 100644 --- a/python/sglang/srt/layers/attention/nsa/utils.py +++ b/python/sglang/srt/layers/attention/nsa/utils.py @@ -1,269 +1,10 @@ -from functools import lru_cache -from typing import TYPE_CHECKING, List, Tuple, Union +# [Deprecated] Re-export shim for backward compatibility. Use dsa.utils instead. +import warnings -import torch -import triton -import triton.language as tl - -from sglang.srt.layers.dp_attention import ( - DpPaddingMode, - get_attention_cp_rank, - get_attention_cp_size, - get_attention_dp_rank, +warnings.warn( + "sglang.srt.layers.attention.nsa.utils is deprecated; " + "use sglang.srt.layers.attention.dsa.utils instead.", + DeprecationWarning, + stacklevel=2, ) -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 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 +from sglang.srt.layers.attention.dsa.utils import * # noqa: F401, F403 diff --git a/python/sglang/srt/layers/attention/nsa_backend.py b/python/sglang/srt/layers/attention/nsa_backend.py index 55a52e991..b2554684a 100644 --- a/python/sglang/srt/layers/attention/nsa_backend.py +++ b/python/sglang/srt/layers/attention/nsa_backend.py @@ -1,2520 +1,23 @@ -from __future__ import annotations +# [Deprecated] nsa_backend.py is a thin re-export shim for backward compatibility. +# Use dsa_backend.py instead. This file will be removed in a future release. +import warnings -from dataclasses import dataclass -from enum import IntEnum, auto -from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Tuple, TypeAlias - -import torch - -from sglang.srt.configs.model_config import get_nsa_index_topk, is_deepseek_nsa -from sglang.srt.environ import envs -from sglang.srt.layers.attention.base_attn_backend import AttentionBackend -from sglang.srt.layers.attention.nsa.dequant_k_cache import dequantize_k_cache_paged -from sglang.srt.layers.attention.nsa.nsa_backend_mtp_precompute import ( - NativeSparseAttnBackendMTPPrecomputeMixin, - PrecomputedMetadata, - compute_cu_seqlens, +warnings.warn( + "sglang.srt.layers.attention.nsa_backend is deprecated; " + "use sglang.srt.layers.attention.dsa_backend instead.", + DeprecationWarning, + stacklevel=2, ) -from sglang.srt.layers.attention.nsa.nsa_indexer import BaseIndexerMetadata -from sglang.srt.layers.attention.nsa.quant_k_cache import quantize_k_cache -from sglang.srt.layers.attention.nsa.transform_index import ( - transform_index_page_table_decode, - transform_index_page_table_prefill, +from sglang.srt.layers.attention.dsa_backend import * # noqa: F401, F403 +from sglang.srt.layers.attention.dsa_backend import ( # noqa: F401 + DeepseekSparseAttnBackend, + DeepseekSparseAttnMultiStepBackend, + DSAFlashMLAMetadata, + DSAIndexerMetadata, + DSAMetadata, + NativeSparseAttnBackend, + NativeSparseAttnMultiStepBackend, + NSAFlashMLAMetadata, + NSAIndexerMetadata, + NSAMetadata, ) -from sglang.srt.layers.attention.nsa.utils import ( - can_nsa_prefill_cp_round_robin_split, - compute_nsa_seqlens, - is_nsa_enable_prefill_cp, - nsa_cp_round_robin_split_data, - nsa_cp_round_robin_split_q_seqs, - pad_nsa_cache_seqlens, -) -from sglang.srt.layers.attention.utils import ( - concat_mla_absorb_q_general, - mla_quantize_and_rope_for_fp8, - seqlens_expand_triton, -) -from sglang.srt.layers.dp_attention import get_attention_tp_size -from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode -from sglang.srt.utils import is_cuda, is_hip - -if TYPE_CHECKING: - from sglang.srt.layers.radix_attention import RadixAttention - from sglang.srt.model_executor.model_runner import ModelRunner - from sglang.srt.speculative.spec_info import SpecInput - - -_is_hip = is_hip() - -if _is_hip: - from sglang.srt.layers.attention.nsa.triton_kernel import get_valid_kv_indices - - try: - from aiter import ( # noqa: F401 - flash_attn_varlen_func, - mha_batch_prefill_func, - paged_attention_ragged, - ) - from aiter.mla import mla_decode_fwd, mla_prefill_fwd # noqa: F401 - except ImportError: - print( - "aiter is AMD specific kernel library. Please make sure aiter is installed on your AMD device." - ) -else: - from sglang.jit_kernel.flash_attention import ( - flash_attn_varlen_func, - flash_attn_with_kvcache, - ) - - -def _to_2d_context_lens(seqlens_32: torch.Tensor, batch_size: int) -> torch.Tensor: - # Always normalize to (N_total, 1) layout, to avoid deadlock at deep_gemm.fp8_paged_mqa_logits - if seqlens_32.dim() == 2: - if seqlens_32.size(1) == 1: - return seqlens_32 - # Fall through and re-flatten if the caller already gave us a (bs, next_n) - # view — we want (N_total, 1) regardless. - seqlens_32 = seqlens_32.reshape(-1) - return seqlens_32.contiguous().view(-1, 1) - - -# Reuse this workspace buffer across all NSA backend instances -global_workspace_buffer = None - -# Control whether to use fused metadata copy kernel for cuda graph replay (default: enabled) -# Set SGLANG_USE_FUSED_METADATA_COPY=0 or false to disable -_USE_FUSED_METADATA_COPY = envs.SGLANG_USE_FUSED_METADATA_COPY.get() and not _is_hip - - -@dataclass(frozen=True) -class NSAFlashMLAMetadata: - """Metadata only needed by FlashMLA""" - - flashmla_metadata: torch.Tensor - num_splits: torch.Tensor - - def slice(self, sli): - return NSAFlashMLAMetadata( - flashmla_metadata=self.flashmla_metadata, - num_splits=self.num_splits[sli], - ) - - def copy_(self, other: "NSAFlashMLAMetadata"): - self.flashmla_metadata.copy_(other.flashmla_metadata) - self.num_splits.copy_(other.num_splits) - - -@dataclass(frozen=True) -class NSAMetadata: - page_size: int - - # Sequence lengths for the forward batch - cache_seqlens_int32: torch.Tensor - # Maximum sequence length for query - max_seq_len_q: int - # Maximum sequence length for key - max_seq_len_k: int - # Cumulative sequence lengths for query - cu_seqlens_q: torch.Tensor - # Cumulative sequence lengths for key - cu_seqlens_k: torch.Tensor - # Page table, the index of KV Cache Tables/Blocks - # this table is always with page_size = 1 - page_table_1: torch.Tensor - - # NOTE(dark): This will property be used in: - # 1. dense decode/prefill, we use paged flash attention, need real_page_table - # 2. sparse decode/prefill, indexer need real_page_table to compute the score - real_page_table: torch.Tensor - - # NSA metadata (nsa prefill are expanded) - nsa_cache_seqlens_int32: torch.Tensor # this seqlens is clipped to `topk` - nsa_cu_seqlens_q: torch.Tensor # must be arange(0, len(nsa_cu_seqlens_k)) - nsa_cu_seqlens_k: torch.Tensor # cumsum of `nsa_cache_seqlens_int32` - nsa_extend_seq_lens_list: List[int] - nsa_seqlens_expanded: torch.Tensor # expanded, unclipped `seqlens` - nsa_max_seqlen_q: Literal[1] = 1 # always 1 for decode, variable for extend - - flashmla_metadata: Optional[NSAFlashMLAMetadata] = None - # DeepGEMM schedule metadata for paged MQA logits (decode/target_verify/draft_extend only). - # Precomputed once per forward batch and reused across layers. - paged_mqa_schedule_metadata: Optional[torch.Tensor] = None - # The sum of sequence lengths for key, prefill only - seq_lens_sum: Optional[int] = None - # The flattened 1D page table with shape (seq_lens_sum,), prefill only - # this table is always with page_size = 1 - page_table_1_flattened: Optional[torch.Tensor] = None - # The offset of topk indices in ragged kv, prefill only - # shape: (seq_lens_sum,) - topk_indices_offset: Optional[torch.Tensor] = None - - # k_start and k_end in kv cache for each token. - indexer_k_start_end: Optional[Tuple[torch.Tensor, torch.Tensor]] = None - # seq lens for each batch. - indexer_seq_lens_cpu: Optional[torch.Tensor] = None - # seq lens for each batch. - indexer_seq_lens: Optional[torch.Tensor] = None - # batch index for each token. - token_to_batch_idx: Optional[torch.Tensor] = None - - -class TopkTransformMethod(IntEnum): - # Transform topk indices to indices to the page table (page_size = 1) - PAGED = auto() - # Transform topk indices to indices to ragged kv (non-paged) - RAGGED = auto() - - -@torch.compile -def _compiled_cat(tensors: list[torch.Tensor], dim: int = -1) -> torch.Tensor: - return torch.cat(tensors, dim=dim) - - -def _cat(tensors: list[torch.Tensor], dim: int = -1) -> torch.Tensor: - """ - Concatenate two tensors along the last dimension. - Use this function to concatenate q_nope and q_rope or k_nope and k_rope. - """ - assert len(tensors) == 2 - - qk_nope, qk_rope = tensors - assert qk_nope.ndim == 3 and qk_rope.ndim == 3 - - torch._dynamo.mark_dynamic(qk_nope, 0) - torch._dynamo.mark_dynamic(qk_rope, 0) - - return _compiled_cat([qk_nope, qk_rope], dim=dim) - - -@dataclass(frozen=True) -class NSAIndexerMetadata(BaseIndexerMetadata): - attn_metadata: NSAMetadata - topk_transform_method: TopkTransformMethod - paged_mqa_schedule_metadata: Optional[torch.Tensor] = None - force_unfused_topk: bool = False - - def get_seqlens_int32(self) -> torch.Tensor: - return self.attn_metadata.cache_seqlens_int32 - - def get_page_table_64(self) -> torch.Tensor: - return self.attn_metadata.real_page_table - - def get_page_table_1(self) -> torch.Tensor: - return self.attn_metadata.page_table_1 - - def get_seqlens_expanded(self) -> torch.Tensor: - return self.attn_metadata.nsa_seqlens_expanded - - def get_cu_seqlens_k(self) -> torch.Tensor: - return self.attn_metadata.cu_seqlens_k - - def get_indexer_kvcache_range(self) -> Tuple[torch.Tensor, torch.Tensor]: - return self.attn_metadata.indexer_k_start_end - - def get_indexer_seq_len(self) -> torch.Tensor: - return self.attn_metadata.indexer_seq_lens - - def get_indexer_seq_len_cpu(self) -> torch.Tensor: - return self.attn_metadata.indexer_seq_lens_cpu - - def get_nsa_extend_len_cpu(self) -> List[int]: - return self.attn_metadata.nsa_extend_seq_lens_list - - def get_token_to_batch_idx(self) -> torch.Tensor: - return self.attn_metadata.token_to_batch_idx - - def topk_transform( - self, - logits: torch.Tensor, - topk: int, - ks: Optional[torch.Tensor] = None, - cu_seqlens_q: torch.Tensor = None, - ke_offset: torch.Tensor = None, - batch_idx_list: List[int] = None, - topk_indices_offset_override: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - from sgl_kernel import ( - fast_topk_transform_fused, - fast_topk_transform_ragged_fused, - fast_topk_v2, - ) - - if topk_indices_offset_override is not None: - cu_topk_indices_offset = topk_indices_offset_override - cu_seqlens_q_topk = None - elif cu_seqlens_q is not None: - cu_seqlens_q = cu_seqlens_q.to(torch.int32) - cu_seqlens_q_topk = compute_cu_seqlens(cu_seqlens_q) - cu_topk_indices_offset = torch.repeat_interleave( - cu_seqlens_q_topk[:-1], - cu_seqlens_q, - ) - else: - cu_seqlens_q_topk = self.attn_metadata.cu_seqlens_q - cu_topk_indices_offset = self.attn_metadata.topk_indices_offset - if ke_offset is not None: - seq_lens_topk = ke_offset - else: - seq_lens_topk = self.get_seqlens_expanded() - if batch_idx_list is not None: - page_table_size_1 = self.attn_metadata.page_table_1[batch_idx_list] - else: - page_table_size_1 = self.attn_metadata.page_table_1 - - if not envs.SGLANG_NSA_FUSE_TOPK.get() or self.force_unfused_topk: - return fast_topk_v2(logits, seq_lens_topk, topk, row_starts=ks) - elif self.topk_transform_method == TopkTransformMethod.PAGED: - # NOTE(dark): if fused, we return a transformed page table directly - return fast_topk_transform_fused( - score=logits, - lengths=seq_lens_topk, - page_table_size_1=page_table_size_1, - cu_seqlens_q=cu_seqlens_q_topk, - topk=topk, - row_starts=ks, - ) - elif self.topk_transform_method == TopkTransformMethod.RAGGED: - if cu_topk_indices_offset is None: - raise RuntimeError( - "RAGGED topk_transform requires topk_indices_offset; " - "expected extend-without-speculative metadata." - ) - return fast_topk_transform_ragged_fused( - score=logits, - lengths=seq_lens_topk, - topk_indices_offset=cu_topk_indices_offset, - topk=topk, - row_starts=ks, - ) - else: - assert False, f"Unsupported {self.topk_transform_method = }" - - -_NSA_IMPL_T: TypeAlias = Literal[ - "flashmla_sparse", "flashmla_kv", "fa3", "tilelang", "trtllm" -] - - -class NativeSparseAttnBackend( - NativeSparseAttnBackendMTPPrecomputeMixin, AttentionBackend -): - def __init__( - self, - model_runner: ModelRunner, - skip_prefill: bool = False, - speculative_step_id=0, - topk=0, - speculative_num_steps=0, - ): - super().__init__() - self.forward_metadata: NSAMetadata - self.device = model_runner.device - assert isinstance(model_runner.page_size, int) - self.real_page_size = model_runner.page_size - self.num_splits = ( - 1 if model_runner.server_args.enable_deterministic_inference else 0 - ) - self.use_nsa = is_deepseek_nsa(model_runner.model_config.hf_config) - assert self.use_nsa, "NSA backend only supports DeepSeek NSA" - self.nsa_kv_cache_store_fp8 = ( - model_runner.token_to_kv_pool.nsa_kv_cache_store_fp8 - ) - self.nsa_index_topk = get_nsa_index_topk(model_runner.model_config.hf_config) - self.max_context_len = model_runner.model_config.context_len - self.num_q_heads = ( - model_runner.model_config.num_attention_heads // get_attention_tp_size() - ) - self.kv_cache_dim = model_runner.token_to_kv_pool.kv_cache_dim - self.qk_nope_head_dim = model_runner.model_config.qk_nope_head_dim - self.kv_lora_rank = model_runner.model_config.kv_lora_rank - self.qk_rope_head_dim = model_runner.model_config.qk_rope_head_dim - - assert model_runner.req_to_token_pool is not None - self.req_to_token = model_runner.req_to_token_pool.req_to_token - - self.use_mha: bool = False - self.nsa_prefill_impl: _NSA_IMPL_T = ( - model_runner.server_args.nsa_prefill_backend - ) - self.nsa_decode_impl: _NSA_IMPL_T = model_runner.server_args.nsa_decode_backend - if self.num_q_heads <= 64: - self.flashmla_kv_num_q_heads = 64 - elif self.num_q_heads <= 128: - self.flashmla_kv_num_q_heads = 128 - else: - # Keep original head count if it exceeds current padded variants. - self.flashmla_kv_num_q_heads = self.num_q_heads - self.enable_auto_select_prefill_impl = self.nsa_prefill_impl == "flashmla_auto" - - self._arange_buf = torch.arange(16384, device=self.device, dtype=torch.int32) - - if _is_hip: - max_bs = model_runner.req_to_token_pool.size - - self.kv_indptr = torch.zeros( - (max_bs + 1,), dtype=torch.int32, device=model_runner.device - ) - - self.kv_indices = torch.zeros( - max_bs * self.nsa_index_topk, - dtype=torch.int32, - device=self.device, - ) - # Aiter mla_decode_fwd supports num_heads multiples of 16 in range [16, 128]. - # For models with fewer heads per GPU (e.g. GLM-5 64 heads / TP8 = 8), need to pad the heads to 16. - self.need_pad_heads = self.num_q_heads < 16 - self.head_repeat_factor = ( - 16 // self.num_q_heads if self.num_q_heads < 16 else 1 - ) - - # Speculative decoding - self.topk = model_runner.server_args.speculative_eagle_topk or 0 - self.speculative_num_steps = speculative_num_steps - self.speculative_num_draft_tokens = ( - model_runner.server_args.speculative_num_draft_tokens - ) - self.speculative_step_id = speculative_step_id - - self.device_capability = torch.cuda.get_device_capability() - self.device_sm_major = self.device_capability[0] - self.kv_cache_dtype = model_runner.kv_cache_dtype - - # Allocate global workspace buffer for TRT-LLM kernels (ragged attention on SM100/B200, or trtllm decode) - if self.device_sm_major >= 10 or self.nsa_decode_impl == "trtllm": - global global_workspace_buffer - if global_workspace_buffer is None: - global_workspace_buffer = torch.empty( - envs.SGLANG_FLASHINFER_WORKSPACE_SIZE.get(), - dtype=torch.uint8, - device=model_runner.device, - ) - self.workspace_buffer = global_workspace_buffer - else: - self.workspace_buffer = None - - def get_device_int32_arange(self, l: int) -> torch.Tensor: - if l > len(self._arange_buf): - next_pow_of_2 = 1 << (l - 1).bit_length() - self._arange_buf = torch.arange( - next_pow_of_2, device=self.device, dtype=torch.int32 - ) - return self._arange_buf[:l] - - def _transform_table_1_to_real(self, page_table: torch.Tensor) -> torch.Tensor: - page_size = self.real_page_size - if page_size == 1: - return page_table - max_seqlen_k = page_table.shape[1] - strided_indices = torch.arange( - 0, max_seqlen_k, page_size, device=page_table.device, dtype=torch.int32 - ) - return page_table[:, strided_indices] // page_size - - def init_forward_metadata(self, forward_batch: ForwardBatch): - """Init the metadata for a forward pass.""" - batch_size = forward_batch.batch_size - device = forward_batch.seq_lens.device - - if forward_batch.forward_mode.is_target_verify(): - draft_token_num = self.speculative_num_draft_tokens - else: - draft_token_num = 0 - - cache_seqlens_int32 = (forward_batch.seq_lens + draft_token_num).to(torch.int32) - cu_seqlens_k = compute_cu_seqlens(cache_seqlens_int32) - assert forward_batch.seq_lens_cpu is not None - max_seqlen_k = int(forward_batch.seq_lens_cpu.max().item() + draft_token_num) - # [b, max_seqlen_k] - page_table = forward_batch.req_to_token_pool.req_to_token[ - forward_batch.req_pool_indices, :max_seqlen_k - ] - - page_table_1_flattened = None - topk_indices_offset = None - - # Centralized dispatch: decide all strategies for this batch - self.set_nsa_prefill_impl(forward_batch) - nsa_impl_for_batch = ( - self.nsa_decode_impl - if ( - forward_batch.forward_mode.is_decode_or_idle() - or forward_batch.forward_mode.is_target_verify() - or forward_batch.forward_mode.is_draft_extend(include_v2=True) - ) - else self.nsa_prefill_impl - ) - use_flashmla_kv = (not self.use_mha) and nsa_impl_for_batch == "flashmla_kv" - topk_transform_method = self.get_topk_transform_method( - forward_batch.forward_mode - ) - # Batch indices selected when cp enabled: After splitting multiple sequences, - # a certain cp rank may not have some of these sequences. - # We use bs_idx_cpu to mark which sequences are finally selected by the current cp rank, - # a default value of None indicates that all sequences are selected. - bs_idx_cpu = None - # seq_len_cpu of selected sequences - indexer_seq_lens_cpu = forward_batch.seq_lens_cpu - indexer_seq_lens = forward_batch.seq_lens - - if forward_batch.forward_mode.is_decode_or_idle(): - extend_seq_lens_cpu = [1] * batch_size - max_seqlen_q = 1 - cu_seqlens_q = self.get_device_int32_arange(batch_size + 1) - seqlens_expanded = cache_seqlens_int32 - elif forward_batch.forward_mode.is_target_verify(): - max_seqlen_q = 1 - cu_seqlens_q = torch.arange( - 0, - batch_size * self.speculative_num_draft_tokens + 1, - 1, - dtype=torch.int32, - device=device, - ) - extend_seq_lens_cpu = [self.speculative_num_draft_tokens] * batch_size - forward_batch.extend_seq_lens_cpu = extend_seq_lens_cpu - - seqlens_expanded = seqlens_expand_triton( - torch.tensor(extend_seq_lens_cpu, dtype=torch.int32, device=device), - cache_seqlens_int32, - self.speculative_num_draft_tokens * batch_size, - self.speculative_num_draft_tokens, - ) - page_table = torch.repeat_interleave( - page_table, repeats=self.speculative_num_draft_tokens, dim=0 - ) - elif forward_batch.forward_mode.is_draft_extend(include_v2=True): - assert ( - forward_batch.extend_seq_lens_cpu is not None - and forward_batch.extend_seq_lens is not None - and forward_batch.extend_prefix_lens_cpu is not None - ), "All of them must not be None" - - extend_seq_lens_cpu = forward_batch.extend_seq_lens_cpu - assert forward_batch.extend_seq_lens is not None - - max_seqlen_q = 1 - cu_seqlens_q = torch.arange( - 0, - forward_batch.extend_num_tokens + 1, - 1, - dtype=torch.int32, - device=device, - ) - - seqlens_expanded = seqlens_expand_triton( - forward_batch.extend_seq_lens, - cache_seqlens_int32, - sum(extend_seq_lens_cpu), - self.speculative_num_draft_tokens, - ) - if forward_batch.forward_mode.is_draft_extend_v2(): - # DRAFT_EXTEND_V2: V2 worker pre-fills draft KV cache with ALL speculated - # tokens upfront. All requests extend by the same fixed - # (speculative_num_draft_tokens). Use scalar to avoid GPU sync. - page_table = torch.repeat_interleave( - page_table, repeats=self.speculative_num_draft_tokens, dim=0 - ) - else: - # DRAFT_EXTEND (v1): V1 worker extends by (num_correct_drafts + 1) per request - # after verification. Lengths vary per request based on how many tokens - # were accepted. - page_table = torch.repeat_interleave( - page_table, repeats=forward_batch.extend_seq_lens, dim=0 - ) - elif forward_batch.forward_mode.is_extend(): - assert ( - forward_batch.extend_seq_lens_cpu is not None - and forward_batch.extend_seq_lens is not None - and forward_batch.extend_prefix_lens_cpu is not None - ), "All of them must not be None" - extend_seq_lens_cpu = forward_batch.extend_seq_lens_cpu - assert forward_batch.extend_seq_lens is not None - extend_seq_lens = forward_batch.extend_seq_lens - - seqlens_expanded = torch.cat( - [ - torch.arange( - kv_len - qo_len + 1, - kv_len + 1, - dtype=torch.int32, - device=device, - ) - for qo_len, kv_len in zip( - forward_batch.extend_seq_lens_cpu, - forward_batch.seq_lens_cpu.tolist(), - strict=True, - ) - ] - ) - - if can_nsa_prefill_cp_round_robin_split(forward_batch): - seqlens_expanded = nsa_cp_round_robin_split_data(seqlens_expanded) - extend_seq_lens_cpu, extend_seq_lens, bs_idx_cpu, bs_idx = ( - nsa_cp_round_robin_split_q_seqs( - extend_seq_lens_cpu, extend_seq_lens - ) - ) - indexer_seq_lens_cpu = indexer_seq_lens_cpu[bs_idx_cpu] - indexer_seq_lens = indexer_seq_lens[bs_idx] - cache_seqlens_int32 = cache_seqlens_int32[bs_idx] - cu_seqlens_k = compute_cu_seqlens(cache_seqlens_int32) - max_seqlen_k = ( - int(indexer_seq_lens_cpu.max().item() + draft_token_num) - if len(indexer_seq_lens_cpu) != 0 - else 0 - ) - page_table = page_table[bs_idx, :max_seqlen_k] - - if ( - any(forward_batch.extend_prefix_lens_cpu) - or forward_batch.forward_mode == ForwardMode.DRAFT_EXTEND - or bs_idx_cpu is not None - ): - max_seqlen_q = ( - max(extend_seq_lens_cpu) if len(extend_seq_lens_cpu) != 0 else 1 - ) - cu_seqlens_q = compute_cu_seqlens(extend_seq_lens.to(torch.int32)) - else: - max_seqlen_q = max_seqlen_k - cu_seqlens_q = cu_seqlens_k - - # Check if MHA FP8 dequantization is needed - mha_dequantize_needed = ( - self.use_mha - and forward_batch.token_to_kv_pool.dtype == torch.float8_e4m3fn - ) - forward_batch.using_mha_one_shot_fp8_dequant = mha_dequantize_needed - - # page_table_1_flattened is only used when prefix sharing is enabled: - has_prefix_sharing = any(forward_batch.extend_prefix_lens_cpu) - if has_prefix_sharing and ( - topk_transform_method == TopkTransformMethod.RAGGED - or mha_dequantize_needed - ): - page_table_1_flattened = torch.cat( - [ - page_table[i, :kv_len] - for i, kv_len in enumerate( - indexer_seq_lens_cpu.tolist(), - ) - ] - ) - assert page_table_1_flattened.shape[0] == sum( - indexer_seq_lens_cpu - ), f"{page_table_1_flattened.shape[0] = } must be the same as {sum(indexer_seq_lens_cpu) = }" - - # Validate indices when logical tokens exceed physical capacity - # This is likely to be triggered by PP with high kv reuse & parallelism - kv_cache_capacity = ( - forward_batch.token_to_kv_pool.size - + forward_batch.token_to_kv_pool.page_size - ) - if forward_batch.seq_lens_sum > kv_cache_capacity: - max_idx = page_table_1_flattened.max().item() - assert max_idx < kv_cache_capacity, ( - f"Invalid page table index: max={max_idx}, " - f"kv_cache_capacity={kv_cache_capacity}" - ) - - if topk_transform_method == TopkTransformMethod.RAGGED: - topk_indices_offset = torch.repeat_interleave( - cu_seqlens_k[:-1], - extend_seq_lens, - ) - else: - assert False, f"Unsupported {forward_batch.forward_mode = }" - - indexer_k_start_end, token_to_batch_idx = self._cal_indexer_k_start_end( - forward_batch, bs_idx_cpu - ) - # 1D, expanded seqlens (1D means cheap to compute, so always compute it) - nsa_cache_seqlens_int32 = compute_nsa_seqlens( - original_seq_lens=seqlens_expanded, - nsa_index_topk=self.nsa_index_topk, - ) - nsa_cache_seqlens_int32 = pad_nsa_cache_seqlens( - forward_batch, nsa_cache_seqlens_int32 - ) - nsa_cu_seqlens_k = compute_cu_seqlens(nsa_cache_seqlens_int32) - nsa_cu_seqlens_q = self.get_device_int32_arange(len(nsa_cu_seqlens_k)) - - paged_mqa_schedule_metadata = None - # DeepGEMM paged MQA logits path needs a schedule metadata tensor. - # Compute it once per forward batch and reuse it across layers. - if is_cuda() and ( - forward_batch.forward_mode.is_decode_or_idle() - or forward_batch.forward_mode.is_target_verify() - or forward_batch.forward_mode.is_draft_extend(include_v2=True) - ): - try: - import deep_gemm - - # NOTE: DeepGEMM paged path uses block_size=64. - seqlens_32 = ( - seqlens_expanded - if ( - forward_batch.forward_mode.is_target_verify() - or forward_batch.forward_mode.is_draft_extend(include_v2=True) - ) - else cache_seqlens_int32 - ) - seqlens_32_2d = _to_2d_context_lens( - seqlens_32, forward_batch.batch_size - ) - paged_mqa_schedule_metadata = deep_gemm.get_paged_mqa_logits_metadata( - seqlens_32_2d, 64, deep_gemm.get_num_sms() - ) - except (ImportError, ModuleNotFoundError): - paged_mqa_schedule_metadata = None - - metadata = NSAMetadata( - page_size=self.real_page_size, - cache_seqlens_int32=cache_seqlens_int32, - max_seq_len_q=max_seqlen_q, - max_seq_len_k=max_seqlen_k, - cu_seqlens_q=cu_seqlens_q, - cu_seqlens_k=cu_seqlens_k, - seq_lens_sum=forward_batch.seq_lens_sum, - page_table_1=page_table, - page_table_1_flattened=page_table_1_flattened, - flashmla_metadata=( - self._compute_flashmla_metadata( - cache_seqlens=nsa_cache_seqlens_int32, - seq_len_q=1, - ) - if use_flashmla_kv - else None - ), - paged_mqa_schedule_metadata=paged_mqa_schedule_metadata, - nsa_cache_seqlens_int32=nsa_cache_seqlens_int32, - nsa_cu_seqlens_q=nsa_cu_seqlens_q, - nsa_cu_seqlens_k=nsa_cu_seqlens_k, - nsa_seqlens_expanded=seqlens_expanded, - nsa_extend_seq_lens_list=extend_seq_lens_cpu, - real_page_table=self._transform_table_1_to_real(page_table), - nsa_max_seqlen_q=1, - topk_indices_offset=topk_indices_offset, - indexer_k_start_end=indexer_k_start_end, - indexer_seq_lens_cpu=indexer_seq_lens_cpu, - indexer_seq_lens=indexer_seq_lens, - token_to_batch_idx=token_to_batch_idx, - ) - self.forward_metadata = metadata - - def _cal_indexer_k_start_end( - self, - forward_batch: ForwardBatch, - bs_idx: Optional[List[int]] = None, - ): - if not forward_batch.forward_mode.is_extend_without_speculative(): - return None, None - if forward_batch.batch_size == 0 or (bs_idx is not None and len(bs_idx) == 0): - empty_t = torch.empty(0, dtype=torch.int32, device=self.device) - return (empty_t, empty_t), empty_t - - # Suppose there are two requests, with extend_seq_len = [3, 2] - # and seq_lens = [10, 4] - # The logits matrix looks like this, with * representing the valid logits - # and - representing the invalid logits: - # - # ********--|---- - # *********-|---- - # **********|---- - # ----------|***- - # ----------|**** - # - # ks = [0, 0, 0, 10, 10] - # ke = [8, 9, 10, 13, 14] - ks_list = [] - ke_list = [] - token_to_batch_idx = [] - - q_offset = 0 - k_offset = 0 - - assert ( - forward_batch.seq_lens_cpu is not None - and forward_batch.extend_seq_lens_cpu is not None - ) - for i in range(forward_batch.batch_size): - seq_len = forward_batch.seq_lens_cpu[i].item() - assert isinstance(seq_len, int) - extend_seq_len = forward_batch.extend_seq_lens_cpu[i] - ks = torch.full( - (extend_seq_len,), k_offset, dtype=torch.int32, device=self.device - ) - kv_len = seq_len - if forward_batch.forward_mode.is_target_verify(): - kv_len += self.speculative_num_draft_tokens - seq_lens_expanded = torch.arange( - kv_len - extend_seq_len + 1, - kv_len + 1, - dtype=torch.int32, - device=self.device, - ) - ke = ks + seq_lens_expanded - ks_list.append(ks) - ke_list.append(ke) - - # bi: The index within the selected batch bs_idx. Entries that were not selected are ignored. - bi = bs_idx.index(i) if (bs_idx is not None and i in bs_idx) else i - tb = torch.full( - (extend_seq_len,), bi, dtype=torch.int32, device=self.device - ) - token_to_batch_idx.append(tb) - - if bs_idx is None or i in bs_idx: # skip batch not included in bs_idx - q_offset += extend_seq_len - k_offset += seq_len - - ks = torch.cat(ks_list, dim=0) - ke = torch.cat(ke_list, dim=0) - token_to_batch_idx = torch.cat(token_to_batch_idx, dim=0) - if bs_idx is not None: - assert can_nsa_prefill_cp_round_robin_split(forward_batch) - ks = nsa_cp_round_robin_split_data(ks) - ke = nsa_cp_round_robin_split_data(ke) - token_to_batch_idx = nsa_cp_round_robin_split_data(token_to_batch_idx) - return (ks, ke), token_to_batch_idx - - def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int): - """Initialize CUDA graph state for the attention backend. - - Args: - max_bs (int): Maximum batch size to support in CUDA graphs - - This creates fixed-size tensors that will be reused during CUDA graph replay - to avoid memory allocations. - """ - self.decode_cuda_graph_metadata: Dict = { - "cache_seqlens": torch.ones( - max_num_tokens, dtype=torch.int32, device=self.device - ), - "cu_seqlens_q": torch.arange( - 0, max_bs + 1, dtype=torch.int32, device=self.device - ), - "cu_seqlens_k": torch.zeros( - max_bs + 1, dtype=torch.int32, device=self.device - ), - # fake page_table for sparse_prefill - # Add extra columns for speculative draft tokens to avoid - # overflow during target_verify when max_seqlen_k = seq_len + num_draft_tokens - "page_table": torch.zeros( - max_num_tokens, - self.max_context_len + (self.speculative_num_draft_tokens or 0), - dtype=torch.int32, - device=self.device, - ), - "flashmla_metadata": ( - self._compute_flashmla_metadata( - cache_seqlens=torch.ones( - max_num_tokens, dtype=torch.int32, device=self.device - ), - seq_len_q=1, - ) - if self.nsa_decode_impl == "flashmla_kv" - else None - ), - } - - def init_forward_metadata_capture_cuda_graph( - self, - bs: int, - num_tokens: int, - req_pool_indices: torch.Tensor, - seq_lens: torch.Tensor, - encoder_lens: Optional[torch.Tensor], - forward_mode: ForwardMode, - spec_info: Optional[SpecInput], - ): - self.set_nsa_prefill_impl(forward_batch=None) - - """Initialize forward metadata for capturing CUDA graph.""" - if forward_mode.is_decode_or_idle(): - # Normal Decode - # Get sequence information - cache_seqlens_int32 = seq_lens.to(torch.int32) - cu_seqlens_k = compute_cu_seqlens(cache_seqlens_int32) - - # Use max context length for seq_len_k - page_table_1 = self.decode_cuda_graph_metadata["page_table"][:bs, :] - max_seqlen_q = 1 - max_seqlen_k = page_table_1.shape[1] - - # Precompute page table - # Precompute cumulative sequence lengths - - # NOTE(dark): this is always arange, since we are decoding - cu_seqlens_q = self.decode_cuda_graph_metadata["cu_seqlens_q"][: bs + 1] - nsa_cache_seqlens_int32 = compute_nsa_seqlens( - cache_seqlens_int32, nsa_index_topk=self.nsa_index_topk - ) - - seqlens_expanded = cache_seqlens_int32 - nsa_extend_seq_lens_list = [1] * num_tokens - if self.nsa_decode_impl == "flashmla_kv": - flashmla_metadata = self.decode_cuda_graph_metadata[ - "flashmla_metadata" - ].slice(slice(0, num_tokens + 1)) - flashmla_metadata.copy_( - self._compute_flashmla_metadata( - cache_seqlens=nsa_cache_seqlens_int32, - seq_len_q=1, - ) - ) - else: - flashmla_metadata = None - elif forward_mode.is_target_verify() or forward_mode.is_draft_extend( - include_v2=True - ): - cache_seqlens_int32 = (seq_lens + self.speculative_num_draft_tokens).to( - torch.int32 - ) - cu_seqlens_k = compute_cu_seqlens(cache_seqlens_int32) - max_seqlen_q = 1 - page_table_1 = self.decode_cuda_graph_metadata["page_table"][ - : bs * self.speculative_num_draft_tokens, : - ] - max_seqlen_k = page_table_1.shape[1] - - cu_seqlens_q = torch.arange( - 0, - bs * self.speculative_num_draft_tokens + 1, - 1, - dtype=torch.int32, - device=self.device, - ) - - 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.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, - ) - ] - ) - nsa_cache_seqlens_int32 = compute_nsa_seqlens( - seqlens_expanded, nsa_index_topk=self.nsa_index_topk - ) - nsa_extend_seq_lens_list = [1] * bs * self.speculative_num_draft_tokens - - if self.nsa_decode_impl == "flashmla_kv": - flashmla_metadata = self.decode_cuda_graph_metadata[ - "flashmla_metadata" - ].slice(slice(0, bs * self.speculative_num_draft_tokens + 1)) - - flashmla_metadata.copy_( - self._compute_flashmla_metadata( - cache_seqlens=nsa_cache_seqlens_int32, - seq_len_q=1, - ) - ) - else: - flashmla_metadata = None - - nsa_cu_seqlens_k = compute_cu_seqlens(nsa_cache_seqlens_int32) - nsa_cu_seqlens_q = self.get_device_int32_arange(len(nsa_cu_seqlens_k)) - real_page_table = self._transform_table_1_to_real(page_table_1) - - paged_mqa_schedule_metadata = None - if is_cuda() and ( - forward_mode.is_decode_or_idle() - or forward_mode.is_target_verify() - or forward_mode.is_draft_extend(include_v2=True) - ): - try: - import deep_gemm - - seqlens_32 = ( - seqlens_expanded - if ( - forward_mode.is_target_verify() - or forward_mode.is_draft_extend(include_v2=True) - ) - else cache_seqlens_int32 - ) - seqlens_32_2d = _to_2d_context_lens(seqlens_32, bs) - paged_mqa_schedule_metadata = deep_gemm.get_paged_mqa_logits_metadata( - seqlens_32_2d, 64, deep_gemm.get_num_sms() - ) - except (ImportError, ModuleNotFoundError): - paged_mqa_schedule_metadata = None - - metadata = NSAMetadata( - page_size=self.real_page_size, - cache_seqlens_int32=cache_seqlens_int32, - max_seq_len_q=max_seqlen_q, - max_seq_len_k=max_seqlen_k, - cu_seqlens_q=cu_seqlens_q, - cu_seqlens_k=cu_seqlens_k, - page_table_1=page_table_1, - flashmla_metadata=flashmla_metadata, - paged_mqa_schedule_metadata=paged_mqa_schedule_metadata, - nsa_cache_seqlens_int32=nsa_cache_seqlens_int32, - nsa_cu_seqlens_q=nsa_cu_seqlens_q, - nsa_cu_seqlens_k=nsa_cu_seqlens_k, - nsa_seqlens_expanded=seqlens_expanded, - real_page_table=real_page_table, - nsa_extend_seq_lens_list=nsa_extend_seq_lens_list, - ) - self.decode_cuda_graph_metadata[bs] = metadata - self.forward_metadata = metadata - - def init_forward_metadata_replay_cuda_graph( - self, - bs: int, - req_pool_indices: torch.Tensor, - seq_lens: torch.Tensor, - seq_lens_sum: int, - encoder_lens: Optional[torch.Tensor], - forward_mode: ForwardMode, - spec_info: Optional[SpecInput], - seq_lens_cpu: Optional[torch.Tensor], - out_cache_loc: Optional[torch.Tensor] = None, - actual_forward_mode: Optional[ForwardMode] = None, - ): - """Initialize forward metadata for replaying CUDA graph.""" - assert seq_lens_cpu is not None - - self.set_nsa_prefill_impl(forward_batch=None) - - seq_lens = seq_lens[:bs] - seq_lens_cpu = seq_lens_cpu[:bs] - req_pool_indices = req_pool_indices[:bs] - - # Normal Decode - metadata: NSAMetadata = self.decode_cuda_graph_metadata[bs] - if forward_mode.is_decode_or_idle(): - # Normal Decode - max_len = int(seq_lens_cpu.max().item()) - - cache_seqlens = seq_lens.to(torch.int32) - metadata.cache_seqlens_int32.copy_(cache_seqlens) - metadata.cu_seqlens_k[1:].copy_( - torch.cumsum(cache_seqlens, dim=0, dtype=torch.int32) - ) - page_indices = self.req_to_token[req_pool_indices, :max_len] - metadata.page_table_1[:, :max_len].copy_(page_indices) - nsa_cache_seqlens = compute_nsa_seqlens( - cache_seqlens, nsa_index_topk=self.nsa_index_topk - ) - metadata.nsa_cache_seqlens_int32.copy_(nsa_cache_seqlens) - seqlens_expanded = cache_seqlens - elif forward_mode.is_target_verify(): - max_seqlen_k = int( - seq_lens_cpu.max().item() + self.speculative_num_draft_tokens - ) - - cache_seqlens = (seq_lens + self.speculative_num_draft_tokens).to( - torch.int32 - ) - metadata.cache_seqlens_int32.copy_(cache_seqlens) - metadata.cu_seqlens_k[1:].copy_( - torch.cumsum(cache_seqlens, dim=0, dtype=torch.int32) - ) - 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 - ) - metadata.page_table_1[:, :max_seqlen_k].copy_(page_indices) - extend_seq_lens_cpu = [self.speculative_num_draft_tokens] * bs - - seqlens_expanded = seqlens_expand_triton( - torch.tensor( - extend_seq_lens_cpu, dtype=torch.int32, device=self.device - ), - cache_seqlens, - self.speculative_num_draft_tokens * bs, - self.speculative_num_draft_tokens, - ) - metadata.nsa_seqlens_expanded.copy_(seqlens_expanded) - nsa_cache_seqlens = compute_nsa_seqlens( - seqlens_expanded, self.nsa_index_topk - ) - metadata.nsa_cache_seqlens_int32.copy_(nsa_cache_seqlens) - elif forward_mode.is_draft_extend(include_v2=True): - max_seqlen_k = int(seq_lens_cpu.max().item()) - cache_seqlens = seq_lens.to(torch.int32) - metadata.cache_seqlens_int32.copy_(cache_seqlens) - metadata.cu_seqlens_k[1:].copy_( - torch.cumsum(cache_seqlens, dim=0, dtype=torch.int32) - ) - - extend_seq_lens = spec_info.num_accept_tokens[:bs] - extend_seq_lens_cpu = extend_seq_lens.tolist() - - 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 - ) - metadata.page_table_1[: page_indices.shape[0], :max_seqlen_k].copy_( - page_indices - ) - - seqlens_expanded = seqlens_expand_triton( - extend_seq_lens, - cache_seqlens, - sum(extend_seq_lens_cpu), - self.speculative_num_draft_tokens, - ) - metadata.nsa_seqlens_expanded[: seqlens_expanded.shape[0]].copy_( - seqlens_expanded - ) - nsa_cache_seqlens = compute_nsa_seqlens( - seqlens_expanded, self.nsa_index_topk - ) - metadata.nsa_cache_seqlens_int32[: seqlens_expanded.shape[0]].copy_( - nsa_cache_seqlens - ) - - # Update DeepGEMM paged MQA schedule metadata outside the captured graph. - if is_cuda() and ( - forward_mode.is_decode_or_idle() - or forward_mode.is_target_verify() - or forward_mode.is_draft_extend(include_v2=True) - ): - try: - import deep_gemm - - seqlens_32 = ( - seqlens_expanded - if ( - forward_mode.is_target_verify() - or forward_mode.is_draft_extend(include_v2=True) - ) - else metadata.cache_seqlens_int32 - ) - seqlens_32_2d = _to_2d_context_lens(seqlens_32, bs) - new_schedule = deep_gemm.get_paged_mqa_logits_metadata( - seqlens_32_2d, 64, deep_gemm.get_num_sms() - ) - if metadata.paged_mqa_schedule_metadata is None: - object.__setattr__( - metadata, "paged_mqa_schedule_metadata", new_schedule - ) - else: - metadata.paged_mqa_schedule_metadata.copy_(new_schedule) - except (ImportError, ModuleNotFoundError): - object.__setattr__(metadata, "paged_mqa_schedule_metadata", None) - seqlens_expanded_size = seqlens_expanded.shape[0] - assert ( - metadata.nsa_cache_seqlens_int32 is not None - and metadata.nsa_cu_seqlens_k is not None - and self.nsa_index_topk is not None - ) - - metadata.nsa_cu_seqlens_k[1 : 1 + seqlens_expanded_size].copy_( - torch.cumsum(nsa_cache_seqlens, dim=0, dtype=torch.int32) - ) - # NOTE(dark): (nsa-) cu_seqlens_q is always arange, no need to copy - - assert self.real_page_size == metadata.page_size - if self.real_page_size > 1: - real_table = self._transform_table_1_to_real(page_indices) - new_rows = real_table.shape[0] - new_cols = real_table.shape[1] - metadata.real_page_table[:new_rows, :new_cols].copy_(real_table) - else: - assert metadata.real_page_table is metadata.page_table_1 - - if self.nsa_decode_impl == "flashmla_kv": - flashmla_metadata = metadata.flashmla_metadata.slice( - slice(0, seqlens_expanded_size + 1) - ) - flashmla_metadata.copy_( - self._compute_flashmla_metadata( - cache_seqlens=nsa_cache_seqlens, - seq_len_q=1, - ) - ) - - self.forward_metadata = metadata - - def init_forward_metadata_replay_cuda_graph_from_precomputed( - self, - bs: int, - precomputed: PrecomputedMetadata, - forward_mode: ForwardMode, - ): - """Fast path: copy precomputed metadata to this backend's metadata. - - This function only performs copy operations, no computation. - - Args: - bs: Batch size - precomputed: Precomputed metadata to copy from - forward_mode: Forward mode - """ - self.set_nsa_prefill_impl(forward_batch=None) - - metadata = self.decode_cuda_graph_metadata[bs] - - # Track whether fused kernel succeeded - fused_kernel_succeeded = False - - # Use fused CUDA kernel for all copy operations - if _USE_FUSED_METADATA_COPY: - try: - from sglang.jit_kernel.fused_metadata_copy import ( - fused_metadata_copy_cuda, - ) - - # Map forward_mode to integer enum - if forward_mode.is_decode_or_idle(): - mode_int = 0 # DECODE - elif forward_mode.is_target_verify(): - mode_int = 1 # TARGET_VERIFY - elif forward_mode.is_draft_extend(): - mode_int = 2 # DRAFT_EXTEND - else: - raise ValueError(f"Unsupported forward_mode: {forward_mode}") - - # Prepare FlashMLA tensors if needed - flashmla_num_splits_src = None - flashmla_num_splits_dst = None - flashmla_metadata_src = None - flashmla_metadata_dst = None - if precomputed.flashmla_metadata is not None: - flashmla_num_splits_src = precomputed.flashmla_metadata.num_splits - flashmla_num_splits_dst = metadata.flashmla_metadata.num_splits - flashmla_metadata_src = ( - precomputed.flashmla_metadata.flashmla_metadata - ) - flashmla_metadata_dst = metadata.flashmla_metadata.flashmla_metadata - - # Call fused kernel - fused_metadata_copy_cuda( - # Source tensors - precomputed.cache_seqlens, - precomputed.cu_seqlens_k, - precomputed.page_indices, - precomputed.nsa_cache_seqlens, - precomputed.seqlens_expanded, - precomputed.nsa_cu_seqlens_k, - precomputed.real_page_table, - flashmla_num_splits_src, - flashmla_metadata_src, - # Destination tensors - metadata.cache_seqlens_int32, - metadata.cu_seqlens_k, - metadata.page_table_1, - metadata.nsa_cache_seqlens_int32, - metadata.nsa_seqlens_expanded, - metadata.nsa_cu_seqlens_k, - ( - metadata.real_page_table - if precomputed.real_page_table is not None - else None - ), - flashmla_num_splits_dst, - flashmla_metadata_dst, - # Parameters - mode_int, - bs, - precomputed.max_len, - precomputed.max_seqlen_k, - precomputed.seqlens_expanded_size, - ) - - # Successfully used fused kernel - fused_kernel_succeeded = True - - except ImportError: - print( - "Warning: Fused metadata copy kernel not available, falling back to individual copies." - ) - except Exception as e: - print( - f"Warning: Fused metadata copy kernel failed with error: {e}, falling back to individual copies." - ) - - # Fallback to individual copy operations if fused kernel disabled or failed - if not fused_kernel_succeeded: - # Copy basic seqlens - metadata.cache_seqlens_int32.copy_(precomputed.cache_seqlens) - metadata.cu_seqlens_k[1:].copy_(precomputed.cu_seqlens_k[1:]) - - # Mode-specific copy logic - if forward_mode.is_decode_or_idle(): - # Decode mode - metadata.page_table_1[:, : precomputed.max_len].copy_( - precomputed.page_indices - ) - metadata.nsa_cache_seqlens_int32.copy_(precomputed.nsa_cache_seqlens) - # seqlens_expanded is same as cache_seqlens (already copied) - - elif forward_mode.is_target_verify(): - # Target verify mode - metadata.page_table_1[:, : precomputed.max_seqlen_k].copy_( - precomputed.page_indices - ) - metadata.nsa_seqlens_expanded.copy_(precomputed.seqlens_expanded) - metadata.nsa_cache_seqlens_int32.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 - metadata.page_table_1[:rows, :cols].copy_(precomputed.page_indices) - - size = precomputed.seqlens_expanded_size - metadata.nsa_seqlens_expanded[:size].copy_(precomputed.seqlens_expanded) - metadata.nsa_cache_seqlens_int32[:size].copy_( - precomputed.nsa_cache_seqlens - ) - - # Copy NSA cu_seqlens - size = precomputed.seqlens_expanded_size - metadata.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 - metadata.real_page_table[:rows, :cols].copy_( - precomputed.real_page_table - ) - - # Copy FlashMLA metadata in fallback path - if precomputed.flashmla_metadata is not None: - size = precomputed.seqlens_expanded_size - flashmla_metadata = metadata.flashmla_metadata.slice(slice(0, size + 1)) - flashmla_metadata.copy_(precomputed.flashmla_metadata) - - # Refresh DeepGEMM paged MQA schedule metadata for the actual seqlens of - # this replay (the captured graph holds stale data otherwise, which can - # deadlock the kernel when the runtime work decomposition diverges from - # the captured one). - if is_cuda(): - try: - import deep_gemm - - if forward_mode.is_decode_or_idle(): - seqlens_32 = metadata.cache_seqlens_int32 - else: - seqlens_32 = metadata.nsa_seqlens_expanded[ - : precomputed.seqlens_expanded_size - ] - seqlens_32_2d = _to_2d_context_lens(seqlens_32, bs) - new_schedule = deep_gemm.get_paged_mqa_logits_metadata( - seqlens_32_2d, 64, deep_gemm.get_num_sms() - ) - if metadata.paged_mqa_schedule_metadata is None: - object.__setattr__( - metadata, "paged_mqa_schedule_metadata", new_schedule - ) - else: - metadata.paged_mqa_schedule_metadata.copy_(new_schedule) - except (ImportError, ModuleNotFoundError): - pass - - self.forward_metadata = metadata - - def forward_extend( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - layer: RadixAttention, - forward_batch: ForwardBatch, - save_kv_cache=True, - # For multi-head latent attention - q_rope: Optional[torch.Tensor] = None, - k_rope: Optional[torch.Tensor] = None, - topk_indices: Optional[torch.Tensor] = None, - cos_sin_cache: Optional[torch.Tensor] = None, - is_neox: Optional[bool] = False, - llama_4_scaling: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - - causal = not layer.is_cross_attention - metadata = self.forward_metadata - assert causal, "NSA is causal only" - - nsa_impl = ( - self.nsa_decode_impl - if ( - forward_batch.forward_mode.is_target_verify() - or forward_batch.forward_mode.is_draft_extend(include_v2=True) - ) - else self.nsa_prefill_impl - ) - - if nsa_impl == "trtllm" and not self.use_mha: - return self._forward_trtllm( - q, - k, - v, - layer, - forward_batch, - metadata.nsa_cache_seqlens_int32, - save_kv_cache, - q_rope, - k_rope, - topk_indices, - cos_sin_cache, - is_neox, - llama_4_scaling, - is_prefill=True, - ) - - if k is not None: - assert v is not None - if save_kv_cache: - cache_loc = ( - forward_batch.out_cache_loc - if not layer.is_cross_attention - else forward_batch.encoder_out_cache_loc - ) - forward_batch.token_to_kv_pool.set_mla_kv_buffer( # type: ignore - layer, - cache_loc, - k, - k_rope, - ) - - # Use MHA kernel if in MHA_ONE_SHOT mode - if self.use_mha: - assert k is not None and v is not None - assert q_rope is None, "MHA_ONE_SHOT path should not pass q_rope" - assert ( - layer.tp_k_head_num == layer.tp_q_head_num > 1 - ), "MHA_ONE_SHOT requires dense multi-head config" - return self._forward_standard_mha( - q=q, - k=k, - v=v, - layer=layer, - forward_batch=forward_batch, - metadata=metadata, - ) - - # Do absorbed multi-latent attention (MLA path) - assert q_rope is not None - kv_cache = forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id) - - if q_rope is not None: - q_nope = q.view(-1, layer.tp_q_head_num, layer.v_head_dim) - q_rope = q_rope.view( - -1, layer.tp_q_head_num, layer.head_dim - layer.v_head_dim - ) - else: - q_all = q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim) - q_nope = q_all[:, :, : layer.v_head_dim] - q_rope = q_all[:, :, layer.v_head_dim :] - - # Align topk_indices with q dimensions - # This handles cases where q is padded (TP + partial DP attention) - if topk_indices is not None: - topk_indices = self._pad_topk_indices(topk_indices, q_nope.shape[0]) - - # NOTE(dark): here, we use page size = 1 - topk_transform_method = self.get_topk_transform_method( - forward_batch.forward_mode - ) - if envs.SGLANG_NSA_FUSE_TOPK.get(): - page_table_1 = topk_indices - else: - if topk_transform_method == TopkTransformMethod.RAGGED: - topk_indices_offset = metadata.topk_indices_offset - assert topk_indices_offset is not None - mask = topk_indices != -1 - topk_indices_offset = ( - topk_indices_offset.unsqueeze(1) - if topk_indices_offset.ndim == 1 - else topk_indices_offset - ) - topk_indices = torch.where( - mask, topk_indices + topk_indices_offset, topk_indices - ) - elif topk_transform_method == TopkTransformMethod.PAGED: - assert metadata.nsa_extend_seq_lens_list is not None - page_table_1 = transform_index_page_table_prefill( - page_table=metadata.page_table_1, - topk_indices=topk_indices, - extend_lens_cpu=metadata.nsa_extend_seq_lens_list, - page_size=1, - ) - - # todo hisparse: to cover more backends - if forward_batch.hisparse_coordinator is not None: - page_table_1 = ( - forward_batch.token_to_kv_pool.translate_loc_to_hisparse_device( - page_table_1 - ) - ) - - if nsa_impl == "tilelang": - if q_rope is not None: - q_all = concat_mla_absorb_q_general(q_nope, q_rope) - return self._forward_tilelang( - q_all=q_all, - kv_cache=kv_cache, - page_table_1=page_table_1, - sm_scale=layer.scaling, - v_head_dim=layer.v_head_dim, - ) - elif nsa_impl == "flashmla_sparse": - if q_rope is not None: - q_all = concat_mla_absorb_q_general(q_nope, q_rope) - - if topk_transform_method == TopkTransformMethod.RAGGED: - if any(forward_batch.extend_prefix_lens_cpu): - page_table_1_flattened = ( - self.forward_metadata.page_table_1_flattened - ) - assert page_table_1_flattened is not None - kv_cache = dequantize_k_cache_paged( - kv_cache, page_table_1_flattened - ) - else: - kv_cache = _cat([k, k_rope], dim=-1) - page_table_1 = topk_indices - - return self._forward_flashmla_sparse( - q_all=q_all, - kv_cache=kv_cache, - page_table_1=page_table_1, - sm_scale=layer.scaling, - v_head_dim=layer.v_head_dim, - ) - elif nsa_impl == "flashmla_kv": - if q_rope is not None: - q_all = concat_mla_absorb_q_general(q_nope, q_rope) - return self._forward_flashmla_kv( - q_all=q_all, - kv_cache=kv_cache, - sm_scale=layer.scaling, - v_head_dim=layer.v_head_dim, - # TODO optimize args - layer=layer, - metadata=metadata, - page_table_1=page_table_1, - ) - elif nsa_impl == "fa3": - return self._forward_fa3( - q_rope=q_rope, - kv_cache=kv_cache, - v_head_dim=layer.v_head_dim, - q_nope=q_nope, - page_table=page_table_1, - cache_seqlens=metadata.nsa_cache_seqlens_int32, - cu_seqlens_q=metadata.nsa_cu_seqlens_q, - cu_seqlens_k=metadata.nsa_cu_seqlens_k, - max_seqlen_q=metadata.nsa_max_seqlen_q, - sm_scale=layer.scaling, - logit_cap=layer.logit_cap, - page_size=1, - ) - elif nsa_impl == "aiter": - if q_rope is not None: - q_all = torch.cat([q_nope, q_rope], dim=-1) - return self._forward_aiter_extend( - q_all=q_all, - kv_cache=kv_cache, - page_table_1=page_table_1, - layer=layer, - ) - else: - raise ValueError( - f"Unsupported {nsa_impl = } for forward_extend. Consider using an other attention backend." - ) - - def forward_decode( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - layer: RadixAttention, - forward_batch: ForwardBatch, - save_kv_cache=True, - # For multi-head latent attention - q_rope: Optional[torch.Tensor] = None, - k_rope: Optional[torch.Tensor] = None, - topk_indices: Optional[torch.Tensor] = None, - cos_sin_cache: Optional[torch.Tensor] = None, - is_neox: Optional[bool] = False, - llama_4_scaling: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - - causal = not layer.is_cross_attention - metadata = self.forward_metadata - assert causal, "NSA is causal only" - - if self.nsa_decode_impl == "trtllm": - return self._forward_trtllm( - q, - k, - v, - layer, - forward_batch, - metadata.cache_seqlens_int32, - save_kv_cache, - q_rope, - k_rope, - topk_indices, - cos_sin_cache, - is_neox, - llama_4_scaling, - ) - - if k is not None: - assert v is not None - if save_kv_cache: - cache_loc = ( - forward_batch.out_cache_loc - if not layer.is_cross_attention - else forward_batch.encoder_out_cache_loc - ) - forward_batch.token_to_kv_pool.set_mla_kv_buffer( # type: ignore - layer, - cache_loc, - k, - k_rope, - ) - - # Do absorbed multi-latent attention - kv_cache = forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id) - if q_rope is not None: - q_nope = q.view(-1, layer.tp_q_head_num, layer.v_head_dim) - q_rope = q_rope.view( - -1, layer.tp_q_head_num, layer.head_dim - layer.v_head_dim - ) - # Caller passed split q_nope / q_rope; we'll need to concat below if - # the chosen impl wants q_all. - q_all = None - else: - # Caller passed already-concatenated q (q_all = q). Reuse it directly - # via a zero-copy view; the impl-specific blocks below will skip the - # otherwise redundant concat_mla_absorb_q_general call. - q_all = q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim) - q_nope = q_all[:, :, : layer.v_head_dim] - q_rope = q_all[:, :, layer.v_head_dim :] - - # Align topk_indices with q dimensions - if topk_indices is not None: - topk_indices = self._pad_topk_indices(topk_indices, q_nope.shape[0]) - - if forward_batch.hisparse_coordinator is not None: - page_table_1 = forward_batch.hisparse_coordinator.swap_in_selected_pages( - forward_batch.req_pool_indices, - forward_batch.seq_lens, - topk_indices, - layer.layer_id, - ) - elif envs.SGLANG_NSA_FUSE_TOPK.get(): - page_table_1 = topk_indices - else: - page_table_1 = transform_index_page_table_decode( - page_table=metadata.page_table_1, - topk_indices=topk_indices, - page_size=1, - ) - - if self.nsa_decode_impl == "flashmla_sparse": - if q_rope is not None: - q_all = concat_mla_absorb_q_general(q_nope, q_rope) - return self._forward_flashmla_sparse( - q_all=q_all, - kv_cache=kv_cache, - page_table_1=page_table_1, - sm_scale=layer.scaling, - v_head_dim=layer.v_head_dim, - ) - elif self.nsa_decode_impl == "flashmla_kv": - if q_rope is not None: - q_all = concat_mla_absorb_q_general(q_nope, q_rope) - return self._forward_flashmla_kv( - q_all=q_all, - kv_cache=kv_cache, - sm_scale=layer.scaling, - v_head_dim=layer.v_head_dim, - # TODO optimize args - layer=layer, - metadata=metadata, - page_table_1=page_table_1, - ) - elif self.nsa_decode_impl == "tilelang": - # Cat-skip (HIP-only): when caller passes q_rope=None on HIP, q_all - # has already been set to a zero-copy view of q in the else branch - # above and we can reuse it directly. The `not _is_hip` clause keeps - # CUDA / MUSA paths byte-identical to pre-patch by always re-cat. - if q_all is None or not _is_hip: - q_all = concat_mla_absorb_q_general(q_nope, q_rope) - return self._forward_tilelang( - q_all=q_all, - kv_cache=kv_cache, - page_table_1=page_table_1, - sm_scale=layer.scaling, - v_head_dim=layer.v_head_dim, - ) - elif self.nsa_decode_impl == "fa3": - return self._forward_fa3( - q_rope=q_rope, - kv_cache=kv_cache, - v_head_dim=layer.v_head_dim, - q_nope=q_nope, - page_table=page_table_1, - cache_seqlens=metadata.nsa_cache_seqlens_int32, - cu_seqlens_q=metadata.nsa_cu_seqlens_q, - cu_seqlens_k=metadata.nsa_cu_seqlens_k, - max_seqlen_q=metadata.nsa_max_seqlen_q, - sm_scale=layer.scaling, - logit_cap=layer.logit_cap, - page_size=1, - ) - elif self.nsa_decode_impl == "aiter": - if q_all is None or not _is_hip: - q_all = torch.cat([q_nope, q_rope], dim=-1) - return self._forward_aiter( - q_all=q_all, - kv_cache=kv_cache, - page_table_1=page_table_1, - layer=layer, - metadata=metadata, - bs=forward_batch.batch_size, - ) - - else: - assert False, f"Unsupported {self.nsa_decode_impl = }" - - def _forward_fa3( - self, - q_rope: torch.Tensor, - kv_cache: torch.Tensor, - v_head_dim: int, - q_nope: torch.Tensor, - page_table: torch.Tensor, - cache_seqlens: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_k: torch.Tensor, - max_seqlen_q: int, - sm_scale: float, - logit_cap: float, - page_size: int, - ) -> torch.Tensor: - k_rope_cache = kv_cache[:, :, v_head_dim:] - c_kv_cache = kv_cache[:, :, :v_head_dim] - qk_rope_dim = k_rope_cache.shape[-1] - k_rope_cache = k_rope_cache.view(-1, page_size, 1, qk_rope_dim) - c_kv_cache = c_kv_cache.view(-1, page_size, 1, v_head_dim) - o = flash_attn_with_kvcache( - q=q_rope, - k_cache=k_rope_cache, - v_cache=c_kv_cache, - qv=q_nope, - page_table=page_table, - cache_seqlens=cache_seqlens, - cu_seqlens_q=cu_seqlens_q, - cu_seqlens_k_new=cu_seqlens_k, - max_seqlen_q=max_seqlen_q, - softmax_scale=sm_scale, - causal=True, - softcap=logit_cap, - return_softmax_lse=False, - num_splits=self.num_splits, - ) - return o # type: ignore - - def _forward_flashmla_sparse( - self, - q_all: torch.Tensor, - kv_cache: torch.Tensor, - v_head_dim: int, - page_table_1: torch.Tensor, - sm_scale: float, - ) -> torch.Tensor: - from sgl_kernel.flash_mla import flash_mla_sparse_fwd - - # FlashMLA sparse kernel requires num_heads to be a multiple of 64 (Hopper) or 128 (Blackwell) - # When using TP, num_heads might be smaller (e.g., 256//8=32) - num_tokens, num_heads, head_dim = q_all.shape - - # Determine required padding based on GPU architecture (use cached value) - required_padding = 128 if self.device_sm_major >= 10 else 64 - - need_padding = num_heads % required_padding != 0 - - if need_padding: - assert required_padding % num_heads == 0, ( - f"num_heads {num_heads} cannot be padded to {required_padding}. " - f"TP size may be too large for this model." - ) - - # Pad q to required size - q_padded = q_all.new_zeros((num_tokens, required_padding, head_dim)) - q_padded[:, :num_heads, :] = q_all - q_input = q_padded - else: - q_input = q_all - - # indices shape must be (s_q, h_kv=1, topk), keep h_kv=1 unchanged - indices_input = page_table_1.unsqueeze(1) - - o, _, _ = flash_mla_sparse_fwd( - q=q_input, - kv=kv_cache, - indices=indices_input, - sm_scale=sm_scale, - d_v=v_head_dim, - ) - - # Trim output back to original num_heads if we padded - if need_padding: - o = o[:, :num_heads, :] - - return o - - def _forward_flashmla_kv( - self, - q_all: torch.Tensor, - kv_cache: torch.Tensor, - v_head_dim: int, - sm_scale: float, - layer, - metadata: NSAMetadata, - page_table_1, - ) -> torch.Tensor: - from sgl_kernel.flash_mla import flash_mla_with_kvcache - - cache_seqlens = metadata.nsa_cache_seqlens_int32 - assert metadata.flashmla_metadata is not None - - # TODO the 2nd dim is seq_len_q, need to be >1 when MTP - q_all = q_all.view(-1, 1, layer.tp_q_head_num, layer.head_dim) - num_q_heads = q_all.shape[2] - target_q_heads = self.flashmla_kv_num_q_heads - if target_q_heads != num_q_heads: - # Pad q heads to match FlashMLA decode supported head-count variants. - q_input = q_all.new_zeros( - q_all.shape[0], q_all.shape[1], target_q_heads, q_all.shape[3] - ) - q_input[:, :, :num_q_heads, :] = q_all - else: - q_input = q_all - - kv_cache = kv_cache.view(-1, self.real_page_size, 1, self.kv_cache_dim) - assert self.real_page_size == 64, "only page size 64 is supported" - - if not self.nsa_kv_cache_store_fp8: - # inefficiently quantize the whole cache - kv_cache = quantize_k_cache(kv_cache) - - indices = page_table_1.unsqueeze(1) - assert ( - indices.shape[-1] == self.nsa_index_topk - ) # requirement of FlashMLA decode kernel - - o, _ = flash_mla_with_kvcache( - q=q_input, - k_cache=kv_cache, - cache_seqlens=cache_seqlens, - head_dim_v=v_head_dim, - tile_scheduler_metadata=metadata.flashmla_metadata.flashmla_metadata, - num_splits=metadata.flashmla_metadata.num_splits, - softmax_scale=sm_scale, - indices=indices, - # doc says it is not used, but if pass in None then error - block_table=torch.empty( - (q_all.shape[0], 0), dtype=torch.int32, device=q_all.device - ), - is_fp8_kvcache=True, - ) - - if target_q_heads != num_q_heads: - o = o[:, :, :num_q_heads, :] - - return o - - def _forward_standard_mha( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - layer: RadixAttention, - forward_batch: ForwardBatch, - metadata: NSAMetadata, - ) -> torch.Tensor: - """Standard MHA using FlashAttention varlen for MHA_ONE_SHOT mode.""" - q = q.view(-1, layer.tp_q_head_num, layer.head_dim) - k = k.view(-1, layer.tp_k_head_num, layer.head_dim) - v = v.view(-1, layer.tp_v_head_num, layer.v_head_dim) - - # MHA_ONE_SHOT: k/v include all tokens (prefix + current) - cu_seqlens_q = metadata.cu_seqlens_q - cu_seqlens_k = metadata.cu_seqlens_k - max_seqlen_k = metadata.max_seq_len_k - causal = True - - # Verify batch sizes match (length of cu_seqlens should be batch_size + 1) - assert len(cu_seqlens_q) == len(cu_seqlens_k), ( - f"batch_size mismatch: cu_seqlens_q has {len(cu_seqlens_q)-1} requests, " - f"cu_seqlens_k has {len(cu_seqlens_k)-1} requests" - ) - - # Use TRTLLm ragged attention for SM100 (Blackwell/B200) to avoid FA4 accuracy issues - if self.device_sm_major >= 10: - import flashinfer - - seq_lens = metadata.cache_seqlens_int32 - return flashinfer.prefill.trtllm_ragged_attention_deepseek( - query=q, - key=k, - value=v, - workspace_buffer=self.workspace_buffer, - seq_lens=seq_lens, - max_q_len=metadata.max_seq_len_q, - max_kv_len=max_seqlen_k, - bmm1_scale=layer.scaling, - bmm2_scale=1.0, - o_sf_scale=1.0, - batch_size=forward_batch.batch_size, - window_left=-1, - cum_seq_lens_q=cu_seqlens_q, - cum_seq_lens_kv=cu_seqlens_k, - enable_pdl=False, - is_causal=causal, - return_lse=False, - skip_softmax_threshold_scale_factor=envs.SGLANG_SKIP_SOFTMAX_PREFILL_THRESHOLD_SCALE_FACTOR.get(), - ) - - # Use FA3 for SM90 (Hopper/H200) - return flash_attn_varlen_func( - q=q, - k=k, - v=v, - cu_seqlens_q=cu_seqlens_q, - cu_seqlens_k=cu_seqlens_k, - max_seqlen_q=metadata.max_seq_len_q, - max_seqlen_k=max_seqlen_k, - softmax_scale=layer.scaling, - causal=causal, - ) - - def _forward_tilelang( - self, - q_all: torch.Tensor, - kv_cache: torch.Tensor, - v_head_dim: int, - page_table_1: torch.Tensor, - sm_scale: float, - ) -> torch.Tensor: - from sglang.srt.layers.attention.nsa.tilelang_kernel import tilelang_sparse_fwd - - return tilelang_sparse_fwd( - q=q_all, - kv=kv_cache, - indices=page_table_1.unsqueeze(1), - sm_scale=sm_scale, - d_v=v_head_dim, - ) - - def _forward_aiter( - self, - q_all: torch.Tensor, - kv_cache: torch.Tensor, - page_table_1: torch.Tensor, - layer: RadixAttention, - metadata: NSAMetadata, - bs: int, - ) -> torch.Tensor: - q = q_all.reshape(-1, layer.tp_q_head_num * layer.head_dim) - - if layer.head_dim != layer.v_head_dim: - o = q.new_empty((q.shape[0], layer.tp_q_head_num * layer.v_head_dim)) - else: - o = torch.empty_like(q) - - if self.need_pad_heads: - q_kernel = q.view( - -1, layer.tp_q_head_num, layer.head_dim - ).repeat_interleave(self.head_repeat_factor, dim=1) - o_kernel = q.new_empty( - ( - q.shape[0], - layer.tp_q_head_num * self.head_repeat_factor, - layer.v_head_dim, - ) - ) - else: - q_kernel = q.view(-1, layer.tp_q_head_num, layer.head_dim) - o_kernel = o.view(-1, layer.tp_q_head_num, layer.v_head_dim) - - kv_indptr = self.kv_indptr - - non_minus1_mask = page_table_1 != -1 - non_minus1_counts = non_minus1_mask.sum(dim=1) - kv_indptr[1 : bs + 1] = torch.cumsum(non_minus1_counts, dim=0) - - kv_indices = self.kv_indices - get_valid_kv_indices(page_table_1, kv_indptr, kv_indices, bs) - - mla_decode_fwd( - q_kernel, - kv_cache.view(-1, 1, 1, layer.head_dim), - o_kernel, - metadata.cu_seqlens_q, - kv_indptr, - kv_indices, - metadata.cu_seqlens_q, - metadata.max_seq_len_q, - sm_scale=layer.scaling, - logit_cap=layer.logit_cap, - ) - - if self.need_pad_heads: - o = o_kernel[:, :: self.head_repeat_factor, :] - - return o - - def _forward_aiter_extend( - self, - q_all: torch.Tensor, - kv_cache: torch.Tensor, - page_table_1: torch.Tensor, - layer: RadixAttention, - ) -> torch.Tensor: - num_tokens = q_all.shape[0] - q = q_all.reshape(-1, layer.tp_q_head_num * layer.head_dim) - - if layer.head_dim != layer.v_head_dim: - o = q.new_empty((num_tokens, layer.tp_q_head_num * layer.v_head_dim)) - else: - o = torch.empty_like(q) - - if self.need_pad_heads: - q_kernel = q.view( - -1, layer.tp_q_head_num, layer.head_dim - ).repeat_interleave(self.head_repeat_factor, dim=1) - o_kernel = q.new_empty( - ( - num_tokens, - layer.tp_q_head_num * self.head_repeat_factor, - layer.v_head_dim, - ) - ) - else: - q_kernel = q.view(-1, layer.tp_q_head_num, layer.head_dim) - o_kernel = o.view(-1, layer.tp_q_head_num, layer.v_head_dim) - - non_minus1_mask = page_table_1 != -1 - non_minus1_counts = non_minus1_mask.sum(dim=1) - - kv_indptr = torch.zeros(num_tokens + 1, dtype=torch.int32, device=self.device) - kv_indptr[1:] = torch.cumsum(non_minus1_counts, dim=0) - - # Allocate kv_indices with upper-bound size (num_tokens * topk) - topk = page_table_1.shape[1] - kv_indices = torch.zeros( - num_tokens * topk, dtype=torch.int32, device=self.device - ) - - # Use get_valid_kv_indices kernel to extract valid indices - get_valid_kv_indices(page_table_1, kv_indptr, kv_indices, num_tokens) - - # Build cu_seqlens_q for extend: each token is treated as seq_len_q=1 - cu_seqlens_q = torch.arange( - 0, num_tokens + 1, dtype=torch.int32, device=self.device - ) - # TODO support more forward_mode - mla_decode_fwd( - q_kernel, - kv_cache.view(-1, 1, 1, layer.head_dim), - o_kernel, - cu_seqlens_q, - kv_indptr, - kv_indices, - cu_seqlens_q, - 1, # max_seq_len_q = 1 for per-token attention - sm_scale=layer.scaling, - logit_cap=layer.logit_cap, - ) - - if self.need_pad_heads: - o = o_kernel[:, :: self.head_repeat_factor, :] - - return o - - def _forward_trtllm( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - layer: RadixAttention, - forward_batch: ForwardBatch, - seq_lens: torch.Tensor, - save_kv_cache=True, - # For multi-head latent attention - q_rope: Optional[torch.Tensor] = None, - k_rope: Optional[torch.Tensor] = None, - topk_indices: Optional[torch.Tensor] = None, - cos_sin_cache: Optional[torch.Tensor] = None, - is_neox: Optional[bool] = False, - llama_4_scaling: Optional[torch.Tensor] = None, - is_prefill: bool = False, - ) -> torch.Tensor: - """Forward using TRT-LLM sparse MLA kernel.""" - import flashinfer.decode - - metadata = self.forward_metadata - - merge_query = q_rope is not None - if self.kv_cache_dtype == torch.float8_e4m3fn: - # For FP8 path, we quantize the query and rope parts and merge them into a single tensor - # Note: rope application in deepseek_v2.py:forward_absorb_prepare is skipped for FP8 decode path of this trtllm_mla backend - assert q_rope is not None, "For FP8 path q_rope should not be None." - assert k_rope is not None, "For FP8 path k_rope should not be None." - assert ( - cos_sin_cache is not None - ), "For FP8 path cos_sin_cache should not be None." - - q, k, k_rope = mla_quantize_and_rope_for_fp8( - q, - q_rope, - k.squeeze(1), - k_rope.squeeze(1), - forward_batch.positions, - cos_sin_cache, - is_neox, - self.kv_lora_rank, - self.qk_rope_head_dim, - ) - merge_query = False - - # Save KV cache if requested - if save_kv_cache: - assert ( - k is not None and k_rope is not None - ), "For populating trtllm_mla kv cache, both k_nope and k_rope should be not None." - cache_loc = ( - forward_batch.out_cache_loc - if not layer.is_cross_attention - else forward_batch.encoder_out_cache_loc - ) - forward_batch.token_to_kv_pool.set_mla_kv_buffer( - layer, cache_loc, k, k_rope - ) - - k_cache = forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id) - kv_cache = k_cache.view(-1, self.real_page_size, self.kv_cache_dim).unsqueeze(1) - - if merge_query: - q_nope = q.view(-1, layer.tp_q_head_num, layer.v_head_dim) - q_rope_reshaped = q_rope.view( - -1, layer.tp_q_head_num, layer.head_dim - layer.v_head_dim - ) - q_all = concat_mla_absorb_q_general(q_nope, q_rope_reshaped) - else: - q_all = q.view(-1, layer.tp_q_head_num, layer.head_dim) - - # Align topk_indices with q dimensions - if topk_indices is not None: - topk_indices = self._pad_topk_indices(topk_indices, q.shape[0]) - - if envs.SGLANG_NSA_FUSE_TOPK.get(): - page_table_1 = topk_indices - elif is_prefill: - page_table_1 = transform_index_page_table_prefill( - page_table=metadata.page_table_1, - topk_indices=topk_indices, - extend_lens_cpu=metadata.nsa_extend_seq_lens_list, - page_size=1, - ) - else: - page_table_1 = transform_index_page_table_decode( - page_table=metadata.page_table_1, - topk_indices=topk_indices, - page_size=1, - ) - - q_scale = 1.0 - k_scale = ( - layer.k_scale_float - if getattr(layer, "k_scale_float", None) is not None - else 1.0 - ) - bmm1_scale = q_scale * k_scale * layer.scaling - - batch_size = page_table_1.shape[0] - _, num_heads, head_dim = q_all.shape - - q = q_all.view(batch_size, 1, num_heads, head_dim) - kv = kv_cache.view(-1, 1, self.real_page_size, self.kv_cache_dim) - block_tables = page_table_1.unsqueeze(1) - seq_lens = metadata.cache_seqlens_int32 if seq_lens is None else seq_lens - - out = flashinfer.decode.trtllm_batch_decode_with_kv_cache_mla( - query=q, - kv_cache=kv, - workspace_buffer=self.workspace_buffer, - qk_nope_head_dim=self.qk_nope_head_dim, - kv_lora_rank=self.kv_lora_rank, - qk_rope_head_dim=self.qk_rope_head_dim, - block_tables=block_tables, - seq_lens=seq_lens, - max_seq_len=metadata.max_seq_len_k, - sparse_mla_top_k=self.nsa_index_topk, - bmm1_scale=bmm1_scale, - backend="trtllm-gen", - skip_softmax_threshold_scale_factor=envs.SGLANG_SKIP_SOFTMAX_DECODE_THRESHOLD_SCALE_FACTOR.get(), - ) - # Output: [batch, q_len=1, heads, v_dim] -> [batch, heads, v_dim] - return out.squeeze(1) - - def _pad_topk_indices( - self, topk_indices: torch.Tensor, num_tokens: int - ) -> torch.Tensor: - current_tokens = topk_indices.shape[0] - if current_tokens == num_tokens: - return topk_indices - - assert current_tokens <= num_tokens, ( - f"topk_indices rows ({current_tokens}) > num_tokens ({num_tokens}); " - "this indicates a mismatch between indexer output and q layout." - ) - - pad_size = num_tokens - current_tokens - padding = torch.full( - (pad_size, topk_indices.shape[1]), - -1, - dtype=topk_indices.dtype, - device=topk_indices.device, - ) - return torch.cat([topk_indices, padding], dim=0) - - def get_cuda_graph_seq_len_fill_value(self): - """Get the fill value for sequence length in CUDA graph.""" - return 1 - - def set_nsa_prefill_impl(self, forward_batch: Optional[ForwardBatch] = None): - """ - Decide all attention prefill dispatch strategies for this batch. - """ - from sglang.srt.utils import get_device_sm, is_blackwell - - # Decide MHA vs MLA - if forward_batch and forward_batch.forward_mode.is_extend_without_speculative(): - # Check if sequence meets criteria for MHA_ONE_SHOT - assert forward_batch.seq_lens_cpu is not None - max_kv_len = forward_batch.seq_lens_cpu.max().item() - sum_seq_lens = sum(forward_batch.seq_lens_cpu) - device_sm = get_device_sm() - - # Requirements: H200/B200, short sequences, supported dtype, fits in chunk - self.use_mha = ( - ( - device_sm == 90 or (device_sm >= 100 and device_sm < 110) - ) # SM90/SM100 only - and max_kv_len - <= envs.SGLANG_NSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD.get() # Short enough for MHA - and forward_batch.token_to_kv_pool.dtype - in [torch.bfloat16, torch.float8_e4m3fn] - and sum_seq_lens - <= forward_batch.get_max_chunk_capacity() # Fits in chunk - and (not is_nsa_enable_prefill_cp()) # CP not enabled - and (forward_batch.hisparse_coordinator is None) - ) - else: - self.use_mha = False # Decode/verify always use MLA - - # Set MLA implementation only if not using MHA - if not self.use_mha and self.enable_auto_select_prefill_impl: - if self.nsa_kv_cache_store_fp8: - if ( - is_blackwell() - and forward_batch is not None - and forward_batch.forward_mode == ForwardMode.EXTEND - ): - total_kv_tokens = forward_batch.seq_lens_sum - total_q_tokens = forward_batch.extend_num_tokens - # Heuristic based on benchmarking flashmla_kv vs flashmla_sparse + dequantize_k_cache_paged - if total_kv_tokens < total_q_tokens * 512: - self.nsa_prefill_impl = "flashmla_sparse" - return - self.nsa_prefill_impl = "flashmla_kv" - else: - # bf16 kv cache - self.nsa_prefill_impl = "flashmla_sparse" - - def get_topk_transform_method( - self, forward_mode: Optional[ForwardMode] = None - ) -> TopkTransformMethod: - """ - SGLANG_NSA_FUSE_TOPK controls whether to fuse the topk transform into the topk kernel. - This method is used to select the topk transform method which can be fused or unfused. - """ - if ( - # disable for MTP - self.nsa_kv_cache_store_fp8 - and self.nsa_prefill_impl == "flashmla_sparse" - and forward_mode == ForwardMode.EXTEND - ): - topk_transform_method = TopkTransformMethod.RAGGED - else: - topk_transform_method = TopkTransformMethod.PAGED - return topk_transform_method - - def get_indexer_metadata( - self, layer_id: int, forward_batch: ForwardBatch - ) -> NSAIndexerMetadata: - force_unfused = ( - forward_batch.hisparse_coordinator is not None - and forward_batch.forward_mode.is_decode_or_idle() - ) - return NSAIndexerMetadata( - attn_metadata=self.forward_metadata, - topk_transform_method=self.get_topk_transform_method( - forward_batch.forward_mode - ), - paged_mqa_schedule_metadata=self.forward_metadata.paged_mqa_schedule_metadata, - force_unfused_topk=force_unfused, - ) - - def _compute_flashmla_metadata(self, cache_seqlens: torch.Tensor, seq_len_q: int): - from sgl_kernel.flash_mla import get_mla_metadata - - num_heads_q = self.flashmla_kv_num_q_heads - - flashmla_metadata, num_splits = get_mla_metadata( - cache_seqlens=cache_seqlens, - # TODO doc says `num_q_tokens_per_q_seq * num_heads_q // num_heads_k` - # but the name looks like need seq_len_q? - num_q_tokens_per_head_k=seq_len_q * num_heads_q // 1, - num_heads_k=1, - num_heads_q=num_heads_q, - is_fp8_kvcache=True, - topk=self.nsa_index_topk, - ) - - return NSAFlashMLAMetadata( - flashmla_metadata=flashmla_metadata, - num_splits=num_splits, - ) - - -class NativeSparseAttnMultiStepBackend: - - def __init__( - self, model_runner: ModelRunner, topk: int, speculative_num_steps: int - ): - self.model_runner = model_runner - self.topk = topk - self.speculative_num_steps = speculative_num_steps - self.attn_backends = [] - for i in range(self.speculative_num_steps - 1): - self.attn_backends.append( - NativeSparseAttnBackend( - model_runner, - speculative_step_id=i, - topk=self.topk, - speculative_num_steps=self.speculative_num_steps, - ) - ) - - def init_forward_metadata(self, forward_batch: ForwardBatch): - for i in range(self.speculative_num_steps - 1): - self.attn_backends[i].init_forward_metadata(forward_batch) - - def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int): - for i in range(self.speculative_num_steps - 1): - self.attn_backends[i].init_cuda_graph_state(max_bs, max_num_tokens) - - def init_forward_metadata_capture_cuda_graph(self, forward_batch: ForwardBatch): - for i in range(self.speculative_num_steps - 1): - self.attn_backends[i].init_forward_metadata_capture_cuda_graph( - forward_batch.batch_size, - forward_batch.batch_size * self.topk, - forward_batch.req_pool_indices, - forward_batch.seq_lens, - encoder_lens=None, - forward_mode=ForwardMode.DECODE, - spec_info=forward_batch.spec_info, - ) - - def init_forward_metadata_replay_cuda_graph( - self, forward_batch: ForwardBatch, bs: int - ): - if envs.SGLANG_NSA_ENABLE_MTP_PRECOMPUTE_METADATA.get(): - # Precompute metadata once (shared across all backends) - precomputed = self.attn_backends[0]._precompute_replay_metadata( - bs=bs, - req_pool_indices=forward_batch.req_pool_indices, - seq_lens=forward_batch.seq_lens, - seq_lens_cpu=forward_batch.seq_lens_cpu, - forward_mode=ForwardMode.DECODE, - spec_info=forward_batch.spec_info, - ) - - # Use multi-backend fused copy when we have 3 or more backends - # This is 3x faster than calling the single-backend copy 3 times - if self.speculative_num_steps > 3: - try: - from sglang.jit_kernel.fused_metadata_copy import ( - fused_metadata_copy_multi_cuda, - ) - - metadata0 = self.attn_backends[0].decode_cuda_graph_metadata[bs] - metadata1 = self.attn_backends[1].decode_cuda_graph_metadata[bs] - metadata2 = self.attn_backends[2].decode_cuda_graph_metadata[bs] - - # Set nsa_prefill_impl for first 3 backends (required by the method) - for i in range(3): - self.attn_backends[i].set_nsa_prefill_impl(forward_batch=None) - - # Prepare FlashMLA tensors if needed - flashmla_num_splits_src = None - flashmla_metadata_src = None - flashmla_num_splits_dst0 = None - flashmla_num_splits_dst1 = None - flashmla_num_splits_dst2 = None - flashmla_metadata_dst0 = None - flashmla_metadata_dst1 = None - flashmla_metadata_dst2 = None - - if precomputed.flashmla_metadata is not None: - flashmla_num_splits_src = ( - precomputed.flashmla_metadata.num_splits - ) - flashmla_metadata_src = ( - precomputed.flashmla_metadata.flashmla_metadata - ) - flashmla_num_splits_dst0 = ( - metadata0.flashmla_metadata.num_splits - ) - flashmla_num_splits_dst1 = ( - metadata1.flashmla_metadata.num_splits - ) - flashmla_num_splits_dst2 = ( - metadata2.flashmla_metadata.num_splits - ) - flashmla_metadata_dst0 = ( - metadata0.flashmla_metadata.flashmla_metadata - ) - flashmla_metadata_dst1 = ( - metadata1.flashmla_metadata.flashmla_metadata - ) - flashmla_metadata_dst2 = ( - metadata2.flashmla_metadata.flashmla_metadata - ) - - # Call the multi-backend fused kernel for first 3 backends - fused_metadata_copy_multi_cuda( - # Source tensors - precomputed.cache_seqlens, - precomputed.cu_seqlens_k, - precomputed.page_indices, - precomputed.nsa_cache_seqlens, - precomputed.nsa_cu_seqlens_k, - precomputed.real_page_table, - flashmla_num_splits_src, - flashmla_metadata_src, - # Destination tensors for backend 0 - metadata0.cache_seqlens_int32, - metadata0.cu_seqlens_k, - metadata0.page_table_1, - metadata0.nsa_cache_seqlens_int32, - metadata0.nsa_cu_seqlens_k, - ( - metadata0.real_page_table - if precomputed.real_page_table is not None - else None - ), - flashmla_num_splits_dst0, - flashmla_metadata_dst0, - # Destination tensors for backend 1 - metadata1.cache_seqlens_int32, - metadata1.cu_seqlens_k, - metadata1.page_table_1, - metadata1.nsa_cache_seqlens_int32, - metadata1.nsa_cu_seqlens_k, - ( - metadata1.real_page_table - if precomputed.real_page_table is not None - else None - ), - flashmla_num_splits_dst1, - flashmla_metadata_dst1, - # Destination tensors for backend 2 - metadata2.cache_seqlens_int32, - metadata2.cu_seqlens_k, - metadata2.page_table_1, - metadata2.nsa_cache_seqlens_int32, - metadata2.nsa_cu_seqlens_k, - ( - metadata2.real_page_table - if precomputed.real_page_table is not None - else None - ), - flashmla_num_splits_dst2, - flashmla_metadata_dst2, - # Parameters - bs, - precomputed.max_len, - precomputed.seqlens_expanded_size, - ) - - # Copy remaining backends one by one (if > 3 backends) - for i in range(3, self.speculative_num_steps - 1): - self.attn_backends[ - i - ].init_forward_metadata_replay_cuda_graph_from_precomputed( - bs=bs, - precomputed=precomputed, - forward_mode=ForwardMode.DECODE, - ) - except (ImportError, Exception) as e: - # Fallback to loop if multi-backend kernel not available or fails - if isinstance(e, ImportError): - print( - "Warning: Multi-backend fused metadata copy kernel not available, falling back to loop." - ) - else: - print( - f"Warning: Multi-backend fused metadata copy kernel failed with error: {e}, falling back to loop." - ) - for i in range(self.speculative_num_steps - 1): - self.attn_backends[ - i - ].init_forward_metadata_replay_cuda_graph_from_precomputed( - bs=bs, - precomputed=precomputed, - forward_mode=ForwardMode.DECODE, - ) - else: - # Less than 3 backends: copy to each backend individually - for i in range(self.speculative_num_steps - 1): - self.attn_backends[ - i - ].init_forward_metadata_replay_cuda_graph_from_precomputed( - bs=bs, - precomputed=precomputed, - forward_mode=ForwardMode.DECODE, - ) - else: - # Fallback: compute metadata separately for each backend - for i in range(self.speculative_num_steps - 1): - self.attn_backends[i].init_forward_metadata_replay_cuda_graph( - bs=bs, - req_pool_indices=forward_batch.req_pool_indices, - seq_lens=forward_batch.seq_lens, - seq_lens_sum=forward_batch.seq_lens_sum, - encoder_lens=None, - forward_mode=ForwardMode.DECODE, - spec_info=forward_batch.spec_info, - seq_lens_cpu=forward_batch.seq_lens_cpu, - out_cache_loc=None, - ) diff --git a/python/sglang/srt/layers/communicator.py b/python/sglang/srt/layers/communicator.py index 591ce2d7f..7c40b0c10 100644 --- a/python/sglang/srt/layers/communicator.py +++ b/python/sglang/srt/layers/communicator.py @@ -33,9 +33,9 @@ from sglang.srt.distributed.device_communicators.pynccl_allocator import ( use_symmetric_memory, ) from sglang.srt.environ import envs -from sglang.srt.layers.attention.nsa.utils import ( - is_nsa_enable_prefill_cp, - nsa_use_prefill_cp, +from sglang.srt.layers.attention.dsa.utils import ( + dsa_use_prefill_cp, + is_dsa_enable_prefill_cp, ) from sglang.srt.layers.dp_attention import ( attn_tp_all_gather_into_tensor, @@ -202,7 +202,7 @@ class ScatterMode(Enum): @staticmethod def model_input_output(): """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.TP_ATTN_FULL @@ -256,15 +256,15 @@ class AttnTpContext: self.allow_input_scattered = False self.input_scattered_ = False self.attn_inputs_: Optional[AttentionInputs] = None - self.is_nsa = False + self.is_dsa = False - def init_context(self, q_lora_rank, is_nsa): - self.is_nsa = is_nsa + def init_context(self, q_lora_rank, is_dsa): + self.is_dsa = is_dsa self.allow_input_scattered = ( get_global_server_args().enable_attn_tp_input_scattered and (_is_cuda or _is_npu) and q_lora_rank is not None - and not is_nsa + and not is_dsa and get_tensor_model_parallel_world_size() > 1 and not is_dp_attention_enabled() and get_moe_a2a_backend().is_none() @@ -379,8 +379,8 @@ class LayerScatterModes: or should_use_flashinfer_cutlass_moe_fp4_allgather() ): return ScatterMode.SCATTERED - # NSA CP doesn't support MOE_FULL yet; fall back to FULL - if is_enable_moe_cp_allgather() and not is_nsa_enable_prefill_cp(): + # DSA CP doesn't support MOE_FULL yet; fall back to FULL + if is_enable_moe_cp_allgather() and not is_dsa_enable_prefill_cp(): return ScatterMode.MOE_FULL return ScatterMode.FULL else: @@ -551,10 +551,10 @@ class LayerCommunicator: ) elif _use_aiter and _is_gfx95_supported and (quant_format == "fp8"): # aiter (ROCm gfx95) fused RMSNorm + FP8 group quant. - # When NSA is active, also preserve the unquantized bf16 - # output as a 3-tuple (fp8, scale, bf16) so the NSA + # When DSA is active, also preserve the unquantized bf16 + # output as a 3-tuple (fp8, scale, bf16) so the DSA # 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, self.input_layernorm.weight, @@ -565,9 +565,9 @@ class LayerCommunicator: group_size=128, dtype_quant=torch.float8_e4m3fn, 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[0], hidden_states[1], @@ -596,9 +596,9 @@ class LayerCommunicator: ) elif _use_aiter and _is_gfx95_supported and (quant_format == "fp8"): # 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). - _nsa_needs_bf16 = get_attn_tp_context().is_nsa + _dsa_needs_bf16 = get_attn_tp_context().is_dsa hidden_states, _unq_bf16, _, residual = ( fused_rms_fp8_group_quant( hidden_states, @@ -610,10 +610,10 @@ class LayerCommunicator: group_size=128, dtype_quant=torch.float8_e4m3fn, 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[0], hidden_states[1], @@ -709,7 +709,7 @@ class LayerCommunicator: return True if forward_batch.dp_padding_mode.is_max_len(): return True - if nsa_use_prefill_cp(forward_batch): + if dsa_use_prefill_cp(forward_batch): return True if get_attn_tp_context().input_scattered and not self.is_last_layer: return True diff --git a/python/sglang/srt/layers/communicator_nsa_cp.py b/python/sglang/srt/layers/communicator_dsa_cp.py similarity index 88% rename from python/sglang/srt/layers/communicator_nsa_cp.py rename to python/sglang/srt/layers/communicator_dsa_cp.py index 2508929f1..20b220278 100644 --- a/python/sglang/srt/layers/communicator_nsa_cp.py +++ b/python/sglang/srt/layers/communicator_dsa_cp.py @@ -18,9 +18,9 @@ from typing import Callable, Optional import torch -from sglang.srt.layers.attention.nsa.utils import ( - is_nsa_enable_prefill_cp, - nsa_use_prefill_cp, +from sglang.srt.layers.attention.dsa.utils import ( + dsa_use_prefill_cp, + is_dsa_enable_prefill_cp, ) from sglang.srt.layers.communicator import ( CommunicateContext, @@ -40,14 +40,14 @@ from sglang.srt.layers.dp_attention import ( 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. # The three parts of prepare_attn, prepare_mlp, and postprocess_layer # 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__( self, layer_scatter_modes: LayerScatterModes, @@ -73,19 +73,19 @@ class NSACPLayerCommunicator(LayerCommunicator): assert ( self._context.attn_dp_size == 1 ), 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, output_mode=ScatterMode.SCATTERED, 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, residual_input_mode=ScatterMode.SCATTERED, hidden_states_output_mode=self.layer_scatter_modes.mlp_mode, # SCATTERED, FULL residual_output_mode=ScatterMode.SCATTERED, 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 residual_input_mode=ScatterMode.SCATTERED, output_mode=ScatterMode.SCATTERED, @@ -93,7 +93,7 @@ class NSACPLayerCommunicator(LayerCommunicator): ) -class NSACPCommunicateSimpleFn(CommunicateSimpleFn): +class DSACPCommunicateSimpleFn(CommunicateSimpleFn): @staticmethod def get_fn( input_mode: ScatterMode, @@ -101,12 +101,12 @@ class NSACPCommunicateSimpleFn(CommunicateSimpleFn): context: CommunicateContext, ): if context.is_same_group_size(input_mode, output_mode): - return NSACPCommunicateSimpleFn._trivial + return DSACPCommunicateSimpleFn._trivial raise NotImplementedError(f"{input_mode=} {output_mode=}") -class NSACPCommunicateWithAllReduceAndLayerNormFn( +class DSACPCommunicateWithAllReduceAndLayerNormFn( CommunicateWithAllReduceAndLayerNormFn ): """Besides communication, needs to @@ -126,11 +126,11 @@ class NSACPCommunicateWithAllReduceAndLayerNormFn( assert residual_input_mode == ScatterMode.SCATTERED assert residual_output_mode == ScatterMode.SCATTERED if hidden_states_output_mode == ScatterMode.SCATTERED: - return NSACPCommunicateWithAllReduceAndLayerNormFn._simple + return DSACPCommunicateWithAllReduceAndLayerNormFn._simple if hidden_states_output_mode == ScatterMode.FULL: return partial( - NSACPCommunicateWithAllReduceAndLayerNormFn._gather_hidden_states_and_residual, + DSACPCommunicateWithAllReduceAndLayerNormFn._gather_hidden_states_and_residual, residual_input_mode=residual_input_mode, ) @@ -152,7 +152,7 @@ class NSACPCommunicateWithAllReduceAndLayerNormFn( hidden_states, residual = layernorm(hidden_states, residual) # for prefill: attn tp scattered -> 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 hidden_states, local_hidden_states = ( get_local_dp_buffer(get_attention_cp_group()), @@ -165,7 +165,7 @@ class NSACPCommunicateWithAllReduceAndLayerNormFn( return hidden_states, residual -class NSACPCommunicateSummableTensorPairFn(CommunicateSummableTensorPairFn): +class DSACPCommunicateSummableTensorPairFn(CommunicateSummableTensorPairFn): """It is allowed to make (hidden_states, residual) := (hidden_states + residual, None) if needed.""" @staticmethod @@ -184,12 +184,12 @@ class NSACPCommunicateSummableTensorPairFn(CommunicateSummableTensorPairFn): and (residual_input_mode == ScatterMode.SCATTERED) and (output_mode == ScatterMode.SCATTERED) ): - return NSACPCommunicateSummableTensorPairFn._scatter_hidden_states + return DSACPCommunicateSummableTensorPairFn._scatter_hidden_states if context.is_same_group_size( hidden_states_input_mode, output_mode ) and context.is_same_group_size(residual_input_mode, output_mode): - return NSACPCommunicateSummableTensorPairFn._trivial + return DSACPCommunicateSummableTensorPairFn._trivial raise NotImplementedError( f"{hidden_states_input_mode=} {residual_input_mode=} {output_mode=}" @@ -205,7 +205,7 @@ class NSACPCommunicateSummableTensorPairFn(CommunicateSummableTensorPairFn): ): # for prefill: full -> attn tp scattered # 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 input_hidden_states = hidden_states hidden_states = hidden_states.tensor_split(context.attn_cp_size)[ diff --git a/python/sglang/srt/layers/dp_attention.py b/python/sglang/srt/layers/dp_attention.py index fc4182b96..eff8ce4f9 100644 --- a/python/sglang/srt/layers/dp_attention.py +++ b/python/sglang/srt/layers/dp_attention.py @@ -72,7 +72,7 @@ class DpPaddingMode(IntEnum): # When is_extend_in_batch and dp_size > 1, use SUM_LEN to avoid padding # overhead from uneven token distribution. # 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: return DpPaddingMode.SUM_LEN diff --git a/python/sglang/srt/layers/mhc.py b/python/sglang/srt/layers/mhc.py index 0be750f1a..414cda649 100644 --- a/python/sglang/srt/layers/mhc.py +++ b/python/sglang/srt/layers/mhc.py @@ -8,7 +8,7 @@ import torch from sglang.jit_kernel.utils import is_arch_support_pdl 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 tilelang.set_log_level("WARNING") @@ -880,7 +880,7 @@ def mhc_post( post_layer_mix: torch.Tensor, comb_res_mix: torch.Tensor, ) -> torch.Tensor: - if is_nsa_prefill_cp_round_robin_split(): + if is_dsa_prefill_cp_round_robin_split(): x = strict_contiguous(x) residual = strict_contiguous(residual) post_layer_mix = strict_contiguous(post_layer_mix) diff --git a/python/sglang/srt/layers/utils/cp_utils.py b/python/sglang/srt/layers/utils/cp_utils.py index f3ee07809..885dfed3b 100644 --- a/python/sglang/srt/layers/utils/cp_utils.py +++ b/python/sglang/srt/layers/utils/cp_utils.py @@ -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): - from sglang.srt.layers.attention.nsa.utils import ( - is_nsa_prefill_cp_round_robin_split, - nsa_cp_round_robin_split_data, + from sglang.srt.layers.attention.dsa.utils import ( + dsa_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() assert ( 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}" - return nsa_cp_round_robin_split_data(input_) + return dsa_cp_round_robin_split_data(input_) input_list = list( 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): - from sglang.srt.layers.attention.nsa.utils import ( - is_nsa_prefill_cp_round_robin_split, - nsa_cp_round_robin_split_data, + from sglang.srt.layers.attention.dsa.utils import ( + dsa_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() assert positions.shape[0] % cp_size == 0, ( f"Expect positions shape 0 can divided by cp size, but got positions shape {positions.shape}, " 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( 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, ... | +-------------------------+ """ - 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, ) - if is_nsa_prefill_cp_round_robin_split(): + if is_dsa_prefill_cp_round_robin_split(): with use_symmetric_memory( get_attention_cp_group(), disabled=not is_allocation_symmetric() ): @@ -395,11 +395,11 @@ def prepare_context_parallel_metadata( cp_size, seqs_len, ): - 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, ) - if is_nsa_prefill_cp_round_robin_split(): + if is_dsa_prefill_cp_round_robin_split(): return ContextParallelMetadata() """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 # 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. - # 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 # would silently drop it whenever the scheduler packs multiple requests # 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. - 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_next = prefix_sum_list[cp_size * 2 - cp_rank - 1] else: diff --git a/python/sglang/srt/managers/hisparse_coordinator.py b/python/sglang/srt/managers/hisparse_coordinator.py index 3f138f3ac..7d5a2dd1c 100644 --- a/python/sglang/srt/managers/hisparse_coordinator.py +++ b/python/sglang/srt/managers/hisparse_coordinator.py @@ -9,7 +9,7 @@ from sglang.srt.managers.schedule_batch import Req from sglang.srt.mem_cache.hisparse_memory_pool import ( DeepSeekV4HiSparseTokenToKVPoolAllocator, DeepSeekV4SingleKVPoolHost, - HiSparseNSATokenToKVPool, + HiSparseDSATokenToKVPool, HiSparseTokenToKVPoolAllocator, ) from sglang.srt.mem_cache.memory_pool_host import MLATokenToKVPoolHost @@ -79,7 +79,7 @@ class HiSparseCoordinator: assert isinstance( 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.mem_pool_host = MLATokenToKVPoolHost( diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index d53dc2db9..933b1ca20 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -34,7 +34,7 @@ from typing import TYPE_CHECKING, Dict, List, Optional, Set, Union import torch 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.managers.schedule_batch import Req, ScheduleBatch from sglang.srt.mem_cache.base_prefix_cache import ( @@ -467,7 +467,7 @@ class PrefillAdder: self.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.prefill_context_parallel_enabled = is_prefill_context_parallel_enabled() self.prefill_max_requests = prefill_max_requests @@ -826,7 +826,7 @@ class PrefillAdder: # Enabling context parallelism currently presents precision issues; # therefore, the prefill-batch setting is temporarily set to 1. 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: return AddReqResult.OTHER diff --git a/python/sglang/srt/managers/scheduler_pp_mixin.py b/python/sglang/srt/managers/scheduler_pp_mixin.py index b2d23ec23..e04ca73a6 100644 --- a/python/sglang/srt/managers/scheduler_pp_mixin.py +++ b/python/sglang/srt/managers/scheduler_pp_mixin.py @@ -523,7 +523,7 @@ class SchedulerPPMixin: 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. 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.last_mbs = [None] * self.pp_loop_size diff --git a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py b/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py index a6b826d7d..82d7366cc 100644 --- a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py +++ b/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py @@ -9,11 +9,11 @@ import torch 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.environ import envs +from sglang.srt.layers.attention.dsa import index_buf_accessor from sglang.srt.layers.attention.dsv4 import ( index_buf_accessor as dsv4_index_buf_accessor, ) 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.deepseek_v4_compress_state import CompressStatePool from sglang.srt.mem_cache.memory_pool import KVCache diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 3416d2a86..136d6e514 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -36,12 +36,12 @@ from sglang.srt.mem_cache.hybrid_cache.hybrid_cache_controller import ( HybridCacheController, ) 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 ( + DSATokenToKVPool, MHATokenToKVPool, MLATokenToKVPool, - NSATokenToKVPool, ) from sglang.srt.mem_cache.memory_pool_host import ( MHATokenToKVPoolHost, @@ -82,8 +82,8 @@ class HiRadixCache(RadixCache): server_args.hicache_mem_layout, allocator_type=server_args.hicache_storage_backend, ) - elif isinstance(self.kv_cache, NSATokenToKVPool): - # Filled by attach_hybrid_nsa_pool_to_hiradix_cache after storage extra_config is parsed. + elif isinstance(self.kv_cache, DSATokenToKVPool): + # Filled by attach_hybrid_dsa_pool_to_hiradix_cache after storage extra_config is parsed. self.token_to_kv_pool_host = None elif isinstance(self.kv_cache, MLATokenToKVPool): self.token_to_kv_pool_host = MLATokenToKVPoolHost( @@ -95,9 +95,7 @@ class HiRadixCache(RadixCache): allocator_type=server_args.hicache_storage_backend, ) else: - raise ValueError( - "HiRadixCache only supports MHA, MLA, and NSA (DSA) models" - ) + raise ValueError("HiRadixCache only supports MHA, MLA, and DSA models") self.tp_group = params.tp_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.load_cache_event = threading.Event() - if isinstance(self.kv_cache, NSATokenToKVPool): - attach_hybrid_nsa_pool_to_hiradix_cache( + if isinstance(self.kv_cache, DSATokenToKVPool): + attach_hybrid_dsa_pool_to_hiradix_cache( self, params, server_args, @@ -643,7 +641,7 @@ class HiRadixCache(RadixCache): def _get_extra_pools(self) -> dict: if not isinstance(self.cache_controller, HybridCacheController): return {} - if isinstance(self.kv_cache, NSATokenToKVPool): + if isinstance(self.kv_cache, DSATokenToKVPool): pool = PoolTransfer( name=PoolName.INDEXER, hit_policy=PoolHitPolicy.ALL_PAGES, diff --git a/python/sglang/srt/mem_cache/hisparse_memory_pool.py b/python/sglang/srt/mem_cache/hisparse_memory_pool.py index 2192381fd..679c0f46d 100644 --- a/python/sglang/srt/mem_cache/hisparse_memory_pool.py +++ b/python/sglang/srt/mem_cache/hisparse_memory_pool.py @@ -16,7 +16,7 @@ from sglang.srt.mem_cache.deepseek_v4_memory_pool import ( DeepSeekV4TokenToKVPool, 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.utils import is_cuda, is_hip from sglang.srt.utils.common import get_num_new_pages @@ -37,7 +37,7 @@ else: ) -class HiSparseNSATokenToKVPool(NSATokenToKVPool): +class HiSparseDSATokenToKVPool(DSATokenToKVPool): def __init__( self, size: int, @@ -143,7 +143,7 @@ class HiSparseTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): page_size: int, dtype: torch.dtype, device: torch.device, - kvcache: HiSparseNSATokenToKVPool, + kvcache: HiSparseDSATokenToKVPool, need_sort: bool, host_to_device_ratio: int = 2, ): diff --git a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py index eb3a896ef..63cb96c4f 100644 --- a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py +++ b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py @@ -10,12 +10,12 @@ from sglang.srt.mem_cache.hybrid_cache.hybrid_cache_controller import ( from sglang.srt.mem_cache.memory_pool_host import ( DeepSeekV4PagedHostPool, DeepSeekV4StateHostPool, + DSAIndexerPoolHost, HostPoolGroup, LogicalHostPool, MambaPoolHost, MHATokenToKVPoolHost, MLATokenToKVPoolHost, - NSAIndexerPoolHost, 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.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool from sglang.srt.mem_cache.memory_pool import ( + DSATokenToKVPool, HybridLinearKVPool, MLATokenToKVPool, - NSATokenToKVPool, ) from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool 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() swa_stack = isinstance(kvcache, SWAKVPool) mamba_stack = isinstance(kvcache, HybridLinearKVPool) - nsa_stack = isinstance(kvcache, NSATokenToKVPool) + dsa_stack = isinstance(kvcache, DSATokenToKVPool) deepseek_v4_stack = isinstance(kvcache, DeepSeekV4TokenToKVPool) if deepseek_v4_stack: @@ -820,7 +820,7 @@ def attach_hybrid_pool_to_unified_cache( cache.swa_kv_pool_host ) transfer_layer_num = len(full_layer_mapping | swa_layer_mapping) - elif nsa_stack: + elif dsa_stack: full_layer_mapping = { 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, use_mla=use_mla, 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, kv_host_pool, server_args.hicache_mem_layout, @@ -897,7 +897,7 @@ def attach_hybrid_pool_to_unified_cache( pools_desc = "KV + MAMBA" elif swa_stack: pools_desc = "KV + SWA" - elif nsa_stack: + elif dsa_stack: pools_desc = "KV + INDEXER" else: pools_desc = "KV" @@ -911,7 +911,7 @@ def attach_hybrid_pool_to_unified_cache( raise -def attach_hybrid_nsa_pool_to_hiradix_cache( +def attach_hybrid_dsa_pool_to_hiradix_cache( radix_cache: HiRadixCache, params: CacheInitParams, server_args: ServerArgs, @@ -925,7 +925,7 @@ def attach_hybrid_nsa_pool_to_hiradix_cache( ) -> None: """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: kv = radix_cache.kv_cache @@ -945,7 +945,7 @@ def attach_hybrid_nsa_pool_to_hiradix_cache( use_mla=True, override_kv_cache_dim=kv.kv_cache_dim, prefetch_threshold=prefetch_threshold, - sidecar_host_pool_factory=lambda kv_host_pool: NSAIndexerPoolHost( + sidecar_host_pool_factory=lambda kv_host_pool: DSAIndexerPoolHost( kv, kv_host_pool, 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.cache_controller = cache_controller 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", len(layer_mapping), ) except Exception: - logger.exception("attach_hybrid_nsa_pool_to_hiradix_cache failed") + logger.exception("attach_hybrid_dsa_pool_to_hiradix_cache failed") raise diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index cedbf4e90..9f06d99c1 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -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.constants import GPU_MEMORY_TYPE_KV_CACHE from sglang.srt.environ import envs -from sglang.srt.layers.attention.nsa import index_buf_accessor -from sglang.srt.layers.attention.nsa.quant_k_cache import ( +from sglang.srt.layers.attention.dsa import index_buf_accessor +from sglang.srt.layers.attention.dsa.quant_k_cache import ( quantize_k_cache, 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.radix_attention import RadixAttention from sglang.srt.mem_cache.utils import ( @@ -1618,7 +1618,7 @@ class MLATokenToKVPool(KVCache): enable_memory_saver: bool, start_layer: Optional[int] = None, end_layer: Optional[int] = None, - use_nsa: bool = False, + use_dsa: bool = False, override_kv_cache_dim: Optional[int] = None, ): super().__init__( @@ -1634,17 +1634,17 @@ class MLATokenToKVPool(KVCache): self.kv_lora_rank = kv_lora_rank self.qk_rope_head_dim = qk_rope_head_dim - self.use_nsa = use_nsa - self.nsa_kv_cache_store_fp8 = ( - use_nsa + self.use_dsa = use_dsa + self.dsa_kv_cache_store_fp8 = ( + use_dsa and dtype == torch.float8_e4m3fn 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. self.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) ) @@ -1655,8 +1655,8 @@ class MLATokenToKVPool(KVCache): dtype=torch.uint64, device=self.device, ) - if not use_nsa: - # NSA will allocate indexer KV cache later and then log the total size + if not use_dsa: + # DSA will allocate indexer KV cache later and then log the total size self._finalize_allocation_log(size) def _create_buffers(self): @@ -1726,7 +1726,7 @@ class MLATokenToKVPool(KVCache): cache_v: torch.Tensor, ): 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: cache_k = cache_k.to(self.dtype) @@ -1746,7 +1746,7 @@ class MLATokenToKVPool(KVCache): ): 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. # Fuse BF16/FP16 -> FP8 cast with paged KV write. set_mla_kv_buffer_triton_fp8_quant( @@ -1756,7 +1756,7 @@ class MLATokenToKVPool(KVCache): cache_k_rope, 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 # 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 @@ -1905,7 +1905,7 @@ class MLATokenToKVPoolFP4(MLATokenToKVPool): cache_v: torch.Tensor, ): 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: from sglang.srt.layers.quantization.kvfp4_tensor import KVFP4QuantizeUtil @@ -1930,7 +1930,7 @@ class MLATokenToKVPoolFP4(MLATokenToKVPool): ): 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 # TODO no need to cat 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 index_k_with_scale_buffer_dtype = torch.uint8 rope_storage_dtype = torch.bfloat16 # rope is always stored in bf16 @@ -2005,7 +2005,7 @@ class NSATokenToKVPool(MLATokenToKVPool): enable_memory_saver, start_layer, end_layer, - use_nsa=True, + use_dsa=True, override_kv_cache_dim=override_dim, ) # self.index_k_dtype = torch.float8_e4m3fn @@ -2013,7 +2013,7 @@ class NSATokenToKVPool(MLATokenToKVPool): self.index_head_dim = index_head_dim if index_buf_size is None: 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 if _is_hip: @@ -2024,7 +2024,7 @@ class NSATokenToKVPool(MLATokenToKVPool): else: assert ( 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: assert self.page_size == 64 with ( @@ -2133,11 +2133,11 @@ class NSATokenToKVPool(MLATokenToKVPool): ) 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' # 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 - # NSA attention reads garbage at those token positions. + # DSA attention reads garbage at those token positions. kv_cache_cpu = super().get_cpu_copy(indices) page_indices = indices[:: self.page_size] // self.page_size diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index 805b137a3..a872e24ea 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/python/sglang/srt/mem_cache/memory_pool_host.py @@ -31,11 +31,11 @@ from sglang.jit_kernel.hicache import ( transfer_hicache_one_layer_mla as jit_transfer_hicache_one_layer_mla, ) from sglang.srt.mem_cache.memory_pool import ( + DSATokenToKVPool, KVCache, MambaPool, MHATokenToKVPool, MLATokenToKVPool, - NSATokenToKVPool, ) from sglang.srt.utils import is_cuda, is_hip, is_mps, is_npu, is_xpu @@ -2608,14 +2608,14 @@ class HostPoolGroup: ) -class NSAIndexerPoolHost(HostKVCache): - """Host-side NSA index buffers only. Slot layout matches the anchor MLA host pool.""" +class DSAIndexerPoolHost(HostKVCache): + """Host-side DSA index buffers only. Slot layout matches the anchor MLA host pool.""" - device_pool: NSATokenToKVPool + device_pool: DSATokenToKVPool def __init__( self, - device_pool: NSATokenToKVPool, + device_pool: DSATokenToKVPool, anchor_host: MLATokenToKVPoolHost, layout: str, pin_memory: bool = True, @@ -2635,7 +2635,7 @@ class NSAIndexerPoolHost(HostKVCache): self.index_head_dim = device_pool.index_head_dim 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.index_head_dim + 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 if requested_bytes > available_bytes: 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"{available_bytes / 1e9:.2f} GB free." ) 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, layout, ) @@ -2726,7 +2726,7 @@ class NSAIndexerPoolHost(HostKVCache): return host_indices, device_indices if host_indices.numel() % self.page_size != 0: raise ValueError( - "Index buffer transfer expects page-aligned indices for NSA." + "Index buffer transfer expects page-aligned indices for DSA." ) host_page_indices = ( host_indices.reshape(-1, self.page_size)[:, 0] // self.page_size diff --git a/python/sglang/srt/mem_cache/sparsity/__init__.py b/python/sglang/srt/mem_cache/sparsity/__init__.py index e226ab5b9..9985b5f1a 100644 --- a/python/sglang/srt/mem_cache/sparsity/__init__.py +++ b/python/sglang/srt/mem_cache/sparsity/__init__.py @@ -1,7 +1,7 @@ from sglang.srt.mem_cache.sparsity.algorithms import ( BaseSparseAlgorithm, BaseSparseAlgorithmImpl, - DeepSeekNSAAlgorithm, + DeepSeekDSAAlgorithm, QuestAlgorithm, ) from sglang.srt.mem_cache.sparsity.backend import BackendAdaptor, FlashAttentionAdaptor @@ -17,7 +17,7 @@ __all__ = [ "BaseSparseAlgorithm", "BaseSparseAlgorithmImpl", "QuestAlgorithm", - "DeepSeekNSAAlgorithm", + "DeepSeekDSAAlgorithm", "BackendAdaptor", "FlashAttentionAdaptor", "SparseConfig", diff --git a/python/sglang/srt/mem_cache/sparsity/algorithms/__init__.py b/python/sglang/srt/mem_cache/sparsity/algorithms/__init__.py index 7e05203af..88d0f0c91 100644 --- a/python/sglang/srt/mem_cache/sparsity/algorithms/__init__.py +++ b/python/sglang/srt/mem_cache/sparsity/algorithms/__init__.py @@ -2,12 +2,12 @@ from sglang.srt.mem_cache.sparsity.algorithms.base_algorithm import ( BaseSparseAlgorithm, BaseSparseAlgorithmImpl, ) -from sglang.srt.mem_cache.sparsity.algorithms.deepseek_nsa import DeepSeekNSAAlgorithm +from sglang.srt.mem_cache.sparsity.algorithms.deepseek_dsa import DeepSeekDSAAlgorithm from sglang.srt.mem_cache.sparsity.algorithms.quest_algorithm import QuestAlgorithm __all__ = [ "BaseSparseAlgorithm", "BaseSparseAlgorithmImpl", - "DeepSeekNSAAlgorithm", + "DeepSeekDSAAlgorithm", "QuestAlgorithm", ] diff --git a/python/sglang/srt/mem_cache/sparsity/algorithms/deepseek_nsa.py b/python/sglang/srt/mem_cache/sparsity/algorithms/deepseek_dsa.py similarity index 88% rename from python/sglang/srt/mem_cache/sparsity/algorithms/deepseek_nsa.py rename to python/sglang/srt/mem_cache/sparsity/algorithms/deepseek_dsa.py index 6d64f10ee..cfa0d35a5 100644 --- a/python/sglang/srt/mem_cache/sparsity/algorithms/deepseek_nsa.py +++ b/python/sglang/srt/mem_cache/sparsity/algorithms/deepseek_dsa.py @@ -7,12 +7,12 @@ from sglang.srt.mem_cache.sparsity.algorithms.base_algorithm import ( ) -class DeepSeekNSAAlgorithm(BaseSparseAlgorithmImpl): +class DeepSeekDSAAlgorithm(BaseSparseAlgorithmImpl): """ - Sparse attention algorithm for DeepSeek NSA. + Sparse attention algorithm for DeepSeek DSA. - This algorithm uses NSA's native indexer for TopK retrieval. - Overrides all parent methods as NSA has its own specialized flow. + This algorithm uses DSA's native indexer for TopK retrieval. + Overrides all parent methods as DSA has its own specialized flow. """ def __init__(self, config, device: torch.device, **kwargs): diff --git a/python/sglang/srt/mem_cache/sparsity/backend/__init__.py b/python/sglang/srt/mem_cache/sparsity/backend/__init__.py index 2a6e1b710..9ac6a9636 100644 --- a/python/sglang/srt/mem_cache/sparsity/backend/__init__.py +++ b/python/sglang/srt/mem_cache/sparsity/backend/__init__.py @@ -1,7 +1,7 @@ from sglang.srt.mem_cache.sparsity.backend.backend_adaptor import ( BackendAdaptor, + DSABackendAdaptor, FlashAttentionAdaptor, - NSABackendAdaptor, ) -__all__ = ["BackendAdaptor", "FlashAttentionAdaptor", "NSABackendAdaptor"] +__all__ = ["BackendAdaptor", "FlashAttentionAdaptor", "DSABackendAdaptor"] diff --git a/python/sglang/srt/mem_cache/sparsity/backend/backend_adaptor.py b/python/sglang/srt/mem_cache/sparsity/backend/backend_adaptor.py index baf6a5c99..c308cc224 100644 --- a/python/sglang/srt/mem_cache/sparsity/backend/backend_adaptor.py +++ b/python/sglang/srt/mem_cache/sparsity/backend/backend_adaptor.py @@ -46,8 +46,8 @@ class BackendAdaptor(ABC): pass -class NSABackendAdaptor(BackendAdaptor): - """Adaptor for NSA (Native Sparse Attention) backend.""" +class DSABackendAdaptor(BackendAdaptor): + """Adaptor for DSA (DeepSeek Sparse Attention) backend.""" def __init__( self, @@ -70,9 +70,9 @@ class NSABackendAdaptor(BackendAdaptor): **kwargs, ) -> Optional[torch.Tensor]: """ - Transform logical page indices to physical device indices for NSA backend. + Transform logical page indices to physical device indices for DSA backend. """ - # TODO: Implement NSA backend adaptor logic + # TODO: Implement DSA backend adaptor logic pass diff --git a/python/sglang/srt/mem_cache/sparsity/factory.py b/python/sglang/srt/mem_cache/sparsity/factory.py index c8b29f041..86804d656 100644 --- a/python/sglang/srt/mem_cache/sparsity/factory.py +++ b/python/sglang/srt/mem_cache/sparsity/factory.py @@ -5,11 +5,11 @@ from typing import Optional import torch from sglang.srt.mem_cache.sparsity.algorithms.base_algorithm import BaseSparseAlgorithm -from sglang.srt.mem_cache.sparsity.algorithms.deepseek_nsa import DeepSeekNSAAlgorithm +from sglang.srt.mem_cache.sparsity.algorithms.deepseek_dsa import DeepSeekDSAAlgorithm from sglang.srt.mem_cache.sparsity.algorithms.quest_algorithm import QuestAlgorithm from sglang.srt.mem_cache.sparsity.backend.backend_adaptor import ( + DSABackendAdaptor, FlashAttentionAdaptor, - NSABackendAdaptor, ) from sglang.srt.mem_cache.sparsity.core.sparse_coordinator import ( SparseConfig, @@ -22,7 +22,7 @@ _global_sparse_coordinator: Optional[SparseCoordinator] = None _ALGORITHM_REGISTRY = { "quest": lambda config, device, **kw: QuestAlgorithm(config, device, **kw), - "deepseek_nsa": lambda config, device, **kw: DeepSeekNSAAlgorithm( + "deepseek_dsa": lambda config, device, **kw: DeepSeekDSAAlgorithm( config, device, **kw ), } @@ -49,8 +49,8 @@ def _create_backend_adaptor( req_to_token_pool, ): """Create backend adaptor.""" - if isinstance(sparse_algorithm, DeepSeekNSAAlgorithm): - return NSABackendAdaptor(device, req_to_token_pool) + if isinstance(sparse_algorithm, DeepSeekDSAAlgorithm): + return DSABackendAdaptor(device, req_to_token_pool) if backend in ["fa3", "flashattention"]: return FlashAttentionAdaptor(device) diff --git a/python/sglang/srt/model_executor/cuda_graph_runner.py b/python/sglang/srt/model_executor/cuda_graph_runner.py index 1afcfb286..918850de8 100644 --- a/python/sglang/srt/model_executor/cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/cuda_graph_runner.py @@ -43,7 +43,7 @@ from sglang.srt.distributed.parallel_state import ( ) from sglang.srt.dllm.config import DllmConfig from sglang.srt.environ import envs -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 ( DpPaddingMode, get_attention_cp_size, @@ -285,7 +285,7 @@ class DecodeInputBuffers(ForwardInputBuffers): seq_len_fill_value: int, require_gathered_buffer: bool, num_tokens_per_bs: int, - nsa_enable_prefill_cp: bool, + dsa_enable_prefill_cp: bool, enable_num_token_non_padded_flag: bool, pp_proxy_tensors: Optional[PPProxyTensors] = None, ): @@ -354,7 +354,7 @@ class DecodeInputBuffers(ForwardInputBuffers): self.global_num_tokens_for_logprob_gpu.fill_(bs * num_tokens_per_bs) if enable_num_token_non_padded_flag: - if require_gathered_buffer and not nsa_enable_prefill_cp: + if require_gathered_buffer and not dsa_enable_prefill_cp: num_tokens_per_dp = bs * num_tokens_per_bs local = compute_local_num_token_non_padded( global_num_token_non_padded=forward_batch.num_token_non_padded, @@ -588,7 +588,7 @@ class CudaGraphRunner: self.attn_tp_size = get_attention_tp_size() self.attn_tp_rank = get_attention_tp_rank() - self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp() + self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp() self.deepep_adapter = DeepEPCudaGraphRunnerAdapter() @@ -946,7 +946,7 @@ class CudaGraphRunner: if ( enable_num_token_non_padded() and self.require_gathered_buffer - and not self.nsa_enable_prefill_cp + and not self.dsa_enable_prefill_cp ): local = compute_local_num_token_non_padded( global_num_token_non_padded=buffers.num_token_non_padded, @@ -1211,7 +1211,7 @@ class CudaGraphRunner: seq_len_fill_value=self.seq_len_fill_value, require_gathered_buffer=self.require_gathered_buffer, num_tokens_per_bs=self.num_tokens_per_bs, - nsa_enable_prefill_cp=self.nsa_enable_prefill_cp, + dsa_enable_prefill_cp=self.dsa_enable_prefill_cp, enable_num_token_non_padded_flag=enable_num_token_non_padded(), pp_proxy_tensors=pp_proxy_tensors, ) diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index bf802312a..9a04a1976 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -110,7 +110,7 @@ from sglang.srt.layers.attention.attention_registry import ( ATTENTION_BACKENDS, attn_backend_wrapper, ) -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.attention.tbo_backend import TboAttnBackend from sglang.srt.layers.dp_attention import ( DpPaddingMode, @@ -246,7 +246,8 @@ MLA_ATTENTION_BACKENDS = [ "trtllm_mla", "tokenspeed_mla", "ascend", - "nsa", + "dsa", + "nsa", # Deprecated alias for "dsa" "intel_xpu", ] @@ -3061,7 +3062,7 @@ class ModelRunner(ModelRunnerKVCacheMixin): ) -> Union[LogitsProcessorOutput, PPProxyTensors]: # In DP Attention, IDLE batches are padded (batch_size > 0) for MLP sync. # in this case, we need to reinit the forward metadata, otherwise the stale - # metadata causes batch_size mismatch in attention kernel(e.g. NSA Indexer). + # metadata causes batch_size mismatch in attention kernel(e.g. DSA Indexer). if forward_batch.batch_size > 0: self.attn_backend.init_forward_metadata(forward_batch) @@ -3237,7 +3238,7 @@ class ModelRunner(ModelRunnerKVCacheMixin): forward_batch.num_token_non_padded is not None and forward_batch.global_num_tokens_gpu is not None and require_gathered_buffer(self.server_args) - and not is_nsa_enable_prefill_cp() + and not is_dsa_enable_prefill_cp() ): forward_batch.adjust_num_token_non_padded_for_attn_tp( server_args=self.server_args, diff --git a/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py b/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py index 79d77032f..fb62c8ebc 100644 --- a/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py +++ b/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py @@ -6,8 +6,8 @@ from typing import TYPE_CHECKING import torch from sglang.srt.configs.model_config import ( - get_nsa_index_head_dim, - is_deepseek_nsa, + get_dsa_index_head_dim, + is_deepseek_dsa, is_deepseek_v4, ) from sglang.srt.distributed.parallel_state import get_world_group @@ -20,10 +20,11 @@ from sglang.srt.mem_cache.allocator import ( from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool from sglang.srt.mem_cache.hisparse_memory_pool import ( DeepSeekV4HiSparseTokenToKVPoolAllocator, - HiSparseNSATokenToKVPool, + HiSparseDSATokenToKVPool, HiSparseTokenToKVPoolAllocator, ) from sglang.srt.mem_cache.memory_pool import ( + DSATokenToKVPool, HybridLinearKVPool, HybridReqToTokenPool, MHATokenToKVPool, @@ -31,7 +32,6 @@ from sglang.srt.mem_cache.memory_pool import ( MLATokenToKVPool, MLATokenToKVPoolFP4, NoOpMHATokenToKVPool, - NSATokenToKVPool, ReqToTokenPool, ) from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool, SWATokenToKVPoolAllocator @@ -137,36 +137,36 @@ class ModelRunnerKVCacheMixin: return total_rest_memory - mamba_state_memory def calculate_mla_kv_cache_dim(self: ModelRunner) -> int: - is_nsa_model = is_deepseek_nsa(self.model_config.hf_config) + is_dsa_model = is_deepseek_dsa(self.model_config.hf_config) kv_cache_dtype = self.kv_cache_dtype kv_lora_rank = self.model_config.kv_lora_rank qk_rope_head_dim = self.model_config.qk_rope_head_dim kv_cache_dim = kv_lora_rank + qk_rope_head_dim # default mla kv cache dim - # For non-NSA models, MLA kv cache dim is simply kv_lora_rank + qk_rope_head_dim - if not is_nsa_model: + # For non-DSA models, MLA kv cache dim is simply kv_lora_rank + qk_rope_head_dim + if not is_dsa_model: return kv_cache_dim # TRTLLM backend does not override kv_cache_dim for MLA kv cache - # Assuming nsa prefill and decode backends are the same when using trtllm MLA backend, + # Assuming dsa prefill and decode backends are the same when using trtllm MLA backend, # since it is not compatible for trtllm and other mla attn backend due to the different # kv cache layout. if ( - self.server_args.nsa_prefill_backend == "trtllm" - or self.server_args.nsa_decode_backend == "trtllm" + self.server_args.dsa_prefill_backend == "trtllm" + or self.server_args.dsa_decode_backend == "trtllm" ): return kv_cache_dim # On HIP with TileLang backend, keep the default MLA KV cache dimension. # FP8 attention uses the nope(512 fp8) + rope(64 fp8) layout, without extra per-block scales. if _is_hip and ( - self.server_args.nsa_prefill_backend == "tilelang" - or self.server_args.nsa_decode_backend == "tilelang" + self.server_args.dsa_prefill_backend == "tilelang" + or self.server_args.dsa_decode_backend == "tilelang" ): return kv_cache_dim - quant_block_size = NSATokenToKVPool.quant_block_size - rope_storage_dtype = NSATokenToKVPool.rope_storage_dtype + quant_block_size = DSATokenToKVPool.quant_block_size + rope_storage_dtype = DSATokenToKVPool.rope_storage_dtype # Calculate override_kv_cache_dim for FP8 storage in backends that use scaled KV layout (excluding TRTLLM and HIP+TileLang). # kv_lora_rank + scale storage (kv_lora_rank // quant_block_size * 4 bytes) + rope dimension storage # Note: rope dimension is stored in original dtype (bf16), not quantized to fp8 @@ -199,7 +199,7 @@ class ModelRunnerKVCacheMixin: def _validate_prefill_only_disable_kv_cache_pool_family( self: ModelRunner, - is_nsa_model: bool, + is_dsa_model: bool, is_dsv4_model: bool, current_platform, ): @@ -215,8 +215,8 @@ class ModelRunnerKVCacheMixin: self.server_args.attention_backend == "ascend" and not self.mambaish_config ): unsupported_pool_family = "NPU/Ascend KV pool" - elif self.use_mla_backend and is_nsa_model: - unsupported_pool_family = "NSA/MLA KV pool" + elif self.use_mla_backend and is_dsa_model: + unsupported_pool_family = "DSA/MLA KV pool" elif self.use_mla_backend and not self.mambaish_config: unsupported_pool_family = "MLA KV pool" elif self.is_hybrid_swa: @@ -328,14 +328,14 @@ class ModelRunnerKVCacheMixin: assert self.is_draft_worker # Initialize token_to_kv_pool - is_nsa_model = is_deepseek_nsa(self.model_config.hf_config) + is_dsa_model = is_deepseek_dsa(self.model_config.hf_config) is_dsv4_model = is_deepseek_v4(self.model_config.hf_config) # Out-of-tree platform plugin system — used by elif below from sglang.srt.platforms import current_platform self._validate_prefill_only_disable_kv_cache_pool_family( - is_nsa_model, is_dsv4_model, current_platform + is_dsa_model, is_dsv4_model, current_platform ) if is_dsv4_model: @@ -375,8 +375,8 @@ class ModelRunnerKVCacheMixin: enable_hisparse=self.enable_hisparse, ) elif current_platform.is_out_of_tree() and not self.mambaish_config: - if self.use_mla_backend and is_nsa_model: - PoolCls = current_platform.get_nsa_kv_pool_cls() + if self.use_mla_backend and is_dsa_model: + PoolCls = current_platform.get_dsa_kv_pool_cls() self.token_to_kv_pool = PoolCls( self.max_total_num_tokens, page_size=self.page_size, @@ -389,7 +389,7 @@ class ModelRunnerKVCacheMixin: enable_memory_saver=self.server_args.enable_memory_saver, start_layer=self.start_layer, end_layer=self.end_layer, - index_head_dim=get_nsa_index_head_dim(self.model_config.hf_config), + index_head_dim=get_dsa_index_head_dim(self.model_config.hf_config), ) elif self.use_mla_backend: PoolCls = current_platform.get_mla_kv_pool_cls() @@ -400,7 +400,7 @@ class ModelRunnerKVCacheMixin: kv_lora_rank=self.model_config.kv_lora_rank, qk_rope_head_dim=self.model_config.qk_rope_head_dim, index_head_dim=( - self.model_config.index_head_dim if is_nsa_model else None + self.model_config.index_head_dim if is_dsa_model else None ), layer_num=self.num_effective_layers, device=self.device, @@ -472,7 +472,7 @@ class ModelRunnerKVCacheMixin: kv_lora_rank=self.model_config.kv_lora_rank, qk_rope_head_dim=self.model_config.qk_rope_head_dim, index_head_dim=( - self.model_config.index_head_dim if is_nsa_model else None + self.model_config.index_head_dim if is_dsa_model else None ), layer_num=self.num_effective_layers, device=self.device, @@ -499,9 +499,9 @@ class ModelRunnerKVCacheMixin: start_layer=self.start_layer, end_layer=self.end_layer, ) - elif self.use_mla_backend and is_nsa_model: + elif self.use_mla_backend and is_dsa_model: PoolCls = ( - HiSparseNSATokenToKVPool if self.enable_hisparse else NSATokenToKVPool + HiSparseDSATokenToKVPool if self.enable_hisparse else DSATokenToKVPool ) pool_kwargs = {} if self.enable_hisparse: @@ -522,11 +522,11 @@ class ModelRunnerKVCacheMixin: enable_memory_saver=self.server_args.enable_memory_saver, start_layer=self.start_layer, end_layer=self.end_layer, - index_head_dim=get_nsa_index_head_dim(self.model_config.hf_config), + index_head_dim=get_dsa_index_head_dim(self.model_config.hf_config), **pool_kwargs, ) elif self.use_mla_backend and not self.mambaish_config: - assert not is_nsa_model + assert not is_dsa_model if is_float4_e2m1fn_x2(self.kv_cache_dtype): self.token_to_kv_pool = MLATokenToKVPoolFP4( self.max_total_num_tokens, diff --git a/python/sglang/srt/model_executor/pool_configurator.py b/python/sglang/srt/model_executor/pool_configurator.py index 1b2cd0f2a..5803e0bcf 100644 --- a/python/sglang/srt/model_executor/pool_configurator.py +++ b/python/sglang/srt/model_executor/pool_configurator.py @@ -20,14 +20,14 @@ from typing import TYPE_CHECKING, Optional import torch from sglang.srt.configs.model_config import ( - get_nsa_index_head_dim, - is_deepseek_nsa, + get_dsa_index_head_dim, + is_deepseek_dsa, is_deepseek_v4, ) from sglang.srt.environ import envs from sglang.srt.layers.dp_attention import get_attention_tp_size from sglang.srt.mem_cache.deepseek_v4_memory_pool import get_compress_state_ring_size -from sglang.srt.mem_cache.memory_pool import NSATokenToKVPool +from sglang.srt.mem_cache.memory_pool import DSATokenToKVPool from sglang.srt.utils.common import is_float4_e2m1fn_x2 @@ -84,7 +84,7 @@ class MemoryPoolConfigurator: class DefaultPoolConfigurator(MemoryPoolConfigurator): - """Configurator for standard models: MHA, MLA, NSA, FP4. + """Configurator for standard models: MHA, MLA, DSA, FP4. coeff = cell_size (bytes per token across all layers) bias = 0 @@ -149,15 +149,15 @@ class DefaultPoolConfigurator(MemoryPoolConfigurator): * kv_size ) - # Add indexer KV cache overhead for NSA models (DeepSeek V3.2) - if is_deepseek_nsa(model_config.hf_config): - index_head_dim = get_nsa_index_head_dim(model_config.hf_config) + # Add indexer KV cache overhead for DSA models (DeepSeek V3.2) + if is_deepseek_dsa(model_config.hf_config): + index_head_dim = get_dsa_index_head_dim(model_config.hf_config) indexer_size_per_token = ( index_head_dim - + index_head_dim // NSATokenToKVPool.quant_block_size * 4 + + index_head_dim // DSATokenToKVPool.quant_block_size * 4 ) element_size = torch._utils._element_size( - NSATokenToKVPool.index_k_with_scale_buffer_dtype + DSATokenToKVPool.index_k_with_scale_buffer_dtype ) cell_size += indexer_size_per_token * num_layers * element_size else: diff --git a/python/sglang/srt/models/deepseek_common/attention_backend_handler.py b/python/sglang/srt/models/deepseek_common/attention_backend_handler.py index 3c93d41a8..de8c6b322 100644 --- a/python/sglang/srt/models/deepseek_common/attention_backend_handler.py +++ b/python/sglang/srt/models/deepseek_common/attention_backend_handler.py @@ -147,9 +147,9 @@ def handle_attention_aiter(attn, forward_batch): return AttnForwardMethod.MLA -def handle_attention_nsa(attn, forward_batch): +def handle_attention_dsa(attn, forward_batch): """ - Dispatch logic is centralized in NativeSparseAttnBackend.set_nsa_prefill_impl and executed + Dispatch logic is centralized in DeepseekSparseAttnBackend.set_dsa_prefill_impl and executed in init_forward_metadata. Read the decision from backend.use_mha. """ @@ -191,6 +191,9 @@ AttentionBackendRegistry.register("fa4", handle_attention_fa4) AttentionBackendRegistry.register("trtllm_mla", handle_attention_trtllm_mla) AttentionBackendRegistry.register("tokenspeed_mla", handle_attention_tokenspeed_mla) AttentionBackendRegistry.register("aiter", handle_attention_aiter) -AttentionBackendRegistry.register("nsa", handle_attention_nsa) +AttentionBackendRegistry.register("dsa", handle_attention_dsa) +AttentionBackendRegistry.register( + "nsa", handle_attention_dsa +) # Deprecated alias; use "dsa" AttentionBackendRegistry.register("triton", handle_attention_triton) AttentionBackendRegistry.register("intel_xpu", handle_attention_intel_xpu) diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py index f317e8fad..fdea0f8f3 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py +++ b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py @@ -5,7 +5,7 @@ from typing import TYPE_CHECKING import torch from sglang.srt.environ import envs -from sglang.srt.layers.attention.nsa.dequant_k_cache import dequantize_k_cache_paged +from sglang.srt.layers.attention.dsa.dequant_k_cache import dequantize_k_cache_paged from sglang.srt.layers.attention.tbo_backend import TboAttnBackend from sglang.srt.layers.attention.utils import concat_and_cast_mha_k_triton from sglang.srt.layers.communicator import get_attn_tp_context @@ -122,10 +122,10 @@ class DeepseekMHAForwardMixin: ) ) - # 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 self.use_nsa: - # NSA requires unquantized q_lora for the indexer. When q_b_proj is FP8 + if self.use_dsa: + # DSA requires unquantized q_lora for the indexer. When q_b_proj is FP8 # on gfx95, we can still use fused RMSNorm+FP8 quant, but MUST request # the unquantized output for q_lora; otherwise q_lora becomes the (fp8,scale) # tuple. @@ -230,15 +230,15 @@ class DeepseekMHAForwardMixin: and sum(forward_batch.extend_prefix_lens_cpu) != 0 ): if ( - self.use_nsa + self.use_dsa and self.kv_cache_dtype == "fp8_e4m3" and ( - not get_global_server_args().nsa_decode_backend == "trtllm" - or not get_global_server_args().nsa_prefill_backend == "trtllm" + not get_global_server_args().dsa_decode_backend == "trtllm" + or not get_global_server_args().dsa_prefill_backend == "trtllm" ) ): - # FP8 path: dequantize NSA-specific FP8 format to BF16 - kv_a, k_pe = self._get_mla_kv_buffer_from_fp8_for_nsa(forward_batch) + # FP8 path: dequantize DSA-specific FP8 format to BF16 + kv_a, k_pe = self._get_mla_kv_buffer_from_fp8_for_dsa(forward_batch) else: # BF16/FP16 path: directly fetch from cache kv_a, k_pe = self._get_mla_kv_buffer( @@ -471,12 +471,12 @@ class DeepseekMHAForwardMixin: kv_a = kv_a.squeeze(1).contiguous() return kv_a, k_pe - def _get_mla_kv_buffer_from_fp8_for_nsa( + def _get_mla_kv_buffer_from_fp8_for_dsa( self: DeepseekV2AttentionMLA, forward_batch: ForwardBatch, ): """ - Dequantize FP8 KV cache to BF16 for MLA attention (NSA-specific format). + Dequantize FP8 KV cache to BF16 for MLA attention (DSA-specific format). Returns: (kv_a, k_pe) both in BF16 """ diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py index 3093c196b..e0ad07511 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py +++ b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py @@ -6,7 +6,7 @@ import torch from sglang.srt.compilation.piecewise_context_manager import is_in_piecewise_cuda_graph from sglang.srt.layers import deep_gemm_wrapper -from sglang.srt.layers.attention.nsa.utils import nsa_use_prefill_cp +from sglang.srt.layers.attention.dsa.utils import dsa_use_prefill_cp from sglang.srt.layers.communicator import get_attn_tp_context from sglang.srt.layers.quantization.fp8_kernel import ( fp8_dtype, @@ -124,7 +124,6 @@ if _use_aiter_gfx95: class DeepseekMLAForwardMixin: - def init_mla_forward(self: DeepseekV2AttentionMLA): self.flashinfer_mla_disable_ragged = ( get_global_server_args().flashinfer_mla_disable_ragged @@ -178,7 +177,7 @@ class DeepseekMLAForwardMixin: _use_aiter_gfx95 and self.q_b_proj.weight.dtype == torch.float8_e4m3fn ): - if self.use_nsa: + if self.use_dsa: q_quanted, q_lora, k_nope, _ = fused_rms_fp8_group_quant( q, self.q_a_layernorm.weight, @@ -220,7 +219,7 @@ class DeepseekMLAForwardMixin: k_nope = self.kv_a_layernorm(k_nope) # q_lora needed by indexer - if self.use_nsa: + if self.use_dsa: if q_lora is None: q_lora = q @@ -281,9 +280,13 @@ class DeepseekMLAForwardMixin: k_pe = latent_cache[..., self.kv_lora_rank :].unsqueeze(1) if self.use_deep_gemm_bmm: - q_nope_val, q_nope_scale, masked_m, expected_m, aligned_m = ( - per_token_group_quant_mla_deep_gemm_masked_fp8(q_nope.transpose(0, 1)) - ) + ( + q_nope_val, + q_nope_scale, + masked_m, + expected_m, + aligned_m, + ) = per_token_group_quant_mla_deep_gemm_masked_fp8(q_nope.transpose(0, 1)) q_nope_out = q_nope.new_empty( (self.num_local_heads, aligned_m, self.kv_lora_rank) ) @@ -362,18 +365,18 @@ class DeepseekMLAForwardMixin: if is_kv_b_lora_active(self): q_nope_out = apply_kv_b_lora_q_correction(self, q_nope, q_nope_out) - skip_rope_for_nsa_tilelang_fused = self._skip_rope_for_nsa_tilelang_fused() + skip_rope_for_dsa_tilelang_fused = self._skip_rope_for_dsa_tilelang_fused() skip_rope_for_aiter_fused_mla = self._skip_rope_for_aiter_fused_mla() if ( self.rotary_emb is not None and (not self._fuse_rope_for_trtllm_mla(forward_batch)) - and (not skip_rope_for_nsa_tilelang_fused) + and (not skip_rope_for_dsa_tilelang_fused) and (not skip_rope_for_aiter_fused_mla) - and (not _use_aiter or not _is_gfx95_supported or self.use_nsa) + and (not _use_aiter or not _is_gfx95_supported or self.use_dsa) ): q_pe, k_pe = self.rotary_emb(positions, q_pe, k_pe) - if nsa_use_prefill_cp(forward_batch): + if dsa_use_prefill_cp(forward_batch): # support allgather+rerrange k_nope, k_pe = self.rebuild_cp_kv_cache( latent_cache, forward_batch, k_nope, k_pe @@ -406,7 +409,7 @@ class DeepseekMLAForwardMixin: save_kv_cache = True if self.current_attention_backend in FORWARD_ABSORB_CORE_ATTENTION_BACKENDS: - if self._skip_rope_for_nsa_tilelang_fused() and self.rotary_emb is not None: + if self._skip_rope_for_dsa_tilelang_fused() and self.rotary_emb is not None: cos = self.rotary_emb.cos_cache sin = self.rotary_emb.sin_cache kv_cache_dtype = ( @@ -430,14 +433,14 @@ class DeepseekMLAForwardMixin: ) save_kv_cache = False # On decode, pass q_cat directly to attn_mqa with q_rope=None so - # nsa_backend.forward_decode reuses q_cat as a zero-copy view + # dsa_backend.forward_decode reuses q_cat as a zero-copy view # (`q.contiguous().view(...)` fast-path) instead of running the # redundant `concat_mla_absorb_q_general(q_nope_fused, q_pe_fused)` # that would otherwise rebuild a tensor byte-identical to q_cat. # On ROCm tilelang decode, this eliminates the # `CatArrayBatchedCopy, ...>` kernel that used to # fire once per layer per decode step (~2.6 us / layer saved). - # Prefill keeps the split form because nsa_backend.forward_extend + # Prefill keeps the split form because dsa_backend.forward_extend # asserts `q_rope is not None`. if forward_batch.forward_mode.is_decode_or_idle(): if llama_4_scaling is not None: @@ -545,10 +548,14 @@ class DeepseekMLAForwardMixin: attn_output = attn_output.view(-1, self.num_local_heads, self.kv_lora_rank) if self.use_deep_gemm_bmm: - attn_output_val, attn_output_scale, masked_m, expected_m, aligned_m = ( - per_token_group_quant_mla_deep_gemm_masked_fp8( - attn_output.transpose(0, 1) - ) + ( + attn_output_val, + attn_output_scale, + masked_m, + expected_m, + aligned_m, + ) = per_token_group_quant_mla_deep_gemm_masked_fp8( + attn_output.transpose(0, 1) ) attn_bmm_output = attn_output.new_empty( (self.num_local_heads, aligned_m, self.v_head_dim) @@ -683,10 +690,10 @@ class DeepseekMLAForwardMixin: """ Check if we should skip rope and do fused rope+quantize for TRTLLM MLA decode in fp8_e4m3 path. """ - if self.current_attention_backend == "nsa": + if self.current_attention_backend in ("dsa", "nsa"): return ( - get_global_server_args().nsa_decode_backend == "trtllm" - or get_global_server_args().nsa_prefill_backend == "trtllm" + get_global_server_args().dsa_decode_backend == "trtllm" + or get_global_server_args().dsa_prefill_backend == "trtllm" ) and forward_batch.attn_backend.kv_cache_dtype == torch.float8_e4m3fn return ( @@ -698,17 +705,17 @@ class DeepseekMLAForwardMixin: and forward_batch.attn_backend.data_type == torch.float8_e4m3fn ) - def _skip_rope_for_nsa_tilelang_fused(self: DeepseekV2AttentionMLA) -> bool: + def _skip_rope_for_dsa_tilelang_fused(self: DeepseekV2AttentionMLA) -> bool: """ - Check if we should skip rope and use fused rope+cache path for TileLang NSA on gfx95. + Check if we should skip rope and use fused rope+cache path for TileLang DSA on gfx95. """ server_args = get_global_server_args() return ( _use_aiter_gfx95 - and self.current_attention_backend == "nsa" + and self.current_attention_backend in ("dsa", "nsa") and ( - server_args.nsa_decode_backend == "tilelang" - or server_args.nsa_prefill_backend == "tilelang" + server_args.dsa_decode_backend == "tilelang" + or server_args.dsa_prefill_backend == "tilelang" ) ) diff --git a/python/sglang/srt/models/deepseek_common/utils.py b/python/sglang/srt/models/deepseek_common/utils.py index 40f552498..60e3b20a0 100644 --- a/python/sglang/srt/models/deepseek_common/utils.py +++ b/python/sglang/srt/models/deepseek_common/utils.py @@ -57,7 +57,8 @@ NVFP4_CKPT_FP8_ATTN_QUANT_MODULES = ["q_b_proj"] FORWARD_ABSORB_CORE_ATTENTION_BACKENDS = [ "fa3", - "nsa", + "dsa", + "nsa", # Deprecated alias for "dsa" "flashinfer", "cutlass_mla", "trtllm_mla", diff --git a/python/sglang/srt/models/deepseek_nextn.py b/python/sglang/srt/models/deepseek_nextn.py index acc893370..59d26282e 100644 --- a/python/sglang/srt/models/deepseek_nextn.py +++ b/python/sglang/srt/models/deepseek_nextn.py @@ -23,14 +23,14 @@ from safetensors.torch import load_file from torch import nn from transformers import PretrainedConfig -from sglang.srt.configs.model_config import is_deepseek_nsa +from sglang.srt.configs.model_config import is_deepseek_dsa from sglang.srt.distributed import get_pp_group, get_tensor_model_parallel_world_size from sglang.srt.environ import envs from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder -from sglang.srt.layers.attention.nsa.utils import ( - can_nsa_cp_split, - is_nsa_enable_prefill_cp, - nsa_use_prefill_cp, +from sglang.srt.layers.attention.dsa.utils import ( + can_dsa_cp_split, + dsa_use_prefill_cp, + is_dsa_enable_prefill_cp, ) from sglang.srt.layers.dp_attention import ( get_attention_cp_rank, @@ -148,8 +148,8 @@ class DeepseekModelNextN(nn.Module): self.shared_head = nn.Module() self.shared_head.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp() - if self.nsa_enable_prefill_cp: + self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp() + if self.dsa_enable_prefill_cp: self.cp_size = get_attention_cp_size() else: self.cp_size = None @@ -193,7 +193,7 @@ class DeepseekModelNextN(nn.Module): else: hidden_states = self.eh_proj(eh_input) - if nsa_use_prefill_cp(forward_batch, self.nsa_enable_prefill_cp): + if dsa_use_prefill_cp(forward_batch, self.dsa_enable_prefill_cp): hidden_states = cp_split_and_rebuild_data(forward_batch, hidden_states) positions = cp_split_and_rebuild_position(forward_batch, positions) residual = None @@ -212,7 +212,7 @@ class DeepseekModelNextN(nn.Module): else: hidden_states = self.shared_head.norm(hidden_states) - if nsa_use_prefill_cp(forward_batch, self.nsa_enable_prefill_cp): + if dsa_use_prefill_cp(forward_batch, self.dsa_enable_prefill_cp): # allgather + rerrange hidden_states = cp_all_gather_rerange_output( hidden_states, @@ -248,9 +248,9 @@ class DeepseekV3ForCausalLMNextN(DeepseekV3ForCausalLM): # if not set, model load will be broken in DeepseekV3ForCausalLM load_weights() self.pp_group = get_pp_group() self.determine_num_fused_shared_experts("DeepseekV3ForCausalLMNextN") - self.use_nsa = is_deepseek_nsa(config) - self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp() - if self.nsa_enable_prefill_cp: + self.use_dsa = is_deepseek_dsa(config) + self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp() + if self.dsa_enable_prefill_cp: self.cp_rank = get_attention_cp_rank() self.cp_size = get_attention_cp_size() else: @@ -289,9 +289,9 @@ class DeepseekV3ForCausalLMNextN(DeepseekV3ForCausalLM): forward_batch: ForwardBatch, ) -> torch.Tensor: # TODO current just support prefill batch=1 and len(input_ids) > self.cp_size * 2 - if self.nsa_enable_prefill_cp: - if can_nsa_cp_split( - len(input_ids), self.cp_size, self.use_nsa, forward_batch + if self.dsa_enable_prefill_cp: + if can_dsa_cp_split( + len(input_ids), self.cp_size, self.use_dsa, forward_batch ): forward_batch.attn_cp_metadata = prepare_context_parallel_metadata( len(input_ids), diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index da76ba597..8e8459eb4 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -40,10 +40,10 @@ from sglang.srt.batch_overlap.two_batch_overlap import ( ) from sglang.srt.configs.model_config import ( compute_mla_mscale_scaling, - get_nsa_index_head_dim, - get_nsa_index_n_heads, - get_nsa_index_topk, - is_deepseek_nsa, + get_dsa_index_head_dim, + get_dsa_index_n_heads, + get_dsa_index_topk, + is_deepseek_dsa, ) from sglang.srt.distributed import ( divide, @@ -59,11 +59,11 @@ from sglang.srt.eplb.expert_location_dispatch import ExpertLocationDispatchInfo from sglang.srt.layers import deep_gemm_wrapper from sglang.srt.layers.activation import SiluAndMul from sglang.srt.layers.amx_utils import PackWeightMethod -from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer -from sglang.srt.layers.attention.nsa.utils import ( - can_nsa_cp_split, - is_nsa_enable_prefill_cp, - nsa_use_prefill_cp, +from sglang.srt.layers.attention.dsa.dsa_indexer import Indexer +from sglang.srt.layers.attention.dsa.utils import ( + can_dsa_cp_split, + dsa_use_prefill_cp, + is_dsa_enable_prefill_cp, ) from sglang.srt.layers.communicator import ( LayerCommunicator, @@ -71,7 +71,7 @@ from sglang.srt.layers.communicator import ( enable_moe_dense_fully_dp, get_attn_tp_context, ) -from sglang.srt.layers.communicator_nsa_cp import NSACPLayerCommunicator +from sglang.srt.layers.communicator_dsa_cp import DSACPLayerCommunicator from sglang.srt.layers.dp_attention import ( get_attention_cp_rank, get_attention_cp_size, @@ -368,7 +368,7 @@ class MoEGate(nn.Module): self.e_score_correction_bias = None if _is_cpu and _is_cpu_amx_available: self.quant_method = PackWeightMethod(weight_names=["weight"]) - self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp() + self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp() def forward( self, @@ -390,7 +390,7 @@ class MoEGate(nn.Module): if ( not self.is_deepseek_v4 and forward_batch is not None - and nsa_use_prefill_cp(forward_batch) + and dsa_use_prefill_cp(forward_batch) ): logits = F.linear(hidden_states, self.weight, None) else: @@ -1352,12 +1352,12 @@ class DeepseekV2AttentionMLA( self.quant_config = quant_config attn_tp_rank = get_attention_tp_rank() attn_tp_size = get_attention_tp_size() - self.use_nsa = is_deepseek_nsa(config) - self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp() - if self.nsa_enable_prefill_cp: - assert self.use_nsa, "CP currently only supports deepseek v3.2 model" + self.use_dsa = is_deepseek_dsa(config) + self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp() + if self.dsa_enable_prefill_cp: + assert self.use_dsa, "CP currently only supports deepseek v3.2 model" # cp reuse the attn_tp comm group but need to duplicate the weights - if self.nsa_enable_prefill_cp and self.use_nsa: + if self.dsa_enable_prefill_cp and self.use_dsa: self.cp_size = get_attention_cp_size() self.num_heads = num_heads assert num_heads % attn_tp_size == 0 @@ -1410,14 +1410,14 @@ class DeepseekV2AttentionMLA( self.skip_topk = None self.next_skip_topk = None - if self.use_nsa: + if self.use_dsa: is_neox_style = not getattr(config, "indexer_rope_interleave", False) self.indexer = Indexer( hidden_size=hidden_size, - index_n_heads=get_nsa_index_n_heads(config), - index_head_dim=get_nsa_index_head_dim(config), + index_n_heads=get_dsa_index_n_heads(config), + index_head_dim=get_dsa_index_head_dim(config), rope_head_dim=qk_rope_head_dim, - index_topk=get_nsa_index_topk(config), + index_topk=get_dsa_index_topk(config), q_lora_rank=q_lora_rank, max_position_embeddings=max_position_embeddings, rope_theta=rope_theta, @@ -1582,7 +1582,7 @@ class DeepseekV2AttentionMLA( def op_core(self, state): result = self.forward_core(state.pop("attn_intermediate_state")) - # forward_core may return (hidden_states, topk_indices) for NSA models + # forward_core may return (hidden_states, topk_indices) for DSA models # with index cache enabled. In the TBO path, topk_indices is not # propagated between layers, so we discard it here. if isinstance(result, tuple): @@ -1808,7 +1808,7 @@ class DeepseekV2DecoderLayer(nn.Module): self.speculative_algorithm = SpeculativeAlgorithm.from_string( get_global_server_args().speculative_algorithm ) - self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp() + self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp() self.layer_id = layer_id self.is_nextn = is_nextn self.self_attn = DeepseekV2AttentionMLA( @@ -1881,8 +1881,8 @@ class DeepseekV2DecoderLayer(nn.Module): self._gfx95_quant_format = self._detect_gfx95_quant_format() - if self.nsa_enable_prefill_cp: - self.layer_communicator = NSACPLayerCommunicator( + if self.dsa_enable_prefill_cp: + self.layer_communicator = DSACPLayerCommunicator( layer_scatter_modes=self.layer_scatter_modes, input_layernorm=self.input_layernorm, post_attention_layernorm=self.post_attention_layernorm, @@ -1996,7 +1996,7 @@ class DeepseekV2DecoderLayer(nn.Module): gemm_output_zero_allocator, ) - if not self.nsa_enable_prefill_cp and should_allreduce_fusion: + if not self.dsa_enable_prefill_cp and should_allreduce_fusion: hidden_states._sglang_needs_allreduce_fusion = True if not should_allreduce_fusion: @@ -2093,8 +2093,8 @@ class DeepseekV2Model(nn.Module): self.vocab_size = config.vocab_size self.first_k_dense_replace = config.first_k_dense_replace self.pp_group = get_pp_group() - self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp() - if self.nsa_enable_prefill_cp: + self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp() + if self.dsa_enable_prefill_cp: self.cp_size = get_attention_cp_size() else: self.cp_size = None @@ -2253,7 +2253,7 @@ class DeepseekV2Model(nn.Module): else None ) - if nsa_use_prefill_cp(forward_batch): + if dsa_use_prefill_cp(forward_batch): if self.pp_group.is_first_rank: hidden_states = cp_split_and_rebuild_data(forward_batch, hidden_states) positions = cp_split_and_rebuild_position(forward_batch, positions) @@ -2338,7 +2338,7 @@ class DeepseekV2Model(nn.Module): else: hidden_states, _ = self.norm(hidden_states, residual) - if self.pp_group.is_last_rank and nsa_use_prefill_cp(forward_batch): + if self.pp_group.is_last_rank and dsa_use_prefill_cp(forward_batch): # allgather + rerrange hidden_states = cp_all_gather_rerange_output( hidden_states, @@ -2385,7 +2385,7 @@ class DeepseekV2ForCausalLM(nn.Module, DeepseekV2WeightLoaderMixin): self.tp_size = get_tensor_model_parallel_world_size() self.quant_config = quant_config self.determine_num_fused_shared_experts() - self.use_nsa = is_deepseek_nsa(config) + self.use_dsa = is_deepseek_dsa(config) self.model = DeepseekV2Model( config, quant_config, prefix=add_prefix("model", prefix) ) @@ -2415,15 +2415,15 @@ class DeepseekV2ForCausalLM(nn.Module, DeepseekV2WeightLoaderMixin): ) self.capture_aux_hidden_states = False - self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp() - if self.nsa_enable_prefill_cp: + self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp() + if self.dsa_enable_prefill_cp: self.cp_rank = get_attention_cp_rank() self.cp_size = get_attention_cp_size() else: self.cp_rank = self.cp_size = None q_lora_rank = config.q_lora_rank if hasattr(config, "q_lora_rank") else None - get_attn_tp_context().init_context(q_lora_rank, is_deepseek_nsa(config)) + get_attn_tp_context().init_context(q_lora_rank, is_deepseek_dsa(config)) @property def routed_experts_weights_of_layer(self): @@ -2498,9 +2498,9 @@ class DeepseekV2ForCausalLM(nn.Module, DeepseekV2WeightLoaderMixin): input_embeds: torch.Tensor = None, pp_proxy_tensors: Optional[PPProxyTensors] = None, ) -> torch.Tensor: - if self.nsa_enable_prefill_cp: - if can_nsa_cp_split( - len(input_ids), self.cp_size, self.use_nsa, forward_batch + if self.dsa_enable_prefill_cp: + if can_dsa_cp_split( + len(input_ids), self.cp_size, self.use_dsa, forward_batch ): forward_batch.attn_cp_metadata = prepare_context_parallel_metadata( len(input_ids), diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index d31370da9..ead96a4a6 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -33,14 +33,14 @@ from sglang.srt.distributed import ( ) from sglang.srt.environ import envs from sglang.srt.eplb.expert_location import ModelConfigForExpertLocation +from sglang.srt.layers.attention.dsa.utils import ( + can_dsa_cp_split, + dsa_use_prefill_cp, + is_dsa_enable_prefill_cp, + is_dsa_prefill_cp_round_robin_split, +) from sglang.srt.layers.attention.dsv4.compressor import Compressor from sglang.srt.layers.attention.dsv4.indexer import C4Indexer -from sglang.srt.layers.attention.nsa.utils import ( - can_nsa_cp_split, - is_nsa_enable_prefill_cp, - is_nsa_prefill_cp_round_robin_split, - nsa_use_prefill_cp, -) from sglang.srt.layers.communicator import get_attn_tp_context from sglang.srt.layers.dp_attention import ( _DpGatheredBufferWrapper, @@ -177,8 +177,8 @@ class MQALayer(nn.Module): super().__init__() self.tp_rank = attn_tp_rank = get_attention_tp_rank() self.tp_size = attn_tp_size = get_attention_tp_size() - self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp() - if self.nsa_enable_prefill_cp: + self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp() + if self.dsa_enable_prefill_cp: self.cp_size = get_attention_cp_size() self.tp_rank = attn_tp_rank = 0 self.tp_size = attn_tp_size = 1 @@ -390,7 +390,7 @@ class MQALayer(nn.Module): ) -> None: """Fused: rmsnorm + RoPE + write directly to FlashMLA paged cache. - Replaces the bf16-kv-intermediate path. Used everywhere except the NSA + Replaces the bf16-kv-intermediate path. Used everywhere except the DSA prefill-CP case (which needs bf16 kv for the cross-rank all-gather). """ if qkv_a is not None: @@ -416,7 +416,7 @@ class MQALayer(nn.Module): positions: torch.Tensor, qkv_a: Optional[torch.Tensor] = None, ) -> torch.Tensor: - """Bf16-kv path used by the NSA prefill-CP case (needs all-gather).""" + """Bf16-kv path used by the DSA prefill-CP case (needs all-gather).""" if qkv_a is not None: kv = qkv_a[..., self.q_lora_rank :] else: @@ -508,10 +508,10 @@ class MQALayer(nn.Module): q_lora = self.q_norm(q_lora) q = self._compute_q_b(q_lora, positions, q_out) - use_cp = self.nsa_enable_prefill_cp and nsa_use_prefill_cp(forward_batch) + use_cp = self.dsa_enable_prefill_cp and dsa_use_prefill_cp(forward_batch) kv: Optional[torch.Tensor] if use_cp: - # NSA CP: keep bf16 kv around for the cross-rank all-gather, then + # DSA CP: keep bf16 kv around for the cross-rank all-gather, then # write to the FlashMLA cache after gather. kv = self._compute_kv_bf16(x, positions, qkv_a=qkv_a) kv = cp_all_gather_rerange_output( @@ -567,7 +567,7 @@ class MQALayer(nn.Module): and self.alt_streams is not None and get_is_capture_mode() and x.shape[0] <= self._multi_stream_bs_limit - and not (self.nsa_enable_prefill_cp and nsa_use_prefill_cp(forward_batch)) + and not (self.dsa_enable_prefill_cp and dsa_use_prefill_cp(forward_batch)) ) tp_slice, q_padded, q_out = slice(None), None, None @@ -591,7 +591,7 @@ class MQALayer(nn.Module): # The cache write is always fused / already done by _forward_prepare* -- # tell the backend to skip its own store_cache. When `kv is None` - # (no NSA-CP), pass `q` as a sentinel for the `k is v` assert; the + # (no DSA-CP), pass `q` as a sentinel for the `k is v` assert; the # attention path doesn't read it once `save_kv_cache=False`. attn_k = kv if kv is not None else q o = attn_backend.forward( @@ -694,7 +694,7 @@ class DeepseekV4DecoderLayer(nn.Module): self.hc_attn_scale = nn.Parameter(torch.empty(3, dtype=torch.float32)) self.hc_ffn_scale = nn.Parameter(torch.empty(3, dtype=torch.float32)) self.rms_norm_eps = config.rms_norm_eps - self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp() + self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp() def hc_pre( self, @@ -869,7 +869,7 @@ class DeepseekV4DecoderLayer(nn.Module): if not norm_fused: hidden_states = self.post_attention_layernorm(hidden_states) - _use_cp = self.nsa_enable_prefill_cp and nsa_use_prefill_cp(forward_batch) + _use_cp = self.dsa_enable_prefill_cp and dsa_use_prefill_cp(forward_batch) _use_tp_moe_gather = ( not _use_cp and get_attention_dp_size() > 1 @@ -979,8 +979,8 @@ class DeepseekV4Model(nn.Module): self.hc_head_base = nn.Parameter(torch.empty(hc_mult, dtype=torch.float32)) self.hc_head_scale = nn.Parameter(torch.empty(1, dtype=torch.float32)) - self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp() - if self.nsa_enable_prefill_cp: + self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp() + if self.dsa_enable_prefill_cp: self.cp_size = get_attention_cp_size() def hc_head( @@ -1040,7 +1040,7 @@ class DeepseekV4Model(nn.Module): else: input_ids_global = input_ids - if nsa_use_prefill_cp(forward_batch): + if dsa_use_prefill_cp(forward_batch): if self.pp_group.is_first_rank: hidden_states = cp_split_and_rebuild_data(forward_batch, hidden_states) positions = cp_split_and_rebuild_position(forward_batch, positions) @@ -1060,7 +1060,7 @@ class DeepseekV4Model(nn.Module): ) # CP all-gather only on the last PP rank; PP IPC carries CP-split tensors. - if self.pp_group.is_last_rank and nsa_use_prefill_cp(forward_batch): + if self.pp_group.is_last_rank and dsa_use_prefill_cp(forward_batch): hidden_states = cp_all_gather_rerange_output( hidden_states, self.cp_size, @@ -1113,7 +1113,7 @@ class DeepseekV4ForCausalLM(nn.Module): self.lm_head = PPMissingLayer() self.logits_processor = LogitsProcessor(config) self.capture_aux_hidden_states = False - get_attn_tp_context().init_context(config.q_lora_rank, is_nsa=True) + get_attn_tp_context().init_context(config.q_lora_rank, is_dsa=True) self._routed_experts_weights_of_layer = LazyValue( lambda: { @@ -1129,8 +1129,8 @@ class DeepseekV4ForCausalLM(nn.Module): self.start_layer = self.model.start_layer self.end_layer = self.model.end_layer - self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp() - if self.nsa_enable_prefill_cp: + self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp() + if self.dsa_enable_prefill_cp: self.cp_rank = get_attention_cp_rank() self.cp_size = get_attention_cp_size() @@ -1159,15 +1159,15 @@ class DeepseekV4ForCausalLM(nn.Module): input_embeds: Optional[torch.Tensor] = None, pp_proxy_tensors: Optional[PPProxyTensors] = None, ) -> torch.Tensor: - if self.nsa_enable_prefill_cp: - if can_nsa_cp_split(len(input_ids), self.cp_size, True, forward_batch): + if self.dsa_enable_prefill_cp: + if can_dsa_cp_split(len(input_ids), self.cp_size, True, forward_batch): forward_batch.attn_cp_metadata = prepare_context_parallel_metadata( len(input_ids), self.cp_rank, self.cp_size, forward_batch.seq_lens_cpu.tolist(), ) - if is_nsa_prefill_cp_round_robin_split(): + if is_dsa_prefill_cp_round_robin_split(): metadata = forward_batch.attn_backend.forward_metadata core_meta = metadata.core_attn_metadata core_meta.apply_cp_reindex() diff --git a/python/sglang/srt/models/deepseek_v4_nextn.py b/python/sglang/srt/models/deepseek_v4_nextn.py index 8d14f7498..bd116c29a 100644 --- a/python/sglang/srt/models/deepseek_v4_nextn.py +++ b/python/sglang/srt/models/deepseek_v4_nextn.py @@ -7,11 +7,11 @@ from torch import nn from transformers import PretrainedConfig from sglang.srt.distributed import get_pp_group, get_tensor_model_parallel_world_size -from sglang.srt.layers.attention.nsa.utils import ( - can_nsa_cp_split, - is_nsa_enable_prefill_cp, - is_nsa_prefill_cp_round_robin_split, - nsa_use_prefill_cp, +from sglang.srt.layers.attention.dsa.utils import ( + can_dsa_cp_split, + dsa_use_prefill_cp, + is_dsa_enable_prefill_cp, + is_dsa_prefill_cp_round_robin_split, ) from sglang.srt.layers.dp_attention import ( _DpGatheredBufferWrapper, @@ -104,8 +104,8 @@ class DeepseekV4ModelNextN(nn.Module): compress_ratio_override=COMPRESS_RATIO_NEXTN_LAYER, ) - self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp() - if self.nsa_enable_prefill_cp: + self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp() + if self.dsa_enable_prefill_cp: self.cp_size = get_attention_cp_size() else: self.cp_size = None @@ -165,7 +165,7 @@ class DeepseekV4ModelNextN(nn.Module): else: input_ids_global = input_ids - if nsa_use_prefill_cp(forward_batch): + if dsa_use_prefill_cp(forward_batch): hidden_states = cp_split_and_rebuild_data(forward_batch, hidden_states) positions = cp_split_and_rebuild_position(forward_batch, positions) @@ -177,7 +177,7 @@ class DeepseekV4ModelNextN(nn.Module): input_ids_global=input_ids_global, ) - if nsa_use_prefill_cp(forward_batch): + if dsa_use_prefill_cp(forward_batch): hidden_states = cp_all_gather_rerange_output( hidden_states, self.cp_size, @@ -209,8 +209,8 @@ class DeepseekV4ForCausalLMNextN(DeepseekV4ForCausalLM): self.pp_group = get_pp_group() self.quant_config = quant_config self.determine_num_fused_shared_experts() - self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp() - if self.nsa_enable_prefill_cp: + self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp() + if self.dsa_enable_prefill_cp: self.cp_rank = get_attention_cp_rank() self.cp_size = get_attention_cp_size() else: @@ -236,15 +236,15 @@ class DeepseekV4ForCausalLMNextN(DeepseekV4ForCausalLM): positions: torch.Tensor, forward_batch: ForwardBatch, ) -> torch.Tensor: - if self.nsa_enable_prefill_cp: - if can_nsa_cp_split(len(input_ids), self.cp_size, True, forward_batch): + if self.dsa_enable_prefill_cp: + if can_dsa_cp_split(len(input_ids), self.cp_size, True, forward_batch): forward_batch.attn_cp_metadata = prepare_context_parallel_metadata( len(input_ids), self.cp_rank, self.cp_size, forward_batch.seq_lens_cpu.tolist(), ) - if is_nsa_prefill_cp_round_robin_split(): + if is_dsa_prefill_cp_round_robin_split(): metadata = forward_batch.attn_backend.forward_metadata core_meta = metadata.core_attn_metadata core_meta.apply_cp_reindex() diff --git a/python/sglang/srt/models/glm4_moe_lite.py b/python/sglang/srt/models/glm4_moe_lite.py index 80a035162..4a58ab697 100644 --- a/python/sglang/srt/models/glm4_moe_lite.py +++ b/python/sglang/srt/models/glm4_moe_lite.py @@ -30,7 +30,7 @@ from sglang.srt.distributed import ( get_tensor_model_parallel_world_size, ) from sglang.srt.layers.activation import SiluAndMul -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.communicator import ( LayerCommunicator, LayerScatterModes, @@ -341,7 +341,7 @@ class Glm4MoeLiteDecoderLayer(DeepseekV2DecoderLayer): nn.Module.__init__(self) self.hidden_size = config.hidden_size self.config = config - self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp() + self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp() rope_theta, rope_scaling = get_rope_config(config) max_position_embeddings = getattr(config, "max_position_embeddings", 202752) self.layer_id = layer_id @@ -433,8 +433,8 @@ class Glm4MoeLiteModel(DeepseekV2Model): self.pp_group = get_pp_group() # DeepseekV2Model.forward expects these attributes to exist. - self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp() - self.cp_size = get_attention_tp_size() if self.nsa_enable_prefill_cp else None + self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp() + self.cp_size = get_attention_tp_size() if self.dsa_enable_prefill_cp else None self.gemm_output_zero_allocator_size = 0 self.llama_4_scaling_config = getattr(config, "llama_4_scaling", None) @@ -503,8 +503,8 @@ class Glm4MoeLiteForCausalLM(DeepseekV2ForCausalLM): ) self.capture_aux_hidden_states = False - self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp() - if self.nsa_enable_prefill_cp: + self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp() + if self.dsa_enable_prefill_cp: self.cp_rank = get_attention_tp_rank() self.cp_size = get_attention_tp_size() else: diff --git a/python/sglang/srt/models/mistral_large_3_eagle.py b/python/sglang/srt/models/mistral_large_3_eagle.py index b0d8956a9..65ffbd820 100644 --- a/python/sglang/srt/models/mistral_large_3_eagle.py +++ b/python/sglang/srt/models/mistral_large_3_eagle.py @@ -8,7 +8,7 @@ from torch import nn from transformers import PretrainedConfig from sglang.srt.distributed import get_pp_group -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.layernorm import RMSNorm from sglang.srt.layers.linear import RowParallelLinear from sglang.srt.layers.quantization.base_config import QuantizationConfig @@ -35,7 +35,7 @@ class MistralLarge3EagleModel(DeepseekV2Model): self.vocab_size = config.vocab_size assert get_pp_group().world_size == 1 self.pp_group = get_pp_group() - self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp() + self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp() self.embed_tokens = VocabParallelEmbedding( config.vocab_size, diff --git a/python/sglang/srt/models/sarvam_moe.py b/python/sglang/srt/models/sarvam_moe.py index 36ead547a..bca26936d 100644 --- a/python/sglang/srt/models/sarvam_moe.py +++ b/python/sglang/srt/models/sarvam_moe.py @@ -104,7 +104,8 @@ class AttnForwardMethod(IntEnum): SEPARATE_ROPE_BACKENDS = frozenset( - ["fa3", "flashinfer", "nsa", "cutlass_mla", "trtllm_mla"] + ["fa3", "flashinfer", "dsa", "nsa", "cutlass_mla", "trtllm_mla"] + # "nsa" is a deprecated alias for "dsa" ) CONCAT_ROPE_BACKENDS = frozenset(["flashmla", "triton"]) @@ -667,7 +668,6 @@ class SarvamMoEMLAAttention(nn.Module): k_pe: torch.Tensor, forward_batch: ForwardBatch, ) -> torch.Tensor: - q_pe, k_pe = self.rotary_emb(positions, q_pe, k_pe) q[..., self.qk_nope_head_dim :] = q_pe @@ -989,7 +989,6 @@ class SarvamMoEMLAAttention(nn.Module): class SarvamMoEMLADecoderLayer(nn.Module): - def __init__( self, config: PretrainedConfig, @@ -1139,7 +1138,6 @@ class SarvamMoEMLADecoderLayer(nn.Module): class SarvamMLAModel(nn.Module): - def __init__( self, config: PretrainedConfig, @@ -1223,7 +1221,6 @@ class SarvamMLAModel(nn.Module): class SarvamMLAForCausalLM(nn.Module): - def __init__( self, config: PretrainedConfig, @@ -1475,7 +1472,6 @@ class SarvamMLAForCausalLM(nn.Module): class SarvamMoEForCausalLM(BailingMoEForCausalLM): - @torch.no_grad() def forward_split_prefill( self, diff --git a/python/sglang/srt/platforms/interface.py b/python/sglang/srt/platforms/interface.py index eb6086d8f..3202136cf 100644 --- a/python/sglang/srt/platforms/interface.py +++ b/python/sglang/srt/platforms/interface.py @@ -61,8 +61,8 @@ class SRTPlatform(DeviceMixin): """Return the MLA KV pool class for this platform.""" raise NotImplementedError - def get_nsa_kv_pool_cls(self) -> type: - """Return the NSA KV pool class for this platform (DeepSeek V3.2).""" + def get_dsa_kv_pool_cls(self) -> type: + """Return the DSA KV pool class for this platform (DeepSeek V3.2).""" raise NotImplementedError def get_paged_allocator_cls(self) -> type: diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 7716921e1..ee751badd 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -155,7 +155,8 @@ ATTENTION_BACKEND_CHOICES = [ "triton", "torch_native", "flex_attention", - "nsa", + "dsa", + "nsa", # Deprecated alias for "dsa" "dsv4", "compressed", # Deprecated alias for "dsv4" # NVIDIA specific @@ -246,13 +247,14 @@ LORA_BACKEND_CHOICES = ["triton", "csgmv", "ascend", "torch_native"] ENCODER_TRANSFER_BACKEND_CHOICES = ["zmq_to_scheduler", "zmq_to_tokenizer", "mooncake"] -NSA_PREFILL_CP_SPLIT_CHOICES = ["in-seq-split", "round-robin-split"] +DSA_PREFILL_CP_SPLIT_CHOICES = ["in-seq-split", "round-robin-split"] +NSA_PREFILL_CP_SPLIT_CHOICES = DSA_PREFILL_CP_SPLIT_CHOICES # deprecated alias PREFILL_CP_SPLIT_CHOICES = ["in-seq-split"] DEFAULT_LORA_EVICTION_POLICY = "lru" -NSA_CHOICES = [ +DSA_CHOICES = [ "flashmla_sparse", "flashmla_kv", "flashmla_auto", @@ -261,6 +263,7 @@ NSA_CHOICES = [ "aiter", "trtllm", ] +NSA_CHOICES = DSA_CHOICES # deprecated alias MAMBA_SCHEDULER_STRATEGY_CHOICES = ["auto", "no_buffer", "extra_buffer"] @@ -535,10 +538,10 @@ class ServerArgs: mm_attention_backend: Optional[str] = None fp8_gemm_runner_backend: str = "auto" fp4_gemm_runner_backend: str = "auto" - nsa_prefill_backend: Optional[str] = ( + dsa_prefill_backend: Optional[str] = ( None # None = auto-detect based on hardware/kv_cache_dtype ) - nsa_decode_backend: Optional[str] = ( + dsa_decode_backend: Optional[str] = ( None # auto-detect based on hardware/kv_cache_dtype ) disable_flashinfer_autotune: bool = False @@ -741,8 +744,8 @@ class ServerArgs: enable_attn_tp_input_scattered: bool = False gc_threshold: Optional[List[int]] = None # Context parallelism used in the long sequence prefill phase of DeepSeek v3.2 - enable_nsa_prefill_context_parallel: bool = False - nsa_prefill_cp_mode: str = "round-robin-split" + enable_dsa_prefill_context_parallel: bool = False + dsa_prefill_cp_mode: str = "round-robin-split" enable_fused_qk_norm_rope: bool = False enable_precise_embedding_interpolation: bool = False enable_fused_moe_sum_all_reduce: bool = False @@ -1647,15 +1650,15 @@ class ServerArgs: return capture_sizes - def _set_default_nsa_kv_cache_dtype(self, major: int, quantization: str) -> str: - user_set_prefill = self.nsa_prefill_backend is not None - user_set_decode = self.nsa_decode_backend is not None + def _set_default_dsa_kv_cache_dtype(self, major: int, quantization: str) -> str: + user_set_prefill = self.dsa_prefill_backend is not None + user_set_decode = self.dsa_decode_backend is not None # If user specified a backend but didn't explicitly set kv_cache_dtype, # suggest them to be explicit about kv_cache_dtype to avoid surprises if (user_set_prefill or user_set_decode) and self.kv_cache_dtype == "auto": logger.warning( - "When specifying --nsa-prefill-backend or --nsa-decode-backend, " + "When specifying --dsa-prefill-backend or --dsa-decode-backend, " "you should also explicitly set --kv-cache-dtype (e.g., 'fp8_e4m3' or 'bfloat16'). " "DeepSeek V3.2 defaults to FP8 KV cache which may not be compatible with all backends." ) @@ -1675,56 +1678,56 @@ class ServerArgs: "fp8_e4m3", ], "DeepSeek DSA only supports bf16/bfloat16 or fp8_e4m3 kv_cache_dtype" - def _set_default_nsa_backends(self, kv_cache_dtype: str, major: int) -> str: + def _set_default_dsa_backends(self, kv_cache_dtype: str, major: int) -> str: from sglang.srt.arg_groups.hisparse_hook import ( - apply_hisparse_nsa_backend_defaults, + apply_hisparse_dsa_backend_defaults, ) - user_set_prefill = self.nsa_prefill_backend is not None - user_set_decode = self.nsa_decode_backend is not None + user_set_prefill = self.dsa_prefill_backend is not None + user_set_decode = self.dsa_decode_backend is not None - if apply_hisparse_nsa_backend_defaults( + if apply_hisparse_dsa_backend_defaults( self, user_set_prefill, user_set_decode, kv_cache_dtype ): return if not user_set_prefill and not user_set_decode and is_hip(): - self.nsa_prefill_backend = "tilelang" - self.nsa_decode_backend = "tilelang" + self.dsa_prefill_backend = "tilelang" + self.dsa_decode_backend = "tilelang" elif kv_cache_dtype == "fp8_e4m3": if major >= 10: if not user_set_prefill: - self.nsa_prefill_backend = "trtllm" + self.dsa_prefill_backend = "trtllm" if not user_set_decode: - self.nsa_decode_backend = "trtllm" + self.dsa_decode_backend = "trtllm" else: # Hopper FP8 defaults to flashmla_kv for both prefill and decode. if not user_set_prefill: - self.nsa_prefill_backend = "flashmla_kv" + self.dsa_prefill_backend = "flashmla_kv" if not user_set_decode: - self.nsa_decode_backend = "flashmla_kv" + self.dsa_decode_backend = "flashmla_kv" else: # set prefill/decode backends based on hardware architecture. if major >= 10: if not user_set_prefill: - self.nsa_prefill_backend = "flashmla_sparse" + self.dsa_prefill_backend = "flashmla_sparse" if not user_set_decode: - self.nsa_decode_backend = "trtllm" + self.dsa_decode_backend = "trtllm" else: # Hopper defaults for bfloat16 if not user_set_prefill: - self.nsa_prefill_backend = "flashmla_sparse" + self.dsa_prefill_backend = "flashmla_sparse" if not user_set_decode: - self.nsa_decode_backend = "fa3" + self.dsa_decode_backend = "fa3" logger.warning( - f"Set NSA backends for {self.kv_cache_dtype} KV Cache: prefill={self.nsa_prefill_backend}, decode={self.nsa_decode_backend}." + f"Set DSA backends for {self.kv_cache_dtype} KV Cache: prefill={self.dsa_prefill_backend}, decode={self.dsa_decode_backend}." ) def _handle_model_specific_adjustments(self): from sglang.srt.configs.model_config import ( get_mimo_v2_fused_qkv_expected_tp_size, - is_deepseek_nsa, + is_deepseek_dsa, ) if parse_connector_type(self.model_path) == ConnectorType.INSTANCE: @@ -1765,37 +1768,37 @@ class ServerArgs: "GlmMoeDsaForCausalLM", ]: # Set attention backend for DeepSeek - if is_deepseek_nsa(hf_config): # DeepSeek 3.2/GLM 5 + if is_deepseek_dsa(hf_config): # DeepSeek 3.2/GLM 5 if model_arch == "GlmMoeDsaForCausalLM" and is_blackwell_supported(): - envs.SGLANG_NSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD.set(0) + envs.SGLANG_DSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD.set(0) logger.warning( - "Force NSA prefill to use sparse MLA (i.e. disable MHA_ONE_SHOT) for GlmMoeDsaForCausalLM on Blackwell." + "Force DSA prefill to use sparse MLA (i.e. disable MHA_ONE_SHOT) for GlmMoeDsaForCausalLM on Blackwell." ) else: - if envs.SGLANG_NSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD.is_set(): + if envs.SGLANG_DSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD.is_set(): logger.warning( - f"Dense attention kv len threshold is manually set to {envs.SGLANG_NSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD.get()} for DSA. Caution: This may cause performance regression if the threshold is larger than the index topk of model." + f"Dense attention kv len threshold is manually set to {envs.SGLANG_DSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD.get()} for DSA. Caution: This may cause performance regression if the threshold is larger than the index topk of model." ) else: # When threshold is not manually set, set it to the index topk of model - from sglang.srt.configs.model_config import get_nsa_index_topk + from sglang.srt.configs.model_config import get_dsa_index_topk - envs.SGLANG_NSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD.set( - get_nsa_index_topk(hf_config) + envs.SGLANG_DSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD.set( + get_dsa_index_topk(hf_config) ) logger.warning( - f"Set dense attention kv len threshold to model index_topk={envs.SGLANG_NSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD.get()} for DeepSeek with DSA." + f"Set dense attention kv len threshold to model index_topk={envs.SGLANG_DSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD.get()} for DeepSeek with DSA." ) if self.is_attention_backend_not_set(): - self.attention_backend = "nsa" - logger.info("Use nsa attention backend for DeepSeek with DSA.") + self.attention_backend = "dsa" + logger.info("Use dsa attention backend for DeepSeek with DSA.") if not is_npu() and not is_xpu(): # CUDA or ROCm GPU - if self.enable_nsa_prefill_context_parallel: + if self.enable_dsa_prefill_context_parallel: logger.warning( "Context parallel feature is still under experiment. It has only been verified on Hopper platform." ) - if self.nsa_prefill_cp_mode == "in-seq-split": + if self.dsa_prefill_cp_mode == "in-seq-split": # TODO Supports moe_dense_tp_size != 1, kv cache dtype = "fp8",moe_a2a_backend non-deepep and cross-machine operation . self.enable_dp_attention = True self.moe_dense_tp_size = 1 @@ -1819,7 +1822,7 @@ class ServerArgs: f"Enable Context Parallel opt for deeeseekv3.2-DSA, Setting dp_size == {self.dp_size} and moe_dense_tp_size == {self.moe_dense_tp_size}, ep_size == {self.ep_size}, tp_size == {self.tp_size}, kv_cache_dtype == {self.kv_cache_dtype}, moe_a2a_backend {self.moe_a2a_backend} " ) else: - # Pure TP and partial DP Attention mode is active for NSA, logging a warning + # Pure TP and partial DP Attention mode is active for DSA, logging a warning if self.dp_size < self.tp_size: logger.warning( f"DSA with TP mode is active, dp_size={self.dp_size}, tp_size={self.tp_size}, " @@ -1827,15 +1830,15 @@ class ServerArgs: ) # Deferred import to avoid a circular import at module-load - # time (nsa.utils imports get_global_server_args). - from sglang.srt.layers.attention.nsa.utils import ( + # time (dsa.utils imports get_global_server_args). + from sglang.srt.layers.attention.dsa.utils import ( aiter_can_use_preshuffle_paged_mqa, ) if is_hip() and not aiter_can_use_preshuffle_paged_mqa(): - # Legacy ROCm NSA path: aiter's gluon paged-MQA kernel is + # Legacy ROCm DSA path: aiter's gluon paged-MQA kernel is # unavailable (Triton<3.5 and AITER_ENABLE_AOT_GLUON_PA_MQA_LOGITS - # not set, or SGLANG_NSA_HIP_DISABLE_PRESHUFFLE=1 / SGLANG_USE_AITER=0). + # not set, or SGLANG_DSA_HIP_DISABLE_PRESHUFFLE=1 / SGLANG_USE_AITER=0). self.page_size = 1 logger.warning( "Setting page size to 1 for DeepSeek DSA on ROCm " @@ -1849,13 +1852,13 @@ class ServerArgs: import torch major, _ = torch.cuda.get_device_capability() - self._set_default_nsa_kv_cache_dtype(major, self.quantization) - self._set_default_nsa_backends(self.kv_cache_dtype, major) + self._set_default_dsa_kv_cache_dtype(major, self.quantization) + self._set_default_dsa_backends(self.kv_cache_dtype, major) - if self.enable_nsa_prefill_context_parallel: + if self.enable_dsa_prefill_context_parallel: assert ( self.disaggregation_mode != "decode" - ), "CP is only supported for prefill when PD disaggregation, please remove --enable-nsa-prefill-context-parallel." + ), "CP is only supported for prefill when PD disaggregation, please remove --enable-dsa-prefill-context-parallel." else: # DeepSeek V3/R1/V3.1 @@ -3408,7 +3411,7 @@ class ServerArgs: "the paged cache, which the no-op pool does not support." ) - # HiSparse selects a different pool class (HiSparseNSATokenToKVPool / + # HiSparse selects a different pool class (HiSparseDSATokenToKVPool / # HiSparseTokenToKVPoolAllocator) that is not the no-op pool. if self.enable_hisparse: raise ValueError( @@ -5349,18 +5352,40 @@ class ServerArgs: help="Set multimodal attention backend.", ) parser.add_argument( - "--nsa-prefill-backend", - default=ServerArgs.nsa_prefill_backend, + "--dsa-prefill-backend", + dest="dsa_prefill_backend", + default=ServerArgs.dsa_prefill_backend, type=str, - choices=NSA_CHOICES, - help="NSA prefill backend. If not specified, auto-detects based on hardware and kv_cache_dtype.", + choices=DSA_CHOICES, + help="DSA (DeepSeek Sparse Attention) prefill backend. If not specified, auto-detects based on hardware and kv_cache_dtype.", + ) + parser.add_argument( + "--nsa-prefill-backend", + dest="dsa_prefill_backend", + action=DeprecatedAliasStoreAction, + new_flag="--dsa-prefill-backend", + default=argparse.SUPPRESS, + type=str, + choices=DSA_CHOICES, + help="[Deprecated] Use --dsa-prefill-backend instead.", + ) + parser.add_argument( + "--dsa-decode-backend", + dest="dsa_decode_backend", + default=ServerArgs.dsa_decode_backend, + type=str, + choices=DSA_CHOICES, + help="DSA (DeepSeek Sparse Attention) decode backend. If not specified, auto-detects based on hardware and kv_cache_dtype.", ) parser.add_argument( "--nsa-decode-backend", - default=ServerArgs.nsa_decode_backend, + dest="dsa_decode_backend", + action=DeprecatedAliasStoreAction, + new_flag="--dsa-decode-backend", + default=argparse.SUPPRESS, type=str, - choices=NSA_CHOICES, - help="NSA decode backend. If not specified, auto-detects based on hardware and kv_cache_dtype.", + choices=DSA_CHOICES, + help="[Deprecated] Use --dsa-decode-backend instead.", ) parser.add_argument( "--fp8-gemm-backend", @@ -6396,15 +6421,34 @@ class ServerArgs: help="Allow input of attention to be scattered when only using tensor parallelism, to reduce the computational load of operations such as qkv latent.", ) parser.add_argument( - "--enable-nsa-prefill-context-parallel", + "--enable-dsa-prefill-context-parallel", + dest="enable_dsa_prefill_context_parallel", action="store_true", help="Enable context parallelism used in the long sequence prefill phase of DeepSeek v3.2.", ) parser.add_argument( - "--nsa-prefill-cp-mode", + "--enable-nsa-prefill-context-parallel", + dest="enable_dsa_prefill_context_parallel", + action=DeprecatedStoreTrueAction, + new_flag="--enable-dsa-prefill-context-parallel", + help="[Deprecated] Use --enable-dsa-prefill-context-parallel instead.", + ) + parser.add_argument( + "--dsa-prefill-cp-mode", + dest="dsa_prefill_cp_mode", type=str, - default=ServerArgs.nsa_prefill_cp_mode, - choices=NSA_PREFILL_CP_SPLIT_CHOICES, + default=ServerArgs.dsa_prefill_cp_mode, + choices=DSA_PREFILL_CP_SPLIT_CHOICES, + help="Token splitting mode for the prefill phase of DeepSeek v3.2 under context parallelism.", + ) + parser.add_argument( + "--nsa-prefill-cp-mode", + dest="dsa_prefill_cp_mode", + action=DeprecatedAliasStoreAction, + new_flag="--dsa-prefill-cp-mode", + default=argparse.SUPPRESS, + type=str, + choices=DSA_PREFILL_CP_SPLIT_CHOICES, help="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.", ) diff --git a/python/sglang/srt/speculative/draft_utils.py b/python/sglang/srt/speculative/draft_utils.py index ce6ac5334..b09e0c6a0 100644 --- a/python/sglang/srt/speculative/draft_utils.py +++ b/python/sglang/srt/speculative/draft_utils.py @@ -54,7 +54,8 @@ class DraftBackendFactory: "trtllm_mha": self._create_trtllm_mha_decode_backend, "trtllm_mla": self._create_trtllm_mla_decode_backend, "tokenspeed_mla": self._create_tokenspeed_mla_decode_backend, - "nsa": self._create_nsa_decode_backend, + "dsa": self._create_dsa_decode_backend, + "nsa": self._create_dsa_decode_backend, # Deprecated alias for "dsa" "ascend": self._create_ascend_decode_backend, "fa4": self._create_fa4_decode_backend, "dsv4": self._create_dsv4_decode_backend, @@ -81,7 +82,8 @@ class DraftBackendFactory: "trtllm_mha": self._create_trtllm_mha_prefill_backend, "trtllm_mla": self._create_trtllm_mla_prefill_backend, "tokenspeed_mla": self._create_tokenspeed_mla_prefill_backend, - "nsa": self._create_nsa_prefill_backend, + "dsa": self._create_dsa_prefill_backend, + "nsa": self._create_dsa_prefill_backend, # Deprecated alias for "dsa" "ascend": self._create_ascend_prefill_backend, "fa4": self._create_fa4_prefill_backend, "dsv4": self._create_dsv4_prefill_backend, @@ -97,19 +99,19 @@ class DraftBackendFactory: "EAGLE is not supported in attention backend {backend_type}", ) - def _create_nsa_decode_backend(self): - from sglang.srt.layers.attention.nsa_backend import ( - NativeSparseAttnMultiStepBackend, + def _create_dsa_decode_backend(self): + from sglang.srt.layers.attention.dsa_backend import ( + DeepseekSparseAttnMultiStepBackend, ) - return NativeSparseAttnMultiStepBackend( + return DeepseekSparseAttnMultiStepBackend( self.draft_model_runner, self.topk, self.speculative_num_steps ) - def _create_nsa_prefill_backend(self): - from sglang.srt.layers.attention.nsa_backend import NativeSparseAttnBackend + def _create_dsa_prefill_backend(self): + from sglang.srt.layers.attention.dsa_backend import DeepseekSparseAttnBackend - return NativeSparseAttnBackend(self.draft_model_runner, skip_prefill=False) + return DeepseekSparseAttnBackend(self.draft_model_runner, skip_prefill=False) def _create_flashinfer_decode_backend(self): if not get_global_server_args().use_mla_backend: diff --git a/python/sglang/test/nightly_utils.py b/python/sglang/test/nightly_utils.py index 2a9d01f2e..a26faea2f 100644 --- a/python/sglang/test/nightly_utils.py +++ b/python/sglang/test/nightly_utils.py @@ -64,7 +64,7 @@ class NightlyBenchmarkRunner: Args: model_path: Path to the model (e.g., "deepseek-ai/DeepSeek-V3.1") - variant: Optional variant suffix (e.g., "basic", "mtp", "nsa") + variant: Optional variant suffix (e.g., "basic", "mtp", "dsa") Returns: Tuple of (profile_path_prefix, json_output_file) diff --git a/test/manual/dsv4/test_b200_flash.py b/test/manual/dsv4/test_b200_flash.py index 05d738bcf..bcf854aba 100644 --- a/test/manual/dsv4/test_b200_flash.py +++ b/test/manual/dsv4/test_b200_flash.py @@ -87,8 +87,8 @@ class TestB200FlashCP(DSV4FlashAime25TestBase): "4", "--moe-a2a-backend", "deepep", - "--enable-nsa-prefill-context-parallel", - "--nsa-prefill-cp-mode", + "--enable-dsa-prefill-context-parallel", + "--dsa-prefill-cp-mode", "round-robin-split", "--chunked-prefill-size", "16384", diff --git a/test/manual/dsv4/test_b200_pro.py b/test/manual/dsv4/test_b200_pro.py index eba9b6546..b6bc279b5 100644 --- a/test/manual/dsv4/test_b200_pro.py +++ b/test/manual/dsv4/test_b200_pro.py @@ -101,8 +101,8 @@ class TestB200ProCP(DSV4ProAime25TestBase): "8", "--moe-a2a-backend", "deepep", - "--enable-nsa-prefill-context-parallel", - "--nsa-prefill-cp-mode", + "--enable-dsa-prefill-context-parallel", + "--dsa-prefill-cp-mode", "round-robin-split", "--chunked-prefill-size", "16384", diff --git a/test/manual/dsv4/test_b300_flash.py b/test/manual/dsv4/test_b300_flash.py index 4e800526d..72279b9f0 100644 --- a/test/manual/dsv4/test_b300_flash.py +++ b/test/manual/dsv4/test_b300_flash.py @@ -89,8 +89,8 @@ class TestB300FlashCP(DSV4FlashAime25TestBase): "4", "--moe-a2a-backend", "deepep", - "--enable-nsa-prefill-context-parallel", - "--nsa-prefill-cp-mode", + "--enable-dsa-prefill-context-parallel", + "--dsa-prefill-cp-mode", "round-robin-split", "--chunked-prefill-size", "16384", diff --git a/test/manual/dsv4/test_b300_pro.py b/test/manual/dsv4/test_b300_pro.py index 701a42a39..4f8618b51 100644 --- a/test/manual/dsv4/test_b300_pro.py +++ b/test/manual/dsv4/test_b300_pro.py @@ -103,8 +103,8 @@ class TestB300ProCP(DSV4ProAime25TestBase): "8", "--moe-a2a-backend", "deepep", - "--enable-nsa-prefill-context-parallel", - "--nsa-prefill-cp-mode", + "--enable-dsa-prefill-context-parallel", + "--dsa-prefill-cp-mode", "round-robin-split", "--chunked-prefill-size", "16384", diff --git a/test/manual/dsv4/test_dsv4_pd_disagg_nixl.py b/test/manual/dsv4/test_dsv4_pd_disagg_nixl.py index 74095c09b..60dcb5431 100644 --- a/test/manual/dsv4/test_dsv4_pd_disagg_nixl.py +++ b/test/manual/dsv4/test_dsv4_pd_disagg_nixl.py @@ -1,7 +1,7 @@ """DSV4 Flash PD-disagg with NIXL backend. Both sides run dp-attention + deepep + EAGLE MTP so attn_tp_size and the V4 state pool layout are fully symmetric: same SWA item_len under matching attn_tp, and same -NSA c4/c128 indexer ring buffer size under matching spec status. nixl +DSA c4/c128 indexer ring buffer size under matching spec status. nixl `send_state` is page-by-index and has no V4 TP-slice / spec-asymmetric path, so any layout mismatch would trip the item_len assert in `nixl/conn.py`.""" diff --git a/test/manual/dsv4/test_gb300_flash.py b/test/manual/dsv4/test_gb300_flash.py index ef997f50e..5553b8c00 100644 --- a/test/manual/dsv4/test_gb300_flash.py +++ b/test/manual/dsv4/test_gb300_flash.py @@ -87,8 +87,8 @@ class TestGB300FlashCP(DSV4FlashAime25TestBase): "4", "--moe-a2a-backend", "deepep", - "--enable-nsa-prefill-context-parallel", - "--nsa-prefill-cp-mode", + "--enable-dsa-prefill-context-parallel", + "--dsa-prefill-cp-mode", "round-robin-split", "--chunked-prefill-size", "16384", diff --git a/test/manual/dsv4/test_gb300_pro.py b/test/manual/dsv4/test_gb300_pro.py index 1c53129b2..54ca95983 100644 --- a/test/manual/dsv4/test_gb300_pro.py +++ b/test/manual/dsv4/test_gb300_pro.py @@ -103,8 +103,8 @@ class TestGB300ProCP(DSV4ProAime25TestBase): "4", "--moe-a2a-backend", "deepep", - "--enable-nsa-prefill-context-parallel", - "--nsa-prefill-cp-mode", + "--enable-dsa-prefill-context-parallel", + "--dsa-prefill-cp-mode", "round-robin-split", "--chunked-prefill-size", "16384", diff --git a/test/manual/dsv4/test_h200_fp8_flash.py b/test/manual/dsv4/test_h200_fp8_flash.py index fe53ded83..b69834f9b 100644 --- a/test/manual/dsv4/test_h200_fp8_flash.py +++ b/test/manual/dsv4/test_h200_fp8_flash.py @@ -98,8 +98,8 @@ class TestH200Fp8FlashCP(DSV4FlashAime25TestBase): "4", "--moe-a2a-backend", "deepep", - "--enable-nsa-prefill-context-parallel", - "--nsa-prefill-cp-mode", + "--enable-dsa-prefill-context-parallel", + "--dsa-prefill-cp-mode", "round-robin-split", "--chunked-prefill-size", "16384", diff --git a/test/manual/layers/attention/nsa/test_act_quant_triton.py b/test/manual/layers/attention/dsa/test_act_quant_triton.py similarity index 98% rename from test/manual/layers/attention/nsa/test_act_quant_triton.py rename to test/manual/layers/attention/dsa/test_act_quant_triton.py index a5257dff6..f24c30285 100644 --- a/test/manual/layers/attention/nsa/test_act_quant_triton.py +++ b/test/manual/layers/attention/dsa/test_act_quant_triton.py @@ -9,8 +9,8 @@ from typing import Tuple import pytest import torch -from sglang.srt.layers.attention.nsa.tilelang_kernel import act_quant -from sglang.srt.layers.attention.nsa.triton_kernel import act_quant as act_quant_triton +from sglang.srt.layers.attention.dsa.tilelang_kernel import act_quant +from sglang.srt.layers.attention.dsa.triton_kernel import act_quant as act_quant_triton def benchmark_kernel( diff --git a/test/manual/layers/attention/nsa/test_get_k_scale_triton_kernel.py b/test/manual/layers/attention/dsa/test_get_k_scale_triton_kernel.py similarity index 98% rename from test/manual/layers/attention/nsa/test_get_k_scale_triton_kernel.py rename to test/manual/layers/attention/dsa/test_get_k_scale_triton_kernel.py index 296567559..a0c5956e1 100644 --- a/test/manual/layers/attention/nsa/test_get_k_scale_triton_kernel.py +++ b/test/manual/layers/attention/dsa/test_get_k_scale_triton_kernel.py @@ -1,6 +1,6 @@ import torch -from sglang.srt.layers.attention.nsa.index_buf_accessor import ( +from sglang.srt.layers.attention.dsa.index_buf_accessor import ( _get_k_and_s_triton_kernel, ) diff --git a/test/manual/layers/attention/nsa/test_index_buf_accessor.py b/test/manual/layers/attention/dsa/test_index_buf_accessor.py similarity index 96% rename from test/manual/layers/attention/nsa/test_index_buf_accessor.py rename to test/manual/layers/attention/dsa/test_index_buf_accessor.py index 49395263d..3e9f9e374 100644 --- a/test/manual/layers/attention/nsa/test_index_buf_accessor.py +++ b/test/manual/layers/attention/dsa/test_index_buf_accessor.py @@ -1,5 +1,5 @@ """ -Correctness tests for NSA Indexer K/S Buffer Access with Fused Triton Kernels. +Correctness tests for DSA Indexer K/S Buffer Access with Fused Triton Kernels. This test verifies that the optimized Triton implementations (GetK, GetS, GetKAndS) produce identical results to the torch_fast baseline implementations. @@ -13,11 +13,11 @@ Test coverage: import pytest import torch -from sglang.srt.layers.attention.nsa.index_buf_accessor import GetK, GetKAndS, GetS +from sglang.srt.layers.attention.dsa.index_buf_accessor import GetK, GetKAndS, GetS -class MockNSATokenToKVPool: - """Mock pool object that mimics NSATokenToKVPool for testing.""" +class MockDSATokenToKVPool: + """Mock pool object that mimics DSATokenToKVPool for testing.""" def __init__( self, @@ -78,7 +78,7 @@ class TestGetK: seq_len = min(seq_len, max_seq_len) # Create mock pool - pool = MockNSATokenToKVPool( + pool = MockDSATokenToKVPool( page_size=page_size, index_head_dim=index_head_dim, device=device ) @@ -119,7 +119,7 @@ class TestGetK: num_pages = 10 seq_len = 320 # 5 pages - pool = MockNSATokenToKVPool( + pool = MockDSATokenToKVPool( page_size=page_size, index_head_dim=index_head_dim, device=device ) buf = create_test_buffer(num_pages, page_size, index_head_dim, device) @@ -140,7 +140,7 @@ class TestGetK: num_pages = 5 seq_len = 192 # 3 pages - pool = MockNSATokenToKVPool( + pool = MockDSATokenToKVPool( page_size=page_size, index_head_dim=index_head_dim, device=device ) buf = create_test_buffer(num_pages, page_size, index_head_dim, device) @@ -171,7 +171,7 @@ class TestGetS: seq_len = min(seq_len, max_seq_len) # Create mock pool - pool = MockNSATokenToKVPool( + pool = MockDSATokenToKVPool( page_size=page_size, index_head_dim=index_head_dim, device=device ) @@ -212,7 +212,7 @@ class TestGetS: num_pages = 10 seq_len = 320 # 5 pages - pool = MockNSATokenToKVPool( + pool = MockDSATokenToKVPool( page_size=page_size, index_head_dim=index_head_dim, device=device ) buf = create_test_buffer(num_pages, page_size, index_head_dim, device) @@ -233,7 +233,7 @@ class TestGetS: num_pages = 5 seq_len = 192 # 3 pages - pool = MockNSATokenToKVPool( + pool = MockDSATokenToKVPool( page_size=page_size, index_head_dim=index_head_dim, device=device ) buf = create_test_buffer(num_pages, page_size, index_head_dim, device) @@ -267,7 +267,7 @@ class TestGetKAndS: seq_len_tensor = torch.tensor([seq_len], dtype=torch.int64, device=device) # Create mock pool - pool = MockNSATokenToKVPool( + pool = MockDSATokenToKVPool( page_size=page_size, index_head_dim=index_head_dim, device=device ) @@ -326,7 +326,7 @@ class TestGetKAndS: seq_len = 320 # 5 pages seq_len_tensor = torch.tensor([seq_len], dtype=torch.int64, device=device) - pool = MockNSATokenToKVPool( + pool = MockDSATokenToKVPool( page_size=page_size, index_head_dim=index_head_dim, device=device ) buf = create_test_buffer(num_pages, page_size, index_head_dim, device) @@ -356,7 +356,7 @@ class TestGetKAndS: seq_len = 192 # 3 pages seq_len_tensor = torch.tensor([seq_len], dtype=torch.int64, device=device) - pool = MockNSATokenToKVPool( + pool = MockDSATokenToKVPool( page_size=page_size, index_head_dim=index_head_dim, device=device ) buf = create_test_buffer(num_pages, page_size, index_head_dim, device) @@ -386,7 +386,7 @@ class TestGetKAndS: seq_len = 100 # Not a multiple of 64 seq_len_tensor = torch.tensor([seq_len], dtype=torch.int64, device=device) - pool = MockNSATokenToKVPool( + pool = MockDSATokenToKVPool( page_size=page_size, index_head_dim=index_head_dim, device=device ) buf = create_test_buffer(num_pages, page_size, index_head_dim, device) @@ -422,7 +422,7 @@ class TestEdgeCases: seq_len = 1 seq_len_tensor = torch.tensor([seq_len], dtype=torch.int64, device=device) - pool = MockNSATokenToKVPool( + pool = MockDSATokenToKVPool( page_size=page_size, index_head_dim=index_head_dim, device=device ) buf = create_test_buffer(num_pages, page_size, index_head_dim, device) @@ -455,7 +455,7 @@ class TestEdgeCases: seq_len = 192 # Exactly 3 pages seq_len_tensor = torch.tensor([seq_len], dtype=torch.int64, device=device) - pool = MockNSATokenToKVPool( + pool = MockDSATokenToKVPool( page_size=page_size, index_head_dim=index_head_dim, device=device ) buf = create_test_buffer(num_pages, page_size, index_head_dim, device) @@ -488,7 +488,7 @@ class TestEdgeCases: seq_len = 4096 # 64 pages seq_len_tensor = torch.tensor([seq_len], dtype=torch.int64, device=device) - pool = MockNSATokenToKVPool( + pool = MockDSATokenToKVPool( page_size=page_size, index_head_dim=index_head_dim, device=device ) buf = create_test_buffer(num_pages, page_size, index_head_dim, device) @@ -520,7 +520,7 @@ class TestEdgeCases: def print_test_summary(): """Print a summary message about the test suite.""" print("\n" + "=" * 80) - print("NSA Indexer K/S Buffer Accessor Correctness Tests") + print("DSA Indexer K/S Buffer Accessor Correctness Tests") print("=" * 80) print("Testing Triton implementations against torch_fast baseline:") print(" - GetK.triton() vs GetK.torch_fast()") diff --git a/test/manual/nightly/test_deepseek_v32_perf.py b/test/manual/nightly/test_deepseek_v32_perf.py index 75b258fdd..b5d6811d2 100644 --- a/test/manual/nightly/test_deepseek_v32_perf.py +++ b/test/manual/nightly/test_deepseek_v32_perf.py @@ -55,7 +55,7 @@ class TestNightlyDeepseekV32Performance(unittest.TestCase): ], }, { - "name": "nsa", + "name": "dsa", "other_args": [ "--trust-remote-code", "--tp", @@ -64,10 +64,10 @@ class TestNightlyDeepseekV32Performance(unittest.TestCase): "8", "--enable-dp-attention", "--attention-backend", - "nsa", - "--nsa-prefill-backend", + "dsa", + "--dsa-prefill-backend", "flashmla_sparse", - "--nsa-decode-backend", + "--dsa-decode-backend", "flashmla_kv", "--model-loader-extra-config", '{"enable_multithread_load": true}', @@ -80,10 +80,10 @@ class TestNightlyDeepseekV32Performance(unittest.TestCase): "--tp", "8", "--attention-backend", - "nsa", - "--nsa-prefill-backend", + "dsa", + "--dsa-prefill-backend", "flashmla_sparse", - "--nsa-decode-backend", + "--dsa-decode-backend", "flashmla_kv", "--model-loader-extra-config", '{"enable_multithread_load": true}', diff --git a/test/manual/test_dsa_alias_cli_registry_env.py b/test/manual/test_dsa_alias_cli_registry_env.py new file mode 100644 index 000000000..859b0a82f --- /dev/null +++ b/test/manual/test_dsa_alias_cli_registry_env.py @@ -0,0 +1,302 @@ +""" +Manual test for step 01: NSA → DSA user-facing alias layer. + +Tests: + 1. CLI: --dsa-* canonical flags write to dsa_* attrs + 2. CLI: --nsa-* deprecated flags write to dsa_* attrs + log deprecation warning + 3. Registry: "dsa" key creates the backend; "nsa" key triggers DeprecationWarning + 4. Env: SGLANG_DSA_* canonical vars work + 5. Env: SGLANG_NSA_* deprecated vars fall back to SGLANG_DSA_* with DeprecationWarning + +Run: + python test/manual/test_dsa_alias_cli_registry_env.py +""" + +import argparse +import os +import sys +import unittest +import warnings + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "../../python")) + + +class TestDSAChoicesAndFields(unittest.TestCase): + """Verify DSA_CHOICES constant and ServerArgs field renaming.""" + + def setUp(self): + from sglang.srt.server_args import ( + DSA_CHOICES, + DSA_PREFILL_CP_SPLIT_CHOICES, + NSA_CHOICES, + NSA_PREFILL_CP_SPLIT_CHOICES, + ServerArgs, + ) + + self.ServerArgs = ServerArgs + self.DSA_CHOICES = DSA_CHOICES + self.NSA_CHOICES = NSA_CHOICES + self.DSA_PREFILL_CP_SPLIT_CHOICES = DSA_PREFILL_CP_SPLIT_CHOICES + self.NSA_PREFILL_CP_SPLIT_CHOICES = NSA_PREFILL_CP_SPLIT_CHOICES + + def test_dsa_choices_is_canonical(self): + self.assertIn("fa3", self.DSA_CHOICES) + self.assertIn("tilelang", self.DSA_CHOICES) + + def test_nsa_choices_is_alias(self): + self.assertIs( + self.NSA_CHOICES, + self.DSA_CHOICES, + "NSA_CHOICES must be the same object as DSA_CHOICES", + ) + + def test_nsa_cp_split_choices_is_alias(self): + self.assertIs( + self.NSA_PREFILL_CP_SPLIT_CHOICES, + self.DSA_PREFILL_CP_SPLIT_CHOICES, + ) + + def test_serverargs_has_dsa_fields(self): + sa = self.ServerArgs + self.assertTrue(hasattr(sa, "dsa_prefill_backend")) + self.assertTrue(hasattr(sa, "dsa_decode_backend")) + self.assertTrue(hasattr(sa, "enable_dsa_prefill_context_parallel")) + self.assertTrue(hasattr(sa, "dsa_prefill_cp_mode")) + + def test_serverargs_no_nsa_fields(self): + """The nsa_* attributes should no longer exist on ServerArgs.""" + sa = self.ServerArgs + self.assertFalse( + hasattr(sa, "nsa_prefill_backend"), + "nsa_prefill_backend should have been renamed", + ) + self.assertFalse( + hasattr(sa, "nsa_decode_backend"), + "nsa_decode_backend should have been renamed", + ) + self.assertFalse(hasattr(sa, "enable_nsa_prefill_context_parallel")) + self.assertFalse(hasattr(sa, "nsa_prefill_cp_mode")) + + +class TestCLICanonicalFlags(unittest.TestCase): + """--dsa-* canonical flags write to dsa_* attributes with no warning.""" + + def setUp(self): + from sglang.srt.server_args import ServerArgs + + self.parser = argparse.ArgumentParser() + ServerArgs.add_cli_args(self.parser) + + def _parse(self, extra_args): + return self.parser.parse_args(["--model", "dummy"] + extra_args) + + def test_dsa_prefill_backend_canonical(self): + args = self._parse(["--dsa-prefill-backend", "fa3"]) + self.assertEqual(args.dsa_prefill_backend, "fa3") + + def test_dsa_decode_backend_canonical(self): + args = self._parse(["--dsa-decode-backend", "tilelang"]) + self.assertEqual(args.dsa_decode_backend, "tilelang") + + def test_enable_dsa_prefill_cp_canonical(self): + args = self._parse(["--enable-dsa-prefill-context-parallel"]) + self.assertTrue(args.enable_dsa_prefill_context_parallel) + + def test_dsa_prefill_cp_mode_canonical(self): + args = self._parse(["--dsa-prefill-cp-mode", "in-seq-split"]) + self.assertEqual(args.dsa_prefill_cp_mode, "in-seq-split") + + def test_defaults_are_none_or_false(self): + args = self._parse([]) + self.assertIsNone(args.dsa_prefill_backend) + self.assertIsNone(args.dsa_decode_backend) + self.assertFalse(args.enable_dsa_prefill_context_parallel) + self.assertEqual(args.dsa_prefill_cp_mode, "round-robin-split") + + def test_attention_backend_dsa_key_in_choices(self): + args = self._parse(["--attention-backend", "dsa"]) + self.assertEqual(args.attention_backend, "dsa") + + +class TestCLIDeprecatedFlags(unittest.TestCase): + """--nsa-* deprecated flags write to dsa_* attributes and emit logger warning.""" + + def setUp(self): + import logging + + from sglang.srt.server_args import ServerArgs + + self.parser = argparse.ArgumentParser() + ServerArgs.add_cli_args(self.parser) + + # Capture log output to detect deprecation warnings + self.log_records = [] + handler = ( + logging.handlers_collector(self.log_records) + if hasattr(logging, "handlers_collector") + else None + ) + + def _parse(self, extra_args): + return self.parser.parse_args(["--model", "dummy"] + extra_args) + + def _parse_capture_warnings(self, extra_args): + """Parse and capture both warnings.warn and logger output.""" + import io + import logging + + log_stream = io.StringIO() + handler = logging.StreamHandler(log_stream) + handler.setLevel(logging.WARNING) + root = logging.getLogger() + root.addHandler(handler) + try: + args = self._parse(extra_args) + finally: + root.removeHandler(handler) + return args, log_stream.getvalue() + + def test_nsa_prefill_backend_deprecated_writes_to_dsa(self): + args, log_output = self._parse_capture_warnings( + ["--nsa-prefill-backend", "fa3"] + ) + self.assertEqual(args.dsa_prefill_backend, "fa3") + self.assertIn( + "deprecated", + log_output.lower(), + f"Expected deprecation warning in log; got: {log_output!r}", + ) + + def test_nsa_decode_backend_deprecated_writes_to_dsa(self): + args, log_output = self._parse_capture_warnings( + ["--nsa-decode-backend", "tilelang"] + ) + self.assertEqual(args.dsa_decode_backend, "tilelang") + self.assertIn("deprecated", log_output.lower()) + + def test_enable_nsa_prefill_cp_deprecated(self): + args, log_output = self._parse_capture_warnings( + ["--enable-nsa-prefill-context-parallel"] + ) + self.assertTrue(args.enable_dsa_prefill_context_parallel) + self.assertIn("deprecated", log_output.lower()) + + def test_nsa_prefill_cp_mode_deprecated(self): + args, log_output = self._parse_capture_warnings( + ["--nsa-prefill-cp-mode", "in-seq-split"] + ) + self.assertEqual(args.dsa_prefill_cp_mode, "in-seq-split") + self.assertIn("deprecated", log_output.lower()) + + def test_attention_backend_nsa_still_accepted(self): + """attention_backend='nsa' still parses without error (registry handles the deprecation).""" + args = self._parse(["--attention-backend", "nsa"]) + self.assertEqual(args.attention_backend, "nsa") + + +class TestAttentionRegistry(unittest.TestCase): + """Registry: 'dsa' key creates backend; 'nsa' key emits DeprecationWarning.""" + + def test_dsa_key_registered(self): + from sglang.srt.layers.attention.attention_registry import ATTENTION_BACKENDS + + self.assertIn("dsa", ATTENTION_BACKENDS) + + def test_nsa_key_still_registered(self): + from sglang.srt.layers.attention.attention_registry import ATTENTION_BACKENDS + + self.assertIn("nsa", ATTENTION_BACKENDS, "nsa must remain as deprecated alias") + + def test_nsa_key_emits_deprecation_warning(self): + """Calling the nsa factory should emit DeprecationWarning.""" + from sglang.srt.layers.attention.attention_registry import ATTENTION_BACKENDS + + nsa_factory = ATTENTION_BACKENDS.get("nsa") + self.assertIsNotNone(nsa_factory) + + class _FakeRunner: + server_args = type("S", (), {"attention_backend": "nsa"})() + + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + try: + nsa_factory(_FakeRunner()) + except Exception: + pass # import errors OK; we only care about DeprecationWarning + dep_warns = [x for x in w if issubclass(x.category, DeprecationWarning)] + self.assertTrue( + len(dep_warns) > 0, + "Expected DeprecationWarning when using 'nsa' registry key", + ) + self.assertIn("deprecated", str(dep_warns[0].message).lower()) + + +class TestEnvVarAliases(unittest.TestCase): + """SGLANG_DSA_* canonical; SGLANG_NSA_* fall back with DeprecationWarning.""" + + def setUp(self): + # Clean state for every test + for key in [ + "SGLANG_DSA_FUSE_TOPK", + "SGLANG_NSA_FUSE_TOPK", + "SGLANG_DSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD", + "SGLANG_NSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD", + "SGLANG_DSA_ENABLE_MTP_PRECOMPUTE_METADATA", + "SGLANG_NSA_ENABLE_MTP_PRECOMPUTE_METADATA", + ]: + os.environ.pop(key, None) + # Re-import to reset descriptor state + from importlib import reload + + import sglang.srt.environ as e + + reload(e) + from sglang.srt.environ import envs + + self.envs = envs + + def tearDown(self): + for key in [ + "SGLANG_DSA_FUSE_TOPK", + "SGLANG_NSA_FUSE_TOPK", + "SGLANG_DSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD", + "SGLANG_NSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD", + ]: + os.environ.pop(key, None) + + def test_dsa_fuse_topk_default(self): + self.assertTrue(self.envs.SGLANG_DSA_FUSE_TOPK.get()) + + def test_dsa_fuse_topk_canonical_set(self): + os.environ["SGLANG_DSA_FUSE_TOPK"] = "0" + self.assertFalse(self.envs.SGLANG_DSA_FUSE_TOPK.get()) + + def test_nsa_fuse_topk_deprecated_fallback(self): + """SGLANG_NSA_FUSE_TOPK=0 should be read by SGLANG_DSA_FUSE_TOPK with DeprecationWarning.""" + os.environ["SGLANG_NSA_FUSE_TOPK"] = "0" + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + val = self.envs.SGLANG_DSA_FUSE_TOPK.get() + self.assertFalse(val) + dep = [x for x in w if issubclass(x.category, DeprecationWarning)] + self.assertTrue( + len(dep) > 0, "Expected DeprecationWarning for SGLANG_NSA_FUSE_TOPK" + ) + + def test_dsa_threshold_default(self): + self.assertEqual( + self.envs.SGLANG_DSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD.get(), 2048 + ) + + def test_nsa_threshold_deprecated_fallback(self): + os.environ["SGLANG_NSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD"] = "1024" + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + val = self.envs.SGLANG_DSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD.get() + self.assertEqual(val, 1024) + dep = [x for x in w if issubclass(x.category, DeprecationWarning)] + self.assertTrue(len(dep) > 0) + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/test/registered/8-gpu-models/test_deepseek_v32.py b/test/registered/8-gpu-models/test_deepseek_v32.py index b79293197..c4b184bac 100644 --- a/test/registered/8-gpu-models/test_deepseek_v32.py +++ b/test/registered/8-gpu-models/test_deepseek_v32.py @@ -102,55 +102,55 @@ class TestDeepseekV32(unittest.TestCase): ) @unittest.skipIf(is_blackwell_system(), "Requires H200 system") - def test_deepseek_v32_nsa_backends(self): - """Test NSA attention backend variants (H200 only). + def test_deepseek_v32_dsa_backends(self): + """Test DSA attention backend variants (H200 only). - Tests three NSA backend configurations: + Tests three DSA backend configurations: - flashmla: flashmla_sparse prefill + flashmla_kv decode - fa3: FA3 prefill + FA3 decode - fp8kvcache: default backends with FP8 KV cache """ - NSA_FLASHMLA_ARGS = [ - "--attention-backend=nsa", - "--nsa-prefill-backend=flashmla_sparse", - "--nsa-decode-backend=flashmla_kv", + DSA_FLASHMLA_ARGS = [ + "--attention-backend=dsa", + "--dsa-prefill-backend=flashmla_sparse", + "--dsa-decode-backend=flashmla_kv", ] - NSA_FA3_ARGS = [ - "--attention-backend=nsa", - "--nsa-prefill-backend=fa3", - "--nsa-decode-backend=fa3", + DSA_FA3_ARGS = [ + "--attention-backend=dsa", + "--dsa-prefill-backend=fa3", + "--dsa-decode-backend=fa3", ] - NSA_FP8KV_ARGS = [ - "--attention-backend=nsa", + DSA_FP8KV_ARGS = [ + "--attention-backend=dsa", "--kv-cache-dtype=fp8_e4m3", ] - nsa_variants = [ + dsa_variants = [ # flashmla backend ModelLaunchSettings( DEEPSEEK_V32_MODEL_PATH, tp_size=8, - extra_args=BASE_ARGS + DP_ARGS + NSA_FLASHMLA_ARGS, + extra_args=BASE_ARGS + DP_ARGS + DSA_FLASHMLA_ARGS, ), # fa3 backend ModelLaunchSettings( DEEPSEEK_V32_MODEL_PATH, tp_size=8, - extra_args=BASE_ARGS + DP_ARGS + NSA_FA3_ARGS, + extra_args=BASE_ARGS + DP_ARGS + DSA_FA3_ARGS, ), # fp8 kv cache ModelLaunchSettings( DEEPSEEK_V32_MODEL_PATH, tp_size=8, - extra_args=BASE_ARGS + DP_ARGS + NSA_FP8KV_ARGS, + extra_args=BASE_ARGS + DP_ARGS + DSA_FP8KV_ARGS, ), ] run_combined_tests( - models=nsa_variants, - test_name="DeepSeek-V3.2 NSA Backends", + models=dsa_variants, + test_name="DeepSeek-V3.2 DSA Backends", accuracy_params=AccuracyTestParams( dataset="gsm8k", baseline_accuracy=GSM8K_BASELINE ), diff --git a/test/registered/8-gpu-models/test_dsa_models_hisparse.py b/test/registered/8-gpu-models/test_dsa_models_hisparse.py index 8bc6575ee..a6a973cf3 100644 --- a/test/registered/8-gpu-models/test_dsa_models_hisparse.py +++ b/test/registered/8-gpu-models/test_dsa_models_hisparse.py @@ -38,7 +38,7 @@ class TestGLM5DPHiSparse(CustomTestCase): "--disable-radix-cache", "--kv-cache-dtype", "bfloat16", - "--nsa-decode-backend", + "--dsa-decode-backend", "flashmla_sparse", "--enable-hisparse", "--hisparse-config", diff --git a/test/registered/8-gpu-models/test_return_indexer_topk.py b/test/registered/8-gpu-models/test_return_indexer_topk.py index 2b503b9d7..5a478928d 100644 --- a/test/registered/8-gpu-models/test_return_indexer_topk.py +++ b/test/registered/8-gpu-models/test_return_indexer_topk.py @@ -33,7 +33,7 @@ logger = logging.getLogger(__name__) class TestReturnIndexerTopk(CustomTestCase): - """Indexer-topk capture e2e test for DSv3.2 (NSA). + """Indexer-topk capture e2e test for DSv3.2 (DSA). Single server with `--enable-return-indexer-topk` and `index_topk_freq=2`. Validates the native `/generate` endpoint only — OpenAI-protocol surface diff --git a/test/registered/amd/accuracy/mi30x/test_glm51_eval_amd.py b/test/registered/amd/accuracy/mi30x/test_glm51_eval_amd.py index 93a4b345a..531d97c38 100644 --- a/test/registered/amd/accuracy/mi30x/test_glm51_eval_amd.py +++ b/test/registered/amd/accuracy/mi30x/test_glm51_eval_amd.py @@ -1,6 +1,6 @@ """AMD GLM-5.1 GSM8K Completion Evaluation Test (8-GPU) -Tests GLM-5.1-FP8 with NSA attention backend using few-shot +Tests GLM-5.1-FP8 with DSA attention backend using few-shot completion benchmark on MI325/MI300X. Registry: nightly-amd-accuracy-8-gpu-glm51 suite @@ -58,16 +58,16 @@ GLM51_MODELS = [ tp_size=8, accuracy_threshold=0.93, timeout=3600, - variant="nsa", + variant="dsa", other_args=[ "--trust-remote-code", "--reasoning-parser", "glm45", "--tool-call-parser", "glm47", - "--nsa-prefill-backend", + "--dsa-prefill-backend", "tilelang", - "--nsa-decode-backend", + "--dsa-decode-backend", "tilelang", "--chunked-prefill-size", "131072", diff --git a/test/registered/amd/accuracy/mi30x/test_glm5_eval_amd.py b/test/registered/amd/accuracy/mi30x/test_glm5_eval_amd.py index 93233439f..3295dce46 100644 --- a/test/registered/amd/accuracy/mi30x/test_glm5_eval_amd.py +++ b/test/registered/amd/accuracy/mi30x/test_glm5_eval_amd.py @@ -1,6 +1,6 @@ """AMD GLM-5 GSM8K Completion Evaluation Test (8-GPU) -Tests GLM-5 with NSA attention backend using few-shot completion +Tests GLM-5 with DSA attention backend using few-shot completion benchmark on MI325/MI300X. Registry: nightly-amd-accuracy-8-gpu-glm5 suite @@ -55,24 +55,24 @@ class ModelConfig: return self.model_path -# GLM-5 models for MI325/MI300X - NSA attention backend +# GLM-5 models for MI325/MI300X - DSA attention backend GLM5_MODELS = [ - # GLM-5 with NSA attention (TP=8) + # GLM-5 with DSA attention (TP=8) ModelConfig( model_path="zai-org/GLM-5-FP8", tp_size=8, accuracy_threshold=0.93, timeout=3600, - variant="nsa", + variant="dsa", other_args=[ "--trust-remote-code", "--reasoning-parser", "glm45", "--tool-call-parser", "glm47", - "--nsa-prefill-backend", + "--dsa-prefill-backend", "tilelang", - "--nsa-decode-backend", + "--dsa-decode-backend", "tilelang", "--chunked-prefill-size", "131072", diff --git a/test/registered/amd/accuracy/mi35x/test_deepseek_v32_dp_eval_mi35x.py b/test/registered/amd/accuracy/mi35x/test_deepseek_v32_dp_eval_mi35x.py index e196a01c0..e4a1e993b 100644 --- a/test/registered/amd/accuracy/mi35x/test_deepseek_v32_dp_eval_mi35x.py +++ b/test/registered/amd/accuracy/mi35x/test_deepseek_v32_dp_eval_mi35x.py @@ -60,9 +60,9 @@ class TestDeepseekV32DP(CustomTestCase): "--enable-dp-attention", "--model-loader-extra-config", '{"enable_multithread_load": true}', - "--nsa-prefill-backend", + "--dsa-prefill-backend", "tilelang", - "--nsa-decode-backend", + "--dsa-decode-backend", "tilelang", ] cls.process = popen_launch_server( diff --git a/test/registered/amd/accuracy/mi35x/test_deepseek_v32_eval_mi35x.py b/test/registered/amd/accuracy/mi35x/test_deepseek_v32_eval_mi35x.py index 0b5a4a71e..1d3e08cc0 100644 --- a/test/registered/amd/accuracy/mi35x/test_deepseek_v32_eval_mi35x.py +++ b/test/registered/amd/accuracy/mi35x/test_deepseek_v32_eval_mi35x.py @@ -78,9 +78,9 @@ MI35X_DEEPSEEK_V32_MODELS = [ variant="basic", other_args=[ "--trust-remote-code", - "--nsa-prefill-backend", + "--dsa-prefill-backend", "tilelang", - "--nsa-decode-backend", + "--dsa-decode-backend", "tilelang", "--mem-fraction-static", "0.85", diff --git a/test/registered/amd/accuracy/mi35x/test_deepseek_v32_mtp_eval_mi35x.py b/test/registered/amd/accuracy/mi35x/test_deepseek_v32_mtp_eval_mi35x.py index dad040a30..cf7aea5a0 100644 --- a/test/registered/amd/accuracy/mi35x/test_deepseek_v32_mtp_eval_mi35x.py +++ b/test/registered/amd/accuracy/mi35x/test_deepseek_v32_mtp_eval_mi35x.py @@ -59,9 +59,9 @@ class TestDeepseekV32TPMTP(CustomTestCase): "--trust-remote-code", "--tp", "8", - "--nsa-prefill-backend", + "--dsa-prefill-backend", "tilelang", - "--nsa-decode-backend", + "--dsa-decode-backend", "tilelang", "--speculative-algorithm", "EAGLE", diff --git a/test/registered/amd/accuracy/mi35x/test_glm51_eval_mi35x.py b/test/registered/amd/accuracy/mi35x/test_glm51_eval_mi35x.py index 3267a0f34..907df4f49 100644 --- a/test/registered/amd/accuracy/mi35x/test_glm51_eval_mi35x.py +++ b/test/registered/amd/accuracy/mi35x/test_glm51_eval_mi35x.py @@ -1,6 +1,6 @@ """MI35x GLM-5.1 GSM8K Completion Evaluation Test (8-GPU) -Tests GLM-5.1-FP8 with NSA attention backend using few-shot +Tests GLM-5.1-FP8 with DSA attention backend using few-shot completion benchmark on MI35x. Registry: nightly-amd-8-gpu-mi35x-glm51 suite @@ -62,16 +62,16 @@ MI35X_GLM51_MODELS = [ tp_size=8, accuracy_threshold=0.93, timeout=5400, - variant="nsa", + variant="dsa", other_args=[ "--trust-remote-code", "--reasoning-parser", "glm45", "--tool-call-parser", "glm47", - "--nsa-prefill-backend", + "--dsa-prefill-backend", "tilelang", - "--nsa-decode-backend", + "--dsa-decode-backend", "tilelang", "--chunked-prefill-size", "131072", diff --git a/test/registered/amd/accuracy/mi35x/test_glm5_eval_mi35x.py b/test/registered/amd/accuracy/mi35x/test_glm5_eval_mi35x.py index 02af23a57..22ca79d00 100644 --- a/test/registered/amd/accuracy/mi35x/test_glm5_eval_mi35x.py +++ b/test/registered/amd/accuracy/mi35x/test_glm5_eval_mi35x.py @@ -1,6 +1,6 @@ """MI35x GLM-5 GSM8K Completion Evaluation Test (8-GPU) -Tests GLM-5 with NSA attention backend using few-shot completion +Tests GLM-5 with DSA attention backend using few-shot completion benchmark on MI35x. Registry: nightly-amd-8-gpu-mi35x-glm5 suite @@ -60,24 +60,24 @@ class ModelConfig: return self.model_path -# GLM-5 models for MI35x - NSA attention backend +# GLM-5 models for MI35x - DSA attention backend MI35X_GLM5_MODELS = [ - # GLM-5 with NSA attention (TP=8) + # GLM-5 with DSA attention (TP=8) ModelConfig( model_path="zai-org/GLM-5-FP8", tp_size=8, accuracy_threshold=0.93, timeout=5400, - variant="nsa", + variant="dsa", other_args=[ "--trust-remote-code", "--reasoning-parser", "glm45", "--tool-call-parser", "glm47", - "--nsa-prefill-backend", + "--dsa-prefill-backend", "tilelang", - "--nsa-decode-backend", + "--dsa-decode-backend", "tilelang", "--chunked-prefill-size", "131072", diff --git a/test/registered/amd/perf/mi30x/test_glm51_perf_amd.py b/test/registered/amd/perf/mi30x/test_glm51_perf_amd.py index 5b2347c9b..4a2d43004 100644 --- a/test/registered/amd/perf/mi30x/test_glm51_perf_amd.py +++ b/test/registered/amd/perf/mi30x/test_glm51_perf_amd.py @@ -1,6 +1,6 @@ """Nightly performance benchmark for GLM-5.1 on MI30x. -Tests GLM-5.1-FP8 with NSA attention backend using bench_one_batch +Tests GLM-5.1-FP8 with DSA attention backend using bench_one_batch on 8 GPUs with TP=8, FP8 KV cache. Model path can be configured via GLM51_MODEL_PATH environment variable. @@ -53,7 +53,7 @@ PROFILE_DIR = "performance_profiles_glm51" class TestNightlyGLM51Performance(unittest.TestCase): """Nightly performance benchmark for GLM-5.1 on MI30x. - Tests GLM-5.1-FP8 with NSA attention backend on TP=8. + Tests GLM-5.1-FP8 with DSA attention backend on TP=8. """ @classmethod @@ -74,9 +74,9 @@ class TestNightlyGLM51Performance(unittest.TestCase): "glm47", "--tp", "8", - "--nsa-prefill-backend", + "--dsa-prefill-backend", "tilelang", - "--nsa-decode-backend", + "--dsa-decode-backend", "tilelang", "--kv-cache-dtype", "fp8_e4m3", diff --git a/test/registered/amd/perf/mi30x/test_glm5_perf_amd.py b/test/registered/amd/perf/mi30x/test_glm5_perf_amd.py index 1cdd8f660..216d8ac21 100644 --- a/test/registered/amd/perf/mi30x/test_glm5_perf_amd.py +++ b/test/registered/amd/perf/mi30x/test_glm5_perf_amd.py @@ -1,6 +1,6 @@ """Nightly performance benchmark for GLM-5 on MI30x. -Tests GLM-5 with NSA attention backend using bench_one_batch on 8 GPUs. +Tests GLM-5 with DSA attention backend using bench_one_batch on 8 GPUs. Model paths can be configured via environment variables: - GLM5_MODEL_PATH: Path to GLM-5 model (default: zai-org/GLM-5-FP8) @@ -54,7 +54,7 @@ PROFILE_DIR = "performance_profiles_glm5" class TestNightlyGLM5Performance(unittest.TestCase): """Nightly performance benchmark for GLM-5. - Tests GLM-5 with NSA attention backend on TP=8. + Tests GLM-5 with DSA attention backend on TP=8. """ @classmethod @@ -75,9 +75,9 @@ class TestNightlyGLM5Performance(unittest.TestCase): "glm47", "--tp", "8", - "--nsa-prefill-backend", + "--dsa-prefill-backend", "tilelang", - "--nsa-decode-backend", + "--dsa-decode-backend", "tilelang", "--kv-cache-dtype", "fp8_e4m3", diff --git a/test/registered/amd/perf/mi35x/test_deepseek_v32_basic_perf_mi35x.py b/test/registered/amd/perf/mi35x/test_deepseek_v32_basic_perf_mi35x.py index 740500e9f..d537e7928 100644 --- a/test/registered/amd/perf/mi35x/test_deepseek_v32_basic_perf_mi35x.py +++ b/test/registered/amd/perf/mi35x/test_deepseek_v32_basic_perf_mi35x.py @@ -78,16 +78,16 @@ class TestNightlyDeepseekV32BasicPerformance(unittest.TestCase): cls.output_lens = tuple(_parse_int_list_env("NIGHTLY_OUTPUT_LENS", "512")) # Basic variant configuration for DeepSeek-V3.2 - # MI35x uses tilelang NSA backends + # MI35x uses tilelang DSA backends cls.variant_config = { "name": "basic", "other_args": [ "--trust-remote-code", "--tp", "8", - "--nsa-prefill-backend", + "--dsa-prefill-backend", "tilelang", - "--nsa-decode-backend", + "--dsa-decode-backend", "tilelang", "--mem-fraction-static", "0.85", diff --git a/test/registered/amd/perf/mi35x/test_deepseek_v32_mtp_perf_mi35x.py b/test/registered/amd/perf/mi35x/test_deepseek_v32_mtp_perf_mi35x.py index 6a0445126..ddfdea753 100644 --- a/test/registered/amd/perf/mi35x/test_deepseek_v32_mtp_perf_mi35x.py +++ b/test/registered/amd/perf/mi35x/test_deepseek_v32_mtp_perf_mi35x.py @@ -132,16 +132,16 @@ class TestNightlyDeepseekV32MTPPerformance(unittest.TestCase): cls.output_lens = tuple(_parse_int_list_env("NIGHTLY_OUTPUT_LENS", "512")) # MTP variant configuration for DeepSeek-V3.2 - # MI35x uses tilelang NSA backends + EAGLE speculative decoding + # MI35x uses tilelang DSA backends + EAGLE speculative decoding cls.variant_config = { "name": "mtp", "other_args": [ "--trust-remote-code", "--tp", "8", - "--nsa-prefill-backend", + "--dsa-prefill-backend", "tilelang", - "--nsa-decode-backend", + "--dsa-decode-backend", "tilelang", "--speculative-algorithm", "EAGLE", diff --git a/test/registered/amd/perf/mi35x/test_glm51_perf_mi35x.py b/test/registered/amd/perf/mi35x/test_glm51_perf_mi35x.py index e4bb32f07..474d6a732 100644 --- a/test/registered/amd/perf/mi35x/test_glm51_perf_mi35x.py +++ b/test/registered/amd/perf/mi35x/test_glm51_perf_mi35x.py @@ -1,6 +1,6 @@ """MI35x Nightly performance benchmark for GLM-5.1. -Tests GLM-5.1-FP8 with NSA attention backend using bench_one_batch +Tests GLM-5.1-FP8 with DSA attention backend using bench_one_batch on 8 GPUs with TP=8, FP8 KV cache. Registry: nightly-perf-8-gpu-mi35x-glm51 suite @@ -55,7 +55,7 @@ PROFILE_DIR = "performance_profiles_glm51_mi35x" class TestGLM51PerfMI35x(unittest.TestCase): """Nightly performance benchmark for GLM-5.1 on MI35x. - Tests GLM-5.1-FP8 with NSA attention backend on TP=8. + Tests GLM-5.1-FP8 with DSA attention backend on TP=8. """ @classmethod @@ -76,9 +76,9 @@ class TestGLM51PerfMI35x(unittest.TestCase): "glm47", "--tp", "8", - "--nsa-prefill-backend", + "--dsa-prefill-backend", "tilelang", - "--nsa-decode-backend", + "--dsa-decode-backend", "tilelang", "--kv-cache-dtype", "fp8_e4m3", diff --git a/test/registered/amd/perf/mi35x/test_glm5_perf_mi35x.py b/test/registered/amd/perf/mi35x/test_glm5_perf_mi35x.py index a742cbc1d..d06e7d6a9 100644 --- a/test/registered/amd/perf/mi35x/test_glm5_perf_mi35x.py +++ b/test/registered/amd/perf/mi35x/test_glm5_perf_mi35x.py @@ -1,6 +1,6 @@ """MI35x Nightly performance benchmark for GLM-5. -Tests GLM-5 with NSA attention backend using bench_one_batch on 8 GPUs. +Tests GLM-5 with DSA attention backend using bench_one_batch on 8 GPUs. Registry: nightly-perf-8-gpu-mi35x-glm5 suite """ @@ -54,7 +54,7 @@ PROFILE_DIR = "performance_profiles_glm5_mi35x" class TestGLM5PerfMI35x(unittest.TestCase): """Nightly performance benchmark for GLM-5 on MI35x. - Tests GLM-5 with NSA attention backend on TP=8. + Tests GLM-5 with DSA attention backend on TP=8. """ @classmethod @@ -75,9 +75,9 @@ class TestGLM5PerfMI35x(unittest.TestCase): "glm47", "--tp", "8", - "--nsa-prefill-backend", + "--dsa-prefill-backend", "tilelang", - "--nsa-decode-backend", + "--dsa-decode-backend", "tilelang", "--kv-cache-dtype", "fp8_e4m3", diff --git a/test/registered/amd/test_deepseek_v32_basic.py b/test/registered/amd/test_deepseek_v32_basic.py index 27d5b0fc1..dbdd4620d 100644 --- a/test/registered/amd/test_deepseek_v32_basic.py +++ b/test/registered/amd/test_deepseek_v32_basic.py @@ -38,9 +38,9 @@ class TestDeepseekV32DP(CustomTestCase): ] if is_in_amd_ci(): other_args += [ - "--nsa-prefill-backend", + "--dsa-prefill-backend", "tilelang", - "--nsa-decode-backend", + "--dsa-decode-backend", "tilelang", ] @@ -106,9 +106,9 @@ class TestDeepseekV32TP(CustomTestCase): ] if is_in_amd_ci(): other_args += [ - "--nsa-prefill-backend", + "--dsa-prefill-backend", "tilelang", - "--nsa-decode-backend", + "--dsa-decode-backend", "tilelang", ] diff --git a/test/registered/amd/test_deepseek_v32_mtp.py b/test/registered/amd/test_deepseek_v32_mtp.py index 69587bdf6..50b8da214 100644 --- a/test/registered/amd/test_deepseek_v32_mtp.py +++ b/test/registered/amd/test_deepseek_v32_mtp.py @@ -54,9 +54,9 @@ class TestDeepseekV32DPMTP(CustomTestCase): ] if is_in_amd_ci(): other_args += [ - "--nsa-prefill-backend", + "--dsa-prefill-backend", "tilelang", - "--nsa-decode-backend", + "--dsa-decode-backend", "tilelang", ] cls.process = popen_launch_server( @@ -146,9 +146,9 @@ class TestDeepseekV32TPMTP(CustomTestCase): ] if is_in_amd_ci(): other_args += [ - "--nsa-prefill-backend", + "--dsa-prefill-backend", "tilelang", - "--nsa-decode-backend", + "--dsa-decode-backend", "tilelang", ] cls.process = popen_launch_server( diff --git a/test/registered/cp/test_deepseek_v32_cp_single_node.py b/test/registered/cp/test_deepseek_v32_cp_single_node.py index fb5abcf1e..5b582e827 100644 --- a/test/registered/cp/test_deepseek_v32_cp_single_node.py +++ b/test/registered/cp/test_deepseek_v32_cp_single_node.py @@ -31,8 +31,8 @@ class TestDeepseekV32CPInSeqSplit(CustomTestCase): "2", "--attn-cp-size", "4", - "--enable-nsa-prefill-context-parallel", - "--nsa-prefill-cp-mode", + "--enable-dsa-prefill-context-parallel", + "--dsa-prefill-cp-mode", "in-seq-split", "--speculative-algorithm", "EAGLE", @@ -97,8 +97,8 @@ class TestDeepseekV32CPRoundRobinSplit(CustomTestCase): "8", "--attn-cp-size", "8", - "--enable-nsa-prefill-context-parallel", - "--nsa-prefill-cp-mode", + "--enable-dsa-prefill-context-parallel", + "--dsa-prefill-cp-mode", "round-robin-split", "--speculative-algorithm", "EAGLE", diff --git a/test/registered/dsv4/test_deepseek_v4_flash_fp4_b200.py b/test/registered/dsv4/test_deepseek_v4_flash_fp4_b200.py index c3500c1bd..60d7bab43 100644 --- a/test/registered/dsv4/test_deepseek_v4_flash_fp4_b200.py +++ b/test/registered/dsv4/test_deepseek_v4_flash_fp4_b200.py @@ -158,8 +158,8 @@ class TestDSV4FlashFP4B200Balanced_CP(ServerSanityMixin, CustomTestCase): "1", "--speculative-num-draft-tokens", "2", - "--enable-nsa-prefill-context-parallel", - "--nsa-prefill-cp-mode", + "--enable-dsa-prefill-context-parallel", + "--dsa-prefill-cp-mode", "round-robin-split", "--deepep-config", DEEPEP_CONFIG, diff --git a/test/registered/kernels/test_nsa_indexer.py b/test/registered/kernels/test_dsa_indexer.py similarity index 95% rename from test/registered/kernels/test_nsa_indexer.py rename to test/registered/kernels/test_dsa_indexer.py index e1d877012..09021180f 100644 --- a/test/registered/kernels/test_nsa_indexer.py +++ b/test/registered/kernels/test_dsa_indexer.py @@ -11,15 +11,15 @@ from sglang.test.ci.ci_register import register_cuda_ci _dp_attn.get_attention_tp_size = lambda: 1 # TP size = 1 for unit test from sglang.srt.configs.model_config import AttentionArch -from sglang.srt.layers.attention.nsa.nsa_indexer import ( +from sglang.srt.layers.attention.dsa.dsa_indexer import ( BaseIndexerMetadata, Indexer, rotate_activation, ) -from sglang.srt.layers.attention.nsa_backend import NativeSparseAttnBackend +from sglang.srt.layers.attention.dsa_backend import DeepseekSparseAttnBackend from sglang.srt.layers.layernorm import LayerNorm from sglang.srt.layers.linear import LinearBase -from sglang.srt.mem_cache.memory_pool import NSATokenToKVPool +from sglang.srt.mem_cache.memory_pool import DSATokenToKVPool from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler from sglang.test.test_utils import CustomTestCase @@ -136,7 +136,7 @@ class MockIndexerMetadata(BaseIndexerMetadata): """Return: seq lens for each batch.""" return torch.tensor(self.seq_lens, dtype=torch.int32, device=self.device) - def get_nsa_extend_len_cpu(self) -> List[int]: + def get_dsa_extend_len_cpu(self) -> List[int]: """ Return: extend seq lens for each batch. """ @@ -179,7 +179,7 @@ class MockModelRunner: max_context_len = self.config["context_len"] max_batch_size = self.config["max_bs"] - # Create mock hf_config for NSA - instantiate it as an object, not a type + # Create mock hf_config for DSA - instantiate it as an object, not a type hf_config = type( "HfConfig", (), @@ -224,9 +224,9 @@ class MockModelRunner: }, )() - # Create NSATokenToKVPool + # Create DSATokenToKVPool max_total_num_tokens = max_batch_size * max_context_len - self.token_to_kv_pool = NSATokenToKVPool( + self.token_to_kv_pool = DSATokenToKVPool( size=max_total_num_tokens, page_size=self.config["page_size"], dtype=self.config["kv_cache_dtype"], @@ -239,7 +239,7 @@ class MockModelRunner: kv_cache_dim=self.config["kv_lora_rank"] + self.config["qk_rope_head_dim"], ) - # Required by backend with NSA-specific attributes + # Required by backend with DSA-specific attributes self.server_args = type( "ServerArgs", (), @@ -248,21 +248,21 @@ class MockModelRunner: "speculative_eagle_topk": None, "speculative_num_draft_tokens": 0, "enable_deterministic_inference": False, - "nsa_prefill_backend": "flashmla_sparse", - "nsa_decode_backend": "fa3", + "dsa_prefill_backend": "flashmla_sparse", + "dsa_decode_backend": "fa3", }, )() @unittest.skipIf(not torch.cuda.is_available(), "Test requires CUDA") -class TestNSAIndexer(CustomTestCase): +class TestDSAIndexer(CustomTestCase): @classmethod def setUpClass(cls): """Set up global server args for testing.""" server_args = ServerArgs(model_path="dummy") server_args.enable_dp_attention = False - server_args.nsa_prefill_backend = "flashmla_sparse" - server_args.nsa_decode_backend = "flashmla_sparse" + server_args.dsa_prefill_backend = "flashmla_sparse" + server_args.dsa_decode_backend = "flashmla_sparse" set_global_server_args_for_scheduler(server_args) # Check GPU capability for FP8 @@ -289,7 +289,7 @@ class TestNSAIndexer(CustomTestCase): if config_override: config.update(config_override) self.model_runner = MockModelRunner(config) - self.backend = NativeSparseAttnBackend(self.model_runner) + self.backend = DeepseekSparseAttnBackend(self.model_runner) def _create_indexer(self, **kwargs): """Create an Indexer instance with default parameters.""" @@ -417,7 +417,7 @@ class TestNSAIndexer(CustomTestCase): "Output should have padding or exact topk size", ) - @patch("sglang.srt.layers.attention.nsa.nsa_indexer.deep_gemm") + @patch("sglang.srt.layers.attention.dsa.dsa_indexer.deep_gemm") def test_indexer_basic_creation(self, mock_deep_gemm): """Test basic indexer creation and initialization.""" mock_deep_gemm.get_num_sms.return_value = 132 @@ -431,8 +431,8 @@ class TestNSAIndexer(CustomTestCase): self.assertEqual(indexer.index_topk, self.config["index_topk"]) self.assertEqual(indexer.layer_id, self.config["layer_id"]) - @patch("sglang.srt.layers.attention.nsa.nsa_indexer.deep_gemm") - @patch("sglang.srt.layers.attention.nsa.triton_kernel.act_quant") + @patch("sglang.srt.layers.attention.dsa.dsa_indexer.deep_gemm") + @patch("sglang.srt.layers.attention.dsa.triton_kernel.act_quant") def test_forward_extend_mode(self, mock_act_quant, mock_deep_gemm): """Test indexer forward pass in extend mode.""" if not self.supports_fp8: @@ -513,8 +513,8 @@ class TestNSAIndexer(CustomTestCase): topk_indices, self.batch_size, self.seq_len, self.config["index_topk"] ) - @patch("sglang.srt.layers.attention.nsa.nsa_indexer.deep_gemm") - @patch("sglang.srt.layers.attention.nsa.triton_kernel.act_quant") + @patch("sglang.srt.layers.attention.dsa.dsa_indexer.deep_gemm") + @patch("sglang.srt.layers.attention.dsa.triton_kernel.act_quant") def test_forward_decode_mode(self, mock_act_quant, mock_deep_gemm): """Test indexer forward pass in decode mode.""" if not self.supports_fp8: @@ -627,7 +627,7 @@ class TestNSAIndexer(CustomTestCase): self.assertEqual(topk_indices.shape, (batch_size, topk)) # TODO: enable this test after indexer accuracy aligned - # @patch("sglang.srt.layers.attention.nsa.nsa_indexer.deep_gemm") + # @patch("sglang.srt.layers.attention.dsa.dsa_indexer.deep_gemm") # def test_indexer_with_different_topk(self, mock_deep_gemm): # """Test indexer with different topk values.""" # mock_deep_gemm.get_num_sms.return_value = 132 @@ -637,7 +637,7 @@ class TestNSAIndexer(CustomTestCase): # indexer = self._create_indexer(index_topk=topk) # self.assertEqual(indexer.index_topk, topk) - @patch("sglang.srt.layers.attention.nsa.nsa_indexer.deep_gemm") + @patch("sglang.srt.layers.attention.dsa.dsa_indexer.deep_gemm") def test_indexer_with_fused_wk(self, mock_deep_gemm): """Test indexer creation with fused wk and weights projection.""" mock_deep_gemm.get_num_sms.return_value = 132 @@ -647,7 +647,7 @@ class TestNSAIndexer(CustomTestCase): indexer = self._create_indexer() self.assertIsNotNone(indexer) - @patch("sglang.srt.layers.attention.nsa.nsa_indexer.deep_gemm") + @patch("sglang.srt.layers.attention.dsa.dsa_indexer.deep_gemm") def test_indexer_with_alt_stream(self, mock_deep_gemm): """Test indexer creation with alternative CUDA stream.""" mock_deep_gemm.get_num_sms.return_value = 132 diff --git a/test/registered/quant/test_deepseek_v32_fp4_mtp_4gpu.py b/test/registered/quant/test_deepseek_v32_fp4_mtp_4gpu.py index 2ef4633e7..77d11b34e 100644 --- a/test/registered/quant/test_deepseek_v32_fp4_mtp_4gpu.py +++ b/test/registered/quant/test_deepseek_v32_fp4_mtp_4gpu.py @@ -37,7 +37,7 @@ class TestDeepseekV32FP4DPSpecV2(CustomTestCase): "4", "--enable-dp-attention", "--attention-backend", - "nsa", + "dsa", "--moe-runner-backend", "flashinfer_trtllm", "--quantization", @@ -127,7 +127,7 @@ class TestDeepseekV32FP4TPSpecV2(CustomTestCase): "--tp", "4", "--attention-backend", - "nsa", + "dsa", "--moe-runner-backend", "flashinfer_trtllm", "--quantization", diff --git a/test/registered/unit/managers/test_hisparse_unit.py b/test/registered/unit/managers/test_hisparse_unit.py index 56fc32a16..22996c54e 100644 --- a/test/registered/unit/managers/test_hisparse_unit.py +++ b/test/registered/unit/managers/test_hisparse_unit.py @@ -94,11 +94,11 @@ class TestHiSparseUnit(unittest.TestCase): global_page_size = 1 if is_hip() else PAGE_SIZE from sglang.srt.mem_cache.hisparse_memory_pool import ( - HiSparseNSATokenToKVPool, + HiSparseDSATokenToKVPool, HiSparseTokenToKVPoolAllocator, ) - cls.device_pool = HiSparseNSATokenToKVPool( + cls.device_pool = HiSparseDSATokenToKVPool( size=SIZE, page_size=global_page_size, kv_lora_rank=KV_LORA_RANK, diff --git a/test/registered/unit/mem_cache/test_nsa_pool_host_unit.py b/test/registered/unit/mem_cache/test_dsa_pool_host_unit.py similarity index 93% rename from test/registered/unit/mem_cache/test_nsa_pool_host_unit.py rename to test/registered/unit/mem_cache/test_dsa_pool_host_unit.py index 17fe9d691..939e40f5e 100644 --- a/test/registered/unit/mem_cache/test_nsa_pool_host_unit.py +++ b/test/registered/unit/mem_cache/test_dsa_pool_host_unit.py @@ -2,11 +2,11 @@ import unittest import torch -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 ( ALLOC_MEMORY_FUNCS, + DSAIndexerPoolHost, MLATokenToKVPoolHost, - NSAIndexerPoolHost, alloc_with_pin_memory, ) from sglang.srt.utils import is_cuda, is_hip, is_npu, is_xpu @@ -15,12 +15,12 @@ from sglang.test.ci.ci_register import register_cuda_ci register_cuda_ci(est_time=9, stage="base-b", runner_config="1-gpu-small") -class TestNSAHiCacheTransfer(unittest.TestCase): +class TestDSAHiCacheTransfer(unittest.TestCase): def setUp(self): if not torch.cuda.is_available(): - self.skipTest("CUDA is required for NSA host transfer tests.") + self.skipTest("CUDA is required for DSA host transfer tests.") if is_npu() or is_xpu(): - self.skipTest("NSA host transfer tests only support CUDA/ROCm.") + self.skipTest("DSA host transfer tests only support CUDA/ROCm.") if not (is_cuda() or is_hip()): self.skipTest("CUDA/ROCm not available.") @@ -42,7 +42,7 @@ class TestNSAHiCacheTransfer(unittest.TestCase): layer_num = 2 size = page_size * 4 - device_pool = NSATokenToKVPool( + device_pool = DSATokenToKVPool( size=size, page_size=page_size, kv_lora_rank=128, @@ -70,7 +70,7 @@ class TestNSAHiCacheTransfer(unittest.TestCase): allocator_type="default", override_kv_cache_dim=device_pool.kv_cache_dim, ) - indexer_host = NSAIndexerPoolHost( + indexer_host = DSAIndexerPoolHost( device_pool=device_pool, anchor_host=mla_host, layout="layer_first", diff --git a/test/registered/unit/platforms/test_platform_interface.py b/test/registered/unit/platforms/test_platform_interface.py index a4b4cd9f7..8b8ef1454 100644 --- a/test/registered/unit/platforms/test_platform_interface.py +++ b/test/registered/unit/platforms/test_platform_interface.py @@ -71,7 +71,7 @@ class _StubPlatform(SRTPlatform): def get_mla_kv_pool_cls(self): return object - def get_nsa_kv_pool_cls(self): + def get_dsa_kv_pool_cls(self): return object def get_paged_allocator_cls(self):