[V32/GLM5] Control the threshold of applying dense attention with an environ (#20062)

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[