diff --git a/docs/references/environment_variables.md b/docs/references/environment_variables.md index 872e7697f..a07a4fcc2 100644 --- a/docs/references/environment_variables.md +++ b/docs/references/environment_variables.md @@ -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 diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index f7f8257a4..d1388d159 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -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) diff --git a/python/sglang/srt/layers/attention/nsa_backend.py b/python/sglang/srt/layers/attention/nsa_backend.py index 2e76cf608..875bb5803 100644 --- a/python/sglang/srt/layers/attention/nsa_backend.py +++ b/python/sglang/srt/layers/attention/nsa_backend.py @@ -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[ diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index b1f39af79..6bfa8ecb3 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -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.") diff --git a/test/registered/quant/test_deepseek_v32_fp4_4gpu.py b/test/registered/quant/test_deepseek_v32_fp4_4gpu.py index dbfb5d641..4bdcd1367 100644 --- a/test/registered/quant/test_deepseek_v32_fp4_4gpu.py +++ b/test/registered/quant/test_deepseek_v32_fp4_4gpu.py @@ -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", diff --git a/test/registered/quant/test_deepseek_v32_fp4_mtp_4gpu.py b/test/registered/quant/test_deepseek_v32_fp4_mtp_4gpu.py index 7950f98f9..3a99f6ad9 100644 --- a/test/registered/quant/test_deepseek_v32_fp4_mtp_4gpu.py +++ b/test/registered/quant/test_deepseek_v32_fp4_mtp_4gpu.py @@ -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",