Reland fix(qsa): clamp the compress gather to the rows (#38346) (#39446)

Co-authored-by: ehuaa <ehuamail1@gmail.com>
2026-09-14 14:05:41 -07:00
parent 9128d57966
commit d72e59508b
1 changed files with 1 additions and 0 deletions
@@ -321,6 +321,7 @@ class QSAIndexer(MultiPlatformOp):
                self.compress_ratio, device=member_rows.device, dtype=torch.long
            )
            source_keys = token_k
+            group_locs = group_locs.clamp_max(source_keys.shape[0] - 1)
            source_rope = metadata.extend_rope_matrix
            if source_rope is None:
                source_rope = build_rope_position_matrix(