[Deepseek V3.2] Clean up MTP (#13236)

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

@@ -290,7 +290,10 @@ class Indexer(CustomOp):
         )
 
         blocksize = page_size
-        if forward_batch.forward_mode.is_target_verify():
+        if (
+            forward_batch.forward_mode.is_target_verify()
+            or forward_batch.forward_mode.is_draft_extend()
+        ):
             seqlens_32 = metadata.get_seqlens_expanded()
         else:
             seqlens_32 = metadata.get_seqlens_int32()
@@ -337,6 +340,8 @@ class Indexer(CustomOp):
         if TYPE_CHECKING:
             assert isinstance(forward_batch.token_to_kv_pool, NSATokenToKVPool)
 
+        assert forward_batch.forward_mode.is_extend_without_speculative()
+
         page_size = forward_batch.token_to_kv_pool.page_size
         assert page_size == 64, "only support page size 64"
         assert len(weights.shape) == 3
@@ -649,6 +654,7 @@ class Indexer(CustomOp):
             if (
                 forward_batch.forward_mode.is_decode_or_idle()
                 or forward_batch.forward_mode.is_target_verify()
+                or forward_batch.forward_mode.is_draft_extend()
             ):
                 topk_result = self._get_topk_paged(
                     forward_batch, layer_id, q_fp8, weights, metadata

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

@@ -137,6 +137,9 @@ class NSAIndexerMetadata(BaseIndexerMetadata):
     def get_seqlens_expanded(self) -> torch.Tensor:
         return self.attn_metadata.nsa_seqlens_expanded
 
+    def get_cu_seqlens_k(self) -> torch.Tensor:
+        return self.attn_metadata.cu_seqlens_k
+
     def topk_transform(
         self,
         logits: torch.Tensor,
@@ -304,8 +307,7 @@ class NativeSparseAttnBackend(AttentionBackend):
             cu_seqlens_q = self.get_device_int32_arange(batch_size + 1)
             seqlens_expanded = cache_seqlens_int32
         elif forward_batch.forward_mode.is_target_verify():
-            max_seqlen_q = self.speculative_num_draft_tokens
+            max_seqlen_q = 1
-            nsa_max_seqlen_q = self.speculative_num_draft_tokens
             cu_seqlens_q = torch.arange(
                 0,
                 batch_size * self.speculative_num_draft_tokens + 1,
@@ -338,6 +340,43 @@ class NativeSparseAttnBackend(AttentionBackend):
             page_table = torch.repeat_interleave(
                 page_table, repeats=self.speculative_num_draft_tokens, dim=0
             )
+        elif forward_batch.forward_mode.is_draft_extend():
+            assert (
+                forward_batch.extend_seq_lens_cpu is not None
+                and forward_batch.extend_seq_lens is not None
+                and forward_batch.extend_prefix_lens_cpu is not None
+            ), "All of them must not be None"
+
+            extend_seq_lens_cpu = forward_batch.extend_seq_lens_cpu
+            assert forward_batch.extend_seq_lens is not None
+
+            max_seqlen_q = 1
+            cu_seqlens_q = torch.arange(
+                0,
+                forward_batch.extend_num_tokens + 1,
+                1,
+                dtype=torch.int32,
+                device=device,
+            )
+            seqlens_expanded = torch.cat(
+                [
+                    torch.arange(
+                        kv_len - qo_len + 1,
+                        kv_len + 1,
+                        dtype=torch.int32,
+                        device=device,
+                    )
+                    for qo_len, kv_len in zip(
+                        forward_batch.extend_seq_lens_cpu,
+                        forward_batch.seq_lens_cpu.tolist(),
+                        strict=True,
+                    )
+                ]
+            )
+            page_table = torch.repeat_interleave(
+                page_table, repeats=forward_batch.extend_seq_lens, dim=0
+            )
+
         elif forward_batch.forward_mode.is_extend():
             assert (
                 forward_batch.extend_seq_lens_cpu is not None
@@ -518,7 +557,7 @@ class NativeSparseAttnBackend(AttentionBackend):
                 )
             else:
                 flashmla_metadata = None
-        elif forward_mode.is_target_verify():
+        elif forward_mode.is_target_verify() or forward_mode.is_draft_extend():
             cache_seqlens_int32 = (seq_lens + self.speculative_num_draft_tokens).to(
                 torch.int32
             )
@@ -576,64 +615,6 @@ class NativeSparseAttnBackend(AttentionBackend):
                 )
             else:
                 flashmla_metadata = None
-        elif forward_mode.is_draft_extend():
-            cache_seqlens_int32 = (seq_lens + self.speculative_num_draft_tokens).to(
-                torch.int32
-            )
-            cu_seqlens_k = compute_cu_seqlens(cache_seqlens_int32)
-            page_table_1 = self.decode_cuda_graph_metadata["page_table"][:bs, :]
-            max_seqlen_k = page_table_1.shape[1]
-
-            extend_seq_lens_cpu = [self.speculative_num_draft_tokens] * bs
-            extend_seq_lens = torch.full(
-                (bs,),
-                self.speculative_num_draft_tokens,
-                device=self.device,
-                dtype=torch.int32,
-            )
-
-            max_seqlen_q = max(extend_seq_lens_cpu)
-            cu_seqlens_q = compute_cu_seqlens(extend_seq_lens.to(torch.int32))
-
-            seqlens_int32_cpu = [
-                self.speculative_num_draft_tokens + kv_len
-                for kv_len in seq_lens.tolist()
-            ]
-            seqlens_expanded = torch.cat(
-                [
-                    torch.arange(
-                        kv_len - qo_len + 1,
-                        kv_len + 1,
-                        dtype=torch.int32,
-                        device=self.device,
-                    )
-                    for qo_len, kv_len in zip(
-                        extend_seq_lens_cpu,
-                        seqlens_int32_cpu,
-                        strict=True,
-                    )
-                ]
-            )
-            nsa_cache_seqlens_int32 = compute_nsa_seqlens(
-                seqlens_expanded, nsa_index_topk=self.nsa_index_topk
-            )
-            nsa_extend_seq_lens_list = [1] * bs
-
-            if NSA_DECODE_IMPL == "flashmla_kv":
-                flashmla_metadata = self.decode_cuda_graph_metadata[
-                    "flashmla_metadata"
-                ].slice(slice(0, bs * self.speculative_num_draft_tokens + 1))
-                # As the DeepGemm is not support for q_len = 3/4 in Indexer and every token has independent topk_indices,
-                # we made the Q shape [bs * speculative_num_draft_tokens, 1, head_nums, dim].
-                # So seq_len_q is 1 for flashmla_metadata in target_verify and draft_extend mode.
-                flashmla_metadata.copy_(
-                    self._compute_flashmla_metadata(
-                        cache_seqlens=nsa_cache_seqlens_int32,
-                        seq_len_q=1,
-                    )
-                )
-            else:
-                flashmla_metadata = None
 
         nsa_cu_seqlens_k = compute_cu_seqlens(nsa_cache_seqlens_int32)
         nsa_cu_seqlens_q = self.get_device_int32_arange(len(nsa_cu_seqlens_k))
@@ -747,14 +728,18 @@ class NativeSparseAttnBackend(AttentionBackend):
             metadata.cu_seqlens_k[1:].copy_(
                 torch.cumsum(cache_seqlens, dim=0, dtype=torch.int32)
             )
-            page_indices = self.req_to_token[req_pool_indices, :max_seqlen_k]
-            metadata.page_table_1[:, :max_seqlen_k].copy_(page_indices)
-            extend_seq_lens_cpu = spec_info.accept_length[:bs].tolist()
 
-            seqlens_int32_cpu = [
+            extend_seq_lens = spec_info.accept_length[:bs]
-                self.speculative_num_draft_tokens + kv_len
+            extend_seq_lens_cpu = extend_seq_lens.tolist()
-                for kv_len in seq_lens_cpu.tolist()
+
-            ]
+            page_indices = self.req_to_token[req_pool_indices, :max_seqlen_k]
+            page_indices = torch.repeat_interleave(
+                page_indices, repeats=extend_seq_lens, dim=0
+            )
+            metadata.page_table_1[: page_indices.shape[0], :max_seqlen_k].copy_(
+                page_indices
+            )
+
             seqlens_expanded = torch.cat(
                 [
                     torch.arange(
@@ -765,21 +750,21 @@ class NativeSparseAttnBackend(AttentionBackend):
                     )
                     for qo_len, kv_len in zip(
                         extend_seq_lens_cpu,
-                        seqlens_int32_cpu,
+                        seq_lens_cpu.tolist(),
                         strict=True,
                     )
                 ]
             )
-            metadata.nsa_seqlens_expanded[: seqlens_expanded.size(0)].copy_(
+            metadata.nsa_seqlens_expanded[: seqlens_expanded.shape[0]].copy_(
                 seqlens_expanded
             )
             nsa_cache_seqlens = compute_nsa_seqlens(
                 seqlens_expanded, self.nsa_index_topk
             )
-            metadata.nsa_cache_seqlens_int32[: seqlens_expanded.size(0)].copy_(
+            metadata.nsa_cache_seqlens_int32[: seqlens_expanded.shape[0]].copy_(
                 nsa_cache_seqlens
             )
-        seqlens_expanded_size = seqlens_expanded.size(0)
+        seqlens_expanded_size = seqlens_expanded.shape[0]
         assert (
             metadata.nsa_cache_seqlens_int32 is not None
             and metadata.nsa_cu_seqlens_k is not None
@@ -794,8 +779,9 @@ class NativeSparseAttnBackend(AttentionBackend):
         assert self.real_page_size == metadata.page_size
         if self.real_page_size > 1:
             real_table = self._transform_table_1_to_real(page_indices)
-            new_len = real_table.shape[1]
+            new_rows = real_table.shape[0]
-            metadata.real_page_table[:, :new_len].copy_(real_table)
+            new_cols = real_table.shape[1]
+            metadata.real_page_table[:new_rows, :new_cols].copy_(real_table)
         else:
             assert metadata.real_page_table is metadata.page_table_1
 

test/srt/test_deepseek_v32_mtp.py

@@ -58,7 +58,7 @@ class TestDeepseekV32MTP(CustomTestCase):
         requests.get(self.base_url + "/flush_cache")
 
         args = SimpleNamespace(
-            num_shots=8,
+            num_shots=20,
             data_path=None,
             num_questions=1400,
             parallel=1400,
@@ -81,7 +81,7 @@ class TestDeepseekV32MTP(CustomTestCase):
                 f'{metrics["accuracy"]=:.3f}\n'
                 f"{avg_spec_accept_length=:.2f}\n"
             )
-            self.assertGreater(metrics["accuracy"], 0.935)
+            self.assertGreater(metrics["accuracy"], 0.94)
             self.assertGreater(avg_spec_accept_length, 2.7)
 
     def test_bs_1_speed(self):