[ROCm][CI] Fix sparse MLA metadata sync fixture (#49128)
Signed-off-by: Andreas Karatzas <Andreas.Karatzas@amd.com> Co-authored-by: OpenAI Codex <noreply@openai.com>
Andreas Karatzas committed
Jul 19, 2026 at 23:02 UTC
752bd106477f1a72385fb38616abfe198dd67815
1 file changed
+2
tests/kernels/attention/test_rocm_aiter_mla_sparse_metadata_sync.py
+2
@@ -58,6 +58,7 @@ def _make_builder():
58
max_num_batched_tokens + 1, dtype=torch.int32, device="cpu"
59
)
60
builder._num_attention_heads = 16
61
+ builder._num_compute_units = current_platform.num_compute_units()
62
builder._mla_work_meta_data = torch.empty(1, dtype=torch.int32, device="cpu")
63
builder._mla_work_indptr = torch.empty(1, dtype=torch.int32, device="cpu")
64
builder._mla_work_info_set = torch.empty(1, dtype=torch.int32, device="cpu")
@@ -116,6 +117,7 @@ def test_sparse_persistent_metadata_syncs_only_after_recompute(monkeypatch):
117
118
assert events == ["metadata", "sync"]
119
assert fake_get_mla_metadata_v1_mock.call_count == 1
120
+ assert fake_get_mla_metadata_v1_mock.call_args.kwargs["max_split_per_batch"] == 1
121
122
events.clear()
123