From 7d757d6f17ef3f48a3eb41257e3baecd89bf6e46 Mon Sep 17 00:00:00 2001 From: Baizhou Zhang Date: Wed, 7 Jan 2026 14:00:16 +0800 Subject: [PATCH] Clean Some Environment Variables for DeepSeek V32 (#15938) --- docs/references/environment_variables.md | 10 ++++ python/sglang/srt/environ.py | 4 ++ .../layers/attention/nsa/dequant_k_cache.py | 9 +--- .../srt/layers/attention/nsa/nsa_indexer.py | 4 +- .../srt/layers/attention/nsa/quant_k_cache.py | 48 +++---------------- .../sglang/srt/layers/attention/nsa/utils.py | 24 ---------- .../srt/layers/attention/nsa_backend.py | 27 ++++------- python/sglang/srt/server_args.py | 21 ++++---- 8 files changed, 39 insertions(+), 108 deletions(-) diff --git a/docs/references/environment_variables.md b/docs/references/environment_variables.md index cac206abc..8621dcc78 100644 --- a/docs/references/environment_variables.md +++ b/docs/references/environment_variables.md @@ -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 | 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 | diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 6d3728ede..e757f9d93 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -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) diff --git a/python/sglang/srt/layers/attention/nsa/dequant_k_cache.py b/python/sglang/srt/layers/attention/nsa/dequant_k_cache.py index b6087eecd..5c69d135b 100644 --- a/python/sglang/srt/layers/attention/nsa/dequant_k_cache.py +++ b/python/sglang/srt/layers/attention/nsa/dequant_k_cache.py @@ -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, diff --git a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py index c9e82e4b1..7d19d0f90 100644 --- a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py +++ b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py @@ -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 diff --git a/python/sglang/srt/layers/attention/nsa/quant_k_cache.py b/python/sglang/srt/layers/attention/nsa/quant_k_cache.py index a7dc745c2..5454071b8 100644 --- a/python/sglang/srt/layers/attention/nsa/quant_k_cache.py +++ b/python/sglang/srt/layers/attention/nsa/quant_k_cache.py @@ -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) diff --git a/python/sglang/srt/layers/attention/nsa/utils.py b/python/sglang/srt/layers/attention/nsa/utils.py index 94cb96e7c..b9a3c7c6b 100644 --- a/python/sglang/srt/layers/attention/nsa/utils.py +++ b/python/sglang/srt/layers/attention/nsa/utils.py @@ -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) diff --git a/python/sglang/srt/layers/attention/nsa_backend.py b/python/sglang/srt/layers/attention/nsa_backend.py index 5f26347c8..b464f3f8d 100644 --- a/python/sglang/srt/layers/attention/nsa_backend.py +++ b/python/sglang/srt/layers/attention/nsa_backend.py @@ -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, diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index a2b26e0e0..bcdad4e3e 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -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"