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:
@@ -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.
|
||||
|
||||
Reference in New Issue
Block a user