Clean Some Environment Variables for DeepSeek V32 (#15938)

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 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 |

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)
 

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)
-        return _dequantize_k_cache_fast_wrapped(quant_k_cache)
-    else:
-        return _dequantize_k_cache_slow(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,

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
+            self.alt_stream is not None
-            and 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

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
+    return _quantize_k_cache_fast_wrapped(cache_k)
-    if NSA_QUANT_K_CACHE_FAST:
-        return _quantize_k_cache_fast_wrapped(cache_k)
-    else:
-        return _quantize_k_cache_slow(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)
+    return _quantize_k_cache_fast_separate(
-    if NSA_QUANT_K_CACHE_FAST:
+        k_nope=k_nope_2d, k_rope=k_rope_2d, group_size=tile_size
-        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
 
 
 # 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)

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)
 

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,

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"