[Deepseek V3.2] Fix accuracy bug in the Indexer (#12583)

Signed-off-by: Hao Lu <14827759+hlu1@users.noreply.github.com>

.github/workflows/pr-test.yml

@@ -868,7 +868,7 @@ jobs:
           CUSTOM_BUILD_SGL_KERNEL=${{needs.check-changes.outputs.sgl_kernel}} bash scripts/ci/ci_install_dependency.sh
 
       - name: Run test
-        timeout-minutes: 20
+        timeout-minutes: 30
         run: |
           cd test/srt
           python3 run_suite.py --suite per-commit-8-gpu-h200-deepseek-v32

python/sglang/srt/layers/attention/nsa/nsa_indexer.py

@@ -345,7 +345,10 @@ class Indexer(CustomOp):
         k_scale_list = []
         ks_list = []
         ke_list = []
-        offset = 0
+
+        q_offset = 0
+        k_offset = 0
+
         seq_lens_expanded = metadata.get_seqlens_expanded()
         block_tables = metadata.get_page_table_64()
 
@@ -368,13 +371,17 @@ class Indexer(CustomOp):
                 block_tables[i],
             )
             extend_seq_len = forward_batch.extend_seq_lens_cpu[i]
-            ks = torch.full((extend_seq_len,), offset, dtype=torch.int32, device="cuda")
+            ks = torch.full(
-            ke = ks + seq_lens_expanded[offset : offset + extend_seq_len]
+                (extend_seq_len,), k_offset, dtype=torch.int32, device="cuda"
+            )
+            ke = ks + seq_lens_expanded[q_offset : q_offset + extend_seq_len]
             k_fp8_list.append(k_fp8)
             k_scale_list.append(k_scale)
             ks_list.append(ks)
             ke_list.append(ke)
-            offset += extend_seq_len
+
+            q_offset += extend_seq_len
+            k_offset += seq_len
 
         k_fp8 = torch.cat(k_fp8_list, dim=0).view(torch.float8_e4m3fn)
         k_scale = torch.cat(k_scale_list, dim=0).view(torch.float32).squeeze(-1)
@@ -382,21 +389,38 @@ class Indexer(CustomOp):
         ks = torch.cat(ks_list, dim=0)
         ke = torch.cat(ke_list, dim=0)
 
+        # Suppose there are two requests, with extend_seq_len = [3, 2]
+        # and seq_lens = [10, 4]
+        # The logits matrix looks like this, with * representing the valid logits
+        # and - representing the invalid logits:
+        #
+        #  ********--|----
+        #  *********-|----
+        #  **********|----
+        #  ----------|***-
+        #  ----------|****
+        #
+        # ks = [0, 0, 0, 10, 10]
+        # ke = [8, 9, 10, 13, 14]
+
         logits = deep_gemm.fp8_mqa_logits(
-            q_fp8[:offset],
+            q_fp8[:q_offset],
             kv_fp8,
-            weights[:offset],
+            weights[:q_offset],
             ks,
             ke,
             clean_logits=False,
         )
+
         token_nums, _, _ = q_fp8.shape
         assert logits.shape[0] == len(seq_lens_expanded)
-        raw_topk_result = metadata.topk_transform(logits, self.index_topk)
+        assert logits.shape[1] == k_offset
+
+        raw_topk_result = metadata.topk_transform(logits, self.index_topk, ks=ks)
         topk_result = torch.full(
             (token_nums, self.index_topk), -1, device=q_fp8.device, dtype=torch.int32
         )
