Preserve CP HiCache page-aligned transfer findings

The LPF direct-kernel work exposed an allocation-dependent performance boundary: layer-page-first reduces descriptors only when host pages have contiguous extents. Capture the benchmark evidence and owner-lane allocator interpretation so future changes do not rediscover the same constraint.

Constraint: Current CP compute-owner allocation selects pages by modulo owner lane, e.g. (page_id - 1) % cp_size.

Rejected: Document LPF as an unconditional production improvement | random and owner-lane same-layout benchmarks are neutral without compact allocation.

Rejected: Treat owner_lane benchmark as full allocator replay | it models the modulo-lane constraint, while free/release history can add fragmentation.

Confidence: medium

Scope-risk: narrow

Directive: Tie any production LPF switch to compact/extent-aware host allocation or gather/staging support.

Tested: Local doc marker check for C115/C116/C117.

Tested: Remote tai-kernel CUDA pytest for the referenced kernel/test changes -> 45 passed in 2.52s.

Not-tested: Full SGLang ETE after changing host allocation policy; no such policy change is included.
This commit is contained in:
laoyao0822
2026-06-02 00:16:03 +08:00
parent 6ef4face89
commit d1627d1da3

View File

@@ -4594,3 +4594,430 @@ C110 verification update:
`test_nsa_cp_utils.py`, `test_cp_shared_kv_layout.py`,
`test_cp_shared_kv_runtime.py`, and `test_cp_hicache_metadata.py`
(`253 passed, 5 warnings, 2 subtests passed`).
### C111 — 2026-06-01 TAI page_first_direct direct transfers must not hide index D2H copies
Finding:
- The TAI `cudaMemcpyBatchAsync` page_first_direct per-layer D2H/H2D paths were
accepting `src_indices` and `dst_indices` tensors and then doing
`src_indices.to(CPU).contiguous()` / `dst_indices.to(CPU).contiguous()` inside
the C++ wrapper before building the memcpy descriptor arrays.
- If a benchmark or future caller passed CUDA indices, the direct path silently
inserted a CUDA-to-host metadata copy and potential host synchronization before
each transfer submission. This made direct-path timing ambiguous and could
hide CPU/control-path regressions.
- SGLang production `CacheController.move_indices()` already uses CPU indices
for `io_backend=direct` + `page_first_direct`; the hidden conversion mainly
affected tai-kernel tests/benchmarks and unsafe future callers.
Correction:
- Tighten the TAI per-layer direct D2H/H2D contract: `src_indices` and
`dst_indices` must be CPU int64 contiguous tensors. CUDA indices now fail
fast with `[CP_HICACHE_FAILFAST][tai_page_first_direct_cpu_indices]` instead
of being copied back to host.
- Update tai-kernel direct tests and benchmarks to pass CPU indices. The
`--indices-on-cuda` benchmark option is kept only as a deprecated error path.
Verification target:
- Build/reinstall tai-kernel remotely, then run the CUDA `test_kvcacheio_lf_pf.py`
direct-path tests on g0034. Local CUDA execution remains intentionally
unsupported.
C111 verification update:
- Local tai-kernel `git diff --check`, Python `py_compile`, and CUDA-skipped
local pytest collection passed.
- Synced tai-kernel direct-path code/tests/benchmarks to g0034 and ran remote
container CUDA tests:
- fail-fast CPU-index contract tests: `4 passed`
- direct LF->PF/PF->LF reference tests: `8 passed`
- full `tests/nsa_prefill/test_kvcacheio_lf_pf.py`: `23 passed`
### C112 — 2026-06-01 Benchmark GFD-style H2D alternatives against current page_first_direct baseline
Finding:
- The GFD article's diagnosis applies to many-small-copy H2D workloads, but the
direct optimization must be compared against our current baseline, not against
naive per-token `cudaMemcpyAsync`.
- While preparing that benchmark, an important layout constraint was identified:
current HiCache `page_first_direct` host layout is `[page, layer, page_size,
...]`. For a fixed `layer_id`, page `N` and page `N+1` are not adjacent in
host memory; they are separated by the payloads of the other layers. Therefore
direct page-run coalescing for the existing per-layer PF->LF H2D path is not
physically valid when `layers > 1`.
- This means a simple "merge consecutive pages into one `cudaMemcpyBatchAsync`
descriptor" optimization is not available for the current per-layer direct
path. A GFD-style staging/gather path remains worth measuring because it can
convert strided/scattered host page-layer fragments into one contiguous H2D DMA,
at the cost of an extra CPU gather.
Correction / benchmark addition:
- Added `benchmark/nsa_prefill/benchmark_hicache_h2d_gfd_style.py` in
tai-kernel. It reports current descriptor count, physically mergeable direct
descriptor count, staging gather fragment count, staging DMA descriptor count,
host fixed-layer gap bytes, submit time, total time, and effective bandwidth.
- Added a pure-CPU unit test for the benchmark helpers to lock the key layout
contract: with `layers=78`, contiguous source/destination page ids still have
`direct_mergeable_descriptors == pages`; only a single-layer layout can merge
fixed-layer page runs directly.
- The current `torch_staging` mode is intentionally a baseline/probe, not a
production GFD implementation. It uses PyTorch CPU gather into pinned staging
plus one non-blocking H2D copy, so it should be interpreted as an upper-bound
control-path prototype rather than as optimized AVX-512 polling-thread code.
Verification update:
- Local tai-kernel `py_compile`, `git diff --check`, and helper tests passed:
`tests/nsa_prefill/test_hicache_h2d_gfd_style_benchmark.py` -> `4 passed`.
- Synced the new benchmark and tests to g0034. Remote container helper tests and
`py_compile` passed: `4 passed`.
- Remote CUDA smoke benchmark ran with `--tokens 1024 --payloads kv index
--src-patterns contiguous random --dst-patterns contiguous --warmup 1
--repeat 2`. It produced CSV rows for `direct_batch` and `torch_staging` and
confirmed the descriptor accounting column stays at `direct_mergeable_descriptors
== current_descriptors` for the 78-layer page_first_direct layout.
Next measurement target:
- Run the benchmark with production-sized tokens (`4k, 8k, 10k, 40k, 120k`) and
both BF16/FP8 payload settings after no inference job is using the remote GPUs.
- If optimized staging is still interesting after the baseline run, implement a
dedicated tai-kernel/GFD-style CPU polling/gather path for index payload first,
because index pages are much smaller and closer to the article's small-fragment
workload than the main KV page payload.
C112 benchmark extension update:
- Extended `benchmark_hicache_h2d_gfd_style.py` with an `owner_lane` page pattern
to model CP owner-lane page distribution. For `cp_size=8`, `cp_rank=3`, the
pattern produces `[3, 11, 19, 27, ...]`.
- Added `benchmark_hicache_h2d_multigpu_submit.py` as a thin multi-GPU entry
point. It can run either Python-thread workers or spawned-process workers,
each bound to a CUDA device, and reports global wall time, per-worker measured
wall time, submit latency, total latency, and aggregate bandwidth.
- The multi-GPU benchmark intentionally separates `global_wall_ms` from
`worker_wall_max_ms`: process/thread startup and CUDA context initialization
can dominate the former, while the latter is the useful transfer-window signal
for driver-lock / submit-scaling analysis.
C112 verification update 2:
- Local tai-kernel helper tests passed: `6 passed`.
- Local `py_compile` passed for both benchmark entry points and the helper test.
- Synced benchmark/test files to g0034; remote container `py_compile` and helper
tests passed: `6 passed`.
- Remote CUDA smoke ran:
- single-GPU production-shape smoke for `4096` logical tokens with `kv/index`,
`owner_lane/random`, `direct_batch/torch_staging`.
- two-GPU thread smoke for `index/direct_batch/owner_lane`: produced CSV with
`workers=2` and separate `global_wall_ms`/`worker_wall_max_ms` columns.
- two-GPU process smoke for the same case: produced CSV and demonstrated that
process startup/context init dominates `global_wall_ms`, so full production
runs should compare `worker_wall_max_ms`, submit latency, and total latency.
C112 layout-compare benchmark update:
- Added `benchmark/nsa_prefill/benchmark_hicache_layout_compare.py` in
tai-kernel to compare the current `page_first_direct` layout against a proposed
`layer_page_first` layout without changing production code.
- The benchmark measures both H2D and D2H for a fixed layer. It reports the
current page_first_direct descriptor count and the descriptor count that would
be possible if the host layout were `[layer, page, page_size, ...]`.
- For `contiguous` logical pages, `layer_page_first` reduces the descriptor count
from pages to one run. For `owner_lane` with CP stride, it does not coalesce
unless the logical owner-lane pages are physically compacted/reordered in the
host layout.
Remote smoke evidence on g0034:
- BF16, 4096 logical tokens, 64 pages:
- contiguous H2D: page_first_direct `0.277 ms`, layer_page_first prototype
`0.146 ms`.
- contiguous D2H: page_first_direct `0.211 ms`, layer_page_first prototype
`0.140 ms`.
- owner_lane H2D/D2H: layer_page_first tensor-copy prototype is slower than
page_first_direct because no page-run coalescing is possible and the
prototype launches many small PyTorch copies.
- FP8 e4m3, 4096 logical tokens, 64 pages:
- contiguous H2D: page_first_direct `0.245 ms`, layer_page_first prototype
`0.103 ms`.
- contiguous D2H: page_first_direct `0.171 ms`, layer_page_first prototype
`0.104 ms`.
- owner_lane again does not benefit without physical compaction.
Interpretation:
- `layer_page_first` can bring clear H2D/D2H speedup when the host pages for a
fixed layer are physically contiguous or can be represented as long runs.
- In the current CP owner-lane distribution, simply changing `[page, layer, ...]`
to `[layer, page, ...]` is not enough if the physical host page ids remain
owner-lane strided. To benefit in CP, we likely need either:
1. compact owner-lane host allocation for each rank's local shard, or
2. a transfer-time staging/gather path for owner-lane strided pages, or
3. a new host layout that stores CP-owned logical page runs compactly per layer.
- Therefore the next decision should not be "switch all HiCache to
layer_page_first" yet. It should be: benchmark/implement a compacted
`layer_page_first` owner-lane mapping, then compare against current direct
page_first_direct under production token sizes.
C112 realistic-scatter benchmark gap:
- The first layout-compare smoke covered contiguous and owner-lane patterns, and
the production-sized sweep covered fully random source pages. Fully random
pages are useful as a worst-case fragmentation bound, but real HiCache host
pages are more likely to be a mixture of short contiguous extents separated by
allocator/eviction gaps.
- To avoid overfitting the layout decision to either ideal contiguous pages or
one-page random fragments, the benchmark should also include an
allocator-fragmented pattern: random short page runs, preserving coalescing
inside each run while scattering runs across the host pool.
- This pattern is especially relevant for comparing `page_first_direct` against
a future `layer_page_first` host allocation. If transferred pages remain
fragmented by physical page id, layout reordering alone can regress; if the
owner/rank-local shard is compacted into long runs, layer-first can reduce
descriptors and improve H2D/D2H.
C112 realistic-scatter implementation / measurement update:
- Added `fragmented` page pattern to the tai-kernel benchmark helpers. It
generates random short contiguous page extents separated by fixed gaps, e.g.
with `--fragment-run-pages 8 --fragment-gap-pages 8` a 3000-page transfer is
represented as 375 coalescible runs instead of either one ideal run or 3000
one-page random fragments.
- Remote g0034 helper verification passed after sync:
`tests/nsa_prefill/test_hicache_h2d_gfd_style_benchmark.py -> 10 passed`.
- Remote 4k-192k sweeps were run for:
- BF16 random source pages:
`/mnt/beegfs/cjy/log/hicache_layout_compare_bf16_random_4k_192k.csv`
- FP8 random source pages:
`/mnt/beegfs/cjy/log/hicache_layout_compare_fp8_random_4k_192k.csv`
- BF16 fragmented source pages:
`/mnt/beegfs/cjy/log/hicache_layout_compare_bf16_fragmented_4k_192k.csv`
- FP8 fragmented source pages:
`/mnt/beegfs/cjy/log/hicache_layout_compare_fp8_fragmented_4k_192k.csv`
Key BF16 observations:
- Fully random pages are worst-case for non-compacted `layer_page_first`: 192k
H2D is about `28.35 ms / 7.80 GB/s` for `layer_page_first same`, worse than
current `page_first_direct` at `7.96 ms / 27.78 GB/s`.
- Fragmented 8-page extents are closer to allocator reality and show the useful
middle ground: 192k BF16 H2D is `6.74 ms / 32.81 GB/s` for current
`page_first_direct`, `5.52 ms / 40.09 GB/s` for `layer_page_first same`, and
`4.50 ms / 49.18 GB/s` for compact `layer_page_first`.
- Compact `layer_page_first` remains the best BF16 path across 4k-192k:
4k H2D improves from current `0.202 ms / 23.38 GB/s` to
`0.132 ms / 35.86 GB/s`; 192k H2D improves from `6.74 ms / 32.81 GB/s` to
`4.50 ms / 49.18 GB/s`.
Key FP8 observations:
- FP8 halves bytes moved, so fixed submit/descriptors matter more. Fully random
non-compacted `layer_page_first same` is very bad (`~3.9-4.0 GB/s`) because it
degenerates into thousands of tiny copies.
- Fragmented FP8 still benefits from layer-first even without full compaction:
192k H2D improves from current `5.43 ms / 20.38 GB/s` to
`4.01 ms / 27.57 GB/s` with `layer_page_first same`, and to
`2.53 ms / 43.71 GB/s` with compact `layer_page_first`.
- Therefore the production decision is not just "change memory dimension order".
The safe performance contract is:
1. if host pages stay fully random, keep current direct path or use a true
gather/staging implementation;
2. if allocator can preserve short extents, `layer_page_first` can already help;
3. if CP owner/rank-local pages can be compacted per layer, `layer_page_first`
gives the clearest H2D/D2H win for both BF16 and FP8.
### C113 — 2026-06-01 Current allocators do not actively optimize for contiguous host page runs
Finding:
- `HostKVCache.alloc()` is a simple FIFO slice of `free_slots`, initialized as a
contiguous token range. It preserves contiguity only while consuming the
never-used tail or when previously freed pages happen to be queued contiguously.
`HostKVCache.free()` appends returned indices to the end of `free_slots` and
does not sort, coalesce, or search for the longest contiguous page run.
- CP device allocation is owner-lane aware, not contiguity aware.
`CPSharedPagedTokenToKVPoolAllocator._select_compute_owner_pages()` selects the
first available pages matching each requested owner lane and then replays the
requested owner sequence. Existing unit tests intentionally cover unsorted
`free_pages` and selection from `release_pages` without sort-merge.
- Therefore the current production allocation behavior cannot guarantee the
short contiguous host extents assumed by the `fragmented` benchmark pattern.
Initial warm state can be contiguous, but after host eviction/release churn the
host page order is driven by radix eviction/release order.
Implication:
- A `layer_page_first` production switch should not rely on today's allocator to
naturally produce compact per-layer runs. To realize the benchmarked compact
path, we need an explicit compact/extent-aware host allocation policy or a
transfer-time gather/staging path.
### C114 — 2026-06-01 Layout benchmark underestimates layer_page_first for fragmented pages
Correction:
- The current layout-compare benchmark is not a fair optimized-kernel comparison
for non-contiguous `layer_page_first` pages. The `page_first_direct` baseline
uses the TAI direct `cudaMemcpyBatchAsync` path, while the `layer_page_first`
prototype currently loops over Python/PyTorch `copy_()` calls for each page
run.
- Therefore poor `layer_page_first same` results on fully random pages should be
interpreted as "many small copies are bad" rather than "the layout is bad".
`page_first_direct` is also affected by page fragmentation, but its benchmark
path is already closer to production direct-batch submission.
- The layout-level signal remains valid: `page_first_direct` cannot make fixed
layer adjacent pages physically contiguous because the host layout is
`[page, layer, page_size, ...]`, while `layer_page_first` can collapse each
contiguous page run to one descriptor because the host layout is
`[layer, page, page_size, ...]`.
Next measurement requirement:
- Before deciding production layout, add an optimized TAI `layer_page_first`
direct-batch path that builds one descriptor per contiguous page run. Compare:
current `page_first_direct`, optimized `layer_page_first same`, and compact
`layer_page_first` under contiguous / fragmented / random page distributions.
### C115 — 2026-06-01 Added optimized TAI layer_page_first direct-batch kernel path
Implementation:
- Added TAI direct-batch per-layer APIs for proposed host layout
`[layer, page, page_size, ...]`:
- `transfer_kv_per_layer_direct_lf_lpf`: device layer-first LF -> host
layer-page-first LPF (D2H backup).
- `transfer_kv_per_layer_direct_lpf_lf`: host LPF -> device LF (H2D load).
- The new path keeps the same fail-fast CPU-index contract as
`page_first_direct`: `src_indices` and `dst_indices` must be CPU int64
contiguous and page-aligned. CUDA indices are rejected to avoid hidden D2H
control-path copies.
- Descriptor construction now coalesces contiguous page runs where both source
and destination token starts advance by exactly `page_size`; each run becomes
one `cudaMemcpyBatchAsync` descriptor per logical KV tensor. This makes the
`fragmented` benchmark contract meaningful: 8-page extents become 8
descriptors instead of 64 for a 4k-token shard.
- `benchmark_hicache_layout_compare.py` now uses the new TAI direct APIs for
`layer_page_first` instead of Python/PyTorch per-run `copy_()` loops.
Verification:
- Remote CUDA unit tests on `g0034` / `sglang-glm5-dev-2`:
`PYTHONPATH=python python -m pytest -q tests/nsa_prefill/test_kvcacheio_lf_pf.py`
passed: `35 passed in 2.44s`.
- Remote benchmark smoke with 4k fragmented pages, 8 layers, item size 96:
- H2D `page_first_direct`: `0.388 ms`, `2.03 GB/s`.
- H2D `layer_page_first same`: `0.104 ms`, `7.54 GB/s`.
- H2D `layer_page_first compact`: `0.088 ms`, `8.90 GB/s`.
- D2H `page_first_direct`: `0.133 ms`, `5.92 GB/s`.
- D2H `layer_page_first same`: `0.100 ms`, `7.89 GB/s`.
- D2H `layer_page_first compact`: `0.086 ms`, `9.12 GB/s`.
- Remote real-shape smoke with 4k fragmented pages, 78 layers, item size 576:
- H2D `page_first_direct`: `0.276 ms`, `17.09 GB/s`.
- H2D `layer_page_first same`: `0.176 ms`, `26.84 GB/s`.
- H2D `layer_page_first compact`: `0.159 ms`, `29.68 GB/s`.
- D2H `page_first_direct`: `0.205 ms`, `22.97 GB/s`.
- D2H `layer_page_first same`: `0.169 ms`, `27.86 GB/s`.
- D2H `layer_page_first compact`: `0.159 ms`, `29.61 GB/s`.
Remaining validation gap:
- This proves the optimized direct path works and removes the unfair Python-loop
benchmark artifact. Full 4k-192k real-shape BF16/FP8 benchmarking is still
needed before making `layer_page_first` the production default.
### C116 — 2026-06-01 Full LPF direct benchmark sweep after optimized direct path
Remote benchmark files on `g0034`:
- `/mnt/beegfs/cjy/log/hicache_layout_compare_lpf_direct_bf16_4k_192k.csv`
- `/mnt/beegfs/cjy/log/hicache_layout_compare_lpf_direct_fp8_4k_192k.csv`
- `/mnt/beegfs/cjy/log/hicache_layout_compare_lpf_direct_bf16_owner_lane_4k_40k.csv`
- `/mnt/beegfs/cjy/log/hicache_layout_compare_lpf_direct_fp8_owner_lane_4k_40k.csv`
Sweep parameters:
- `tokens=4k,10k,40k,120k,192k`, `layers=78`, `item_size=576`,
`page_size=64`, `warmup=2`, `repeat=5`.
- Patterns: `contiguous`, `fragmented(run=8,gap=8)`, `random`.
- Additional CP owner-lane pattern was measured up to 40k. Extending
owner-lane to 192k requires physical page ids up to `3000*8`, which made the
benchmark allocate about hundreds of GB of pinned host layout memory for both
page-first and layer-page-first layouts. That path was terminated and should
be re-measured only after adding sparse/used-page host allocation to the
benchmark.
Key 192k BF16 results, total median time:
| pattern | direction | page_first | LPF same | LPF compact | result |
|---|---:|---:|---:|---:|---|
| contiguous | H2D | 6.292 ms | 4.408 ms | 4.419 ms | LPF same 1.43x |
| contiguous | D2H | 6.196 ms | 4.408 ms | 4.408 ms | LPF same 1.41x |
| fragmented | H2D | 6.290 ms | 4.739 ms | 4.431 ms | LPF same 1.33x, compact 1.42x |
| fragmented | D2H | 6.226 ms | 4.724 ms | 4.428 ms | LPF same 1.32x, compact 1.41x |
| random | H2D | 6.315 ms | 6.326 ms | 4.398 ms | same neutral, compact 1.44x |
| random | D2H | 6.371 ms | 6.201 ms | 4.461 ms | same near-neutral, compact 1.43x |
Key 192k FP8 results, total median time:
| pattern | direction | page_first | LPF same | LPF compact | result |
|---|---:|---:|---:|---:|---|
| contiguous | H2D | 4.607 ms | 2.460 ms | 2.453 ms | LPF same 1.87x |
| contiguous | D2H | 4.520 ms | 2.616 ms | 2.451 ms | LPF compact 1.84x |
| fragmented | H2D | 4.612 ms | 2.737 ms | 2.450 ms | LPF same 1.68x, compact 1.88x |
| fragmented | D2H | 4.560 ms | 2.709 ms | 2.455 ms | LPF same 1.68x, compact 1.86x |
| random | H2D | 4.780 ms | 4.614 ms | 2.452 ms | same near-neutral, compact 1.95x |
| random | D2H | 4.857 ms | 4.506 ms | 2.447 ms | same near-neutral, compact 1.98x |
Owner-lane result at 40k:
- BF16 owner-lane `same` is neutral (`~1.0x`) because pages are strided by
`cp_size=8`, so LPF cannot coalesce descriptors unless allocation is compacted.
LPF compact is `~1.44-1.45x` faster.
- FP8 owner-lane `same` is also neutral (`~1.0x`); LPF compact is
`~1.75-1.78x` faster.
Conclusion:
- Optimized LPF direct path is consistently better when physical host pages have
contiguous extents. The win is larger for FP8 because descriptor overhead is
a larger fraction of total transfer time.
- LPF by itself does not fix fully random or owner-lane-strided host page ids;
in those cases it is neutral unless we also compact/extent-aware allocate host
pages or add a gather/staging path.
- Production switch should therefore be tied to host allocation policy: either
compact CP owner-lane pages per layer, or keep page_first_direct for random
residency and use LPF only when compact extents are guaranteed.
### C117 — 2026-06-02 Owner-lane benchmark models the CP allocator modulo-owner constraint, not full allocation history
Clarification:
- `owner_lane` in the benchmark is not a complete replay of production
allocation history. It isolates one important current allocator constraint:
CP compute-owner pages are selected from modulo lanes.
- In production, `CPSharedPagedTokenToKVPoolAllocator._select_compute_owner_pages()`
chooses pages for an owner where `(page_id - 1) % cp_size == owner`.
Therefore one owner lane is physically strided by `cp_size`, for example
`1, 9, 17, 25, ...` when `cp_size=8`.
- `build_in_seq_page_compute_owners()` generates a zigzag owner sequence for
newly allocated pages. The allocator then maps each owner in that sequence to
the next available page from that owner's modulo lane. Reuse from
`free_pages` and `release_pages` can make the real layout more fragmented or
random over time, but the modulo-lane stride is already enough to prevent
LPF-same descriptor coalescing.
Implication:
- `layer_page_first same` is neutral on `owner_lane` because LPF only changes the
per-page memory order; it does not make owner-lane physical page ids
contiguous. A fixed layer still sees many small descriptors when host page ids
are `owner, owner+cp_size, owner+2*cp_size, ...`.
- Production LPF speedup requires changing allocation/residency policy, not only
tensor layout: allocate compact/extent-aware host pages for each backed node or
add a transfer-time gather/staging path.