[HiCache] Back up MXFP8 KV scales in the host pool

[alphabetc1]

Motivation

--kv-cache-dtype mxfp8 stores the KV payload as fp8 e4m3 plus a UE8M0 scale per 32 elements, kept in a separate k_scale_buffer / v_scale_buffer. The HiCache host pool does not know about them: MHATokenToKVPoolHost only holds pointers to k_buffer / v_buffer, neither backup_from_device_all_layer nor load_to_device_per_layer touches the scale buffers, and get_size_per_token counts one byte per element. A page restored from the host therefore comes back with whatever scales happen to be in the device rows it landed on, and the host pool is sized as if the scales did not exist.

Modifications

• New MHATokenToKVPoolMXFP8Host, selected by get_mha_host_pool_cls when the device pool is MHATokenToKVPoolMXFP8.
• The scales of one page occupy a contiguous byte range in the device buffer, so the host mirror stores one row per page and reuses the existing transfer kernels with page_size=1 and an element size of "one page of scales": the staged all-layer write-back for D2H, and the single-layer gather for H2D. No new kernel.
• get_size_per_token includes the scale bytes, so --hicache-ratio sizing accounts for them.
• The class requires the production combination (page_first layout, io_backend=kernel, interleaved scale layout) and raises NotImplementedError otherwise; the L3 flat-page interface also raises, since a flat page of an MXFP8 pool is not defined yet.

Accuracy Tests

Round trip with the real JIT kernels is bit exact: payload and scales read back from the host pool equal the device rows they were copied from. End to end on 4 ranks with both host pools, payload and scales were byte-compared against the device copy after backup and again after load. A shape-matched probe (same request served once as a device hit and once as a host restore) returns identical text and a maximum per-token logprob difference of 0.

Speed Tests and Profiling

No throughput claim: this adds the missing scale traffic, so the host pool moves slightly more bytes per token. It is independent of the write-back batching change in #39050.

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.

New unit test file test/registered/unit/mem_cache/test_mxfp8_mha_pool_host_unit.py: factory dispatch, size_per_token accounting, a CPU round trip of payload and scales, and rejection of the storage interface.

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 #35082102036
Latest PR Test (Extra): ❌ Run #35082101775
Latest PR Test (AMD ROCm 10): ❌ Run #35082101983

[alphabetc1]

/rerun-test -c

[github-actions]

⚠️ Rebase Required Before Re-run

A major update has landed on main. Your PR is diverged relative to required base commit ddd4600197a4.

Re-run was not dispatched. What to do:

• Rebase your branch onto the latest main and push again
• Follow issue #21065 for context
• CI-fix PRs may request the bypass-maintenance label to skip this check

[alphabetc1]

/rerun-test -c

[github-actions]

Results for /rerun-test -c:

🚀 ubuntu-latest (1 test): ✅ View workflow run

cd test/ && python3 registered/unit/mem_cache/test_mxfp8_mha_pool_host_unit.py

[alphabetc1]

/rerun-test -c

[github-actions]

Results for /rerun-test -c:

🚀 ubuntu-latest (1 test): ✅ View workflow run

cd test/ && python3 registered/unit/mem_cache/test_mxfp8_mha_pool_host_unit.py