Fix MLA decode path returning unwritten (padded) rows (#19902)

2026-03-06 00:54:29 -08:00
parent 0252ca8255
commit 27053aa5ed
1 changed files with 2 additions and 1 deletions
@@ -1704,7 +1704,8 @@ class AiterAttnBackend(AttentionBackend):
                        num_kv_splits=num_kv_splits,
                    )
-                    return o[q_mask]
+                    total_valid_q = int(qo_indptr[-1].item())
                    return o[:total_valid_q]
                else:
                    o = q.new_empty(
                        (q.shape[0], layer.tp_q_head_num, layer.v_head_dim),