From f78c414b64cf1938a51e5f7171e123f0ece5c968 Mon Sep 17 00:00:00 2001 From: laoyao0822 Date: Wed, 10 Jun 2026 07:00:43 +0800 Subject: [PATCH] Skip unused FP8 NSA K-cache dequant work High cache-hit NSA sparse attention only gathers rows selected by topk_indices, but the FP8 RAGGED path dequantized the whole materialized K buffer first. Wire the syh in-place topk-only dequant path behind an explicit env gate so the sparse path can dequant referenced rows at their original row ids while keeping topk_indices unchanged. The old full-dequant path remains the default and the verification gate stays available for small correctness checks only. Constraint: Production runs should enable only SGLANG_NSA_DEQUANT_ONLY_TOPK=1; VERIFY also runs the full-dequant reference and is too expensive for perf tests. Rejected: Compact/remap topk buffer | data-dependent shape and remap add synchronization/aliasing risk; syh final in-place contract avoids both. Confidence: medium Scope-risk: moderate Directive: Do not enable SGLANG_NSA_DEQUANT_ONLY_TOPK_VERIFY in throughput runs; use it only for small correctness probes. Tested: python -m py_compile python/sglang/srt/environ.py python/sglang/srt/layers/attention/nsa_backend.py test/registered/unit/layers/test_nsa_dequant_only_topk.py Tested: git diff --check Tested: remote g0034 container PYTHONPATH=python python -m pytest -q test/registered/unit/layers/test_nsa_dequant_only_topk.py -> 2 passed Tested: remote g0034 CUDA smoke for tai_kernel.nsa_prefill.nsa_dequant_topk_inplace topk rows matched torch reference Not-tested: Full GSM8K or replay ETE with SGLANG_NSA_DEQUANT_ONLY_TOPK=1 --- python/sglang/srt/environ.py | 5 ++ .../srt/layers/attention/nsa_backend.py | 55 ++++++++++++++++++- .../unit/layers/test_nsa_dequant_only_topk.py | 44 +++++++++++++++ 3 files changed, 101 insertions(+), 3 deletions(-) create mode 100644 test/registered/unit/layers/test_nsa_dequant_only_topk.py diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 17c837ec5..47b6c8f6e 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -422,6 +422,11 @@ class Envs: # NSA Backend SGLANG_NSA_FUSE_TOPK = EnvBool(True) + # Dequantize only the KV rows referenced by sparse-attention topk_indices + # instead of full materialized FP8 K cache. Default OFF; VERIFY runs the full + # dequant reference and compares gathered rows, so it is for validation only. + SGLANG_NSA_DEQUANT_ONLY_TOPK = EnvBool(False) + SGLANG_NSA_DEQUANT_ONLY_TOPK_VERIFY = EnvBool(False) SGLANG_NSA_ENABLE_MTP_PRECOMPUTE_METADATA = EnvBool(True) SGLANG_USE_FUSED_METADATA_COPY = EnvBool(True) SGLANG_NSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD = EnvInt(2048) diff --git a/python/sglang/srt/layers/attention/nsa_backend.py b/python/sglang/srt/layers/attention/nsa_backend.py index b0c0af69c..ae34b2f4a 100644 --- a/python/sglang/srt/layers/attention/nsa_backend.py +++ b/python/sglang/srt/layers/attention/nsa_backend.py @@ -2641,9 +2641,17 @@ class NativeSparseAttnBackend( nvtx_source="mla.ragged_partial_current_sync", ) ) - kv_cache = dequantize_k_cache_paged( - kv_cache, page_table_1_flattened - ) + if envs.SGLANG_NSA_DEQUANT_ONLY_TOPK.get(): + kv_cache = self._dequant_topk_inplace( + kv_cache, + page_table_1_flattened, + topk_indices, + verify=envs.SGLANG_NSA_DEQUANT_ONLY_TOPK_VERIFY.get(), + ) + else: + kv_cache = dequantize_k_cache_paged( + kv_cache, page_table_1_flattened + ) else: kv_cache = _cat([k, k_rope], dim=-1) page_table_1 = topk_indices @@ -2720,6 +2728,47 @@ class NativeSparseAttnBackend( index_prefetcher.wait_attention_window() return attn_output + def _dequant_topk_inplace( + self, + kv_cache_fp8: torch.Tensor, + page_table_1_flattened: torch.Tensor, + topk_indices: torch.Tensor, + verify: bool = False, + ) -> torch.Tensor: + """Dequantize only topk-referenced FP8 KV rows at their original row ids. + + The output keeps the full logical row count ``[num_kv_rows, 1, 576]`` so + sparse attention can keep using the unchanged ``topk_indices``. This is + the final syh in-place contract, not the earlier compact/remap variant: + no data-dependent output shape, no topk remap, and no CPU synchronization. + """ + try: + from tai_kernel.nsa_prefill import nsa_dequant_topk_inplace + except ImportError: + nsa_dequant_topk_inplace = None + + if nsa_dequant_topk_inplace is not None: + out = nsa_dequant_topk_inplace( + kv_cache_fp8, + page_table_1_flattened, + topk_indices, + ) + else: + out = dequantize_k_cache_paged(kv_cache_fp8, page_table_1_flattened) + + if verify: + kv_ref = dequantize_k_cache_paged(kv_cache_fp8, page_table_1_flattened) + flat = topk_indices.reshape(-1) + valid = (flat >= 0) & (flat < int(page_table_1_flattened.shape[0])) + rows = flat[valid].long() + ref_rows = kv_ref.reshape(-1, kv_ref.shape[-1])[rows] + new_rows = out.reshape(-1, out.shape[-1])[rows] + assert torch.equal(ref_rows, new_rows), ( + "[SGLANG_NSA_DEQUANT_ONLY_TOPK_VERIFY] in-place dequant != full " + f"dequant (N={int(page_table_1_flattened.shape[0])})" + ) + return out + def forward_decode( self, q: torch.Tensor, diff --git a/test/registered/unit/layers/test_nsa_dequant_only_topk.py b/test/registered/unit/layers/test_nsa_dequant_only_topk.py new file mode 100644 index 000000000..9dcdde0a4 --- /dev/null +++ b/test/registered/unit/layers/test_nsa_dequant_only_topk.py @@ -0,0 +1,44 @@ +import builtins +import unittest +from unittest.mock import patch + +import torch + +from sglang.srt.environ import envs +from sglang.srt.layers.attention.nsa_backend import NativeSparseAttnBackend +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=1, suite="stage-a-test-cpu") + + +class TestNSADequantOnlyTopK(unittest.TestCase): + def test_env_gate_is_defined_and_default_off(self): + self.assertFalse(envs.SGLANG_NSA_DEQUANT_ONLY_TOPK.get()) + self.assertFalse(envs.SGLANG_NSA_DEQUANT_ONLY_TOPK_VERIFY.get()) + + def test_dequant_topk_inplace_falls_back_to_full_dequant_when_kernel_missing(self): + backend = object.__new__(NativeSparseAttnBackend) + kv_cache = torch.empty((4, 656), dtype=torch.uint8) + page_table = torch.tensor([3, 1, 2, 0], dtype=torch.int32) + topk_indices = torch.tensor([[0, 2, -1], [1, 3, 2]], dtype=torch.int32) + expected = torch.empty((4, 1, 576), dtype=torch.bfloat16) + + real_import = builtins.__import__ + + def fake_import(name, *args, **kwargs): + if name == "tai_kernel.nsa_prefill": + raise ImportError("tai kernel intentionally unavailable") + return real_import(name, *args, **kwargs) + + with patch("builtins.__import__", side_effect=fake_import), patch( + "sglang.srt.layers.attention.nsa_backend.dequantize_k_cache_paged", + return_value=expected, + ) as dequant: + out = backend._dequant_topk_inplace(kv_cache, page_table, topk_indices) + + self.assertIs(out, expected) + dequant.assert_called_once_with(kv_cache, page_table) + + +if __name__ == "__main__": + unittest.main()