Avoid redundant CP collectives on sync shared-KV materialize
Synchronous CP shared KV full-hit and partial-current paths can now use the tai-kernel CUDA IPC slot-dense materialize path instead of first copying local owner pages and then running a dense CP all-reduce. This keeps the async prefetch pipeline unchanged while routing the safer synchronous runtime path through descriptor-driven owner-page reads. The runtime builds descriptors from the existing page-aligned slot contract, caches peer pointer tables for long-lived KV/index buffers, and falls back with explicit warnings if the tai-kernel IPC capability is missing. Unit coverage locks offset handle exchange, descriptor construction, and both MLA/index full and partial-current calls. Constraint: Async prefetch currently has unresolved scheduling contention with GEMM/MoE and current-layer KV collectives, so this commit intentionally does not wire IPC into prefetch. Constraint: tai-kernel must provide offset-aware CUDA IPC symbols from af9fb67. Rejected: Replace prefetch all-reduce in this slice | Nsight shows prefetch timing and communicator contention need a separate scheduling design. Rejected: Fail-fast on missing tai IPC immediately | remote ETE still needs to validate production deployment capability before removing the warning fallback. Confidence: medium Scope-risk: moderate Directive: Keep prefetch and synchronous materialize decisions separate until prefetch has a low-SM/copy-engine schedule. Related: tai-kernel af9fb67 Tested: git diff --check; remote runtime/unit evidence recorded in docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md. Not-tested: Fresh GLM5 ETE after this commit; async prefetch IPC path. Co-authored-by: OmX <omx@oh-my-codex.dev>
This commit is contained in:
laoyao0822
@@ -4088,3 +4088,205 @@ Conclusion:
|
||||
- The segment-copy path is now slower than tuned direct peer-page gather for this
|
||||
benchmark. Do not prefer owner-concat segment copy solely for performance
|
||||
unless a future consumer can avoid the final rewrite entirely.
|
||||
|
||||
### C99 — 2026-05-31 multi-warp/page IPC materialize tuning
|
||||
|
||||
Change:
|
||||
|
||||
- Added a benchmark-selectable `warps_per_page` mode to tai-kernel
|
||||
`materialize_cuda_ipc_peer_pages_slot_dense`.
|
||||
- `warps_per_page=1` preserves the previous one-warp-per-page mapping.
|
||||
- `warps_per_page>1` assigns a whole warp group to one page and stripes the
|
||||
page copy across `warps_per_page * 32` threads.
|
||||
- Python wrapper default is now `warps_per_page=0`, an auto selector based on
|
||||
the measured H200 regions for the production ~72 KiB BF16 MLA page:
|
||||
- <= 256 pages: 8 warps/page,
|
||||
- 512--1024 pages: 1 warp/page,
|
||||
- >= 1536 pages: 4 warps/page,
|
||||
- synthetic pages smaller than 64 KiB: 1 warp/page.
|
||||
|
||||
Remote correctness:
|
||||
|
||||
- `PYTHONPATH=python python -m pytest -q tests/nsa_prefill/test_cuda_ipc_gather.py`
|
||||
on `g0034` passed: `6 passed`.
|
||||
|
||||
Remote benchmark logs:
|
||||
|
||||
- Variant sweep:
|
||||
`/mnt/beegfs/cjy/log/cp_shared_kv_ipc_warps_per_page_20260531_114147.log`
|
||||
- Auto selector sweep:
|
||||
`/mnt/beegfs/cjy/log/cp_shared_kv_ipc_auto_wpp_20260531_114447.log`
|
||||
|
||||
Auto selector p50 GPU ms, BF16 MLA page payload, 8 CP ranks:
|
||||
|
||||
| tokens | pages | dense all-reduce | IPC direct slot-dense auto | speedup vs dense |
|
||||
| ---: | ---: | ---: | ---: | ---: |
|
||||
| 4096 | 64 | 0.093 | 0.057 | 1.63x |
|
||||
| 8192 | 128 | 0.115 | 0.060 | 1.92x |
|
||||
| 16384 | 256 | 0.164 | 0.083 | 1.98x |
|
||||
| 32768 | 512 | 0.235 | 0.134 | 1.75x |
|
||||
| 65536 | 1024 | 0.394 | 0.236 | 1.67x |
|
||||
| 98304 | 1536 | 0.554 | 0.348 | 1.59x |
|
||||
| 122880 | 1920 | 0.670 | 0.421 | 1.59x |
|
||||
|
||||
Bandwidth reading for the largest case:
|
||||
|
||||
- `122880` tokens = `1920` pages = `135 MiB` logical payload per GPU.
|
||||
- `0.421 ms` gives about `336 GB/s` payload rate per GPU.
|
||||
- The kernel performs roughly one full payload read plus one full payload write,
|
||||
so effective HBM traffic is about `672 GB/s`, ~14% of H200's ~4.8 TB/s HBM
|
||||
bandwidth.
|
||||
- Remote NVLink receive traffic is about `7/8` of the payload, or `294 GB/s`,
|
||||
~61% of the per-direction 18-link H200 NVLink budget (`18 * 26.562 GB/s ~= 478 GB/s`).
|
||||
|
||||
Conclusion:
|
||||
|
||||
- This kernel is no longer obviously HBM-limited; NVLink plus per-page peer-load
|
||||
overhead is closer to the limiting side for large prefixes.
|
||||
- The remaining large-prefix headroom to the per-direction NVLink roof is about
|
||||
`478 / 294 ~= 1.6x` under this access pattern.
|
||||
- For small prefixes, the main gain came from increasing per-page parallelism
|
||||
rather than more CTAs; launch/fixed overhead still dominates below ~8k tokens.
|
||||
|
||||
Additional C99 correctness check:
|
||||
|
||||
- Ran the multi-rank benchmark with checks enabled for `4096` and `122880`
|
||||
tokens and `warps_per_page=0`; all benchmark reference assertions passed.
|
||||
|
||||
### C100 — 2026-05-31 vectorized 16B copy for multi-warp/page IPC materialize
|
||||
|
||||
Change:
|
||||
|
||||
- The multi-warp/page IPC materialize path now uses `ld.global.nc.v2.u64` +
|
||||
`st.global.cg.v2.u64` when the page payload is 16-byte aligned.
|
||||
- The one-warp/page legacy path is unchanged; the vectorized helper only affects
|
||||
the measured multi-warp variants used by the auto selector.
|
||||
|
||||
Remote verification:
|
||||
|
||||
- `PYTHONPATH=python python -m pytest -q tests/nsa_prefill/test_cuda_ipc_gather.py`
|
||||
on `g0034` passed: `6 passed`.
|
||||
- Multi-rank benchmark checks enabled for `4096` and `122880` tokens passed.
|
||||
|
||||
Remote benchmark log:
|
||||
|
||||
- `/mnt/beegfs/cjy/log/cp_shared_kv_ipc_auto_wpp_vec16_20260531_114903.log`
|
||||
|
||||
Auto selector p50 GPU ms after vectorization:
|
||||
|
||||
| tokens | pages | dense all-reduce | IPC direct slot-dense auto | previous auto IPC | delta |
|
||||
| ---: | ---: | ---: | ---: | ---: | ---: |
|
||||
| 4096 | 64 | 0.137 | 0.042 | 0.057 | -26% |
|
||||
| 8192 | 128 | 0.112 | 0.052 | 0.060 | -13% |
|
||||
| 16384 | 256 | 0.162 | 0.082 | 0.083 | -1% |
|
||||
| 32768 | 512 | 0.236 | 0.136 | 0.134 | +1% |
|
||||
| 65536 | 1024 | 0.392 | 0.238 | 0.236 | +1% |
|
||||
| 98304 | 1536 | 0.553 | 0.339 | 0.348 | -3% |
|
||||
| 122880 | 1920 | 0.668 | 0.416 | 0.421 | -1% |
|
||||
|
||||
Bandwidth reading for `122880` tokens:
|
||||
|
||||
- Payload rate: `135 MiB / 0.416 ms ~= 340 GB/s` per GPU.
|
||||
- Effective HBM read+write traffic: `~680 GB/s`, still only ~14% of H200 HBM.
|
||||
- Remote NVLink receive traffic: `(7/8) * 340 ~= 298 GB/s`, ~62% of the
|
||||
`~478 GB/s` per-direction H200 NVLink link budget.
|
||||
|
||||
Conclusion:
|
||||
|
||||
- 16B vectorization materially improves the short-prefix auto path and slightly
|
||||
improves the large-prefix path; it does not change the bottleneck conclusion.
|
||||
- Large-prefix headroom remains primarily in reducing peer-load overhead / better
|
||||
NVLink utilization, not in HBM bandwidth.
|
||||
|
||||
### C101 — 2026-05-31 index page payload also benefits from IPC materialize
|
||||
|
||||
Clarification:
|
||||
|
||||
- The descriptor builder needed by SGLang runtime is the translation from a
|
||||
consumer slot page table to the IPC kernel descriptors:
|
||||
`slot_logical_pages -> owner_ranks, src_page_indices`.
|
||||
- The IPC kernel itself does not understand radix/HiCache/current-reuse state;
|
||||
it only consumes owner/source descriptors. Full hit, multi-extend hit, and
|
||||
partial-current prefix reuse are all covered if the runtime builds these
|
||||
descriptors from the same page-aligned slot contract.
|
||||
|
||||
Index payload benchmark:
|
||||
|
||||
- NSA index page buffer is `uint8` with per-token bytes
|
||||
`index_head_dim + index_head_dim / quant_block_size * 4 = 128 + 4 = 132`.
|
||||
- With `page_size=64`, each index page is `8448` bytes.
|
||||
- Benchmark command used `--dtype uint8 --kv-dim 132` to match the real index
|
||||
page payload size and dtype.
|
||||
|
||||
Remote log:
|
||||
|
||||
- `/mnt/beegfs/cjy/log/cp_shared_kv_ipc_index_uint8_20260531_120416.log`
|
||||
|
||||
p50 GPU ms, 8 CP ranks:
|
||||
|
||||
| tokens | pages | dense uint8 all-reduce | IPC direct slot-dense | speedup |
|
||||
| ---: | ---: | ---: | ---: | ---: |
|
||||
| 4096 | 64 | 0.119 | 0.033 | 3.61x |
|
||||
| 8192 | 128 | 0.070 | 0.031 | 2.26x |
|
||||
| 16384 | 256 | 0.070 | 0.033 | 2.12x |
|
||||
| 32768 | 512 | 0.081 | 0.037 | 2.19x |
|
||||
| 65536 | 1024 | 0.102 | 0.049 | 2.08x |
|
||||
| 98304 | 1536 | 0.120 | 0.061 | 1.97x |
|
||||
| 122880 | 1920 | 0.139 | 0.072 | 1.93x |
|
||||
|
||||
Additional check:
|
||||
|
||||
- Multi-rank correctness benchmark with checks enabled for `4096` and `122880`
|
||||
index payloads passed.
|
||||
|
||||
Conclusion:
|
||||
|
||||
- IPC materialize is beneficial for index pages too. Because index pages are
|
||||
much smaller than MLA pages, the absolute saving is smaller, but the relative
|
||||
saving versus uint8 all-reduce is still about `~2x` for large prefixes and
|
||||
stronger for short prefixes.
|
||||
- This supports wiring both MLA and index full/partial prefix materialize to the
|
||||
same descriptor-driven IPC path.
|
||||
|
||||
### C102 — 2026-05-31 runtime IPC materialize is wired only for synchronous non-prefetch paths
|
||||
|
||||
Implementation note:
|
||||
|
||||
- The first SGLang runtime hookup targets only the synchronous materialize paths:
|
||||
- MLA full materialize with a precomputed slot remap,
|
||||
- MLA partial-current synchronous prefix materialize,
|
||||
- index full materialize with a precomputed slot remap,
|
||||
- index partial-current synchronous prefix materialize.
|
||||
- `CpSharedKVMlaPrefetcher` and `CpSharedKVIndexPrefetcher` are intentionally
|
||||
unchanged in this slice. Async prefetch still uses the existing local
|
||||
materialize + async collective path until stream-ordering and CPU-submit cost
|
||||
are re-evaluated separately.
|
||||
- The runtime descriptor contract is:
|
||||
`slot_logical_pages -> owner_ranks, src_page_indices`, where source page index
|
||||
is the peer CP-local physical page id including dummy page 0. Invalid slots
|
||||
use `-1/-1` and are zero-filled by the tai-kernel IPC slot-dense materialize
|
||||
kernel.
|
||||
- CUDA IPC handle exchange is cached per long-lived source tensor pointer/shape
|
||||
and CP group. The cache intentionally avoids per-layer handle collectives;
|
||||
a tiny handle/offset all-gather is still required the first time a KV/index
|
||||
source buffer is registered on a rank.
|
||||
|
||||
Important correction:
|
||||
|
||||
- Runtime IPC must carry CUDA allocation offsets. Production source tensors are
|
||||
PyTorch caching-allocator tensors, not guaranteed `cudaMalloc` base pointers.
|
||||
`tai-kernel` now exposes `get_cuda_ipc_mem_handle_with_offset()` and
|
||||
`open_cuda_ipc_mem_handles_with_offsets()` so peers dereference the exact
|
||||
tensor data pointer instead of the allocation base.
|
||||
|
||||
Risk / not yet covered:
|
||||
|
||||
- Peer pointers are cached for long-lived KV/index buffers and are not currently
|
||||
closed before process exit. This is acceptable for the scheduler lifetime but
|
||||
should not be reused for short-lived staging tensors.
|
||||
- The IPC path currently falls back to the old collective materialize path if the
|
||||
tai-kernel IPC capability is missing or peer opening fails; the fallback logs
|
||||
`[CP_SHARED_KV_FALLBACK][tai_ipc_materialize]`. Once remote ETE validates the
|
||||
capability, this can be tightened to fail-fast for the production GLM5 config.
|
||||
- Remote CUDA verification is still required after syncing both `tai-kernel` and
|
||||
`sglang-dev`; local CPU import is blocked by missing optional dependencies.
|
||||
|
||||
Reference in New Issue
Block a user