[ROCm][CI] Add MTP and other spec-decode acceptance coverage (#53399)

Signed-off-by: Andreas Karatzas <Andreas.Karatzas@amd.com>
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Co-authored-by: OpenAI Codex <codex@openai.com>

.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

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,