[AMD] WIP - Clamp disagg max_total_num_tokens to MORI's single-MR size limit

[yctseng0211]

Motivation

The Nightly Test (AMD MI355X 2N 1P1D Disagg) has been failing for
DeepSeek-V4-Flash since the 2026-07-02 scheduled run
(https://github.com/sgl-project/sglang/actions/runs/28565894749); the last
green scheduled run was 2026-07-01
(https://github.com/sgl-project/sglang/actions/runs/28494350421). Every
dsv4flash job (fp8/fp4, base/mtp/dp8ep8) fails while every dsv4pro job
stays green.
Root cause (as investigated in 30313): DSV4-Flash disagg auto-sizing raises max_total_num_tokens from the 07-01 value 8,551,168 to 23,448,064. Each per-layer unified C4 KV buffer is 256 bytes/token, so the largest single KV region grows from ~2.19 GB to ~6.0 GB. MORI registers each KV buffer as one RDMA memory region, and RegisterRdmaMemoryRegion fails with errno=22 (EINVAL) once a single region crosses the ~4 GiB (2**32) single-MR ceiling.
The PD path then never serves (disagg warmup returns near-all-zero output_ids; /generate -> 500 -> router 502 -> PD path not serving; aborting), and the GSM8K gate reports no accuracy. 30313 mitigates this by hardcoding max_total_tokens: 8551168 in the MTP recipe.
This PR replaces that recipe-scoped magic number with a general, limit-aware clamp, so all DSV4 disagg recipes — and real deployments, not just CI — are protected without a per-recipe constant.

Modifications

• environ.py: add SGLANG_MORI_MAX_MR_BYTES (default 4 GiB = 2**32,
  <= 0 disables) — the max byte size of a single MORI-registered KV region.
• pool_configurator.py: add largest_registered_kv_region_bytes_per_token().
  The base returns 0 (no known bound, so callers skip clamping).
  DSV4PoolConfigurator returns
  (qk_nope_head_dim + qk_rope_head_dim) * 2 / min(compress_ratios), which is
  exactly the per-token size of the largest registered region (a C4 layer)
  under the unified bf16 layout and an upper bound for the packed-FP8 layout.
  It reads the same model_config dims the pool uses to allocate the buffers,
  so it cannot drift from get_contiguous_buf_infos.
• model_runner_kv_cache_mixin.py: in _resolve_memory_pool_config, after the
  existing user-cap / PP constraints, _apply_mori_mr_limit() clamps
  max_total_num_tokens so the largest region stays strictly below the limit
  (minus one page to absorb get_contiguous_buf_infos page rounding). No-op
  unless disaggregation_mode != "null" and
  disaggregation_transfer_backend == "mori"; other backends and models are
  unchanged.
  With the 4 GiB default and DSV4-Flash ((448+64)*2/4 = 256 bytes/token) the
  cap is ~16.78M tokens — above the 8,551,168 that was green on 07-01 and below
  the 23,448,064 that failed on 07-02. If a NIC/driver caps MRs smaller, lower
  SGLANG_MORI_MAX_MR_BYTES (no code change).

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 #28857410081
Latest PR Test (Extra): ❌ Run #28857409874

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[yctseng0211]

Triggered a MI355X disagg workflow_dispatch on this branch (amd_test_disag_0707, checkout-runtime enabled so the clamp actually runs, not the image-baked package) to validate all the dsv4flash configs that were failing on the scheduled nightly:

• Run: https://github.com/sgl-project/sglang/actions/runs/28841147584

Configs under test (all 8 flash variants, fp8/fp4 × base/mtp/dp8ep8):
dsv4flash-fp8-1k1k-1p1d, dsv4flash-fp8-1k1k-1p1d-mtp, dsv4flash-fp8-1k1k-1p1d-dp8ep8, dsv4flash-fp8-1k1k-1p1d-dp8ep8-mtp, dsv4flash-fp4-1k1k-1p1d, dsv4flash-fp4-1k1k-1p1d-mtp, dsv4flash-fp4-1k1k-1p1d-dp8ep8, dsv4flash-fp4-1k1k-1p1d-dp8ep8-mtp.

Will update once it completes.

[yctseng0211]

Re-triggered the MI355X disagg workflow_dispatch — the previous run (28841147584) was evicted from the single-slot nightly-test-mi355x concurrency group (a scheduled run queued behind it and bumped the older pending one).

• New run: https://github.com/sgl-project/sglang/actions/runs/28843631030 (currently pending, queued behind an in-progress run)

Same 8 dsv4flash configs (fp8/fp4 × base/mtp/dp8ep8), checkout-runtime enabled.

[yctseng0211]

Rebased/merged latest main into the branch (now at 8a0b1f7) and re-triggered. The prior run failed several configs with scheduler died during initialization (exit code: -3) — that crash happens before pool sizing (so it's unrelated to this PR's clamp) and was already fixed on main after the branch's old base (#30237, 07-06); #30313 passed because it sat on newer main. The branch now includes those fixes.

• New run (updated branch): https://github.com/sgl-project/sglang/actions/runs/28844720096
• Cancelled the stale run (28843631030) to free the MI355X slot.

Same 8 dsv4flash configs, checkout-runtime enabled.

[yctseng0211]

Fixed a ZeroDivisionError in the clamp: DSV4PoolConfigurator.largest_registered_kv_region_bytes_per_token() took min(compression_ratios), but dense layers carry ratio 0, so the divisor became 0. Now filtered to positive (compressed) ratios. Note the earlier scheduler died during initialization crash is gone after merging latest main, so the clamp path is now actually exercised.

• New run (branch 6ef68f9): https://github.com/sgl-project/sglang/actions/runs/28845190610

Same 8 dsv4flash configs, checkout-runtime enabled.

[yctseng0211]

Rebased the branch onto 3cbb756 (the exact commit #30313 passed 8/8 on) + the clamp commits, dropping the latest-main merge. The prior run's base-config garbage output (output_ids: [0, 0, 201, 0, ...]) was a fresh main regression unrelated to this PR — prime suspect 27867 ([DSv4] Loading Time Weight Dequant, landed 07-07 after #30313's base), which the merge had pulled in. This isolates the clamp for a clean validation vs main churn.

• New run (branch 9df498e, based on 3cbb756): https://github.com/sgl-project/sglang/actions/runs/28846089164

Same 8 dsv4flash configs, checkout-runtime enabled.

[yctseng0211]

Re-based onto the exact sglang commit baked into the MI355X CI image (5f98f62a = 0.5.14.dev20260706+g5f98f62a8a, what #30313 actually ran) + #30237 (HIP sparse-off) + checkout-runtime + this PR's clamp. This aligns the runtime with #30313's green while still exercising the clamp via checkout-runtime, avoiding the post-image main regressions (worker-crash + #27867) that are unrelated to this change.

• Aligned validation run: https://github.com/sgl-project/sglang/actions/runs/28848738618

Same 8 dsv4flash configs.

[yctseng0211]

Prior aligned run gave direct evidence that the MI355X RDMA single-MR ceiling is below 4 GiB: after the clamp reduced the region to size=4294836224 (~3.9999 GiB), RegisterRdmaMemoryRegion still failed with errno=22. Since ~2.19 GB registers fine (the value #30313's launcher cap used) while ~4 GiB does not, lowered the default SGLANG_MORI_MAX_MR_BYTES from 4 GiB to 2 GiB (cap ~8.39M tokens -> ~2.0 GiB region, safely below the proven-good 2.038 GiB). The clamp mechanism itself is confirmed working; this only tunes the default threshold.

• New run (branch 3bee4a2, 2 GiB default): https://github.com/sgl-project/sglang/actions/runs/28849981143

[yctseng0211]

[yctseng0211]

https://github.com/sgl-project/sglang/actions/runs/28855170861

[yctseng0211]

Side investigation on this branch (separate from the clamp): the kimik26-fp8-1k1k-1p1d disagg non-MTP GSM8K only reaches ~0.88 (< 0.92) while the MTP variant passes at ~0.95, and single-node non-MTP passes at 0.944.

Findings so far:

• page_size ruled out: disagg non-MTP scores 0.887 @ page_size=256 and 0.880 @ page_size=1 — no meaningful difference.
• All disagg runtime args (kv_cache_dtype=auto, backends, disable_radix_cache, chunked_prefill, mem_fraction) are identical between MTP (pass) and non-MTP (fail), so they're ruled out.
• The only real difference is the attention path: dispatch_attn_forward_method routes non-MTP decode → triton (decode backend), but MTP verify → aiter (prefill backend, since speculative_attention_mode=prefill). Single-node triton decode reads self-computed KV and is fine (0.944); disagg triton decode reads MORI-transferred KV and degrades.

Experiment now running: kimik26 non-MTP with decode_attention_backend: aiter (page_size back to 256) to test whether the triton MLA decode on transferred KV is the culprit.

• Run: https://github.com/sgl-project/sglang/actions/runs/28857548173

If accuracy returns to ~0.94, the fix is to use aiter decode for Kimi disagg (matching the backend MTP verify already uses).

[yctseng0211]

closed as #30313 being merged.