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

2026-03-06 02:54:29 -06:00
parent 0252ca8255
commit 27053aa5ed
1 changed files with 2 additions and 1 deletions
--- a/python/sglang/srt/layers/attention/aiter_backend.py
+++ b/python/sglang/srt/layers/attention/aiter_backend.py
@@ -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),