chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 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

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.")