[Bugfix][ROCm][CI] Give the AITER MLA decode metadata stub its MLA dims (#52139)
Signed-off-by: Stefan Koncarevic <Stefan.Koncarevic@amd.com>
stefankoncarevic committed
Aug 13, 2026 at 20:35 UTC
11c3fa4adc8105bd1dbad888d6df81617ee7c8fc
1 file changed
+13
-3
tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py
+13
-3
@@ -82,6 +82,7 @@ def _build_decode_metadata():
82
create_vllm_config,
83
)
84
from vllm.config.vllm import set_current_vllm_config
85
+ from vllm.model_executor.layers.attention.mla_attention import get_mla_dims
86
from vllm.v1.attention.backends.registry import AttentionBackendEnum
87
from vllm.v1.kv_cache_interface import MLAAttentionSpec
88
from vllm.v1.worker.workspace import init_workspace_manager
@@ -110,11 +111,20 @@ def _build_decode_metadata():
111
112
builder_cls = AttentionBackendEnum.ROCM_AITER_MLA.get_class().get_builder_cls()
113
113
- # The builder reads layer.prefill_backend from static_forward_context; a
114
- # stub with the attribute is enough for metadata construction.
114
+ # The builder reads prefill_backend and the MLA latent dims off the layer,
115
+ # so the stub carries both; the dims come from the model config to stay in
116
+ # step with the checkpoint.
117
+ mla_dims = get_mla_dims(vllm_config.model_config)
118
layer_name = "placeholder"
119
vllm_config.compilation_config.static_forward_context[layer_name] = (
117
- types.SimpleNamespace(prefill_backend=torch.empty((1,)))
120
+ types.SimpleNamespace(
121
+ prefill_backend=torch.empty((1,)),
122
+ q_lora_rank=mla_dims.q_lora_rank,
123
+ kv_lora_rank=mla_dims.kv_lora_rank,
124
+ qk_nope_head_dim=mla_dims.qk_nope_head_dim,
125
+ qk_rope_head_dim=mla_dims.qk_rope_head_dim,
126
+ v_head_dim=mla_dims.v_head_dim,
127
+ )
128
)
129
130
init_workspace_manager(device)