[Bugfix] Fix speculative decoding for short_conv (LFM2) models (#50272)
Signed-off-by: jens.luecke <jens.luecke@liquid.ai> Co-authored-by: Tonoken3 <sakamakismile@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Jens committed
Aug 22, 2026 at 15:34 UTC
da329cc303a5233e17fa3d553ce0a3d6ceea87a8
4 files changed
+43
-11
vllm/model_executor/layers/mamba/mamba_utils.py
+2
-1
@@ -235,9 +235,10 @@ class MambaStateShapeCalculator:
235
tp_world_size: int,
236
intermediate_size: int,
237
conv_kernel: int,
238
+ num_spec: int = 0,
239
) -> tuple[tuple[int, int]]:
240
conv_dim = divide(intermediate_size, tp_world_size)
240
- conv_state_shape = cls._orient_conv_shape(conv_dim, conv_kernel - 1)
241
+ conv_state_shape = cls._orient_conv_shape(conv_dim, conv_kernel - 1 + num_spec)
242
return (conv_state_shape,)
243
244
@classmethod
vllm/model_executor/layers/mamba/short_conv.py
+33
-9
@@ -80,7 +80,9 @@ class ShortConv(MambaBase, PluggableLayer):
80
prefix=f"{prefix}.out_proj",
81
)
82
83
- compilation_config = get_current_vllm_config().compilation_config
83
+ vllm_config = get_current_vllm_config()
84
+ self.num_spec = vllm_config.num_speculative_tokens
85
+ compilation_config = vllm_config.compilation_config
86
if prefix in compilation_config.static_forward_context:
87
raise ValueError(f"Duplicate layer name: {prefix}")
88
compilation_config.static_forward_context[prefix] = self
@@ -234,6 +236,8 @@ class ShortConv(MambaBase, PluggableLayer):
236
state_indices_tensor_d = attn_metadata.state_indices_tensor_d
237
has_initial_states_p = attn_metadata.has_initial_states_p
238
query_start_loc_p = attn_metadata.query_start_loc_p
239
+ num_accepted_tokens = attn_metadata.num_accepted_tokens
240
+ query_start_loc_d = attn_metadata.query_start_loc_d
241
242
BCx, _ = self.in_proj(hidden_states)
243
@@ -296,14 +300,33 @@ class ShortConv(MambaBase, PluggableLayer):
300
301
if has_decode:
302
Bx_d = (B_d * x_d).contiguous()
299
- Bx = causal_conv1d_update(
300
- Bx_d,
301
- conv_state,
302
- conv_weights,
303
- self.conv.bias,
304
- activation=None,
305
- conv_state_indices=state_indices_tensor_d,
306
- )
303
+ if num_accepted_tokens is not None:
304
+ # Speculative decode: the verify step feeds >1 query token per
305
+ # decode request, so use the spec-aware conv update path
306
+ # (mirrors mamba_mixer2). state_indices_tensor_d is
307
+ # (num_decodes, 1 + num_spec_tokens) here.
308
+ assert state_indices_tensor_d is not None
309
+ Bx = causal_conv1d_update(
310
+ Bx_d,
311
+ conv_state,
312
+ conv_weights,
313
+ self.conv.bias,
314
+ activation=None,
315
+ conv_state_indices=state_indices_tensor_d,
316
+ num_accepted_tokens=num_accepted_tokens,
317
+ query_start_loc=query_start_loc_d,
318
+ max_query_len=state_indices_tensor_d.size(-1),
319
+ )
320
+ else:
321
+ # Non-spec decode
322
+ Bx = causal_conv1d_update(
323
+ Bx_d,
324
+ conv_state,
325
+ conv_weights,
326
+ self.conv.bias,
327
+ activation=None,
328
+ conv_state_indices=state_indices_tensor_d,
329
+ )
330
y = C_d * Bx
331
conv_output_list.insert(0, y)
332
@@ -326,6 +349,7 @@ class ShortConv(MambaBase, PluggableLayer):
349
tp_world_size=get_tensor_model_parallel_world_size(),
350
intermediate_size=self.conv_dim,
351
conv_kernel=self.L_cache,
352
+ num_spec=self.num_spec,
353
)
354
355
@property
vllm/v1/worker/gpu/model_states/mamba_hybrid.py
+6
-1
@@ -12,6 +12,7 @@ from vllm.config.compilation import CUDAGraphMode
12
from vllm.triton_utils import tl, triton
13
from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadataBuilder
14
from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionMetadataBuilder
15
+from vllm.v1.attention.backends.short_conv_attn import ShortConvAttentionMetadataBuilder
16
from vllm.v1.core.sched.output import NewRequestData
17
from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec
18
from vllm.v1.utils import CpuGpuBuffer
@@ -48,7 +49,11 @@ class MambaHybridAttnMetadata(ModelSpecificAttnMetadata):
49
) -> dict[str, Any]:
50
if not isinstance(
51
attn_metadata_builder,
51
- (Mamba2AttentionMetadataBuilder, GDNAttentionMetadataBuilder),
52
+ (
53
+ Mamba2AttentionMetadataBuilder,
54
+ GDNAttentionMetadataBuilder,
55
+ ShortConvAttentionMetadataBuilder,
56
+ ),
57
):
58
return {}
59
return {
vllm/v1/worker/gpu_model_runner.py
+2
@@ -148,6 +148,7 @@ from vllm.v1.attention.backends.linear_attn import (
148
BailingLinearAttentionMetadataBuilder,
149
)
150
from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionMetadataBuilder
151
+from vllm.v1.attention.backends.short_conv_attn import ShortConvAttentionMetadataBuilder
152
from vllm.v1.attention.backends.utils import (
153
NULL_BLOCK_ID,
154
create_fast_prefill_custom_backend,
@@ -2572,6 +2573,7 @@ class GPUModelRunner(
2573
Mamba2AttentionMetadataBuilder,
2574
GDNAttentionMetadataBuilder,
2575
BailingLinearAttentionMetadataBuilder,
2576
+ ShortConvAttentionMetadataBuilder,
2577
),
2578
):
2579
assert ubid is None, (