From 8bbfc17cf99ef440eed2a7264e79b425fbaf4ff6 Mon Sep 17 00:00:00 2001
From: Roger Wang <hey@rogerw.io>
Date: Wed, 11 Feb 2026 20:59:28 +0000
Subject: [PATCH] add

Signed-off-by: Roger Wang <hey@rogerw.io>
---
 vllm/v1/attention/backends/mla/indexer.py | 43 ++++++++++++++++++++---
 vllm/v1/spec_decode/eagle.py              | 23 +++++++++---
 2 files changed, 58 insertions(+), 8 deletions(-)

diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py
index 368b217f0ba..32dad91cfbb 100644
--- a/vllm/v1/attention/backends/mla/indexer.py
+++ b/vllm/v1/attention/backends/mla/indexer.py
@@ -268,11 +268,39 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
             num_reqs=reqs_end - reqs_start,
         )
 
+    def build_for_drafting(
+        self,
+        common_attn_metadata: CommonAttentionMetadata,
+        draft_index: int,
+    ) -> DeepseekV32IndexerMetadata:
+        """Build indexer metadata for draft model.
+
+        During drafting, each forward pass generates only 1 token per request,
+        so we set offsets=None (no speculative decoding offsets needed).
+        """
+        return self._build_impl(
+            common_attn_metadata=common_attn_metadata,
+            fast_build=True,
+            drafting=True,
+        )
+
     def build(
         self,
         common_prefix_len: int,
         common_attn_metadata: CommonAttentionMetadata,
         fast_build: bool = False,
+    ) -> DeepseekV32IndexerMetadata:
+        return self._build_impl(
+            common_attn_metadata=common_attn_metadata,
+            fast_build=fast_build,
+            drafting=False,
+        )
+
+    def _build_impl(
+        self,
+        common_attn_metadata: CommonAttentionMetadata,
+        fast_build: bool = False,
+        drafting: bool = False,
     ) -> DeepseekV32IndexerMetadata:
         num_reqs = common_attn_metadata.num_reqs
         num_tokens = common_attn_metadata.num_actual_tokens
@@ -331,11 +359,18 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
             # - top_k_per_row_decode wins for batch > 128 or seq_len <= 8K
             use_large_context_topk = batch_size <= 128 and _is_large_context
 
-            next_n = 1 + self.num_speculative_tokens
-            if next_n > 1:
-                offsets = torch.arange(next_n, device=self.device, dtype=torch.int32)
-            else:
+            if drafting:
+                # During drafting, each forward generates 1 token per request,
+                # so no speculative offsets are needed.
                 offsets = None
+            else:
+                next_n = 1 + self.num_speculative_tokens
+                if next_n > 1:
+                    offsets = torch.arange(
+                        next_n, device=self.device, dtype=torch.int32
+                    )
+                else:
+                    offsets = None
 
             seq_lens = common_attn_metadata.seq_lens[:num_decodes]
             if is_deep_gemm_supported():
diff --git a/vllm/v1/spec_decode/eagle.py b/vllm/v1/spec_decode/eagle.py
index b5532d65261..2df40f33e47 100644
--- a/vllm/v1/spec_decode/eagle.py
+++ b/vllm/v1/spec_decode/eagle.py
@@ -547,10 +547,14 @@ class SpecDecodeBaseProposer:
             num_tokens_unpadded=batch_size, num_tokens_padded=batch_size
         )
 
-        cudagraph_runtime_mode, batch_desc = self.cudagraph_dispatcher.dispatch(
-            batch_size_dp_padded
-        )
-        input_batch_size = batch_desc.num_tokens
+        # Disable CUDA graphs for draft loop iterations.
+        # When CUDA graphs pad the batch (e.g., 100 → 104), the MLA attention
+        # backends call reshape_query_for_spec_decode(q, num_decodes) which
+        # fails when padded_size % num_decodes != 0. Even when divisible,
+        # padding positions contain invalid topk indices that cause CUDA
+        # illegal memory access in the FlashMLA kernel.
+        cudagraph_runtime_mode = CUDAGraphMode.NONE
+        input_batch_size = batch_size_dp_padded
         if batch_size_across_dp is not None:
             batch_size_across_dp[self.dp_rank] = input_batch_size
 
@@ -653,6 +657,17 @@ class SpecDecodeBaseProposer:
             for layer_name in self.attn_layer_names:
                 per_layer_attn_metadata[layer_name] = attn_metadata
 
+            # Rebuild indexer/DSA attention metadata for models like GLM5
+            if self.draft_indexer_metadata_builder:
+                draft_indexer_metadata = (
+                    self.draft_indexer_metadata_builder.build_for_drafting(
+                        common_attn_metadata=common_attn_metadata,
+                        draft_index=token_index + 1,
+                    )
+                )
+                for layer_name in self.indexer_layer_names:
+                    per_layer_attn_metadata[layer_name] = draft_indexer_metadata
+
             # copy inputs to buffer for cudagraph
             self.input_ids[:batch_size] = input_ids
             self._set_positions(batch_size, clamped_positions)