From 11c3fa4adc8105bd1dbad888d6df81617ee7c8fc Mon Sep 17 00:00:00 2001 From: stefankoncarevic Date: Thu, 13 Aug 2026 20:35:00 +0200 Subject: [PATCH] [Bugfix][ROCm][CI] Give the AITER MLA decode metadata stub its MLA dims (#52139) Signed-off-by: Stefan Koncarevic --- .../test_rocm_aiter_mla_decode_metadata.py | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py b/tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py index a384af6f13..0c15cb99f1 100644 --- a/tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py +++ b/tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py @@ -82,6 +82,7 @@ def _build_decode_metadata(): create_vllm_config, ) from vllm.config.vllm import set_current_vllm_config + from vllm.model_executor.layers.attention.mla_attention import get_mla_dims from vllm.v1.attention.backends.registry import AttentionBackendEnum from vllm.v1.kv_cache_interface import MLAAttentionSpec from vllm.v1.worker.workspace import init_workspace_manager @@ -110,11 +111,20 @@ def _build_decode_metadata(): builder_cls = AttentionBackendEnum.ROCM_AITER_MLA.get_class().get_builder_cls() - # The builder reads layer.prefill_backend from static_forward_context; a - # stub with the attribute is enough for metadata construction. + # The builder reads prefill_backend and the MLA latent dims off the layer, + # so the stub carries both; the dims come from the model config to stay in + # step with the checkpoint. + mla_dims = get_mla_dims(vllm_config.model_config) layer_name = "placeholder" vllm_config.compilation_config.static_forward_context[layer_name] = ( - types.SimpleNamespace(prefill_backend=torch.empty((1,))) + types.SimpleNamespace( + prefill_backend=torch.empty((1,)), + q_lora_rank=mla_dims.q_lora_rank, + kv_lora_rank=mla_dims.kv_lora_rank, + qk_nope_head_dim=mla_dims.qk_nope_head_dim, + qk_rope_head_dim=mla_dims.qk_rope_head_dim, + v_head_dim=mla_dims.v_head_dim, + ) ) init_workspace_manager(device)