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

Clean Some Environment Variables for DeepSeek V32 (#15938)

Browse Source
This commit is contained in:
Baizhou Zhang
2026-01-07 14:00:16 +08:00
committed by GitHub
parent 5c04088b3a
commit 7d757d6f17
8 changed files with 39 additions and 108 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
docs/references/environment_variables.md
View File

@@ -60,6 +60,16 @@ SGLang supports various environment variables that can be used to configure its
| `SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK` | The maximum number of dispatched tokens on each GPU | `"128"` |
| `SGLANG_DEEPEP_LL_COMBINE_SEND_NUM_SMS` | Number of SMs used for DeepEP combine when single batch overlap is enabled | `"32"` |
## NSA 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` |
## Memory Management
| Environment Variable | Description | Default Value |

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

@@ -323,6 +323,10 @@ class Envs:
SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK = EnvInt(128)
SGLANG_DEEPEP_LL_COMBINE_SEND_NUM_SMS = EnvInt(32)
# NSA Backend
SGLANG_NSA_FUSE_TOPK = EnvBool(True)
SGLANG_NSA_ENABLE_MTP_PRECOMPUTE_METADATA = EnvBool(True)
# sgl-kernel
SGLANG_SKIP_SGL_KERNEL_VERSION_CHECK = EnvBool(False)

9
python/sglang/srt/layers/attention/nsa/dequant_k_cache.py
View File

@@ -2,17 +2,12 @@ import torch
import triton
import triton.language as tl
from sglang.srt.layers.attention.nsa.utils import NSA_DEQUANT_K_CACHE_FAST
def dequantize_k_cache(quant_k_cache):
if NSA_DEQUANT_K_CACHE_FAST:
return _dequantize_k_cache_fast_wrapped(quant_k_cache)
else:
return _dequantize_k_cache_slow(quant_k_cache)
return _dequantize_k_cache_fast_wrapped(quant_k_cache)
def _dequantize_k_cache_slow(
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,

4
python/sglang/srt/layers/attention/nsa/nsa_indexer.py
View File

@@ -25,7 +25,6 @@ if is_npu():
from sglang.srt.layers import deep_gemm_wrapper
from sglang.srt.layers.attention.nsa.utils import (
NSA_DUAL_STREAM,
cp_all_gather_rerange_output,
is_nsa_enable_prefill_cp,
is_nsa_prefill_cp_in_seq_split,
@@ -802,8 +801,7 @@ class Indexer(MultiPlatformOp):
)
enable_dual_stream = (
NSA_DUAL_STREAM
and self.alt_stream is not None
self.alt_stream is not None
and get_is_capture_mode()
and q_lora.shape[0] > 0
and q_lora.shape[0] <= DUAL_STREAM_TOKEN_THRESHOLD

48
python/sglang/srt/layers/attention/nsa/quant_k_cache.py
View File

@@ -2,15 +2,9 @@ import torch
import triton
import triton.language as tl
from sglang.srt.layers.attention.nsa.utils import NSA_QUANT_K_CACHE_FAST
def quantize_k_cache(cache_k):
# TODO upstream can skip concat([k_nope, k_pe]) since we split them here
if NSA_QUANT_K_CACHE_FAST:
return _quantize_k_cache_fast_wrapped(cache_k)
else:
return _quantize_k_cache_slow(cache_k)
return _quantize_k_cache_fast_wrapped(cache_k)
def quantize_k_cache_separate(
@@ -56,43 +50,13 @@ def quantize_k_cache_separate(
f"k_nope and k_rope must have same num_tokens, got {num_tokens} vs {k_rope_2d.shape[0]}"
)
# Call fast kernel that directly produces two separate outputs (single Triton kernel)
if NSA_QUANT_K_CACHE_FAST:
nope_part, rope_part = _quantize_k_cache_fast_separate(
k_nope=k_nope_2d, k_rope=k_rope_2d, group_size=tile_size
)
else:
# Fallback: use existing slow path with post-processing
cache_k_concat = torch.cat([k_nope_2d, k_rope_2d], dim=-1)
packed_output_4d = quantize_k_cache(cache_k_concat.unsqueeze(1).unsqueeze(1))
packed_output = packed_output_4d.squeeze(1).squeeze(1)
# Convert to uint8 bytes view
packed_bytes = packed_output.contiguous().view(torch.uint8)
# Strict byte-size validation
expected_total_bytes = 656 # 512 (nope_fp8) + 16 (scales) + 128 (rope_bf16)
if packed_bytes.shape[1] != expected_total_bytes:
raise ValueError(
f"Packed output has {packed_bytes.shape[1]} bytes, expected {expected_total_bytes}. "
f"Original dtype: {packed_output.dtype}, shape: {packed_output.shape}"
)
# Split into nope and rope parts
num_tiles = dim_nope // tile_size # 4
nope_part_bytes = dim_nope + num_tiles * 4 # 512 + 16 = 528
rope_part_bytes = 128
nope_part = packed_bytes[:, :nope_part_bytes].unsqueeze(1)
rope_part = packed_bytes[
:, nope_part_bytes : nope_part_bytes + rope_part_bytes
].unsqueeze(1)
return nope_part, rope_part
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_slow(
def _quantize_k_cache_ref(
input_k_cache: torch.Tensor, # (num_blocks, block_size, h_k, d)
dv: int = 512,
tile_size: int = 128,
@@ -375,7 +339,7 @@ if __name__ == "__main__":
device="cuda",
)
ref_quant = _quantize_k_cache_slow(input_k_cache)
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)

24
python/sglang/srt/layers/attention/nsa/utils.py
View File

@@ -15,35 +15,11 @@ from sglang.srt.layers.dp_attention import (
get_attention_tp_size,
)
from sglang.srt.server_args import get_global_server_args
from sglang.srt.utils import get_bool_env_var
if TYPE_CHECKING:
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
NSA_DUAL_STREAM = get_bool_env_var("SGLANG_NSA_DUAL_STREAM", "true")
NSA_FUSE_TOPK = get_bool_env_var("SGLANG_NSA_FUSE_TOPK", "true")
NSA_FLASHMLA_BACKEND_DECODE_COMPUTE_FP8 = get_bool_env_var(
"SGLANG_NSA_FLASHMLA_BACKEND_DECODE_COMPUTE_FP8", "true"
)
NSA_QUANT_K_CACHE_FAST = get_bool_env_var("SGLANG_NSA_QUANT_K_CACHE_FAST", "true")
NSA_DEQUANT_K_CACHE_FAST = get_bool_env_var("SGLANG_NSA_DEQUANT_K_CACHE_FAST", "true")
# Environment variable to control mtp precomputing of metadata for multi-step speculative decoding
NSA_ENABLE_MTP_PRECOMPUTE_METADATA = get_bool_env_var(
"SGLANG_NSA_ENABLE_MTP_PRECOMPUTE_METADATA", "true"
)
def print_nsa_bool_env_vars():
msg = ""
for k, v in globals().items():
if k.startswith("NSA_") and isinstance(v, bool):
msg += f"{k}={v} "
print(msg, flush=True)
def compute_nsa_seqlens(original_seq_lens, nsa_index_topk: int):
return original_seq_lens.clamp(max=nsa_index_topk)

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

@@ -22,9 +22,6 @@ from sglang.srt.layers.attention.nsa.transform_index import (
transform_index_page_table_prefill,
)
from sglang.srt.layers.attention.nsa.utils import (
NSA_ENABLE_MTP_PRECOMPUTE_METADATA,
NSA_FLASHMLA_BACKEND_DECODE_COMPUTE_FP8,
NSA_FUSE_TOPK,
can_nsa_prefill_cp_round_robin_split,
compute_nsa_seqlens,
is_nsa_enable_prefill_cp,
@@ -230,7 +227,7 @@ class NSAIndexerMetadata(BaseIndexerMetadata):
else:
page_table_size_1 = self.attn_metadata.page_table_1
if not NSA_FUSE_TOPK:
if not envs.SGLANG_NSA_FUSE_TOPK.get():
return fast_topk_v2(logits, seq_lens_topk, topk, row_starts=ks)
elif self.topk_transform_method == TopkTransformMethod.PAGED:
# NOTE(dark): if fused, we return a transformed page table directly
@@ -1222,12 +1219,6 @@ class NativeSparseAttnBackend(
assert q_rope is not None
kv_cache = forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id)
# when store in fp8 and compute in fp8, no need to convert dtype
if not (
NSA_FLASHMLA_BACKEND_DECODE_COMPUTE_FP8 and self.nsa_kv_cache_store_fp8
):
kv_cache = kv_cache.to(q.dtype)
if q_rope is not None:
q_nope = q.view(-1, layer.tp_q_head_num, layer.v_head_dim)
q_rope = q_rope.view(
@@ -1245,7 +1236,7 @@ class NativeSparseAttnBackend(
# NOTE(dark): here, we use page size = 1
topk_transform_method = self.get_topk_transform_method()
if NSA_FUSE_TOPK:
if envs.SGLANG_NSA_FUSE_TOPK.get():
page_table_1 = topk_indices
else:
if topk_transform_method == TopkTransformMethod.RAGGED:
@@ -1283,8 +1274,6 @@ class NativeSparseAttnBackend(
if q_rope is not None:
q_all = _concat_mla_absorb_q_general(q_nope, q_rope)
# NSA_FLASHMLA_BACKEND_DECODE_COMPUTE_FP8 has no effect here,
# because the flashmla_sparse kernel doesn't support fp8 compute
if topk_transform_method == TopkTransformMethod.RAGGED:
if any(forward_batch.extend_prefix_lens_cpu):
page_table_1_flattened = (
@@ -1384,7 +1373,7 @@ class NativeSparseAttnBackend(
if topk_indices is not None:
topk_indices = self._pad_topk_indices(topk_indices, q_nope.shape[0])
if NSA_FUSE_TOPK:
if envs.SGLANG_NSA_FUSE_TOPK.get():
page_table_1 = topk_indices
else:
page_table_1 = transform_index_page_table_decode(
@@ -1562,7 +1551,7 @@ class NativeSparseAttnBackend(
kv_cache = kv_cache.view(-1, self.real_page_size, 1, self.kv_cache_dim)
assert self.real_page_size == 64, "only page size 64 is supported"
if NSA_FLASHMLA_BACKEND_DECODE_COMPUTE_FP8 and not self.nsa_kv_cache_store_fp8:
if not self.nsa_kv_cache_store_fp8:
# inefficiently quantize the whole cache
kv_cache = quantize_k_cache(kv_cache)
@@ -1584,7 +1573,7 @@ class NativeSparseAttnBackend(
block_table=torch.empty(
(q_all.shape[0], 0), dtype=torch.int32, device=q_all.device
),
is_fp8_kvcache=NSA_FLASHMLA_BACKEND_DECODE_COMPUTE_FP8,
is_fp8_kvcache=True,
)
return o
@@ -1787,7 +1776,7 @@ class NativeSparseAttnBackend(
def get_topk_transform_method(self) -> TopkTransformMethod:
"""
NSA_FUSE_TOPK controls whether to fuse the topk transform into the topk kernel.
SGLANG_NSA_FUSE_TOPK controls whether to fuse the topk transform into the topk kernel.
This method is used to select the topk transform method which can be fused or unfused.
"""
if (
@@ -1819,7 +1808,7 @@ class NativeSparseAttnBackend(
num_q_tokens_per_head_k=seq_len_q * self.num_q_heads // 1,
num_heads_k=1,
num_heads_q=self.num_q_heads,
is_fp8_kvcache=NSA_FLASHMLA_BACKEND_DECODE_COMPUTE_FP8,
is_fp8_kvcache=True,
topk=self.nsa_index_topk,
)
@@ -1871,7 +1860,7 @@ class NativeSparseAttnMultiStepBackend:
def init_forward_metadata_replay_cuda_graph(
self, forward_batch: ForwardBatch, bs: int
):
if NSA_ENABLE_MTP_PRECOMPUTE_METADATA:
if envs.SGLANG_NSA_ENABLE_MTP_PRECOMPUTE_METADATA.get():
# Precompute metadata once (shared across all backends)
precomputed = self.attn_backends[0]._precompute_replay_metadata(
bs=bs,

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

@@ -1066,7 +1066,7 @@ class ServerArgs:
if is_deepseek_nsa(hf_config): # DeepSeek 3.2
if self.is_attention_backend_not_set():
self.attention_backend = "nsa"
logger.info("Use nsa attention backend for DeepSeek NSA.")
logger.info("Use nsa attention backend for DeepSeek with DSA.")
if not is_npu(): # CUDA GPU
if self.enable_nsa_prefill_context_parallel:
@@ -1100,12 +1100,12 @@ class ServerArgs:
# Pure TP and partial DP Attention mode is active for NSA, logging a warning
if self.dp_size < self.tp_size:
logger.warning(
f"NSA with TP mode is active, dp_size={self.dp_size}, tp_size={self.tp_size}, "
f"DSA with TP mode is active, dp_size={self.dp_size}, tp_size={self.tp_size}, "
f"attn_tp_size={self.tp_size}, attention weights will be sharded across {self.tp_size} ranks."
)
self.page_size = 64
logger.warning("Setting page size to 64 for DeepSeek NSA.")
logger.warning("Setting page size to 64 for DeepSeek DSA.")
# For Hopper, we support both bf16 and fp8 kv cache; for Blackwell, we support fp8 only currently
import torch
@@ -1114,34 +1114,29 @@ class ServerArgs:
if self.kv_cache_dtype == "auto":
self.kv_cache_dtype = "fp8_e4m3" if major >= 10 else "bfloat16"
logger.warning(
f"Setting KV cache dtype to {self.kv_cache_dtype} for DeepSeek NSA."
f"Setting KV cache dtype to {self.kv_cache_dtype} for DeepSeek DSA on SM{major} device."
)
if self.kv_cache_dtype == "bf16":
self.kv_cache_dtype = "bfloat16"
assert self.kv_cache_dtype in [
"bfloat16",
"fp8_e4m3",
], "DeepSeek NSA only supports bf16/bfloat16 or fp8_e4m3 kv_cache_dtype"
], "DeepSeek DSA only supports bf16/bfloat16 or fp8_e4m3 kv_cache_dtype"
if self.kv_cache_dtype == "fp8_e4m3":
# flashmla_auto dispatches to flashmla_sparse/flashmla_kv based on hardware and heuristics
self.nsa_prefill_backend = "flashmla_auto"
self.nsa_decode_backend = "flashmla_kv"
logger.warning(
"Setting NSA backend to flashmla_auto for prefill and flashmla_kv for decode for FP8 KV Cache."
"Setting DSA backend to flashmla_auto for prefill and flashmla_kv for decode for FP8 KV Cache."
)
else:
# set prefill/decode backends for Blackwell. The default settings are for Hopper.
# set prefill/decode backends to flashmla_sparse for Blackwell.
# The default settings (P=flashmla_sparse, D=fa3) are for Hopper.
if major >= 10:
self.nsa_prefill_backend = "flashmla_sparse"
self.nsa_decode_backend = "flashmla_sparse"
# Logging env vars for NSA
from sglang.srt.layers.attention.nsa.utils import (
print_nsa_bool_env_vars,
)
print_nsa_bool_env_vars()
if self.enable_nsa_prefill_context_parallel:
assert (
self.disaggregation_mode != "decode"