[ROCm][Spec Decode] Fix probabilistic draft probs test attention backend (#45706)
Signed-off-by: Stefan Koncarevic <stefan.koncarevic@amd.com>
stefankoncarevic committed
Jun 18, 2026 at 18:59 UTC
e2352c29743aeec4a2dafc66c4fdd0e10b37072e
2 files changed
+12
-2
.buildkite/test_areas/misc.yaml
+6
@@ -21,6 +21,12 @@ steps:
21
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
22
# TODO: create another `optional` test group for slow tests
23
- pytest -v -s -m 'not slow_test' v1/spec_decode
24
+ mirror:
25
+ amd:
26
+ device: mi300_1
27
+ timeout_in_minutes: 65
28
+ depends_on:
29
+ - image-build-amd
30
31
- label: V1 Sample + Logits
32
key: v1-sample-logits
tests/v1/spec_decode/test_eagle.py
+6
-2
@@ -1002,7 +1002,11 @@ def test_propose(method, attn_backend, num_speculative_tokens, monkeypatch):
1002
assert torch.equal(result, expected_tokens)
1003
1004
1005
-def test_propose_stores_probabilistic_draft_probs(monkeypatch):
1005
+@pytest.mark.parametrize(
1006
+ "attn_backend",
1007
+ ["ROCM_ATTN", "TRITON_ATTN"] if current_platform.is_rocm() else ["FLASH_ATTN"],
1008
+)
1009
+def test_propose_stores_probabilistic_draft_probs(attn_backend, monkeypatch):
1010
device = torch.device(DEVICE_TYPE)
1011
batch_size = 2
1012
seq_lens = [5, 3]
@@ -1053,7 +1057,7 @@ def test_propose_stores_probabilistic_draft_probs(monkeypatch):
1057
)
1058
1059
attn_metadata_builder_cls, _ = try_get_attention_backend(
1056
- AttentionBackendEnum.FLASH_ATTN
1060
+ AttentionBackendEnum[attn_backend]
1061
)
1062
attn_metadata_builder = attn_metadata_builder_cls(
1063
kv_cache_spec=create_standard_kv_cache_spec(proposer.vllm_config),