-        topk_result[:offset] = raw_topk_result
+        topk_result[:q_offset] = raw_topk_result
         return topk_result
 
     def _forward_cuda_k_only(

python/sglang/srt/layers/attention/nsa_backend.py

@@ -140,6 +140,7 @@ class NSAIndexerMetadata(BaseIndexerMetadata):
         self,
         logits: torch.Tensor,
         topk: int,
+        ks: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
         from sgl_kernel import (
             fast_topk_transform_fused,
@@ -148,7 +149,9 @@ class NSAIndexerMetadata(BaseIndexerMetadata):
         )
 
         if not NSA_FUSE_TOPK:
-            return fast_topk_v2(logits, self.get_seqlens_expanded(), topk)
+            return fast_topk_v2(
+                logits, self.get_seqlens_expanded(), topk, row_starts=ks
+            )
         elif self.topk_transform_method == TopkTransformMethod.PAGED:
             # NOTE(dark): if fused, we return a transformed page table directly
             return fast_topk_transform_fused(
@@ -157,6 +160,7 @@ class NSAIndexerMetadata(BaseIndexerMetadata):
                 page_table_size_1=self.attn_metadata.page_table_1,
                 cu_seqlens_q=self.attn_metadata.cu_seqlens_q,
                 topk=topk,
+                row_starts=ks,
             )
         elif self.topk_transform_method == TopkTransformMethod.RAGGED:
             return fast_topk_transform_ragged_fused(
@@ -164,6 +168,7 @@ class NSAIndexerMetadata(BaseIndexerMetadata):
                 lengths=self.get_seqlens_expanded(),
                 topk_indices_offset=self.attn_metadata.topk_indices_offset,
                 topk=topk,
+                row_starts=ks,
             )
         else:
             assert False, f"Unsupported {self.topk_transform_method = }"

test/srt/test_deepseek_v32_basic.py

@@ -44,7 +44,7 @@ class TestDeepseekV32Basic(CustomTestCase):
         self,
     ):  # Append an "a" to make this test run first (alphabetically) to warm up the server
         args = SimpleNamespace(
-            num_shots=8,
+            num_shots=20,
             data_path=None,
             num_questions=1400,
             parallel=1400,

test/srt/test_deepseek_v32_mtp.py

@@ -15,13 +15,13 @@ from sglang.test.test_utils import (
     write_github_step_summary,
 )
 
-FULL_DEEPSEEK_V3_MODEL_PATH = "deepseek-ai/DeepSeek-V3.2-Exp"
+FULL_DEEPSEEK_V32_MODEL_PATH = "deepseek-ai/DeepSeek-V3.2-Exp"
 
 
 class TestDeepseekV32MTP(CustomTestCase):
     @classmethod
     def setUpClass(cls):
-        cls.model = FULL_DEEPSEEK_V3_MODEL_PATH
+        cls.model = FULL_DEEPSEEK_V32_MODEL_PATH
         cls.base_url = DEFAULT_URL_FOR_TEST
         other_args = [
             "--trust-remote-code",
@@ -82,7 +82,7 @@ class TestDeepseekV32MTP(CustomTestCase):
                 f"{avg_spec_accept_length=:.2f}\n"
             )
             self.assertGreater(metrics["accuracy"], 0.935)
-            self.assertGreater(avg_spec_accept_length, 2.9)
+            self.assertGreater(avg_spec_accept_length, 2.7)
 
     def test_bs_1_speed(self):
         args = BenchArgs(port=int(self.base_url.split(":")[-1]), max_new_tokens=2048)
@@ -97,7 +97,7 @@ class TestDeepseekV32MTP(CustomTestCase):
                 f"{speed=:.2f} token/s\n"
             )
 
-            self.assertGreater(acc_length, 2.9)
+            self.assertGreater(acc_length, 2.7)
             self.assertGreater(speed, 75)
 
 

test/srt/test_deepseek_v32_nsabackend.py

@@ -50,7 +50,7 @@ class TestDeepseekV32NasBackend_flashmla(CustomTestCase):
         self,
     ):  # Append an "a" to make this test run first (alphabetically) to warm up the server
         args = SimpleNamespace(
-            num_shots=8,
+            num_shots=20,
             data_path=None,
             num_questions=1400,
             parallel=1400,
@@ -102,7 +102,57 @@ class TestDeepseekV32NasBackend_fa3(CustomTestCase):
         self,
     ):  # Append an "a" to make this test run first (alphabetically) to warm up the server
         args = SimpleNamespace(
-            num_shots=8,
+            num_shots=20,
+            data_path=None,
+            num_questions=1400,
+            parallel=1400,
+            max_new_tokens=512,
+            host="http://127.0.0.1",
+            port=int(self.base_url.split(":")[-1]),
+        )
+        metrics = run_eval_few_shot_gsm8k(args)
+        print(f"{metrics=}")
+
+        if is_in_ci():
+            write_github_step_summary(
+                f"### test_gsm8k (deepseek-v3)\n" f'{metrics["accuracy"]=:.3f}\n'
+            )
+            self.assertGreater(metrics["accuracy"], 0.935)
+
+
+class TestDeepseekV32NasBackend_fp8kvcache(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.model = DEEPSEEK_V32_MODEL_PATH
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        other_args = [
+            "--trust-remote-code",
+            "--attention-backend",
+            "nsa",
+            "--kv-cache-dtype",
+            "fp8_e4m3",
+            "--tp",
+            "8",
+            "--dp",
+            "8",
+            "--enable-dp-attention",
+        ]
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=other_args,
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+
+    def test_a_gsm8k(
+        self,
+    ):  # Append an "a" to make this test run first (alphabetically) to warm up the server
+        args = SimpleNamespace(
+            num_shots=20,
             data_path=None,
             num_questions=1400,
             parallel=1400,