[ROCm][V1][MLA] Clone prefill backend state per metadata builder (#46993)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Andreas Karatzas committed
Jun 30, 2026 at 13:43 UTC
c8f9c156a5f823209f7373b2e842f82335cbbd08
3 files changed
+37
-1
tests/v1/attention/test_mla_prefill_registry.py
+22
@@ -28,6 +28,28 @@ class CustomMLAPrefillBackend(MLAPrefillBackend):
28
raise NotImplementedError
29
30
31
+def test_prefill_backend_clone_has_isolated_metadata():
32
+ backend = CustomMLAPrefillBackend(
33
+ num_heads=4,
34
+ scale=0.5,
35
+ kv_lora_rank=8,
36
+ qk_nope_head_dim=16,
37
+ qk_rope_head_dim=8,
38
+ v_head_dim=32,
39
+ vllm_config=object(),
40
+ )
41
+
42
+ clone = backend.clone()
43
+
44
+ assert isinstance(clone, CustomMLAPrefillBackend)
45
+ assert clone is not backend
46
+ assert clone.num_heads == backend.num_heads
47
+ assert clone.scale == backend.scale
48
+ backend._prefill_metadata = object()
49
+ clone._prefill_metadata = object()
50
+ assert clone._prefill_metadata is not backend._prefill_metadata
51
+
52
+
53
@pytest.fixture(autouse=True)
54
def cleanup_overrides():
55
"""Clear any overrides after each test."""
vllm/model_executor/layers/attention/mla_attention.py
+4
-1
@@ -1563,9 +1563,12 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
1563
device=device,
1564
)
1565
1566
+ # Metadata builders are created per ubatch when DBO is enabled. MLA
1567
+ # prefill backends keep the prepared metadata on the backend object, so
1568
+ # each builder needs its own backend instance to avoid cross-ubatch races.
1569
self._prefill_backend = self.compilation_config.static_forward_context[
1570
layer_names[0]
1568
- ].prefill_backend
1571
+ ].prefill_backend.clone()
1572
1573
supports_spec_decode = self.query_len_support != QueryLenSupport.SINGLE_ONLY
1574
self._init_reorder_batch_threshold(
vllm/v1/attention/backends/mla/prefill/base.py
+11
@@ -116,6 +116,17 @@ class MLAPrefillBackend(ABC):
116
self.v_head_dim = v_head_dim
117
self.vllm_config = vllm_config
118
119
+ def clone(self) -> "MLAPrefillBackend":
120
+ return self.__class__(
121
+ num_heads=self.num_heads,
122
+ scale=self.scale,
123
+ kv_lora_rank=self.kv_lora_rank,
124
+ qk_nope_head_dim=self.qk_nope_head_dim,
125
+ qk_rope_head_dim=self.qk_rope_head_dim,
126
+ v_head_dim=self.v_head_dim,
127
+ vllm_config=self.vllm_config,
128
+ )
129
+
130
def prepare_metadata( # noqa: B027
131
self,
132
prefill_metadata: "MLACommonPrefillMetadata",