From e52407ef4ce8400bbbb6d0a34d26e0456ce72cb3 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Wed, 2 Sep 2026 03:27:05 -0500 Subject: [PATCH] [ROCm][CI] Add MTP and other spec-decode acceptance coverage (#53399) Signed-off-by: Andreas Karatzas Signed-off-by: Andreas Karatzas Co-authored-by: OpenAI Codex --- .buildkite/test-amd.yaml | 28 ++++++++++++++++++++++++++++ vllm/v1/spec_decode/gemma4.py | 3 ++- 2 files changed, 30 insertions(+), 1 deletion(-) diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index edcf975c0e..79dbbf4e5f 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -3929,6 +3929,34 @@ steps: commands: - pytest -v -s -m 'not slow_test' v1/spec_decode +- label: ":amd: (MI355) Spec Decode AL MTP + Other Acceptance" # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false + agent_pool: mi355_1 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/v1/spec_decode/ + - vllm/v1/sample/ + - vllm/v1/worker/gpu/ + - vllm/v1/worker/gpu_model_runner.py + - vllm/v1/worker/gpu_worker.py + - vllm/model_executor/model_loader/ + - vllm/model_executor/models/gemma4.py + - vllm/model_executor/models/gemma4_mm.py + - vllm/model_executor/models/gemma4_mtp.py + - vllm/model_executor/models/llama.py + - vllm/model_executor/models/llama_eagle3.py + - vllm/model_executor/models/medusa.py + - vllm/model_executor/models/registry.py + - tests/evals/gsm8k/ + - tests/utils.py + - tests/v1/e2e/spec_decode/ + - vllm/platforms/rocm.py + commands: + - pytest -v -s v1/e2e/spec_decode/acceptance_rates/mtp_other/ + #------------------------------------------------------ mi355 ยท weight_loading -------------------------------------------------------# - label: ":amd: (MI355) Weight Loading Multi-GPU" # TBD diff --git a/vllm/v1/spec_decode/gemma4.py b/vllm/v1/spec_decode/gemma4.py index b94fa6b867..bea677dde8 100644 --- a/vllm/v1/spec_decode/gemma4.py +++ b/vllm/v1/spec_decode/gemma4.py @@ -118,6 +118,7 @@ class Gemma4Proposer(SpecDecodeBaseProposer): """Capture CUDA graphs for centroids get_top_tokens at key sizes.""" masked_emb = self.model.masked_embedding lm_head_weight = self.model._get_full_lm_head_weight() + capture_stream = current_stream() for size in [1, 2, 4, 8, 16, 32, 64]: static_input = torch.zeros( @@ -131,7 +132,7 @@ class Gemma4Proposer(SpecDecodeBaseProposer): torch.accelerator.synchronize() g = torch.cuda.CUDAGraph() - with torch.cuda.graph(g, stream=current_stream()): + with torch.cuda.graph(g, stream=capture_stream): static_output = masked_emb.get_top_tokens( static_input, lm_head_weight,