chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[V32/GLM5] Control the threshold of applying dense attention with an environ (#20062)

Browse Source
This commit is contained in:
Baizhou Zhang
2026-03-09 14:36:10 -07:00
committed by GitHub
parent d39ed074cf
commit be63f982b7
6 changed files with 32 additions and 59 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
docs/references/environment_variables.md
View File

@@ -91,6 +91,8 @@ SGLang supports various environment variables that can be used to configure its
| --- | --- | --- |
| `SGLANG_NSA_FUSE_TOPK` | Fuse the operation of picking topk logits and picking topk indices from page table | `true` |
| `SGLANG_NSA_ENABLE_MTP_PRECOMPUTE_METADATA` | Precompute metadata that can be shared among different draft steps when MTP is enabled | `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` |
## Memory Management

3
python/sglang/srt/environ.py
View File

@@ -378,8 +378,7 @@ class Envs:
SGLANG_NSA_FUSE_TOPK = EnvBool(True)
SGLANG_NSA_ENABLE_MTP_PRECOMPUTE_METADATA = EnvBool(True)
SGLANG_USE_FUSED_METADATA_COPY = EnvBool(True)
SGLANG_VERIFY_FUSED_METADATA_COPY = EnvBool(False)
SGLANG_NSA_FORCE_MLA = EnvBool(False)
SGLANG_NSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD = EnvInt(2048)
# sgl-kernel
SGLANG_SKIP_SGL_KERNEL_VERSION_CHECK = EnvBool(False)

49
python/sglang/srt/layers/attention/nsa_backend.py
View File

@@ -16,10 +16,6 @@ from sglang.srt.layers.attention.nsa.nsa_backend_mtp_precompute import (
compute_cu_seqlens,
)
from sglang.srt.layers.attention.nsa.nsa_indexer import BaseIndexerMetadata
from sglang.srt.layers.attention.nsa.nsa_mtp_verification import (
verify_multi_backend_fused_metadata_copy,
verify_single_backend_fused_metadata_copy,
)
from sglang.srt.layers.attention.nsa.quant_k_cache import quantize_k_cache
from sglang.srt.layers.attention.nsa.transform_index import (
transform_index_page_table_decode,
@@ -71,15 +67,10 @@ else:
# Reuse this workspace buffer across all NSA backend instances
global_workspace_buffer = None
# Control whether to use fused metadata copy kernel (default: enabled)
# Control whether to use fused metadata copy kernel for cuda graph replay (default: enabled)
# Set SGLANG_USE_FUSED_METADATA_COPY=0 or false to disable
_USE_FUSED_METADATA_COPY = envs.SGLANG_USE_FUSED_METADATA_COPY.get() and not _is_hip
# Control whether to verify fused metadata copy against individual copies (default: disabled)
# Set SGLANG_VERIFY_FUSED_METADATA_COPY=1 or true to enable verification
# This will crash with detailed error message if any inconsistency is detected
_VERIFY_FUSED_METADATA_COPY = envs.SGLANG_VERIFY_FUSED_METADATA_COPY.get()
@dataclass(frozen=True)
class NSAFlashMLAMetadata:
@@ -317,8 +308,6 @@ class NativeSparseAttnBackend(
self.req_to_token = model_runner.req_to_token_pool.req_to_token
self.use_mha: bool = False
# Force NSA prefill to use MLA (i.e. disable MHA_ONE_SHOT), controlled by env var.
self._force_attn_forward_mla: bool = envs.SGLANG_NSA_FORCE_MLA.get()
self.nsa_prefill_impl: _NSA_IMPL_T = (
model_runner.server_args.nsa_prefill_backend
)
@@ -1182,18 +1171,6 @@ class NativeSparseAttnBackend(
# Successfully used fused kernel
fused_kernel_succeeded = True
# Verification: compare fused kernel results against individual copies
if _VERIFY_FUSED_METADATA_COPY:
verify_single_backend_fused_metadata_copy(
metadata=metadata,
precomputed=precomputed,
forward_mode=forward_mode,
bs=bs,
flashmla_num_splits_src=flashmla_num_splits_src,
flashmla_metadata_src=flashmla_metadata_src,
flashmla_num_splits_dst=flashmla_num_splits_dst,
flashmla_metadata_dst=flashmla_metadata_dst,
)
except ImportError:
print(
"Warning: Fused metadata copy kernel not available, falling back to individual copies."
@@ -2058,19 +2035,13 @@ class NativeSparseAttnBackend(
sum_seq_lens = sum(forward_batch.seq_lens_cpu)
device_sm = get_device_sm()
# when nsa prefill impl is trtllm, use its max chunk capacity as mha max kv len
mha_max_kv_len = (
forward_batch.get_max_chunk_capacity()
if self.nsa_prefill_impl == "trtllm"
else self.nsa_index_topk
)
# Requirements: H200/B200, short sequences, supported dtype, fits in chunk
self.use_mha = (
(
device_sm == 90 or (device_sm >= 100 and device_sm < 110)
) # SM90/SM100 only
and max_kv_len <= mha_max_kv_len # Short enough for MHA
and max_kv_len
<= envs.SGLANG_NSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD.get() # Short enough for MHA
and forward_batch.token_to_kv_pool.dtype
in [torch.bfloat16, torch.float8_e4m3fn]
and sum_seq_lens
@@ -2079,8 +2050,6 @@ class NativeSparseAttnBackend(
)
else:
self.use_mha = False # Decode/verify always use MLA
if self._force_attn_forward_mla:
self.use_mha = False
# Set MLA implementation only if not using MHA
if not self.use_mha and self.enable_auto_select_prefill_impl:
@@ -2306,18 +2275,6 @@ class NativeSparseAttnMultiStepBackend:
precomputed.seqlens_expanded_size,
)
# Verification: compare fused kernel results against individual copies
if _VERIFY_FUSED_METADATA_COPY:
verify_multi_backend_fused_metadata_copy(
metadata0=metadata0,
metadata1=metadata1,
metadata2=metadata2,
precomputed=precomputed,
bs=bs,
flashmla_num_splits_src=flashmla_num_splits_src,
flashmla_metadata_src=flashmla_metadata_src,
)
# Copy remaining backends one by one (if > 3 backends)
for i in range(3, self.speculative_num_steps - 1):
self.attn_backends[

29
python/sglang/srt/server_args.py
View File

@@ -1406,12 +1406,35 @@ class ServerArgs:
"GlmMoeDsaForCausalLM",
]:
# Set attention backend for DeepSeek
if is_deepseek_nsa(hf_config): # DeepSeek 3.2, GlmMoeDsaForCausalLM
if is_deepseek_nsa(hf_config): # DeepSeek 3.2/GLM 5
if model_arch == "GlmMoeDsaForCausalLM" and is_blackwell_supported():
envs.SGLANG_NSA_FORCE_MLA.set(True)
envs.SGLANG_NSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD.set(0)
logger.warning(
"Force NSA prefill to use MLA (i.e. disable MHA_ONE_SHOT) for GlmMoeDsaForCausalLM on Blackwell."
"Force NSA prefill to use sparse MLA (i.e. disable MHA_ONE_SHOT) for GlmMoeDsaForCausalLM on Blackwell."
)
else:
if envs.SGLANG_NSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD.is_set():
logger.warning(
f"Dense attention kv len threshold is manually set to {envs.SGLANG_NSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD.get()} for DSA. Caution: This may cause performance regression if the threshold is larger than the index topk of model."
)
else:
# When threshold is not manually set, set it to the index topk of model
from sglang.srt.configs.model_config import get_nsa_index_topk
envs.SGLANG_NSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD.set(
get_nsa_index_topk(hf_config)
)
logger.warning(
f"Set dense attention kv len threshold to model index_topk={envs.SGLANG_NSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD.get()} for DeepSeek with DSA."
)
if self.nsa_prefill_backend == "trtllm":
# We temporarily set the threshold to 128k to avoid IMA error. Should be removed after supporting flashmla prefill impl with trtllm decode impl.
envs.SGLANG_NSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD.set(
128 * 1024
)
logger.warning(
"TRTLLM sparse MLA kernel requires MHA as prefill impl, the threshold for dense attention is overridden. This will be fixed in the future."
)
if self.is_attention_backend_not_set():
self.attention_backend = "nsa"
logger.info("Use nsa attention backend for DeepSeek with DSA.")

4
test/registered/quant/test_deepseek_v32_fp4_4gpu.py
View File

@@ -34,8 +34,6 @@ class TestDeepseekV32FP4DP(CustomTestCase):
"flashinfer_trtllm",
"--quantization",
"modelopt_fp4",
"--kv-cache-dtype",
"fp8_e4m3",
"--tool-call-parser",
"deepseekv32",
"--reasoning-parser",
@@ -103,8 +101,6 @@ class TestDeepseekV32FP4TP(CustomTestCase):
"flashinfer_trtllm",
"--quantization",
"modelopt_fp4",
"--kv-cache-dtype",
"fp8_e4m3",
"--tool-call-parser",
"deepseekv32",
"--reasoning-parser",

4
test/registered/quant/test_deepseek_v32_fp4_mtp_4gpu.py
View File

@@ -39,8 +39,6 @@ class TestDeepseekV32FP4DPSpecV2(CustomTestCase):
"flashinfer_trtllm",
"--quantization",
"modelopt_fp4",
"--kv-cache-dtype",
"fp8_e4m3",
"--tool-call-parser",
"deepseekv32",
"--reasoning-parser",
@@ -131,8 +129,6 @@ class TestDeepseekV32FP4TPSpecV2(CustomTestCase):
"flashinfer_trtllm",
"--quantization",
"modelopt_fp4",
"--kv-cache-dtype",
"fp8_e4m3",
"--tool-call-parser",
"deepseekv32",
"--reasoning-parser",