fix(qsa): clamp the compress gather to the rows this forward has

[ehuaa]

Motivation

Related to the IMA error described in #37633.
In the indexer every group spanscompress_ratio rows, so a padded entry gathers rows 0..ratio-1 of token_k.
A forward that carries fewer than compress_ratio(4) token rows — the 1..3 token chunk that ends a request under chunked prefill — therefore reads past token_k.
The read almost always lands in neighbouring allocated memory and is absorbed silently, it faults only when it happens to cross an unmapped page.

Modifications

Clamp the gather row index to the rows this forward actually has:

group_locs = group_locs.clamp_max(source_keys.shape[0] - 1)

This is exact rather than a mask over the symptom: padded entries write the
reserved slot 0 and their gathered values are discarded, and real entries are
provably in range, so the clamp cannot change any real computation.

Accuracy Tests

Speed Tests and Profiling

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

---

CI States

Latest PR Test (Base): ❌ Run #34582374121
Latest PR Test (Extra): ❌ Run #34582373804
Latest PR Test (AMD ROCm 10): ❌ Run #34582374085

[ehuaa]

Please help review this PR, thanks! @Qiaolin-Yu @yhyang201

[ykamakazi]

Independent confirmation on ROCm — 1× MI355X, rocm/sgl-dev:v0.5.19-rocm724-mi35x-20260909, Qwen/Qwen3.8-Flash-Next-FP8, TP1.

Without this change, any /generate prefill that adds 1–3 tokens aborts the HSA queue deterministically (HSA_STATUS_ERROR_EXCEPTION 0x1016, scheduler SIGABRT) — the out-of-bounds read is not absorbed on ROCm. Minimal reproducer:

input_ids=[1000,1001,1002]

Instrumenting the gather shows exactly the padded entry described here, on every full-attention layer: token_k rows=3 groups=1 member_rows=[0] group_locs.max=3.

With this one-line clamp applied: the reproducer passes, 1,500 varied-length requests (5 seeds, three of which crashed within 6–98 requests before) and a 300-prompt ShareGPT run complete with no faults.

Same conclusion applies to #37786.

[YAMY1234]

LGTM, but "Fixes #37633" is likely overstated. The crashing batch there was a single request with 64 new tokens, which doesn't hit the <4-row trigger, and the reporter later attributed it to an overlap-scheduler race. Suggest "Related to #37633" so it isn't auto-closed.

[ehuaa]

LGTM, but "Fixes #37633" is likely overstated. The crashing batch there was a single request with 64 new tokens, which doesn't hit the <4-row trigger, and the reporter later attributed it to an overlap-scheduler race. Suggest "Related to #37633" so it isn't auto-closed.

@YAMY1234 You're right. What we fixed only shares the same IMA stack trace as #37633 —
it isn't necessarily the same root cause. I'll fix the comment above.