Repository navigation
[feat]Implement depth-aware KV and asynchronous decode scheduling - #30
Conversation
Add automatic KV sizing, bounded admission bypass, chunked prefill controls, shared KV, reusable buffers, padding, and trace replay. Support random lookahead and delayed exits using the checkpoint Ouro gate in synchronous and asynchronous execution. Add runtime regression tests and local pre-commit checks. Document the unresolved BF16 multi-stream output discrepancy observed with real Ouro weights. Signed-off-by: bjf-frz <frz123db@gmail.com>
hsliuustc0106
left a comment
hsliuustc0106
left a comment
There was a problem hiding this comment.
Reviewed e525f64 against 6e8abbc. Found two P2 issues, detailed inline. Synchronous execution remains a useful fallback while the async path is validated.
Validation: Ruff lint, Python formatting, syntax parsing of all 43 Python files, and git diff --check passed. Ruff 0.16.3 also flagged formatting in one Markdown example in docs/cdb_runtime.md. Runtime tests were not run because isolation for executing the fork's code was unavailable on the reviewer host. The disclosed BF16 multi-stream discrepancy remains unverified.
| * math.ceil(config.max_position_embeddings / cache.block_size) | ||
| * 4 | ||
| ) | ||
| budget = min(free_before, int(total * cache.gpu_memory_utilization) - reserved) |
There was a problem hiding this comment.
[P2] Include reusable allocator memory in the KV budget
After an engine is deleted, PyTorch can retain its freed KV storage in the caching allocator. Both free_before and the utilization limit minus reserved count that storage as unavailable. Recreating an engine in the same process can therefore produce a much smaller cache or fail with "no memory remains" despite sufficient reusable memory. Account for reserved-but-unallocated storage when calculating capacity, while preserving the utilization limit and runtime headroom. See PyTorch's allocator semantics.
| engine = LLMEngine( | ||
| model, | ||
| cache_config=CacheConfig(num_blocks=args.num_blocks, block_size=args.block_size), | ||
| **runtime_configs(args), |
There was a problem hiding this comment.
[P2] Provide a request-ID mapping before enabling HTTP trace replay
Passing the shared runtime configuration here enables --exit-mode trace --exit-trace ... in the HTTP server, but RequestChannel assigns a fresh cmpl-<uuid> ID to every request. LLMEngine.add_request() requires that exact ID in the startup trace, and the HTTP request schema provides no way to select one. Normal replay requests therefore fail trace validation with HTTP 400. Add a replay-ID mapping, or reject trace mode at server startup until the serving interface supports it.
|
run the A/B test |
Owner reviewReviewed Verdict: no new P1/P2 findings; 4×P3. The two P2s from the earlier review remain the open blockers — both re-verified. Earlier P2s — re-verified, still open
New findingsP3 — auto-sizing probe ignores P3 — P3 — cross-module private reach-ins: P3 — docs evidence inconsistencies ( Verified sound (full-diff read + tracing)Async event chains (core round → finalize KV copies → coda gather), readback slot leases (bounded at Validation (reviewed SHA
|
Signed-off-by: bjf-frz <frz123db@gmail.com>
Signed-off-by: bjf-frz <frz123db@gmail.com>
|
Signed-off-by: bjf-frz <frz123db@gmail.com>
hsliuustc0106
left a comment
hsliuustc0106
left a comment
There was a problem hiding this comment.
Re-reviewed c0d7a5a against base 6e8abbc, focusing on changes since e525f64. Found one new P2, detailed inline.
Both earlier P2 findings remain open after rechecking this head:
- KV memory accounting: reusable allocator memory remains excluded, potentially breaking engine recreation.
- HTTP trace replay: fresh UUIDs still cannot match predefined trace IDs.
Validation: Ruff lint, Python formatting, syntax parsing across all 49 Python files, and git diff --check passed. Runtime tests were not run because isolation for executing the fork's code remains unavailable on the reviewer host. BF16 equivalence remains unresolved; campaign artifacts and before/after single-request profiling are not attached.
| major, minor = capability | ||
| automatic = 4 if major in (10, 12) else 3 if major == 9 else 2 | ||
| selected = automatic if backend == "flash_attn" else int(backend.rsplit("_", 1)[1]) | ||
| supported = {2: major in (8, 9), 3: major == 9, 4: major in (9, 10, 12)} |
There was a problem hiding this comment.
[P2] Reject SM12 for the pinned FA4 paged backend
The selector accepts SM12 and automatically chooses FA4, but the pinned flash-attn-4==4.0.0b30 implementation explicitly rejects page tables on SM12. In flash_attn/cute/interface.py:1368, its SM12 branch checks assert page_table is None with the message "Paged KV not supported on SM 12.0 in this PR". Both attention paths in this adapter always supply a page table, so this configuration fails during automatic KV sizing or first inference. Reject this combination with guidance to use Triton until a compatible paged implementation is available, and update the architecture-selection test accordingly.
Verified by static inspection of the exact pinned FA4 wheel, with its SHA-256 checked against PyPI metadata.
There was a problem hiding this comment.
[P2] Reject SM12 for the pinned FA4 paged backend
The selector accepts SM12 and automatically chooses FA4, but the pinned
flash-attn-4==4.0.0b30implementation explicitly rejects page tables on SM12. Inflash_attn/cute/interface.py:1368, its SM12 branch checksassert page_table is Nonewith the message "Paged KV not supported on SM 12.0 in this PR". Both attention paths in this adapter always supply a page table, so this configuration fails during automatic KV sizing or first inference. Reject this combination with guidance to use Triton until a compatible paged implementation is available, and update the architecture-selection test accordingly.Verified by static inspection of the exact pinned FA4 wheel, with its SHA-256 checked against PyPI metadata.
resolved
Signed-off-by: bjf-frz <frz123db@gmail.com>
Signed-off-by: bjf-frz <frz123db@gmail.com>
Signed-off-by: bjf-frz <frz123db@gmail.com>
Signed-off-by: bjf-frz <frz123db@gmail.com>
2 participants
Purpose
Implement depth-aware KV, delayed exit routing and asynchronous decode scheduling for Ouro. The latest update fixes long-context FA4 prefill by packing contiguous query chunks with request-level page tables, and adds bounded CUDA Graph replay for decode recurrent cores.
The table below distinguishes existing functionality from changes introduced by this PR and features that remain deferred.
early_exit_gate.ouro_delayed)flash_attnselects FA2/FA3/FA4 by GPU architecture; explicit versions are available. B300 validation uses FA4 4.0.0b30; FA2/FA3 hardware paths remain unverified.Known limitations: The final 612-trial campaign recorded differences in 196 trials against the respective policy reference: 182 had token differences and 168 had exit-depth differences (overlapping categories). Exact real-model BF16 equivalence and model quality are not established. Graphs help small batches substantially, but high-concurrency E2E can regress because graph coverage and scheduling/batch fragmentation remain limiting factors. See the measured results below. The latest commit contains runtime code and regression tests only; local experiment scripts and analysis documents are excluded.
Test Plan
Runtime parameters and activation
The following options are available in both the CLI and HTTP serving entrypoints.
--attention-backendtriton,flash_attn, or explicitflash_attn_2/3/4flash_attnselects FA4 on B300; no silent backend fallback.--exit-modeouro,ouro_delayed,random_lookahead, ortrace; defaults toouro--lookahead-seed7; defaults to0--async-schedulingouro_delayed,random_lookahead, ortrace. CUDA execution supportstritonandflash_attnbackends. Static buffers are independently opt-in.--single-stream--async-scheduling--cuda-graphs--static-buffers; does not capture prefill, sampling, or Python request bookkeeping.--cuda-graph-max-batch-size128--cuda-graph-max-graphs16--cuda-graph-memory-reserve-bytes1073741824--static-buffers--pad-to-power-of-two--static-buffers--kv-layoutlast_exitedorshared; defaults tolast_exited--gpu-memory-utilization0.9--memory-reserve-bytes268435456--kv-cache-memory-bytes--num-blocks.--num-blocks--prefill-chunk-size128--max-prefill-batches-before-decode1--admission-scan-limit64--max-admission-bypasses8--min-coda-batch-size1--exit-tracedepths_by_requestwith--exit-mode traceExisting request parameters
min_loops,max_loops, andexit_thresholdalso control the new exit modes. Setexit_threshold=1.0to disable adaptive early exit.Start the asynchronous runtime
Activate the existing environment in fish and start the server with the local checkpoint:
Use an available GPU. For a synchronous comparison, remove only
--async-scheduling, keeping the exit mode, cache layout, padding, and request parameters unchanged. For single-stream async, add--single-stream. For SHARED KV, replace--kv-layout last_exitedwith--kv-layout shared.To exercise the random Lookahead path, replace
--exit-mode ouro_delayedwith--exit-mode random_lookahead --lookahead-seed 7.Lookahead weights:
random_lookaheadcurrently uses random, untrained weights. A meaningful Lookahead quality evaluation requires trained weights, such as weights obtained through distillation. Those weights are not available in this implementation yet.ouro_delayedreuses the trained original Ouro gate, but delays its decision by one round; it is not a distilled Lookahead predictor.Send a request
Wait for
/healthto return HTTP 200, then send a completion request:A single request checks generation functionality. Concurrent requests are needed to exercise cross-request overlap. Earlier FP32 comparisons passed for their recorded workloads; exact real-model BF16 multi-stream equivalence remains unresolved. FA4 requires FP16/BF16; use Triton for FP32 checks.
Test Results
Inputs: frozen independent prompts, 128 greedy output tokens, ignore EOS,
max_num_batched_tokens=2048,prefill_chunk_size=2048, LAST_EXITED KV, page size 16 and identical explicit KV capacity per context across variants. Contexts/concurrencies: 1K/[1,32,128], 4K/[1,16,64], 16K/[1,4,16], 64K/[1,2,4]; the 64K prompt is 65,409 tokens. E2E uses 2*C requests with completion-driven refill; decode uses C independent real-prefilled prefixes with prefill excluded. Each context uses one B300 for sequential P8/P9 runs, with order reversed in round two; four contexts run concurrently. Every measurement follows warmup, and graph capture is frozen during measurement.All three rounds are retained, including suspected slow outliers. The new throughput results are appended to the original E2E and decode stacking tables below, using medians of three trials. Latency changes discussed here use medians of per-trial quantiles, not pooled samples. Throughput increases are better; latency decreases are better. Timing is engine host delivery, excluding HTTP/network/tokenization. Small-C tails are not SLO estimates. CPU affinity was not pinned, and external-contention attribution is not established.
Interpretation and remaining issues
walkthrough_artifacts/async_state_p9_20260916/; they are not attached to this PR. This commit contains runtime/dependency changes and regression tests only, with no benchmark scripts or analysis documents.Historical update (
c0d7a5a).b_rdmapre-commitRuntime Python source in this commit matches the frozen source used by the completed campaign. The earlier report that FA4 prefill was still wired exclusively as length-one query sequences is superseded by this fix. Full CUDA Graph execution of all stages is not claimed.
Final protocol and stage definitions
B300, Ouro-1.4B BF16, FA4 4.0.0b30, LAST_EXITED KV, block size 16. Inputs: 1,024 / 4,096 / 16,384 / 65,409 tokens. Output: 128 tokens, greedy, ignore EOS.
max_num_batched_tokens=2048,prefill_chunk_size=2048,max_num_seqs=C. Adaptive threshold 0.2, min/max rounds 2/4; fixed-depth threshold 1.0. E2E uses 2C requests with closed-loop replacement; decode uses C restored independent prefixes. Each configuration has full warmup and three measured trials. Input IDs are frozen LongBench-v2 head/tail-truncated questions; this is not an accuracy evaluation.Decode stages share a common corrected-FA4 prefix. The earlier campaign used Triton prefixes, so cross-campaign decode comparisons are not strict same-prefix A/Bs. E2E includes actual prefill and drain time, excluding HTTP, tokenization, loading and warmup. Stage order is interleaved on the same GPU per context/scope. Percentages are ratios of unrounded trial medians; small changes are not established statistically significant improvements.
P2→P3 changes exit semantics; it is not a pure scheduling change. P8 currently caches actual valid row counts (up to 16 shapes, batch size up to 128) and strips inactive padded rows in its graph path; it is not a power-of-two graph-bucket implementation. Its incremental result includes these implementation choices.
Complete E2E stacking table
Output tokens/s, median of three trials. Each cell after P0 shows
throughput (change vs previous stage); the final columns showP9 / P0 − 1andP9 / P1 − 1against historical baselines. Percentages are calculated from unrounded medians.P0 is historical, not rerun.* Consequently P1 vs P0* and the final baseline comparison are historical comparisons, not matched same-campaign A/B results. Rows without a complete historical P0 baseline show N/A for those comparisons. All P1–P8 values are from the new campaign.
The two columns following historical P8 append the new async-state experiment: P8 rerun with fixed FA4 split=1, then P9 with resident async state and its change against that rerun P8. The original P0–P8 columns retain the historical automatic-split results. The final P9/P0 and P9/P1 columns are cumulative cross-campaign comparisons, not controlled single-feature A/Bs: P0–P7 were not rerun and the FA4 split policy changed. These percentages use the displayed historical baseline values and the unrounded P9 medians; missing historical P0 stays N/A. The split=1 P8 rerun is not an isolated split-policy A/B. All three new repetitions are retained. High-concurrency output/depth differences remain unresolved.
E2E latency — fixed-split P8 → P9
Each cell shows P8 → P9 (percentage change); lower is better. All values are milliseconds, using medians of three per-trial percentiles (not pooled percentiles). Both variants use FA4 split=1; these are the same runs as the P8 split=1 rerun and P9 throughput columns above. All repetitions, including slow outliers, are retained. Small-concurrency p95 values are descriptive, not SLO estimates. TTFT is time to first delivered token; TPOT is each request's average interval after its first token; ITL describes individual delivered-token intervals; E2E is request completion latency. Timing excludes HTTP/network/tokenization overhead.
Complete decode-only stacking table
Output tokens/s, median of three trials. Each cell after P0 shows
throughput (change vs previous stage); the final columns showP9 / P0 − 1andP9 / P1 − 1against historical baselines. Percentages are calculated from unrounded medians.All stages were rerun with a common corrected-FA4 prefix; prefill time is excluded.
The two columns following historical P8 append the new async-state experiment: P8 rerun with fixed FA4 split=1, then P9 with resident async state and its change against that rerun P8. The original P0–P8 columns retain the historical automatic-split results. The final P9/P0 and P9/P1 columns are cumulative cross-campaign comparisons, not controlled single-feature A/Bs: P0–P7 were not rerun and the FA4 split policy changed. These percentages use the displayed historical baseline values and the unrounded P9 medians; missing historical P0 stays N/A. The split=1 P8 rerun is not an isolated split-policy A/B. All three new repetitions are retained. High-concurrency output/depth differences remain unresolved.
Decode-only latency — fixed-split P8 → P9
Each cell shows P8 → P9 (percentage change); lower is better. All values are milliseconds, using medians of three per-trial percentiles (not pooled percentiles). Both variants use FA4 split=1; these are the same runs as the P8 split=1 rerun and P9 throughput columns above. All repetitions, including slow outliers, are retained. Small-concurrency p95 values are descriptive, not SLO estimates. TTFT and request E2E are not reported for decode-only runs because prefix construction occurs outside the measured interval.
Incremental E2E gain or regression
Each column compares with the immediately preceding stage. P1/P0 is omitted because E2E P0 is historical.
Updated analysis
FA4 prefill is fixed, not merely relabeled. The previous adapter represented a chunk as independent length-one query sequences and duplicated the page table per token. The new LAST_EXITED FA4 path packs each contiguous request chunk, supplies query boundaries and the final visible prefix length, and uses request-level page tables with bottom-right causal alignment. It bypasses decode static metadata and does not reduce input length, prefill rounds or token budget. Historical P1→current P1 comparisons show the size of the improvement, but are not interleaved cross-version A/Bs:
Original early exit remains valuable, especially at high concurrency. At the maximum tested concurrency for 1K/4K/16K/64K, E2E P2/P1 gains are about +11.9%/+15.9%/+18.7%/+14.2%. Lower prefill cost makes decode savings more visible. Mean decode rounds are approximately 2.4–2.5 for P2 and 3.3–3.4 for delayed P3. The extra round still consumes a substantial part of the original early-exit benefit; policy-changing comparisons do not establish equal model quality.
Async and multiple streams remain workload-dependent. Maximum-concurrency decode P4/P3 improves about +10.5% at 4K/C64, +9.1% at 16K/C16 and +8.4% at 64K/C4, but loses about 1.1% at 1K/C128. Multiple streams can reverse gains. In this campaign 4K/C64 decode P4/P5 has essentially unchanged recurrent batch size and call count, so its regression cannot be attributed to the old campaign's fragmentation pattern. Other long-context cases still exhibit smaller batches and more calls. Profiling is needed to apportion event, launch and resource-contention costs.
Static E2E regressions are reduced, not eliminated. At 1K/C128 the static increment changes from approximately −14.8% in the old campaign to −1.8% now; at 4K/C64 from −9.6% to −0.6%. This is consistent with removing the old prefill workspace overhead. Decode static regressions remain: approximately −6.4% at 1K/C128 and −12.0% at 4K/C64. Padding has small, mixed effects and is not a uniformly beneficial default.
Graph replay strongly helps small batches. C1 decode gains over P7 are +162.1%/+253.7%/+131.3%/+176.8% for 1K/4K/16K/64K; E2E gains are +151.2%/+142.9%/+93.0%/+32.3%. This is consistent with reducing recurrent-core host/launch overhead, but is not a measurement of GPU idle-time percentage. Only warmed core shapes are graphed; first-use capture latency is excluded from these steady-state numbers.
High-concurrency Graph regressions involve both coverage and batch fragmentation. The table below uses E2E data. Hit rate is the median per-trial fraction of recurrent calls replayed, not token coverage.
16K/C16 and 4K/C16 still regress despite 100% graph hits; increasing cache capacity alone cannot resolve this. Faster core execution changes refill/exit-retirement timing, and the measured batches become smaller with more calls at similar average depth. Stable-input copies, output clones and event protection also cost time. Batch statistics support the fragmentation finding; they do not establish each cost's causal contribution without profiling.
Full stacking is not universally fastest. P8 is fastest in all measured E2E C1 cases and slightly ahead at 16K/C4. P2 is fastest in the other E2E cases. Comparing these policies does not establish answer-quality equivalence. Investigate same-policy output differences, graph batching/bucketing and high-concurrency scheduling before expanding the matrix again.
Correctness and measurement limitations
Code pointers