fix(mla): plan persistent decode over the rows q/o have

atom/model_ops/attentions/aiter_mla.py

@@ -2830,7 +2830,14 @@ def prepare_decode(
             # that padded nothing reaches here with the two equal, and then this
             # is a no-op -- but eager no longer implies that, since a DP-unified
             # batch pads off the graph path too.
-            meta_bs = scheduled_bs if cp_round_robin else running_bs
+            # Otherwise plan the sequences q/o have rows for: running_tokens is
+            # running_bs * max_seqlen_q on a DP-unified step, which pads the
+            # activations as well, but only this rank's own count while a peer
+            # is prefilling. The persistent kernel loads q[qo_start] and stores
+            # o[qo_start] for every work item, and the planner hands a zero-KV
+            # sequence one whenever its CU has payload left.
+            q_rows_bs = min(running_bs, running_tokens // max_seqlen_q)
+            meta_bs = scheduled_bs if cp_round_robin else q_rows_bs
             ctx_mla_ps = self.set_mla_persistent_worker_buffers(
                 meta_bs, max_seqlen_q, is_cp_round_robin=cp_round_robin
             )

tests/test_h2d_attention_publication.py

@@ -642,3 +642,21 @@ def test_shared_maps_skip_h2d_per_slot_and_tbo_buffer(monkeypatch, kind, transpo
     torch.cuda.synchronize()
     for actual, expected in observed:
         assert torch.equal(actual.cpu(), expected)
+
+
+@pytest.mark.parametrize("q", [1, 3])
+def test_mla_persistent_plan_spans_only_rows_q_has(monkeypatch, q):
+    # running_bs is the group's batch either way; running_tokens is the
+    # group's height only on a DP-unified step, and this rank's own count
+    # while a peer is prefilling -- the height q and o are allocated at.
+    runner, builder = make_runner(monkeypatch, "mla", q=q, tbo=False)
+    planned = []
+    builder.set_mla_persistent_worker_buffers = (
+        lambda bs, *a, **kw: planned.append(bs) or {}
+    )
+    for count, running_tokens in ((1, q), (2, 2 * q), (2, 4 * q), (4, 4 * q)):
+        batch = decode_batch(count, q)
+        start_decode(runner, builder, batch, q)
+        builder.prepare_decode(batch, 4, running_tokens, q)
+        finish(runner)
+    assert planned == [1, 2, 4, 4]