[HiCache] TMA-staged host<->device KV transfer kernel (sm_90+) - #40278

Motivation

HiCache brings cached prefixes back from the pinned host pool into the device pool layer by layer, on a side stream, while the forward pass runs on the same GPU. Two properties of the transfer kernel therefore matter: how fast it moves a page, because the load gates the prefix hit, and how many SMs it takes from the forward pass to do so.

The register-staging JIT kernel in hicache.cuh is bounded by bytes in flight, not by the link. Each of its 1024 threads holds one 64 B slice of a row in registers between its load and its store, so a CTA has at most 64 KB outstanding at any time. With

bandwidth = bytes in flight / round trip

the measured ~48 GB/s per CTA on a GB300 (NVLink-C2C host link) corresponds to an effective host round trip of ~1.3 us. The copy-engine ceiling of the same link is ~210 GB/s H2D and ~192 GB/s D2H, so the default 2-block launch reaches ~90 GB/s, and saturating the link takes 6+ blocks of 1024 threads, every one of them issuing dependent loads for the whole transfer. Making the register pipeline deeper does not help: two rows in flight per thread measured 40 GB/s instead of 48, because at 1024 threads the extra registers cost more occupancy than the extra outstanding bytes buy.

TMA bulk copies remove that trade-off. One thread issues a 32 KB cp.async.bulk and the data lands in shared memory with no registers involved, so a single warp keeps a 192 KB ring in flight: three times the register kernel's outstanding bytes with nine warps per CTA instead of 32. Two hardware limits, measured with raw probe kernels on the same GPU, then set what any staged copy can reach:

The SM write port (32 B/clk) caps one CTA around 66 GB/s in either design, and the TMA unit handles roughly one bulk op per ~50 cycles, which makes per-row bulk copies a trap. The achievable goal is therefore to run each CTA at its write port instead of at its load latency: 4 blocks saturate the link instead of 6+, and 1 or 2 blocks move more per SM than today.

Modifications

New JIT kernel hicache_tma.cuh, routed from the existing transfer_hicache_* entry points.

CTA = 1 loader warp + 8 store warps, smem ring of 6 x 32 KB stages
chunk = (K|V, layer, 2^k consecutive index positions)   # never straddles a page

loader warp, per chunk:
  prefetch next chunk's src/dst indices          # overlaps the ring wait
  if src rows are one run:
      contiguous (stride == row bytes) -> 1 x cp.async.bulk of the whole span
      strided (page-first host view)    -> 1 x cp.async.bulk.tensor 2D box
  else                                  -> 1 bulk op per row (lanes issue in parallel)
  stash dst indices in smem, arm mbarrier(expect_tx)

store warps, per chunk:
  wait mbarrier(full)
  if dst rows are one contiguous run: warp 0 bulk-stores half (2 groups in flight),
                                      warps 1..7 vector-store the other half
  else:                               vector stores with per-row addressing
  release stage (mbarrier empty, count = store warps)

1. Deep loads without registers. The loader keeps the whole ring in flight with cp.async.bulk; completion is counted on per-stage mbarriers, so no thread holds data.
2. One op per run. Contiguous page runs are one 1D bulk copy. A page run in a page-first host view (rows strided by layers * row_bytes) is one 2D tensor-map box (rows mapped as 8-byte elements, box [rows_per_chunk][row_bytes/8]), encoded per call with cuTensorMapEncodeTiled via the driver entry point (no -lcuda). Per-row ops remain only for scattered tokens.
3. Index prefetch. Loading the next chunk's indices before blocking on the ring removed a ~0.5 us dependent-load bubble per chunk (3x on the loader loop).
4. Store side at the port. Bulk stores and vector stores split contiguous spans so the TMA engine's request generation and the LSU share the write port; destination indices come from smem and the unit-to-row division is a magic multiply, keeping the store loop branch-free.
5. Runtime row size. One compiled module serves every KV shape (multiple of 16 B), including rows the register kernel rejects (e.g. 1152 B MLA rows).
6. Gate. SGLANG_HICACHE_TMA_TRANSFER (default on) selects the kernel when the GPU is sm_90+, the row size is a multiple of 16 B, and page_size % rows_per_chunk == 0 (pools now pass their page size); otherwise the register kernel is used. ROCm and pre-Hopper are unaffected.

Invariants: correctness never depends on the run detection (it only picks the op shape); every stage is reused only after its readers release it through the empty mbarrier; bulk stores are waited with cp.async.bulk.wait_group 0 before the grid completes.

Accuracy Tests

• test/registered/kernels/ops/kvcache/test_hicache_tma.py: H2D then D2H round trips through layer-first and page-first host views, int32 and int64 indices, whole pages and fully scattered rows, a partial tail chunk, untouched rows verified unchanged; all-layer pointer-table D2H into a page-first host; MLA single-buffer one-layer and all-layer. 10 cases pass (GB300).
• test/registered/kernels/ops/kvcache/test_hicache.py and test_hicache_page_first_write_back.py pass through the pool API with the new default (59) and with SGLANG_HICACHE_TMA_TRANSFER=0 (47).

Speed Tests and Profiling

H2D, 512 B rows (bf16, 4 KV heads x 64), 128-token pages, 32768 tokens, page-first host, GB/s (TMA / register kernel), GB300:

At 4 blocks the TMA kernel reaches the copy-engine ceiling of the link; the register kernel needs 6+. D2H to a layer-first host is at the host write cap for both kernels (~45 GB/s per block). For scattered single-token indices (page_size=1) the register kernel stays faster for small rows, which is why the gate keeps it for pages that do not tile a chunk. bench_hicache.py gained a tma provider line.

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.

CI States

Latest PR Test (Base): ❌ Run #35412153100
Latest PR Test (Extra): ❌ Run #35412152921
Latest PR Test (AMD ROCm 10): ❌ Run #35412153051