From b4f5c3bc0e5c37ffbc062aeee45af3f29d5e09e5 Mon Sep 17 00:00:00 2001 From: laoyao0822 Date: Sun, 31 May 2026 18:28:59 +0800 Subject: [PATCH] Record CP shared KV IPC materialize evidence before runtime wiring The materialize transport direction changed after measurement: direct CUDA IPC peer-page materialize was initially slow because of an under-parallel launch policy, not because slot-dense materialize was inherently too expensive. Capture the corrected evidence and the remaining production constraints before SGLang runtime wiring starts. Constraint: Current near-term goal accepts SM-consuming kernels for low latency/high throughput; low-SM prefetch friendliness is deferred. Rejected: Treat the first slow IPC result as a design blocker | tuned launch parameters beat dense all-reduce across the measured 4k-120k prefix range. Rejected: Switch consumers to owner-concat immediately | slot-dense fused materialize is competitive enough to preserve the current consumer contract for the next integration step. Confidence: medium Scope-risk: narrow Directive: Keep this document updated with every benchmark/result correction to avoid re-litigating stale conclusions. Tested: Documentation update based on remote g0034 tai-kernel CUDA tests and cp_shared_kv_ipc_sm_tuned_20260531_181926 benchmark log Not-tested: SGLang runtime ETE serving with IPC materialize enabled --- ..._prefill_cp_page_aligned_cache_contract.md | 199 ++++++++++++++++++ 1 file changed, 199 insertions(+) diff --git a/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md b/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md index 4c96191e6..d903be5bd 100644 --- a/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md +++ b/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md @@ -3889,3 +3889,202 @@ Design implication: - Expose handles as fixed-size uint8 tensors or keep opened peer pointers in a C++ registry keyed by an integer transport id. Avoid Python object handles in production hot paths. + +### C94 — 2026-05-31 CUDA IPC handle offset issue with PyTorch allocator tensors + +Finding: + +- A first distributed CUDA IPC logical-dense gather using PyTorch-allocated + staging tensors produced correct self-owned pages but corrupted remote-owned + pages. The failure pattern was roughly `1 / cp_size` correct data and the rest + invalid, which pointed to remote IPC pointer interpretation rather than owner + metadata. +- The root cause is that CUDA IPC exports an allocation handle, while PyTorch's + caching allocator may place a tensor data pointer at an internal offset within + a larger allocation. Opening the IPC handle on a peer can return the allocation + base, not the tensor's exact `data_ptr()`. Existing SGLang custom all-reduce + code avoids this by using long-lived `cudaMalloc` IPC buffers or by carrying + `(handle, offset)` metadata. + +Resolution in the prototype: + +- `tai-kernel` now has a CUDA IPC-exportable staging allocation helper backed by + `cudaMalloc`. The benchmark copies compact owner pages into this pointer-stable + byte buffer before exporting the handle. +- Production integration must either use long-lived `cudaMalloc` staging buffers + or carry and apply explicit IPC handle offsets. Do not export arbitrary + PyTorch cache-allocator tensors and assume the opened peer pointer equals + `data_ptr()`. + +### C95 — 2026-05-31 CUDA IPC P2P owner-page benchmark evidence + +Remote benchmark log: + +- `/mnt/beegfs/cjy/log/cp_shared_kv_ipc_gather_segments_20260531_092319.log` + +Configuration: + +- 8 local CP ranks on g0034, BF16 MLA payload, `page_size=64`, `kv_dim=576`. +- Prefix tokens: `4096 8192 16384 32768 65536 98304 122880`. +- Paths: + - `dense_all_reduce`: local dense clone + NCCL all-reduce. + - `cuda_ipc_logical_dense_gather`: one SM kernel directly loads peer pages into + logical-dense order. + - `cuda_ipc_owner_concat_segments`: one `cudaMemcpyAsync` segment copy per peer + into owner-concat layout; not final consumer layout. + - `cuda_ipc_segments_reorder_logical`: segment copy plus logical-dense unpack. + +p50 GPU ms: + +| tokens | dense all-reduce | IPC direct logical | IPC owner-concat | IPC segment+unpack | +| ---: | ---: | ---: | ---: | ---: | +| 4096 | 0.091 | 0.092 | 0.065 | 0.100 | +| 8192 | 0.115 | 0.160 | 0.093 | 0.108 | +| 16384 | 0.156 | 0.297 | 0.189 | 0.198 | +| 32768 | 0.234 | 0.573 | 0.337 | 0.371 | +| 65536 | 0.392 | 1.177 | 0.662 | 0.725 | +| 98304 | 0.553 | 1.828 | 0.983 | 1.091 | +| 122880 | 0.666 | 2.316 | 1.261 | 1.347 | + +Conclusion: + +- The naive CUDA IPC direct logical-dense kernel is not a production replacement + for NCCL all-reduce. It uses SM peer loads and becomes much slower as prefix + grows. +- Segment-level `cudaMemcpyAsync` P2P is better than direct page-gather but is + still slower than NCCL dense all-reduce for large MLA prefixes once the final + logical-dense layout is required. +- The useful takeaway is architectural, not performance-ready: CUDA IPC is viable + as a transport substrate, but a production win likely requires a more fused + layout contract that avoids a full logical-dense rewrite or uses a consumer + that can read owner-concat/segmented layout directly. Do not replace the current + materialize collective with this prototype path. + +### C96 — 2026-05-31 attention idx/page correspondence is the real owner-concat contract + +Finding: + +- The attention-side requirement is not that the materialized KV buffer must be + physically sorted by logical page. The hard requirement is that every attention + idx/loc points to the page slot that contains the corresponding token/page. +- Current CP shared KV already has this abstraction through `page_inverse` and + `remap_logical_locs_to_slot_dense_locs`: logical token locs are rewritten to + dense slot locs before consumers read the materialized buffer. + +Implication: + +- An owner-concat or segmented buffer can be a valid consumer layout if we build + `page_inverse[logical_page] = owner_concat_slot` and remap attention locs to + `owner_concat_slot * page_size + offset`. +- This avoids requiring a full owner-concat -> logical-dense rewrite, as long as + every downstream consumer uses remapped locs/page ids and no kernel assumes + `physical_slot == logical_page` or monotonic logical ordering. +- The production audit should therefore focus on consumers that bypass the + `page_inverse` remap or derive page slots directly from logical positions. + +### C97 — 2026-05-31 CUDA IPC fused materialize must live in CUDA extension, not Triton + +Finding: + +- Triton kernels can remap/materialize a local tensor that already exists in the + current process, but they cannot own CUDA IPC handle creation/opening or safely + dereference peer IPC pointer tables by themselves. +- The first slot-dense materialize helper was useful for local owner-concat + tensors, but it is not the correct primitive for a fused CUDA IPC transport. + The fused IPC path belongs in `tai-kernel`'s C++/CUDA extension. + +Implemented prototype: + +- Added a tai-kernel CUDA extension op: + `materialize_cuda_ipc_peer_pages_slot_dense(peer_ptrs, dst, owner_ranks, + src_page_indices, page_nbytes, ...)`. +- The op writes the current SGLang consumer contract directly: + - dense page 0 is zeroed as dummy, + - valid slot `i` copies + `peer_ptrs[owner_ranks[i]][src_page_indices[i]] -> dst[i + 1]`, + - negative owner/source entries are zero-filled sentinel slots. +- This keeps attention/index consumers on the existing slot-dense contract while + allowing the transport side to experiment with owner-packed IPC staging. + +Remote verification: + +- Unit: + `PYTHONPATH=python python -m pytest -q tests/nsa_prefill/test_cuda_ipc_gather.py` + on `g0034` passed: `4 passed`. +- Benchmark log: + `/mnt/beegfs/cjy/log/cp_shared_kv_ipc_fused_materialize_20260531_175520.log`. + +Benchmark result summary, BF16 MLA page payload, 8 CP ranks: + +| tokens | dense all-reduce | IPC segment+slot-dense | IPC direct peer slot-dense | +| ---: | ---: | ---: | ---: | +| 4096 | 0.100 ms | 0.094 ms | 0.092 ms | +| 8192 | 0.122 ms | 0.108 ms | 0.160 ms | +| 16384 | 0.158 ms | 0.207 ms | 0.297 ms | +| 32768 | 0.234 ms | 0.402 ms | 0.578 ms | +| 65536 | 0.390 ms | 0.752 ms | 1.196 ms | +| 98304 | 0.551 ms | 1.143 ms | 1.812 ms | +| 122880 | 0.668 ms | 1.366 ms | 2.319 ms | + +Conclusion: + +- The direct SM CUDA IPC peer-load fused materialize is correctness-useful but + was not a production replacement for NCCL all-reduce on large prefixes under + the initial under-parallel launch configuration. +- The segment-copy + local materialize path is better than per-page peer loads, + but still loses to dense all-reduce for large MLA prefixes in this benchmark. +- Do not wire this prototype into production hot paths as a default. The next + performance direction should be either a copy-engine/contiguous segment + transport that avoids SM peer loads, or a consumer contract that can consume + owner-concat/segmented layout without a full rewrite. + +### C98 — 2026-05-31 SM fused IPC materialize launch tuning changes the conclusion + +Correction: + +- C97's negative performance conclusion was caused by a bad launch policy, not + by the fused materialize mapping itself. The default wrapper used + `block_quota=2, num_warps_per_block=32`, which effectively capped the copy + kernel around two CTAs. For a 72 KiB BF16 MLA page, each warp then copied many + pages serially and underutilized the GPU badly. +- The current short-term goal is explicitly SM-consuming low latency/high + throughput. Low-SM/copy-engine prefetch friendliness is deferred. + +Change: + +- The CUDA IPC gather/materialize Python wrappers and benchmark now default to + `block_quota=256, num_warps_per_block=8`. +- This keeps enough CTAs active for large prefixes while avoiding 1024-thread + CTAs from the old `32`-warp default. + +Remote benchmark log: + +- `/mnt/beegfs/cjy/log/cp_shared_kv_ipc_sm_tuned_20260531_181926.log` + +Configuration: + +- 8 local CP ranks on g0034, BF16 MLA payload, `page_size=64`, `kv_dim=576`. +- Prefix tokens: `4096 8192 16384 32768 65536 98304 122880`. +- `block_quota=256`, `num_warps_per_block=8`. + +p50 GPU ms: + +| tokens | dense all-reduce | IPC direct logical | IPC direct slot-dense | +| ---: | ---: | ---: | ---: | +| 4096 | 0.100 | 0.090 | 0.090 | +| 8192 | 0.118 | 0.091 | 0.091 | +| 16384 | 0.160 | 0.096 | 0.096 | +| 32768 | 0.236 | 0.138 | 0.138 | +| 65536 | 0.393 | 0.237 | 0.240 | +| 98304 | 0.554 | 0.347 | 0.347 | +| 122880 | 0.668 | 0.421 | 0.421 | + +Conclusion: + +- Under the tuned SM launch, direct IPC peer-page fused materialize beats dense + all-reduce across the tested 4k--120k MLA prefix range. +- The slot-dense variant has essentially the same cost as logical-dense gather, + so preserving the existing SGLang consumer contract is not the dominant cost. +- 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.