[Bugfix] Fix mamba+dflash for MRV2 (#47698)
Benjamin Chislett committed
Jul 8, 2026 at 04:59 UTC
7bd154375dc505046a6e59e6d8c884a9c6b8fc0f
1 file changed
+1
-4
vllm/v1/worker/gpu/spec_decode/dflash/speculator.py
+1
-4
@@ -135,9 +135,6 @@ class DFlashSpeculator(DraftModelSpeculator):
135
]
136
assert self.draft_kv_cache_group_ids, "No draft attention groups found."
137
self.draft_kv_cache_group_id = self.draft_kv_cache_group_ids[0]
138
- self.draft_block_size = self.block_tables.block_sizes[
139
- self.draft_kv_cache_group_id
140
- ]
138
139
# Per-group context slot buffers for the precompute (one row per group).
140
self._context_slot_mappings = torch.zeros(
@@ -337,7 +334,7 @@ class DFlashSpeculator(DraftModelSpeculator):
334
last_sampled,
335
next_prefill_tokens,
336
self.block_tables.input_block_tables[gid],
340
- self.block_tables.block_sizes[gid],
337
+ self.block_tables.kernel_block_sizes[gid],
338
self.parallel_drafting_token_id,
339
self.num_query_per_req,
340
self.num_speculative_steps,