[Refactor] Rename NSA → DSA: user-facing aliases, file/class/import rename (#25821)
Co-authored-by: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Cheng Wan
co-authored by
Claude Sonnet 4.6
Claude Sonnet 4.6
parent
da6d549ab2
commit
8131641bc6
@@ -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*`<br>RoPE followed by reshape / cache DtoD | `python/sglang/srt/layers/attention/utils.py::fused_qk_rope_reshape_and_cache` | One Triton kernel applies RoPE to Q / K, reshapes cache layout, and writes K / V directly to paged cache | Treat separate RoPE + reshape + cache-write ladders as an existing attention-prep fusion family. |
|
||||
| Fused RoPE + KV cache store | `fused_set_kv_buffer`<br>RoPE followed by KV-store, DtoD, or cache-write kernels | `python/sglang/jit_kernel/rope.py`<br>`python/sglang/srt/models/utils.py::enable_fused_set_kv_buffer` | Shared entrypoints can route to fused RoPE + KV-store or model-side `fused_set_kv_buffer` fast paths | Compare against the fused cache-store path before proposing a new KV rewrite. |
|
||||
| Fused decode metadata setup | `normal_decode_set_metadata`<br>`cache_seqlens_int32`<br>`cu_seqlens_k`<br>`page_table`<br>`swa_page_table` | `python/sglang/srt/layers/attention/flashattention_backend.py::normal_decode_set_metadata` | Triton decode path fuses seq-len cast/add, prefix-sum, req-to-token gather, page-table divide, and optional SWA metadata build into 1-2 kernels | If decode exposes multiple tiny metadata kernels before attention, first compare against this existing fused metadata-prep path. |
|
||||
| NSA fused metadata copy for graph replay | `fused_metadata_copy`<br>`fused_metadata_copy_multi`<br>`fused_nsa_cache_seqlens`<br>`fused_flashmla_metadata` | `python/sglang/jit_kernel/fused_metadata_copy.py` | CUDA graph replay path fuses multiple metadata copies into one kernel or one multi-destination kernel | Treat bursts of tiny metadata-copy kernels around NSA replay as a missed existing replay fusion. |
|
||||
| DSA fused metadata copy for graph replay | `fused_metadata_copy`<br>`fused_metadata_copy_multi`<br>`fused_dsa_cache_seqlens`<br>`fused_flashmla_metadata` | `python/sglang/jit_kernel/fused_metadata_copy.py` | CUDA graph replay path fuses multiple metadata copies into one kernel or one multi-destination kernel | Treat bursts of tiny metadata-copy kernels around DSA replay as a missed existing replay fusion. |
|
||||
| DeepSeek MLA fused projection + norm + RoPE | `qkv_proj_with_rope_fused_weight`<br>`fused_qkv_a_proj_with_mqa`<br>`forward_absorb_fused_mla_rope*` | `python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_fused_rope_cpu.py`<br>`python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_fused_rope_rocm.py`<br>`python/sglang/srt/models/deepseek_v2.py` | CPU / ROCm paths fuse DeepSeek MLA projection packing with q / k norm, RoPE, and cache-oriented MLA prep | For DeepSeek MLA, split proj / norm / rope prep is usually an existing backend-specific fuse that did not fire. |
|
||||
| Fused QK RoPE concat + MLA cache write | `fused_qk_rope_cat_and_cache_mla`<br>`set_mla_kv_buffer` | `python/sglang/srt/layers/rocm_linear_utils.py`<br>`python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py` | ROCm MLA path can fuse Q / K RoPE packing, concat, and MLA cache write in one backend-specific op | On DeepSeek / MLA traces, separate RoPE-cat-cache steps are not automatically novel. |
|
||||
| Qwen3 decode fused QK norm + 3D mRoPE + KV cache write | `fused_qk_norm_mrope_3d_cache_pts_quant_shuffle`<br>`mrope`<br>decode cache write | `python/sglang/srt/models/qwen3.py` | ROCm / AITER decode path fuses QK norm, 3D mRoPE, and paged KV cache write | On Qwen3-style decode, separate norm + mRoPE + cache-store kernels are not a novel opportunity. |
|
||||
| NPU fused split-QKV + RMSNorm + RoPE | `split_qkv_rmsnorm_rope` | `python/sglang/srt/models/llama.py`<br>`python/sglang/srt/models/qwen3.py`<br>`python/sglang/srt/models/qwen3_moe.py`<br>`python/sglang/srt/models/glm4_moe.py` | Ascend path fuses QKV split, Q / K RMSNorm, and RoPE in one op | On NPU traces, separate split / norm / rope kernels usually mean the fused path is unavailable or bypassed. |
|
||||
| Fused FP8 quantize + paged KV cache write | `trtllm_fp8_kv_kernel`<br>`fp8 kv cache write`<br>`paged KV cache write` | `python/sglang/srt/layers/attention/triton_ops/trtllm_fp8_kv_kernel.py` | TRTLLM MHA path fuses FP8 quantization, scale computation, and paged K / V cache write | If FP8 KV cache traces show standalone quant plus write kernels, first compare against this existing Triton fuse. |
|
||||
| Fused MLA KV cache write + FP8 quant | `set_mla_kv_buffer_fp8_quant*`<br>`set_mla_kv_buffer_triton_fp8_quant` | `python/sglang/srt/mem_cache/utils.py`<br>`python/sglang/srt/mem_cache/memory_pool.py` | MLA / NSA KV pool path can quantize K and write directly into KV storage without a separate concat-and-quant chain | Treat standalone quant + KV-buffer write on MLA paths as missing existing fusion first. |
|
||||
| Fused MLA KV cache write + FP8 quant | `set_mla_kv_buffer_fp8_quant*`<br>`set_mla_kv_buffer_triton_fp8_quant` | `python/sglang/srt/mem_cache/utils.py`<br>`python/sglang/srt/mem_cache/memory_pool.py` | MLA / DSA KV pool path can quantize K and write directly into KV storage without a separate concat-and-quant chain | Treat standalone quant + KV-buffer write on MLA paths as missing existing fusion first. |
|
||||
| Fused MoE router / top-k / softcapping | `FusedMoeRouter`<br>`fused_moe_router*`<br>router GEMM + `topk` + `tanh` | `python/sglang/srt/layers/moe/router.py` | Single fused router kernel covers router matmul, softcapping, and top-k selection | Treat exposed router matmul + softcap + top-k chains as an existing MoE fusion family. |
|
||||
| Fused MoE grouped-topk / gate kernels | `fused_topk_deepseek`<br>`moe_fused_gate`<br>`aiter_fused_topk`<br>`kimi_k2_moe_fused_gate` | `python/sglang/srt/layers/moe/topk.py` | CUDA / ROCm / FlashInfer kernels fuse bias, grouped-topk, renorm, and routed scaling into one gate op | Check backend / model eligibility before proposing a novel router-gate fusion. |
|
||||
| Qwen-style shared-expert append into routed top-k output | `_append_shared_to_topk_output`<br>`fused_append_shared_experts_with_weights`<br>`num_fused_shared_experts` | `python/sglang/srt/models/qwen2_moe.py`<br>`python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe_triton_kernels.py` | Qwen-style MoE paths can append shared-expert ids and sigmoid gate weights to routed top-k output in one Triton kernel so the shared experts execute inside the fused MoE path | Treat routed top-k plus shared-expert pad / concat ladders as an existing MoE-prep fusion family first. |
|
||||
@@ -67,8 +67,8 @@ work. Recheck PR state before treating an in-flight row as shipped.
|
||||
| Fused MoE sum + all-reduce | routed MoE followed by explicit sum-reduce kernels | `python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py`<br>`python/sglang/srt/layers/moe/fused_moe_triton/fused_moe_triton_kernels.py` | `fuse_sum_all_reduce=True` path in the second MoE GEMM | Before inventing a new MoE reduction fuse, check whether `enable_fused_moe_sum_all_reduce` is simply off or the quant path is incompatible. |
|
||||
| Fused MoE activation + quant / re-quant | `silu_and_mul_*quant*`<br>`npu_dequant_swiglu_quant`<br>`swiglu_quant` | `python/sglang/srt/layers/moe/ep_moe/kernels.py`<br>`python/sglang/jit_kernel/nvfp4.py`<br>`python/sglang/srt/layers/moe/cutlass_w4a8_moe.py`<br>`python/sglang/srt/hardware_backend/npu/quantization/fused_moe_method_npu.py` | Quantized MoE backends fuse SwiGLU / SiLU-and-mul with FP8 / FP4 / NPU re-quant before the second expert GEMM | If MoE traces show standalone activation then quant kernels, first check whether the quantized fused path is missing. |
|
||||
| DeepSeek comm-prep fused RMSNorm + quant / flatten-quant | `fused_rms_fp8_group_quant`<br>`fused_rms_mxfp4_quant`<br>`fused_flatten_fp8_group_quant`<br>`fused_flatten_mxfp4_quant` | `python/sglang/srt/layers/communicator.py`<br>`python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py`<br>`python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py` | DeepSeek MLA / MHA ROCm paths fuse RMSNorm or flatten with FP8 / MXFP4 quantization for comm / attention prep | On DeepSeek quant traces, split norm + quant or flatten + quant is an existing family, not a new idea. |
|
||||
| NSA fused top-k transform / page-table build | `fast_topk_transform_fused`<br>`fast_topk_transform_ragged_fused` | `python/sglang/srt/layers/attention/nsa_backend.py` | NSA can fuse top-k selection with paged / ragged index transform instead of separate top-k plus metadata scatter | If NSA top-k metadata work is split, check `SGLANG_NSA_FUSE_TOPK` and backend support first. |
|
||||
| NSA fused quantize + indexed K-cache store | `fused_store_index_k_cache`<br>`act_quant`<br>`index_k_with_scale_buffer` | `python/sglang/jit_kernel/fused_store_index_cache.py`<br>`python/sglang/srt/layers/attention/nsa/nsa_indexer.py` | Single JIT kernel quantizes bf16 K to fp8 + scale and writes directly into NSA index cache | Treat split `act_quant` + buffer-store on CUDA as missing an existing fused store path. |
|
||||
| DSA fused top-k transform / page-table build | `fast_topk_transform_fused`<br>`fast_topk_transform_ragged_fused` | `python/sglang/srt/layers/attention/dsa_backend.py` | DSA can fuse top-k selection with paged / ragged index transform instead of separate top-k plus metadata scatter | If DSA top-k metadata work is split, check `SGLANG_DSA_FUSE_TOPK` and backend support first. |
|
||||
| DSA fused quantize + indexed K-cache store | `fused_store_index_k_cache`<br>`act_quant`<br>`index_k_with_scale_buffer` | `python/sglang/jit_kernel/fused_store_index_cache.py`<br>`python/sglang/srt/layers/attention/dsa/dsa_indexer.py` | Single JIT kernel quantizes bf16 K to fp8 + scale and writes directly into DSA index cache | Treat split `act_quant` + buffer-store on CUDA as missing an existing fused store path. |
|
||||
| Fused sampling temperature + softmax | `fused_temperature_softmax*` | `python/sglang/srt/layers/fused_sampling.py`<br>`python/sglang/srt/layers/sampler.py` | Triton single-pass / multi-pass kernels fuse temperature scaling and softmax during decode | Separate temp-divide + softmax at decode batch sizes is often a missed existing fusion. |
|
||||
| Fused logit softcap | `fused_softcap`<br>`final_logit_softcapping` | `python/sglang/srt/layers/elementwise.py`<br>`python/sglang/srt/layers/logits_processor.py` | Triton kernels fuse cast-to-float and softcap / tanh math for logits or generic elementwise softcapping | Treat exposed cast + softcap ladders as an existing Triton fuse family. |
|
||||
| Linear-attention packed projection reshuffle | `fused_qkvzba_split_reshape_cat*`<br>`qkvz_proj`<br>`ba_proj`<br>`qkvabz_proj`<br>`fused_qkvbfg_a_proj` | `python/sglang/jit_kernel/triton/gdn_fused_proj.py`<br>`python/sglang/srt/models/qwen3_next.py`<br>`python/sglang/srt/models/qwen3_5.py`<br>`python/sglang/srt/models/kimi_linear.py`<br>`python/sglang/srt/models/jet_nemotron.py` | GDN / Kimi / Jet-style linear-attn models pack multiple projections, then fuse split / reshape / cat into one kernel | Treat split reshape / transpose / cat ladders as an existing linear-attention fusion family. |
|
||||
@@ -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`<br>`_comm_stream`<br>`dispatch`<br>`combine`<br>`done_event` | `python/sglang/srt/layers/moe/token_dispatcher/moriep.py` | MoriEP can submit dispatch and combine onto a dedicated communication stream and synchronize only through events | Treat MoriEP comm / compute interleave as an existing MoE overlap family. |
|
||||
| Heterogeneous-TP staging scatter overlap | `scatter_stream`<br>`_scatter_stream`<br>`staging` | `python/sglang/srt/disaggregation/common/staging_handler.py`<br>`python/sglang/srt/disaggregation/common/staging_buffer.py` | decode-side staging scatter kernels can run on a dedicated stream while forward continues on the main stream | If decode traces show staging scatter kernels adjacent to forward kernels, classify them against this existing overlap family first. |
|
||||
| Generic `alt_stream` overlap families | `alt_stream` plus explicit `wait_stream` / `with torch.cuda.stream(...)` | `qwen2_moe.py`<br>`qwen3_moe.py`<br>`glm4_moe.py`<br>`bailing_moe.py`<br>`llada2.py`<br>`grok.py`<br>`olmo2.py`<br>`step3p5.py`<br>`longcat_flash.py`<br>`falcon_h1.py` | model-specific overlap on attention prep, MoE branches, or cache-store | Search these families before designing a new overlap scheme from scratch. |
|
||||
@@ -134,16 +134,16 @@ Stable entries should be folded into the mainline family rows above.
|
||||
| Pattern | Trace keywords | Primary code | Existing path | Skill should conclude |
|
||||
| --- | --- | --- | --- | --- |
|
||||
| PR `#21877` fused grouped down-GEMM + combine | `grouped_gemm_nt_masked`<br>`combine`<br>`fused grouped gemm combine` | `PR #21877`<br>`python/sglang/srt/layers/moe/ep_moe/flashinfer_cutedsl_moe.py`<br>`python/sglang/srt/layers/moe/token_dispatcher/deepep.py` | FlashInfer CuTeDSL kernel fuses the second expert GEMM with DeepEP low-latency combine | Treat this as a concrete upstream MoE fuse / overlap family, not a new thought experiment. |
|
||||
| PR `#21889` fused BF16 to FP4 quant + paged KV write | `set_mla_kv_buffer_fp4_quant_kernel`<br>`fp4 kv cache` | `PR #21889`<br>`python/sglang/srt/mem_cache/utils.py` | Triton kernel writes FP4 NSA KV pages directly while quantizing BF16 input | If NSA FP4 KV paths are split into quant plus store, classify them as an in-flight upstream fuse family. |
|
||||
| PR `#21889` fused FP4 paged dequant to FP8 + page-table remap | `_dequant_fp4_to_fp8_paged_kernel`<br>`WRITE_PT`<br>`dequant_fp4_paged_decode` | `PR #21889`<br>`python/sglang/srt/layers/attention/nsa/dequant_fp4_to_fp8.py` | Triton kernel reads FP4 pages, writes FP8 directly, and can fuse decode-side page-table remap | Treat this as an upstream in-flight decode-prep fusion family. |
|
||||
| PR `#21889` fused BF16 to FP4 quant + paged KV write | `set_mla_kv_buffer_fp4_quant_kernel`<br>`fp4 kv cache` | `PR #21889`<br>`python/sglang/srt/mem_cache/utils.py` | Triton kernel writes FP4 DSA KV pages directly while quantizing BF16 input | If DSA FP4 KV paths are split into quant plus store, classify them as an in-flight upstream fuse family. |
|
||||
| PR `#21889` fused FP4 paged dequant to FP8 + page-table remap | `_dequant_fp4_to_fp8_paged_kernel`<br>`WRITE_PT`<br>`dequant_fp4_paged_decode` | `PR #21889`<br>`python/sglang/srt/layers/attention/dsa/dequant_fp4_to_fp8.py` | Triton kernel reads FP4 pages, writes FP8 directly, and can fuse decode-side page-table remap | Treat this as an upstream in-flight decode-prep fusion family. |
|
||||
| PR `#21491` FlashInfer TRTLLM FP8 MoE with fused shared experts | `num_fused_shared_experts`<br>`trtllm_fp8_block_scale_moe` | `PR #21491`<br>`python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py`<br>`python/sglang/srt/models/deepseek_v2.py` | FlashInfer TRTLLM FP8 MoE path can fuse shared experts inside the routed MoE kernel | On FP8 TRTLLM MoE discussions, treat fused shared experts as an upstream pattern that already has a concrete PR. |
|
||||
| PR `#22005` fused add + RMSNorm + per-token FP8 quant | `fused_add_rmsnorm_per_token_quant`<br>`per_token_quant_fp8` | `PR #22005`<br>`python/sglang/jit_kernel/csrc/elementwise/fused_add_rmsnorm_per_token_quant.cuh`<br>`python/sglang/jit_kernel/fused_add_rmsnorm_per_token_quant.py` | CUDA JIT kernel keeps normed values in registers and emits BF16 + FP8 outputs plus per-token scales | If FP8 online-quant traces show add+norm followed by per-token quant, treat this as an in-flight upstream CUDA fuse family. |
|
||||
| PR `#20667` Qwen3.5 fused QK norm + RoPE + KV cache write | `fused_qk_norm_rope_cache_pts_quant_shuffle`<br>`fused_qk_norm_mrope_3d_cache_pts_quant_shuffle`<br>`rotary_dim` | `PR #20667`<br>`python/sglang/srt/models/qwen3_5.py`<br>`python/sglang/srt/models/utils.py` | ROCm / AITER path fuses Q / K RMSNorm, partial or 3D RoPE, and direct KV cache write for Qwen3.5 attention | Treat split QK-norm + RoPE + cache-store on Qwen3.5 as a concrete in-flight upstream family, not a novel idea. |
|
||||
| PR `#22392` CUTLASS FP8 GEMM replacing nvjet | `cutlass_scaled_mm`<br>`fp8_scaled_mm`<br>`nvjet`<br>`cudaMemsetAsync` | `PR #22392`<br>`sgl-kernel/python/sgl_kernel/gemm.py`<br>`python/sglang/srt/layers/quantization/fp8_utils.py` | Runtime replacement swaps nvjet FP8 GEMMs for CUTLASS kernels, removing per-launch memset bubbles and extra output-copy kernels | Treat nvjet GEMM + memset bubble ladders as an in-flight SGLang linear-kernel family before calling them novel. |
|
||||
| PR `#18612` NVFP4 CUTLASS MoE fused SiLU+Mul+quant | `silu_and_mul_scaled_nvfp4`<br>`nvfp4 expert quant`<br>`cutlass moe` | `PR #18612`<br>`python/sglang/srt/layers/moe/cutlass_w4a8_moe.py`<br>`python/sglang/jit_kernel/nvfp4.py` | Fuses MoE activation epilogue and NVFP4 expert quantization before the CUTLASS MoE second GEMM | Treat split SiLU+Mul then NVFP4 expert quant in CUTLASS MoE traces as an in-flight upstream SGLang family. |
|
||||
| PR `#22918` FlashInfer per-token NVFP4 MoE | `per_token_nvfp4`<br>`trtllm_fp4_block_scale_moe`<br>`FlashInfer MoE` | `PR #22918`<br>`python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py` | Adds FlashInfer-backed per-token NVFP4 MoE execution so expert quant/dequant work can move into the fused MoE backend | Treat standalone per-token NVFP4 MoE support kernels as a candidate missing backend-selection path, not an automatically novel kernel idea. |
|
||||
| PR `#22851` NSA top-k backend and FlashInfer / PyTorch top-k split | `nsa topk`<br>`flashinfer_topk`<br>`pytorch_topk`<br>`fast_topk_transform` | `PR #22851`<br>`python/sglang/srt/layers/attention/nsa_backend.py` | Makes NSA top-k backend selection explicit and aligns fused top-k transform with FlashInfer / PyTorch fallbacks | When NSA top-k dominates decode, first classify it as backend selection or fused-transform eligibility work. |
|
||||
| PR `#24125` GLM5 NSA decode CatArrayBatchedCopy removal | `CatArrayBatchedCopy`<br>`GLM-5`<br>`NSA`<br>`TileLang decode` | `PR #24125`<br>`python/sglang/srt/layers/attention/nsa_backend.py` | Skips redundant cat/copy work in the GLM5 NSA TileLang decode path | Treat cat/copy bursts in GLM5 NSA decode as a concrete in-flight cleanup opportunity. |
|
||||
| PR `#22851` DSA top-k backend and FlashInfer / PyTorch top-k split | `dsa topk`<br>`flashinfer_topk`<br>`pytorch_topk`<br>`fast_topk_transform` | `PR #22851`<br>`python/sglang/srt/layers/attention/dsa_backend.py` | Makes DSA top-k backend selection explicit and aligns fused top-k transform with FlashInfer / PyTorch fallbacks | When DSA top-k dominates decode, first classify it as backend selection or fused-transform eligibility work. |
|
||||
| PR `#24125` GLM5 DSA decode CatArrayBatchedCopy removal | `CatArrayBatchedCopy`<br>`GLM-5`<br>`DSA`<br>`TileLang decode` | `PR #24125`<br>`python/sglang/srt/layers/attention/dsa_backend.py` | Skips redundant cat/copy work in the GLM5 DSA TileLang decode path | Treat cat/copy bursts in GLM5 DSA decode as a concrete in-flight cleanup opportunity. |
|
||||
| PR `#24007` MoE LoRA virtual experts for csgmv backend | `csgmv`<br>`virtual experts`<br>`MoE LoRA`<br>`fused_moe_lora` | `PR #24007`<br>`python/sglang/srt/layers/lora_backend.py`<br>`python/sglang/srt/layers/moe` | Routes MoE LoRA adapter work through virtual experts so csgmv-style kernels can batch it instead of launching fragmented adapter work | Treat MoE-LoRA tiny-kernel ladders as an in-flight batching/fusion family. |
|
||||
| PR `#24150` torch.compile local decode support | `enable_torch_compile`<br>`local compile`<br>`decode compile`<br>`torchinductor` | `PR #24150`<br>`python/sglang/srt` | Extends SGLang torch.compile coverage to local decode regions, so Inductor-generated fusion may replace hand-authored tiny kernels | When decode traces show compiler-generated kernels or missing named fused kernels, check this in-flight compile path before calling the shape unsupported. |
|
||||
|
||||
@@ -294,7 +294,7 @@ contain the same implementation.
|
||||
| `enable_single_batch_overlap` | `python/sglang/srt/server_args.py` | Enables the SBO family. |
|
||||
| `enable_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. |
|
||||
|
||||
@@ -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`<br>`_comm_stream`<br>`dispatch`<br>`combine`<br>`done_event` | `python/sglang/srt/layers/moe/token_dispatcher/moriep.py` | MoriEP can submit dispatch and combine onto a dedicated communication stream and synchronize only through events | Treat MoriEP comm / compute interleave as an existing MoE overlap family. |
|
||||
| Generic `alt_stream` overlap families | `alt_stream` plus explicit `wait_stream` / `with torch.cuda.stream(...)` | `qwen2_moe.py`<br>`qwen3_moe.py`<br>`glm4_moe.py`<br>`bailing_moe.py`<br>`llada2.py`<br>`grok.py`<br>`olmo2.py`<br>`step3p5.py`<br>`longcat_flash.py`<br>`falcon_h1.py` | model-specific overlap on attention prep, MoE branches, or cache-store | Search these families before designing a new overlap scheme from scratch. |
|
||||
|
||||
|
||||
@@ -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"
|
||||
"<br>python/sglang/srt/layers/attention/nsa/nsa_indexer.py"
|
||||
"<br>python/sglang/srt/layers/attention/dsa/dsa_indexer.py"
|
||||
),
|
||||
active_keywords=("fused_store_index_k_cache",),
|
||||
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,
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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** |
|
||||
|-----------------------|-------------|------------|-----------------------------------------------|
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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` |
|
||||
|
||||
@@ -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 <dp-size>` in command. DP Attention is better for large concurrency scenarios.
|
||||
- **TP Attention**: Launching with TP attention is also supported. TP attention is better for low latency scenarios.
|
||||
- **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 <HEAD_NODE_IP>: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 <HEAD_NODE_IP>: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 <PREFILL_HEAD_IP>: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 <PREFILL_HEAD_IP>: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 \
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -9,7 +9,7 @@ This document provides a list of commonly used environment variables and aims to
|
||||
|--------------------------------------------------|-------------------------------------------------------------------------------------------------------------------------------------------------------------|---------------|
|
||||
| `SGLANG_NPU_USE_MLAPO` | Adopts the `MLAPO` fusion operator in attention <br/> preprocessing stage of the MLA model. | `false` |
|
||||
| `SGLANG_USE_FIA_NZ` | Reshapes KV Cache for FIA NZ format.<br/> `SGLANG_USE_FIA_NZ` must be enabled with `SGLANG_NPU_USE_MLAPO` | `false` |
|
||||
| `SGLANG_NPU_USE_MULTI_STREAM` | Enable dual-stream computation of shared experts <br/> and routing experts in DeepSeek models.<br/> Enable dual-stream computation in DeepSeek NSA Indexer. | `false` |
|
||||
| `SGLANG_NPU_USE_MULTI_STREAM` | Enable dual-stream computation of shared experts <br/> and routing experts in DeepSeek models.<br/> Enable dual-stream computation in DeepSeek DSA Indexer. | `false` |
|
||||
| `SGLANG_NPU_DISABLE_ACL_FORMAT_WEIGHT` | Disable cast model weight tensor to a specific NPU <br/> ACL format. | `false` |
|
||||
| `SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK` | The maximum number of dispatched tokens on each rank. | `128` |
|
||||
|
||||
|
||||
@@ -204,8 +204,8 @@ click [Server Arguments](https://docs.sglang.io/advanced_features/server_argumen
|
||||
| `--sampling-backend` | `None` | `pytorch`,<br/>`ascend` | A2, A3 |
|
||||
| `--grammar-backend` | `None` | `xgrammar` | A2, A3 |
|
||||
| `--mm-attention-backend` | `None` | `ascend_attn` | A2, A3 |
|
||||
| `--nsa-prefill-backend` | `flashmla_sparse` | `flashmla_sparse`,<br/> `flashmla_decode`,<br/>`fa3`,<br/> `tilelang`,<br/> `aiter` | Special for GPU |
|
||||
| `--nsa-decode-backend` | `fa3` | `flashmla_prefill`,<br/> `flashmla_kv`,<br/> `fa3`,<br/>`tilelang`,<br/> `aiter` | Special for GPU |
|
||||
| `--dsa-prefill-backend` | `flashmla_sparse` | `flashmla_sparse`,<br/> `flashmla_decode`,<br/>`fa3`,<br/> `tilelang`,<br/> `aiter` | Special for GPU |
|
||||
| `--dsa-decode-backend` | `fa3` | `flashmla_prefill`,<br/> `flashmla_kv`,<br/> `fa3`,<br/>`tilelang`,<br/> `aiter` | Special for GPU |
|
||||
| `--fp8-gemm-backend` | `auto` | `auto`,<br/> `deep_gemm`,<br/> `flashinfer_trtllm`,<br/>`flashinfer_cutlass`,<br/>`flashinfer_deepgemm`,<br/>`cutlass`,<br/> `triton`,<br/> `aiter` | Special for GPU |
|
||||
| `--disable-flashinfer-`<br/>`autotune` | `False` | bool flag<br/> (set to enable) | Special for GPU |
|
||||
|
||||
@@ -371,7 +371,7 @@ click [Server Arguments](https://docs.sglang.io/advanced_features/server_argumen
|
||||
| `--rl-on-policy-target` | `None` | `fsdp` | Planned |
|
||||
| `--enable-layerwise-`<br/>`nvtx-marker` | `False` | bool flag<br/> (set to enable) | Special for GPU |
|
||||
| `--enable-attn-tp-`<br/>`input-scattered` | `False` | bool flag<br/> (set to enable) | Experimental |
|
||||
| `--enable-nsa-prefill-`<br/>`context-parallel` | `False` | bool flag<br/> (set to enable) | A2, A3 |
|
||||
| `--enable-dsa-prefill-`<br/>`context-parallel` | `False` | bool flag<br/> (set to enable) | A2, A3 |
|
||||
| `--enable-fused-qk-`<br/>`norm-rope` | `False` | bool flag<br/> (set to enable) | Special for GPU |
|
||||
|
||||
## Dynamic batch tokenizer
|
||||
|
||||
@@ -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 |
|
||||
|
||||
@@ -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 to control mtp precomputing of metadata for multi-step speculative decoding -->
|
||||
|
||||
| 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
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -79,7 +79,7 @@ import { GLM51Deployment } from '/src/snippets/autoregressive/glm-51-deployment.
|
||||
</tbody>
|
||||
</table>
|
||||
|
||||
- **AMD GPUs**: Both BF16 and FP8 checkpoints are supported on MI300X/MI325X/MI355X at tp=8. Use `--nsa-prefill-backend tilelang --nsa-decode-backend tilelang` for the NSA attention backend. Add `--chunked-prefill-size 131072` and `--watchdog-timeout 1200` (20 minutes for weight loading). FP8 uses approximately half the memory of BF16 (~89 GB/GPU vs ~175 GB/GPU). EAGLE speculative decoding is not currently supported on AMD for GLM-5.1.
|
||||
- **AMD GPUs**: Both BF16 and FP8 checkpoints are supported on MI300X/MI325X/MI355X at tp=8. Use `--dsa-prefill-backend tilelang --dsa-decode-backend tilelang` for the DSA attention backend. Add `--chunked-prefill-size 131072` and `--watchdog-timeout 1200` (20 minutes for weight loading). FP8 uses approximately half the memory of BF16 (~89 GB/GPU vs ~175 GB/GPU). EAGLE speculative decoding is not currently supported on AMD for GLM-5.1.
|
||||
- **GB300**: Only the FP8 checkpoint is recommended on GB300, with `tp=4`. For high-throughput DP attention on GB300, use `--dp 4`.
|
||||
- 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
|
||||
|
||||
@@ -85,9 +85,9 @@ import { GLM5Deployment } from '/src/snippets/autoregressive/glm-5-deployment.js
|
||||
</tbody>
|
||||
</table>
|
||||
|
||||
- **B200 (FP8)**: Use `--ep 1 --attention-backend nsa --nsa-decode-backend trtllm --nsa-prefill-backend trtllm --moe-runner-backend flashinfer_trtllm --enable-flashinfer-allreduce-fusion` for optimized NSA and MoE backends on Blackwell. Also add `--quantization fp8` for FP8 weight quantization.
|
||||
- **B200 (FP8)**: Use `--ep 1 --attention-backend dsa --dsa-decode-backend trtllm --dsa-prefill-backend trtllm --moe-runner-backend flashinfer_trtllm --enable-flashinfer-allreduce-fusion` for optimized DSA and MoE backends on Blackwell. Also add `--quantization fp8` for FP8 weight quantization.
|
||||
|
||||
- **AMD GPUs**: Use `--nsa-prefill-backend tilelang --nsa-decode-backend tilelang` for the NSA attention backend. Add `--chunked-prefill-size 131072` and `--watchdog-timeout 1200` (20 minutes for weight loading). EAGLE speculative decoding is not currently supported on AMD for GLM-5.
|
||||
- **AMD GPUs**: Use `--dsa-prefill-backend tilelang --dsa-decode-backend tilelang` for the DSA attention backend. Add `--chunked-prefill-size 131072` and `--watchdog-timeout 1200` (20 minutes for weight loading). EAGLE speculative decoding is not currently supported on AMD for GLM-5.
|
||||
- For other configuration tips, please refer to [DeepSeek V3.2 documentation](../../../docs/basic_usage/deepseek_v32). GLM-5 and DeepSeek V3.2 share the same model structure, so the optimization techniques between these two models are also common (MTP, DSA kernel, Context Parallel...).
|
||||
- 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
|
||||
|
||||
@@ -280,7 +280,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) for details.
|
||||
- DSA is specifically designed for [DeepSeek V3.2](https://lmsys.org/blog/2025-09-29-deepseek-V32/). See the [DSA Attention Backend](#dsa-attention-backend) section and [DeepSeek V3.2 deployment guide](../basic_usage/deepseek_v32) for details.
|
||||
</Note>
|
||||
|
||||
<Warning>
|
||||
@@ -378,11 +378,11 @@ GDN models are hybrid: the full-attention layers still require a standard `--att
|
||||
- **Other CUDA (Hopper, Ampere, etc.)**: auto-selection works; no special constraints.
|
||||
</Warning>
|
||||
|
||||
### DSA Attention Backend (NSA)
|
||||
### DSA Attention Backend
|
||||
|
||||
DSA (Deepseek Sparse Attention) is a native sparse attention mechanism used by [DeepSeek V3.2](https://lmsys.org/blog/2025-09-29-deepseek-V32/). It is activated automatically when the model architecture requires it and is selected via `--attention-backend nsa`.
|
||||
DSA (DeepSeek Sparse Attention) is a native sparse attention mechanism used by [DeepSeek V3.2](https://lmsys.org/blog/2025-09-29-deepseek-V32/). It is activated automatically when the model architecture requires it and is selected via `--attention-backend dsa` (deprecated alias: `nsa`).
|
||||
|
||||
Internally, the NSA backend dispatches to different sub-backends for prefill and decode phases. You can override these with `--nsa-prefill-backend` and `--nsa-decode-backend`:
|
||||
Internally, the DSA backend dispatches to different sub-backends for prefill and decode phases. You can override these with `--dsa-prefill-backend` and `--dsa-decode-backend`:
|
||||
|
||||
<table style={{width: "100%", borderCollapse: "collapse", tableLayout: "fixed"}}>
|
||||
<colgroup>
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -1162,19 +1162,19 @@ Please consult the documentation below and [server_args.py](https://github.com/s
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--attention-backend`</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Choose the kernels for attention layers.</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`None`</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>triton</code>, <code>torch_native</code>, <code>flex_attention</code>, <code>nsa</code>, <code>cutlass_mla</code>, <code>fa3</code>, <code>fa4</code>, <code>flashinfer</code>, <code>flashmla</code>, <code>trtllm_mla</code>, <code>trtllm_mha</code>, <code>dual_chunk_flash_attn</code>, <code>aiter</code>, <code>wave</code>, <code>intel_amx</code>, <code>ascend</code></td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>triton</code>, <code>torch_native</code>, <code>flex_attention</code>, <code>dsa</code>, <code>cutlass_mla</code>, <code>fa3</code>, <code>fa4</code>, <code>flashinfer</code>, <code>flashmla</code>, <code>trtllm_mla</code>, <code>trtllm_mha</code>, <code>dual_chunk_flash_attn</code>, <code>aiter</code>, <code>wave</code>, <code>intel_amx</code>, <code>ascend</code></td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--prefill-attention-backend`</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Choose the kernels for prefill attention layers (have priority over --attention-backend).</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`None`</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>triton</code>, <code>torch_native</code>, <code>flex_attention</code>, <code>nsa</code>, <code>cutlass_mla</code>, <code>fa3</code>, <code>fa4</code>, <code>flashinfer</code>, <code>flashmla</code>, <code>trtllm_mla</code>, <code>trtllm_mha</code>, <code>dual_chunk_flash_attn</code>, <code>aiter</code>, <code>wave</code>, <code>intel_amx</code>, <code>ascend</code></td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>triton</code>, <code>torch_native</code>, <code>flex_attention</code>, <code>dsa</code>, <code>cutlass_mla</code>, <code>fa3</code>, <code>fa4</code>, <code>flashinfer</code>, <code>flashmla</code>, <code>trtllm_mla</code>, <code>trtllm_mha</code>, <code>dual_chunk_flash_attn</code>, <code>aiter</code>, <code>wave</code>, <code>intel_amx</code>, <code>ascend</code></td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--decode-attention-backend`</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Choose the kernels for decode attention layers (have priority over --attention-backend).</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`None`</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>triton</code>, <code>torch_native</code>, <code>flex_attention</code>, <code>nsa</code>, <code>cutlass_mla</code>, <code>fa3</code>, <code>fa4</code>, <code>flashinfer</code>, <code>flashmla</code>, <code>trtllm_mla</code>, <code>trtllm_mha</code>, <code>dual_chunk_flash_attn</code>, <code>aiter</code>, <code>wave</code>, <code>intel_amx</code>, <code>ascend</code></td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>triton</code>, <code>torch_native</code>, <code>flex_attention</code>, <code>dsa</code>, <code>cutlass_mla</code>, <code>fa3</code>, <code>fa4</code>, <code>flashinfer</code>, <code>flashmla</code>, <code>trtllm_mla</code>, <code>trtllm_mha</code>, <code>dual_chunk_flash_attn</code>, <code>aiter</code>, <code>wave</code>, <code>intel_amx</code>, <code>ascend</code></td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--sampling-backend`</td>
|
||||
@@ -1195,14 +1195,14 @@ Please consult the documentation below and [server_args.py](https://github.com/s
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>sdpa</code>, <code>fa3</code>, <code>fa4</code>, <code>triton_attn</code>, <code>ascend_attn</code>, <code>aiter_attn</code></td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--nsa-prefill-backend`</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Choose the NSA backend for the prefill stage (overrides `--attention-backend` when running DeepSeek NSA-style attention).</td>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--dsa-prefill-backend`</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>DSA backend for the prefill stage (overrides `--attention-backend` when running DeepSeek DSA-style attention).</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`flashmla_sparse`</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>flashmla_sparse</code>, <code>flashmla_kv</code>, <code>flashmla_auto</code>, <code>fa3</code>, <code>tilelang</code>, <code>aiter</code>, <code>trtllm</code></td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--nsa-decode-backend`</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Choose the NSA backend for the decode stage when running DeepSeek NSA-style attention. Overrides `--attention-backend` for decoding.</td>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--dsa-decode-backend`</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>DSA backend for the decode stage when running DeepSeek DSA-style attention. Overrides `--attention-backend` for decoding.</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`fa3`</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>flashmla_sparse</code>, <code>flashmla_kv</code>, <code>fa3</code>, <code>tilelang</code>, <code>aiter</code>, <code>trtllm</code></td>
|
||||
</tr>
|
||||
@@ -2278,13 +2278,13 @@ Please consult the documentation below and [server_args.py](https://github.com/s
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>bool flag (set to enable)</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>--enable-nsa-prefill-context-parallel</code></td>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>--enable-dsa-prefill-context-parallel</code></td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Enable context parallelism used in the long sequence prefill phase of DeepSeek v3.2.</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>False</code></td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>bool flag (set to enable)</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>--nsa-prefill-cp-mode</code></td>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>--dsa-prefill-cp-mode</code></td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Token splitting mode for the prefill phase of DeepSeek v3.2 under context parallelism. Optional values: <code>round-robin-split</code>(default),<code>in-seq-split</code>. <code>round-robin-split</code> distributes tokens across ranks based on <code>token_idx % cp_size</code>. It supports multi-batch prefill, fused MoE, and FP8 KV cache.</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>in-seq-split</code></td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>in-seq-split</code>, <code>round-robin-split</code></td>
|
||||
@@ -2841,13 +2841,13 @@ Please consult the documentation below and [server_args.py](https://github.com/s
|
||||
</tr>
|
||||
<tr>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--nsa-prefill`</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Choose the NSA backend for the prefill stage (overrides `--attention-backend` when running DeepSeek NSA-style attention).</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Deprecated alias for `--dsa-prefill-backend`.</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`flashmla_sparse`</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>`flashmla_sparse`, `flashmla_decode`, `fa3`, `tilelang`, `aiter`</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--nsa-decode`</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Choose the NSA backend for the decode stage when running DeepSeek NSA-style attention. Overrides `--attention-backend` for decoding.</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Deprecated alias for `--dsa-decode-backend`.</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`flashmla_kv`</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>`flashmla_prefill`, `flashmla_kv`, `fa3`, `tilelang`, `aiter`</td>
|
||||
</tr>
|
||||
|
||||
@@ -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 <dp-size>` in command. DP Attention is better for large concurrency scenarios.
|
||||
- **TP Attention**: Launching with TP attention is also supported. TP attention is better for low latency scenarios.
|
||||
- **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 <HEAD_NODE_IP>: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 <HEAD_NODE_IP>: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 <PREFILL_HEAD_IP>: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 <PREFILL_HEAD_IP>: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 \
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -29,7 +29,7 @@ This document provides a list of commonly used environment variables and aims to
|
||||
</tr>
|
||||
<tr>
|
||||
<td><code>SGLANG_NPU_USE_MULTI_STREAM</code></td>
|
||||
<td>Enable dual-stream computation of shared experts <br/> and routing experts in DeepSeek models.<br/> Enable dual-stream computation in DeepSeek NSA Indexer.</td>
|
||||
<td>Enable dual-stream computation of shared experts <br/> and routing experts in DeepSeek models.<br/> Enable dual-stream computation in DeepSeek DSA Indexer.</td>
|
||||
<td><code>false</code></td>
|
||||
</tr>
|
||||
<tr>
|
||||
|
||||
@@ -187,18 +187,18 @@ These arguments and environment variables are critical for tuning prefill perfor
|
||||
<td>true</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td>`--enable-nsa-prefill-context-parallel`</td>
|
||||
<td><strong>(DeepSeek V3.2 NSA-specific)</strong> Enables context parallelism for the long-sequence prefill phase of DeepSeek V3.2 with NSA (Native Sparse Attention). Distributes the sequence across CP ranks to parallelize the computationally expensive NSA prefill for ultra-long contexts.</td>
|
||||
<td>`--enable-dsa-prefill-context-parallel`</td>
|
||||
<td><strong>(DeepSeek V3.2 DSA-specific)</strong> Enables context parallelism for the long-sequence prefill phase of DeepSeek V3.2 with DSA (DeepSeek Sparse Attention). Distributes the sequence across CP ranks to parallelize the computationally expensive DSA prefill for ultra-long contexts.</td>
|
||||
<td>Enabled</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td>`--nsa-prefill-cp-mode`</td>
|
||||
<td><strong>(DeepSeek V3.2 NSA-specific)</strong> Controls how the long sequence is split across context parallel ranks: `in-seq-split` divides each sequence uniformly across CP ranks, optimal for single-request prefill. `round-robin-split` (code default) distributes tokens by index mod CP size, supporting multi-batch prefill. Only effective when `--enable-nsa-prefill-context-parallel` is enabled.</td>
|
||||
<td>`--dsa-prefill-cp-mode`</td>
|
||||
<td><strong>(DeepSeek V3.2 DSA-specific)</strong> Controls how the long sequence is split across context parallel ranks: `in-seq-split` divides each sequence uniformly across CP ranks, optimal for single-request prefill. `round-robin-split` (code default) distributes tokens by index mod CP size, supporting multi-batch prefill. Only effective when `--enable-dsa-prefill-context-parallel` is enabled.</td>
|
||||
<td>`in-seq-split`</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td>`--attn-cp-size`</td>
|
||||
<td>Specifies the context parallelism group size for attention computation. Larger values distribute the sequence across more ranks, reducing per-rank memory and compute at the cost of increased communication. For models with NSA, this controls the CP size for sparse attention prefill. Set to the number of available devices for maximum parallelization.</td>
|
||||
<td>Specifies the context parallelism group size for attention computation. Larger values distribute the sequence across more ranks, reducing per-rank memory and compute at the cost of increased communication. For models with DSA, this controls the CP size for sparse attention prefill. Set to the number of available devices for maximum parallelization.</td>
|
||||
<td>`32`</td>
|
||||
</tr>
|
||||
</tbody>
|
||||
@@ -360,7 +360,7 @@ The following environment variables are used in other best practice configuratio
|
||||
<tr>
|
||||
<td>`HCCL_OP_EXPANSION_MODE=AIV`</td>
|
||||
<td>Configures the HCCL communication algorithm scheduling to use AIV (Ascend Intelligent Vision) expansion mode, which can improve communication efficiency for certain collective operations.</td>
|
||||
<td>Used in Qwen MoE and R1 non-NSA configurations</td>
|
||||
<td>Used in Qwen MoE and R1 non-DSA configurations</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td>`SGLANG_NPU_FUSED_MOE_MODE`</td>
|
||||
|
||||
@@ -1225,13 +1225,13 @@ click [Server Arguments](../../advanced_features/server_arguments).
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>A2, A3</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--nsa-prefill-backend`</td>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--dsa-prefill-backend`</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>`flashmla_sparse`</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>flashmla_sparse</code>,<br/> <code>flashmla_decode</code>,<br/><code>fa3</code>,<br/> <code>tilelang</code>,<br/> <code>aiter</code></td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Special for GPU</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--nsa-decode-backend`</td>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--dsa-decode-backend`</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>`fa3`</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>flashmla_prefill</code>,<br/> <code>flashmla_kv</code>,<br/> <code>fa3</code>,<br/><code>tilelang</code>,<br/> <code>aiter</code></td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Special for GPU</td>
|
||||
@@ -2226,7 +2226,7 @@ click [Server Arguments](../../advanced_features/server_arguments).
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Experimental</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--enable-nsa-prefill-`<br/>`context-parallel`</td>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--enable-dsa-prefill-`<br/>`context-parallel`</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>`False`</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>bool flag<br/> (set to enable)</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>A2, A3</td>
|
||||
|
||||
@@ -512,9 +512,9 @@ python -c "from sglang.srt.platforms import current_platform; print(current_plat
|
||||
<td>MLA KV cache pool class</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td><code>get_nsa_kv_pool_cls()</code></td>
|
||||
<td><code>get_dsa_kv_pool_cls()</code></td>
|
||||
<td><code>raise NotImplementedError</code></td>
|
||||
<td>NSA KV cache pool class (DeepSeek V3.2)</td>
|
||||
<td>DSA KV cache pool class (DeepSeek V3.2)</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td><code>get_paged_allocator_cls()</code></td>
|
||||
|
||||
@@ -393,7 +393,7 @@ SGLang supports various environment variables that can be used to configure its
|
||||
</tbody>
|
||||
</table>
|
||||
|
||||
## NSA Backend Configuration (For DeepSeek V3.2)
|
||||
## DSA Backend Configuration (For DeepSeek V3.2)
|
||||
|
||||
{/* # Environment variable to control mtp precomputing of metadata for multi-step speculative decoding */}
|
||||
|
||||
@@ -412,13 +412,13 @@ SGLang supports various environment variables that can be used to configure its
|
||||
</thead>
|
||||
<tbody>
|
||||
<tr>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_NSA_FUSE_TOPK</code></td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Fuse the operation of picking topk logits and picking topk indices from page table</td>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_DSA_FUSE_TOPK</code></td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Fuse the operation of picking topk logits and picking topk indices from page table. <code>SGLANG_NSA_FUSE_TOPK</code> is a deprecated alias.</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>true</code></td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_NSA_ENABLE_MTP_PRECOMPUTE_METADATA</code></td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Precompute metadata that can be shared among different draft steps when MTP is enabled</td>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_DSA_ENABLE_MTP_PRECOMPUTE_METADATA</code></td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Precompute metadata that can be shared among different draft steps when MTP is enabled. <code>SGLANG_NSA_ENABLE_MTP_PRECOMPUTE_METADATA</code> is a deprecated alias.</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>true</code></td>
|
||||
</tr>
|
||||
<tr>
|
||||
@@ -427,8 +427,8 @@ SGLang supports various environment variables that can be used to configure its
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>true</code></td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_NSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD</code></td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>When the maximum kv len in current prefill batch exceeds this value, the sparse mla kernel will be applied, else it falls back to dense MHA implementation. Default to the index topk of model (2048 for DeepSeek V3.2)</td>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_DSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD</code></td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>When the maximum kv len in current prefill batch exceeds this value, the sparse mla kernel will be applied, else it falls back to dense MHA implementation. Default to the index topk of model (2048 for DeepSeek V3.2). <code>SGLANG_NSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD</code> is a deprecated alias.</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>2048</code></td>
|
||||
</tr>
|
||||
</tbody>
|
||||
|
||||
@@ -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';
|
||||
}
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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';
|
||||
}
|
||||
|
||||
@@ -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';
|
||||
}
|
||||
|
||||
@@ -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 <algorithm> // for std::min
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
// Forward mode enum (must match Python ForwardMode in sglang/srt/layers/attention/nsa_backend.py)
|
||||
// Forward mode enum (must match Python ForwardMode in sglang/srt/layers/attention/dsa_backend.py)
|
||||
enum ForwardModeEnum { DECODE = 0, TARGET_VERIFY = 1, DRAFT_EXTEND = 2 };
|
||||
|
||||
/**
|
||||
@@ -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<tvm::ffi::TensorView> seqlens_expanded_src,
|
||||
const tvm::ffi::TensorView nsa_cu_seqlens_k_src,
|
||||
const tvm::ffi::TensorView dsa_cu_seqlens_k_src,
|
||||
const tvm::ffi::Optional<tvm::ffi::TensorView> real_page_table_src,
|
||||
const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_num_splits_src,
|
||||
const tvm::ffi::Optional<tvm::ffi::TensorView> 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<tvm::ffi::TensorView> seqlens_expanded_dst,
|
||||
const tvm::ffi::TensorView nsa_cu_seqlens_k_dst,
|
||||
const tvm::ffi::TensorView dsa_cu_seqlens_k_dst,
|
||||
const tvm::ffi::Optional<tvm::ffi::TensorView> real_page_table_dst,
|
||||
const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_num_splits_dst,
|
||||
const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_metadata_dst,
|
||||
@@ -522,9 +522,9 @@ struct FusedMetadataCopyKernel {
|
||||
.cache_seqlens = unwrap_data_ptr<int32_t>(cache_seqlens_src, "cache_seqlens_src"),
|
||||
.cu_seqlens_k = unwrap_data_ptr<int32_t>(cu_seqlens_k_src, "cu_seqlens_k_src"),
|
||||
.page_indices = unwrap_data_ptr<int32_t>(page_indices_src, "page_indices_src"),
|
||||
.nsa_cache_seqlens = unwrap_data_ptr<int32_t>(nsa_cache_seqlens_src, "nsa_cache_seqlens_src"),
|
||||
.dsa_cache_seqlens = unwrap_data_ptr<int32_t>(dsa_cache_seqlens_src, "dsa_cache_seqlens_src"),
|
||||
.seqlens_expanded = unwrap_optional_data_ptr<int32_t>(seqlens_expanded_src, "seqlens_expanded_src"),
|
||||
.nsa_cu_seqlens_k = unwrap_data_ptr<int32_t>(nsa_cu_seqlens_k_src, "nsa_cu_seqlens_k_src"),
|
||||
.dsa_cu_seqlens_k = unwrap_data_ptr<int32_t>(dsa_cu_seqlens_k_src, "dsa_cu_seqlens_k_src"),
|
||||
.real_page_table = unwrap_optional_data_ptr<int32_t>(real_page_table_src, "real_page_table_src"),
|
||||
.flashmla_num_splits =
|
||||
unwrap_optional_data_ptr<int32_t>(flashmla_num_splits_src, "flashmla_num_splits_src"),
|
||||
@@ -535,9 +535,9 @@ struct FusedMetadataCopyKernel {
|
||||
.cache_seqlens = unwrap_data_ptr_mut<int32_t>(cache_seqlens_dst, "cache_seqlens_dst"),
|
||||
.cu_seqlens_k = unwrap_data_ptr_mut<int32_t>(cu_seqlens_k_dst, "cu_seqlens_k_dst"),
|
||||
.page_table_1 = unwrap_data_ptr_mut<int32_t>(page_table_1_dst, "page_table_1_dst"),
|
||||
.nsa_cache_seqlens = unwrap_data_ptr_mut<int32_t>(nsa_cache_seqlens_dst, "nsa_cache_seqlens_dst"),
|
||||
.dsa_cache_seqlens = unwrap_data_ptr_mut<int32_t>(dsa_cache_seqlens_dst, "dsa_cache_seqlens_dst"),
|
||||
.seqlens_expanded = unwrap_optional_data_ptr_mut<int32_t>(seqlens_expanded_dst, "seqlens_expanded_dst"),
|
||||
.nsa_cu_seqlens_k = unwrap_data_ptr_mut<int32_t>(nsa_cu_seqlens_k_dst, "nsa_cu_seqlens_k_dst"),
|
||||
.dsa_cu_seqlens_k = unwrap_data_ptr_mut<int32_t>(dsa_cu_seqlens_k_dst, "dsa_cu_seqlens_k_dst"),
|
||||
.real_page_table = unwrap_optional_data_ptr_mut<int32_t>(real_page_table_dst, "real_page_table_dst"),
|
||||
.flashmla_num_splits =
|
||||
unwrap_optional_data_ptr_mut<int32_t>(flashmla_num_splits_dst, "flashmla_num_splits_dst"),
|
||||
@@ -605,32 +605,32 @@ struct FusedMetadataCopyMultiKernel {
|
||||
run(const tvm::ffi::TensorView cache_seqlens_src,
|
||||
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<tvm::ffi::TensorView> real_page_table_src,
|
||||
const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_num_splits_src,
|
||||
const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_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<tvm::ffi::TensorView> real_page_table_dst0,
|
||||
const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_num_splits_dst0,
|
||||
const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_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<tvm::ffi::TensorView> real_page_table_dst1,
|
||||
const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_num_splits_dst1,
|
||||
const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_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<tvm::ffi::TensorView> real_page_table_dst2,
|
||||
const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_num_splits_dst2,
|
||||
const tvm::ffi::Optional<tvm::ffi::TensorView> flashmla_metadata_dst2,
|
||||
@@ -647,9 +647,9 @@ struct FusedMetadataCopyMultiKernel {
|
||||
.cache_seqlens = unwrap_data_ptr<int32_t>(cache_seqlens_src, "cache_seqlens_src"),
|
||||
.cu_seqlens_k = unwrap_data_ptr<int32_t>(cu_seqlens_k_src, "cu_seqlens_k_src"),
|
||||
.page_indices = unwrap_data_ptr<int32_t>(page_indices_src, "page_indices_src"),
|
||||
.nsa_cache_seqlens = unwrap_data_ptr<int32_t>(nsa_cache_seqlens_src, "nsa_cache_seqlens_src"),
|
||||
.dsa_cache_seqlens = unwrap_data_ptr<int32_t>(dsa_cache_seqlens_src, "dsa_cache_seqlens_src"),
|
||||
.seqlens_expanded = nullptr, // Not used in multi-backend DECODE mode
|
||||
.nsa_cu_seqlens_k = unwrap_data_ptr<int32_t>(nsa_cu_seqlens_k_src, "nsa_cu_seqlens_k_src"),
|
||||
.dsa_cu_seqlens_k = unwrap_data_ptr<int32_t>(dsa_cu_seqlens_k_src, "dsa_cu_seqlens_k_src"),
|
||||
.real_page_table = unwrap_optional_data_ptr<int32_t>(real_page_table_src, "real_page_table_src"),
|
||||
.flashmla_num_splits =
|
||||
unwrap_optional_data_ptr<int32_t>(flashmla_num_splits_src, "flashmla_num_splits_src"),
|
||||
@@ -660,9 +660,9 @@ struct FusedMetadataCopyMultiKernel {
|
||||
.cache_seqlens = unwrap_data_ptr_mut<int32_t>(cache_seqlens_dst0, "cache_seqlens_dst0"),
|
||||
.cu_seqlens_k = unwrap_data_ptr_mut<int32_t>(cu_seqlens_k_dst0, "cu_seqlens_k_dst0"),
|
||||
.page_table_1 = unwrap_data_ptr_mut<int32_t>(page_table_1_dst0, "page_table_1_dst0"),
|
||||
.nsa_cache_seqlens = unwrap_data_ptr_mut<int32_t>(nsa_cache_seqlens_dst0, "nsa_cache_seqlens_dst0"),
|
||||
.dsa_cache_seqlens = unwrap_data_ptr_mut<int32_t>(dsa_cache_seqlens_dst0, "dsa_cache_seqlens_dst0"),
|
||||
.seqlens_expanded = nullptr,
|
||||
.nsa_cu_seqlens_k = unwrap_data_ptr_mut<int32_t>(nsa_cu_seqlens_k_dst0, "nsa_cu_seqlens_k_dst0"),
|
||||
.dsa_cu_seqlens_k = unwrap_data_ptr_mut<int32_t>(dsa_cu_seqlens_k_dst0, "dsa_cu_seqlens_k_dst0"),
|
||||
.real_page_table = unwrap_optional_data_ptr_mut<int32_t>(real_page_table_dst0, "real_page_table_dst0"),
|
||||
.flashmla_num_splits =
|
||||
unwrap_optional_data_ptr_mut<int32_t>(flashmla_num_splits_dst0, "flashmla_num_splits_dst0"),
|
||||
@@ -674,9 +674,9 @@ struct FusedMetadataCopyMultiKernel {
|
||||
.cache_seqlens = unwrap_data_ptr_mut<int32_t>(cache_seqlens_dst1, "cache_seqlens_dst1"),
|
||||
.cu_seqlens_k = unwrap_data_ptr_mut<int32_t>(cu_seqlens_k_dst1, "cu_seqlens_k_dst1"),
|
||||
.page_table_1 = unwrap_data_ptr_mut<int32_t>(page_table_1_dst1, "page_table_1_dst1"),
|
||||
.nsa_cache_seqlens = unwrap_data_ptr_mut<int32_t>(nsa_cache_seqlens_dst1, "nsa_cache_seqlens_dst1"),
|
||||
.dsa_cache_seqlens = unwrap_data_ptr_mut<int32_t>(dsa_cache_seqlens_dst1, "dsa_cache_seqlens_dst1"),
|
||||
.seqlens_expanded = nullptr,
|
||||
.nsa_cu_seqlens_k = unwrap_data_ptr_mut<int32_t>(nsa_cu_seqlens_k_dst1, "nsa_cu_seqlens_k_dst1"),
|
||||
.dsa_cu_seqlens_k = unwrap_data_ptr_mut<int32_t>(dsa_cu_seqlens_k_dst1, "dsa_cu_seqlens_k_dst1"),
|
||||
.real_page_table = unwrap_optional_data_ptr_mut<int32_t>(real_page_table_dst1, "real_page_table_dst1"),
|
||||
.flashmla_num_splits =
|
||||
unwrap_optional_data_ptr_mut<int32_t>(flashmla_num_splits_dst1, "flashmla_num_splits_dst1"),
|
||||
@@ -688,9 +688,9 @@ struct FusedMetadataCopyMultiKernel {
|
||||
.cache_seqlens = unwrap_data_ptr_mut<int32_t>(cache_seqlens_dst2, "cache_seqlens_dst2"),
|
||||
.cu_seqlens_k = unwrap_data_ptr_mut<int32_t>(cu_seqlens_k_dst2, "cu_seqlens_k_dst2"),
|
||||
.page_table_1 = unwrap_data_ptr_mut<int32_t>(page_table_1_dst2, "page_table_1_dst2"),
|
||||
.nsa_cache_seqlens = unwrap_data_ptr_mut<int32_t>(nsa_cache_seqlens_dst2, "nsa_cache_seqlens_dst2"),
|
||||
.dsa_cache_seqlens = unwrap_data_ptr_mut<int32_t>(dsa_cache_seqlens_dst2, "dsa_cache_seqlens_dst2"),
|
||||
.seqlens_expanded = nullptr,
|
||||
.nsa_cu_seqlens_k = unwrap_data_ptr_mut<int32_t>(nsa_cu_seqlens_k_dst2, "nsa_cu_seqlens_k_dst2"),
|
||||
.dsa_cu_seqlens_k = unwrap_data_ptr_mut<int32_t>(dsa_cu_seqlens_k_dst2, "dsa_cu_seqlens_k_dst2"),
|
||||
.real_page_table = unwrap_optional_data_ptr_mut<int32_t>(real_page_table_dst2, "real_page_table_dst2"),
|
||||
.flashmla_num_splits =
|
||||
unwrap_optional_data_ptr_mut<int32_t>(flashmla_num_splits_dst2, "flashmla_num_splits_dst2"),
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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"],
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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}."
|
||||
)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -17,7 +17,7 @@ if TYPE_CHECKING:
|
||||
class StateType(str, enum.Enum):
|
||||
MAMBA = "mamba"
|
||||
SWA = "swa"
|
||||
NSA = "nsa"
|
||||
DSA = "dsa"
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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],
|
||||
|
||||
@@ -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."
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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 (
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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(
|
||||
|
||||
+2
-2
@@ -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
|
||||
)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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")
|
||||
@@ -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
|
||||
)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -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"],
|
||||
)
|
||||
@@ -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)
|
||||
@@ -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")
|
||||
File diff suppressed because it is too large
Load Diff
@@ -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")
|
||||
@@ -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,
|
||||
)
|
||||
@@ -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
|
||||
File diff suppressed because it is too large
Load Diff
@@ -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,
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -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():
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -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
|
||||
|
||||
+19
-19
@@ -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)[
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
):
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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",
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user