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
This commit is contained in:
laoyao0822
@@ -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.
|
||||
|
||||
Reference in New Issue
Block a user