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 5, 19:45 UTC
[ROCm][CI] Pin test_hybrid test to TRITON_ATTN on ROCm (#38381)
Signed-off-by: Micah Williamson <micah.williamson@amd.com>
Micah Williamson committed
Mar 30, 2026 at 15:26 UTC
d9c7db18da98acff14f01edec69f3090b953e21b
1 file changed
+11
-2
tests/models/language/generation/test_hybrid.py
+11
-2
@@ -57,6 +57,8 @@ FP32_STATE_MODELS = [
57
# Avoid OOM
58
MAX_NUM_SEQS = 4
59
60
+ATTN_BACKEND = "TRITON_ATTN" if current_platform.is_rocm() else "auto"
61
+
62
63
@pytest.mark.parametrize("model", SSM_MODELS + HYBRID_MODELS)
64
@pytest.mark.parametrize("max_tokens", [64])
@@ -82,7 +84,9 @@ def test_models(
84
example_prompts, max_tokens, num_logprobs
85
)
86
85
- with vllm_runner(model, max_num_seqs=MAX_NUM_SEQS) as vllm_model:
87
+ with vllm_runner(
88
+ model, max_num_seqs=MAX_NUM_SEQS, attention_backend=ATTN_BACKEND
89
+ ) as vllm_model:
90
vllm_outputs = vllm_model.generate_greedy_logprobs(
91
example_prompts, max_tokens, num_logprobs
92
)
@@ -157,6 +161,7 @@ def test_chunked_prefill_with_parallel_sampling(
161
# forces prefill chunks with decoding
162
max_num_batched_tokens=MAX_NUM_SEQS * 3,
163
max_num_seqs=MAX_NUM_SEQS,
164
+ attention_backend=ATTN_BACKEND,
165
) as vllm_model:
166
vllm_model.generate(example_prompts, sampling_params)
167
@@ -301,7 +306,9 @@ def test_full_cuda_graph(
306
example_prompts, max_tokens, num_logprobs
307
)
308
304
- with vllm_runner(model, max_num_seqs=MAX_NUM_SEQS) as vllm_model:
309
+ with vllm_runner(
310
+ model, max_num_seqs=MAX_NUM_SEQS, attention_backend=ATTN_BACKEND
311
+ ) as vllm_model:
312
vllm_outputs = vllm_model.generate_greedy_logprobs(
313
example_prompts, max_tokens, num_logprobs
314
)
@@ -370,6 +377,7 @@ def _get_vllm_runner_params(
377
"max_model_len": max_model_len,
378
"tensor_parallel_size": tensor_parallel_size,
379
"gpu_memory_utilization": 0.4,
380
+ "attention_backend": ATTN_BACKEND,
381
}
382
383
@@ -844,6 +852,7 @@ def test_apc_common_prefix_same_batch(
852
mamba_block_size=16,
853
enable_prefix_caching=True,
854
seed=42,
855
+ attention_backend=ATTN_BACKEND,
856
)
857
prompts = [
858
"hello what is one plus one what is one plus one what is one plus one the answer is", # noqa: E501