A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 10, 06:30 UTC
[ROCm][Bugfix] Fix Mamba batched decode producing incorrect output (#32099)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Andreas Karatzas committed
Jan 12, 2026 at 23:46 UTC
11b6af5280d6d6dfb8953af16e67b25f819b3be9
2 files changed
+15
-5
vllm/model_executor/layers/mamba/mamba_mixer.py
+7
-5
@@ -34,6 +34,7 @@ from vllm.model_executor.layers.mamba.ops.mamba_ssm import (
34
selective_state_update,
35
)
36
from vllm.model_executor.utils import set_weight_attrs
37
+from vllm.platforms import current_platform
38
from vllm.utils.torch_utils import direct_register_custom_op
39
from vllm.v1.attention.backends.mamba1_attn import Mamba1AttentionMetadata
40
@@ -195,11 +196,12 @@ class MambaMixer(MambaBase, CustomOp):
196
def _ssm_transform(
197
self, x: torch.Tensor
198
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
198
- if self.is_lora_enabled:
199
- # Lora kernel requires contiguous tensor.
200
- ssm_params = self.x_proj(x.contiguous())[0]
201
- else:
202
- ssm_params = self.x_proj(x)[0]
199
+ # LoRA kernel requires contiguous tensor.
200
+ # ROCm: Non-contiguous tensors cause incorrect GEMM
201
+ # results when batch > 1.
202
+ if self.is_lora_enabled or current_platform.is_rocm():
203
+ x = x.contiguous()
204
+ ssm_params = self.x_proj(x)[0]
205
time_step, B, C = torch.split(
206
ssm_params,
207
[self.time_step_rank, self.ssm_state_size, self.ssm_state_size],
vllm/model_executor/models/plamo2.py
+8
@@ -63,6 +63,7 @@ from vllm.model_executor.models.utils import (
63
maybe_prefix,
64
)
65
from vllm.model_executor.utils import set_weight_attrs
66
+from vllm.platforms import current_platform
67
from vllm.sequence import IntermediateTensors
68
from vllm.utils.torch_utils import direct_register_custom_op
69
from vllm.v1.attention.backend import AttentionMetadata
@@ -414,6 +415,13 @@ class Plamo2MambaMixer(MambaBase, CustomOp):
415
conv_state_indices=state_indices_tensor_d,
416
)
417
418
+ # ROCm: Ensure contiguous tensor for bcdt_proj linear layer.
419
+ # causal_conv1d_update returns a non-contiguous view (stride 8192
420
+ # instead of 4096 for shape [batch, 4096]), causing incorrect GEMM
421
+ # results when batch > 1 on ROCm.
422
+ if current_platform.is_rocm():
423
+ hidden_states_d = hidden_states_d.contiguous()
424
+
425
B, C, dt = self._project_ssm_parameters(hidden_states_d)
426
427
# 3. State Space Model sequence transformation