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 8, 21:30 UTC
[ROCm][Bugfix][CI] Fix hybrid models and their tests (Mamba/Jamba/Bamba) (#32710)
Signed-off-by: Andreas Karatzas <akaratza@amd.com> Signed-off-by: Matthew Wong <Matthew.Wong2@amd.com> Co-authored-by: Matthew Wong <Matthew.Wong2@amd.com>
Andreas Karatzas committed
Feb 5, 2026 at 04:01 UTC
3e472e81f99b5bcf494369ee2d26ee9d6ceeffe3
2 files changed
+11
tests/models/language/generation/test_hybrid.py
+5
@@ -8,6 +8,7 @@ import pytest
8
from tests.models.registry import HF_EXAMPLE_MODELS
9
from tests.utils import multi_gpu_test
10
from vllm.engine.arg_utils import EngineArgs
11
+from vllm.platforms import current_platform
12
from vllm.sampling_params import SamplingParams
13
from vllm.v1.cudagraph_dispatcher import CudagraphDispatcher
14
@@ -577,6 +578,10 @@ def test_apc_multiple_prompts_all_cached_outputs(
578
model, max_model_len, tensor_parallel_size=tensor_parallel_size
579
)
580
vllm_runner_kwargs["mamba_ssm_cache_dtype"] = "float32"
581
+ # Reduce the effects of batch variance on ROCm since batch invariance is not
582
+ # yet supported. See: https://github.com/vllm-project/vllm/issues/27433
583
+ if current_platform.is_rocm():
584
+ vllm_runner_kwargs["max_num_seqs"] = 4
585
586
vllm_outputs_no_cache, _ = _get_vLLM_output(
587
vllm_runner, vllm_runner_kwargs, generated_prompts, max_tokens, num_logprobs
vllm/model_executor/layers/mamba/mamba_mixer.py
+6
@@ -214,6 +214,12 @@ class MambaMixer(MambaBase, CustomOp):
214
time_step = self.dt_layernorm(time_step.contiguous())
215
B = self.b_layernorm(B.contiguous())
216
C = self.c_layernorm(C.contiguous())
217
+
218
+ # ROCm: tensor from split is non-contiguous, causing incorrect
219
+ # GEMM results in dt_proj.
220
+ if current_platform.is_rocm():
221
+ time_step = time_step.contiguous()
222
+
223
discrete_time_step = self.dt_proj(time_step)[0].transpose(-2, -1)
224
return discrete_time_step, B, C
225