Enable compact GLM NoPE KV storage with native FlashInfer

[ormandj]

Motivation

With compact-aware FlashInfer kernels, the same memory budget for GLM NoPE KV rows can hold 24.2% more token slots: each row uses 528 rather than 656 bytes, with the same FP8 values and FP32 scales. That provides more room for long prompts or cached request context without reducing precision. This is the row-capacity gain before allocator padding; other memory pools and configured context limits still apply.

This PR makes that capacity available to SGLang by integrating FlashInfer's native GLM NoPE attention backend on SM120 and selecting compact rows when supported. The existing RoPE entry point cannot serve GLM's NoPE query and sparse-candidate geometry. The required FlashInfer #5075 provides Gabriel Wu’s canonical compact payload, runtime row strides and masked-read correction; David Orman’s follow-up supplies the eight-head decode fix and additional validation. No isolated serving speedup is claimed.

Related to #36830 for the SM120 native-backend path; this PR does not add SM90 support.

Modifications

Use a persistent native FlashInfer wrapper for GLM NoPE and recognize the Glm5Next architectures. Pad 2,051 candidates to 2,176 columns and bound each row by its last valid column, preserving valid tail entries after missing groups. Size the runner buffers for mixed prefill/decode/verify batches and the scheduler's final whole-prompt admission when chunking is disabled. Allow this backend with multiple KPool index groups. When either attention phase selects the native backend, require both phases to select it; preserve other backend pairs for their existing validation paths.

Select compact 528-byte FP8 KV rows from FlashInfer’s model-specific canonical payload metadata. Native GLM NoPE runner construction requires this capability; the earlier padded-only native API is rejected. Both layouts contain the same 512 FP8 values and four FP32 scales. Compact rows remove 128 unused RoPE bytes per cached token per affected layer without requantization. This applies to GLM NoPE, not models that store RoPE values in those bytes.

This backend requires the complete changes in FlashInfer #5075 and lucifer1004/flashinfer#3, including the masked-read correction and compatible eight-head decode scratch. Until these changes land upstream, the complete FlashInfer dependency is available at ormandj/flashinfer commit e7cad40ef1c1c279186c92a22656cdf42325d31c. Runner construction requires supported_sparse_mla_sm120_configs()["glm53_nope"].bytes_per_token == 528. That check excludes older payload metadata, but does not independently verify correctness fixes; partial cherry-picks of compact-row support are unsupported. The compact-capable kernels accept both 528- and 656-byte inputs. The separate open PR #37625 addresses finite top-k selection; this adapter also handles -1 holes without dropping subsequent valid candidates. The existing RoPE API remains available. The native GLM NoPE path does not support skip-softmax thresholds: leave SGLANG_SKIP_SOFTMAX_DECODE_THRESHOLD_SCALE_FACTOR and SGLANG_SKIP_SOFTMAX_PREFILL_THRESHOLD_SCALE_FACTOR unset.

Accuracy Tests

Author CPU validation at 28379a7c55, based on SGLang main dc5f59c3a2: 18 tests and 21 subtests passed with CUDA hidden, using FlashInfer with #5075 and lucifer1004/flashinfer#3. The tests cover canonical model-specific payload metadata, rejecting older or incompatible APIs, preserving legacy RoPE entry points, and the layout/capacity cases below. Configured hooks passed on the changed files. Earlier GPU and sanitizer results retain their original source scope.

Earlier author-run GPU and sanitizer results used SGLang source 6ccd5e37e4 with FlashInfer source c35c4d12a6. These are separate from the current CPU validation above.

Test scope	Result
Adapter with the companion FlashInfer changes, RTX PRO 6000 Blackwell SM120	6 GPU cases passed
Compute Sanitizer on those 6 adapter cases and 64 FlashInfer masked-entry cases	70 cases passed; 0 errors

The CPU tests cover backend selection, layout capability, unchunked scheduler overshoot, mixed-batch capacity, holes and empty rows, and exact versus one-byte-short workspace at 8/24/32 heads. The GPU cases use (tokens, heads) of (1, 32), (4, 8) and (65, 32) with both 528- and 656-byte rows. They check finite, expected outputs with NaNs planted in unused cache slot 0, missing candidates and valid tail entries. These adapter/kernel checks do not establish whole-model accuracy.

python -m pytest -q test/registered/unit/test_flashinfer_sparse_mla.py

Speed Tests and Profiling

Per-row storage decreases from 656 to 528 bytes, a 128-byte reduction (19.5%). Total storage saved is 128 bytes × cached token slots × affected layers, before allocator padding. The FP8 values and FP32 scales are unchanged. No isolated serving throughput result is available for this adapter; smaller cache rows do not by themselves establish faster attention.

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.

Developed with AI assistance.

---

CI States

Latest PR Test (Base): ❌ Run #34542801921
Latest PR Test (Extra): ❌ Run #34542801577
Latest PR Test (AMD ROCm 10): ❌ Run #34542801955

[lucifer1004]

DSv4.1 part has been separated out and will be resubmitted as a new PR. I will merge your parts into current 5075 branch and keep you as coauthor. Thanks a lot!

[ormandj]

DSv4.1 part has been separated out and will be resubmitted as a new PR. I will merge your parts into current 5075 branch and keep you as coauthor. Thanks a lot!

I'll clean it up for you here in a jiffy, just saw your changes, will update shortly!