Skip to content

[Kernel] Split-KV decode: one launch per call in eager mode (stacked on #128) - #130

Open
0z5a wants to merge 5 commits into
InferMatrixLab:mainfrom
0z5a:kernel/split-kv-eager-parity
Open

0z5a wants to merge 5 commits into
InferMatrixLab:mainfrom
0z5a:kernel/split-kv-eager-parity

Conversation

@0z5a

@0z5a 0z5a commented Oct 16, 2026 •

Copy link
Copy Markdown
Contributor

Required human contributor answers

The human contributor must answer all four questions before PR submission.
An explicit answer already recorded elsewhere may be referenced with a precise,
accessible source; a bare link without an identifiable answer is insufficient.
Agents must not invent or fill in the human's answers. Blank, placeholder, or
unresolved answers do not complete this gate. See the
contributor-intent gate
and contributing guide.

1. What outcome are you trying to achieve?

Prefill/decode disaggregation at large scale, for users with larger demand.

2. What is the highest-leverage variable for that outcome?

The split-KV decode kernel and the router, which this work optimizes.

3. How will this work move that variable rather than create busywork?

High-concurrency end-to-end testing and reuse of existing interfaces.

4. Who is the customer or user, and what do they actually need?

Users of other GPUs, i.e. those without Hopper FA4 support.

PR necessity and bugfix evidence

A PR is needed. The human contributor's reason is answer 1: prefill/decode
disaggregation at large scale, for users with larger demand. This follow-up
fixes the eager-mode regression that #128 disclosed.

Observed regression this fixes, disclosed in #128: in eager mode (sync, no
CUDA graphs) the split-KV decode kernel of #128 is about 9% slower per token
than main at small batch (TPOT B/A 1.087 at 512 × 1, 1.090 at 512 × 4,
1.098 at 2K × 1), and the GSM8K-87 harness took 466 s vs 434 s. It was
reproduced on H100 in separate ABBA runs. #128's in-engine attribution: the
whole gap is inside the attention call, from three partial-state allocations
and the second (merge) launch, whose Triton launcher cost inside the engine is
several times what a hot microbenchmark loop shows.

Purpose

Keep the split-KV eager decode optimization and make Triton packed prefill an explicit prefill_uva opt-in. At merged head f17578b, globally advertising packed prefill changed existing mixed execution: four GPU cases produced different tokens from their atomic-prefill oracle. The default capability is now false; requesting packed prefill changes only that selected backend instance.

This follows the review of f17578b. The existing split-prefill and Huginn CODA paths retain their prior kernel selection when UVA is off. A real Triton prefill_uva=True engine case covers metadata growth and reuse.

Final tree e4cb75f — 2026-10-25

main (c6730cc, with #128 and the #141 capture fix) merged into this branch. Diff against main: vllm_rlt/kernels/triton_attention.py and tests/test_attention.py only. All results below are for this tree, on one GH200 (SM90), torch 2.14.1, Triton 3.8.0; see the validation comment for the tables.

  • GSM8K-87 (frozen protocol, Triton, batch 1, eager, 4 loops, greedy): main 59/87, this PR 59/87, token-identical on 87/87, so there is no first divergence. main is also 87/87 identical to [Kernel] Triton packed paged prefill and split-KV decode attention #128 6b02350.
  • Core GPU suite: 395 passed, 3 skipped (FA4-only cases), no capture abort. CPU: 806 passed, 35 skipped.
  • Eager ABBA TPOT, this PR / main: 0.92–0.96 in all four cells and both rounds. The run-to-run spread within a tree was 15–18% on a shared node, so this shows no eager regression rather than a measured speedup.
  • Not rerun: the four-GPU PD suite.

Review follow-up at 9874287 — 2026-10-22

This addresses the non-blocking notes in the review of f17578b and the self-review. Both @triton.jit kernel bodies are unchanged from 50e8f7e; only host code, comments and one test change.

  • Bounded eager workspaces: at most 16 (device, stream) entries keep a partial-state buffer. A new stream evicts the oldest entry, and that bookkeeping runs only on the miss/grow path, so a hit is still one dict lookup. Eviction is safe while the evicted stream's kernels still run: the caching allocator reuses a freed block only on the stream that allocated it, which is the evicted key's own stream. A new CPU test, test_split_kv_eager_workspace_is_bounded_per_stream, covers eviction order, hit reuse and growth.
  • Private stream API: the raw handle is read exactly as Triton's own launcher reads it (torch._C._cuda_getCurrentRawStream, falling back to torch.cuda.current_stream(index).cuda_stream if that import fails).
  • Shared grid dimension 0: a comment at the launch site states that the merge grid's dimension 0 must match the split grid's, because both kernels size PARTS from tl.num_programs(0).

Validation of the exact 9874287 tree on one GH200 (SM90), torch 2.14.0+cu130, Triton 3.8.0, transformers 4.55.0, Python 3.12:

  • Core GPU suite (the command below): 339 passed (338 + the new eviction test) in 6 of 7 full runs, 4 of them with an empty Triton cache. The first run, also with an empty cache, aborted inside a CUDA graph capture in test_mixed_execution while the Python garbage collector was running. 50e8f7e in the same environment passed 338/338 in 6 of 6 runs.
  • That abort is a pre-existing bug, present on main as well: torch.cuda.graph no longer collects garbage before capture (torch.compiler.config.force_cudagraph_gc defaults to False), so reference cycles left by an earlier test's engine (CUDA graphs, streams, pinned buffers) can be freed by an automatic collection in the middle of a later capture. Forcing one collection inside every capture produces the same 13 failures on 50e8f7e and 9874287 (12 cudaErrorStreamCaptureInvalidated, 1 invalid argument). This PR only shifts when the collector happens to run. The fix, gc.collect() before capture, is in a separate PR against main, [Bugfix] Collect garbage before CUDA graph capture #141, with a GPU regression test.
  • Decode numerics vs [Kernel] Triton packed paged prefill and split-KV decode attention #128: 36/36 bitwise equal, max absolute difference 0 (same matrix as below). Packed BF16 prefill: 3/3 shapes within rtol=atol=0.02, max absolute differences 0.0078125 / 0.015625 / 0.015625, not bitwise equal.
  • CPU: 783 passed / 35 skipped locally without Triton (the new eviction test skips there; it ran on the GPU host). Ruff 0.16.8 check and format pass.
  • Not rerun at this head: the four-GPU PD suite and any performance measurement. GSM8K-87 was rerun as a 2×2 attribution (see GSM8K attribution below): the earlier −1.15 pp comes from the packed-prefill selection, and this PR changes no tokens against the corrected [Kernel] Triton packed paged prefill and split-KV decode attention #128.

Validation of 50e8f7e — 2026-10-21

Source: 50e8f7e8b598322a289d63a06f4dfe15b81a22f6, based on f17578bca05934fa2b3d4cb8b841d730d8d08794. The cluster ran the parent plus this exact patch; SHA256 verification of all 186 tracked files binds the tested tree to this commit. Reference decode/legacy GSM8K source: #128 43aa99eb808d66b545d5f6bc5c592d8041fb3183; additional matched-main GSM8K reference: current #128 8a19a9a1d298cdc35765882c29c0fddf2e17664d.

Hardware/runtime: Arrhenius ARM64 NVIDIA GH200 (SM90, 95 GiB reported by CUDA), torch 2.14.0+cu130, Triton 3.8.0, transformers 4.55.0. Model: ByteDance/Ouro-1.4B at 574fa66cb8bf5abdc979642d01cf2b79b16bfab1, BF16. Model, runtime and GSM8K data were staged and verified before GPU submission.

  • CPU suite: 783 passed, 365 skipped (GPU cases are not qualified by this result).

  • GitHub CPU CI run 39643672703 also passes Ruff lint, format and 803 CPU tests / 365 skipped. It checked out the integration merge 91f46f7162e9f7712e3b5687c10902bbe82ef7f6 (50e8f7e merged into main 4a350344027960a48384626aa646f1765fda6d4f); this is separate from the exact-head CPU/GPU evidence above.

  • GPU regression: 338 passed, 0 failed, across attention/64-bit indexing, speculative/async speculative, CUDA graphs, KV, split prefill, mixed execution, Huginn, input preparation and backend selection. 31 FA4 cases deselected because FA4 is unavailable; no selected case was skipped. All four mixed-execution failures from the parent pass with the fix.

  • Decode numerics vs [Kernel] Triton packed paged prefill and split-KV decode attention #128: 36/36 bitwise equal, max absolute difference 0, across FP32/FP16/BF16, head dimensions 32/64/128, four seeds, ragged lengths and noncontiguous physical pages. This claim is about decode, not packed prefill.

  • Packed BF16 numerics: 3/3 shapes × 298 query rows satisfy the declared rtol=atol=0.02 against both per-row decode and a dense FP32 oracle. Maximum absolute differences are 0.0078125, 0.015625, 0.015625. Packed and per-row outputs are not bitwise equal; speculative verification keeps packed_matches_decode=False.

  • Lint: ruff check passes with ruff 0.16.8; all three changed files pass ruff format --check. Unchanged baseline formatting differences are not included in this fix.

  • PD suite: 37 passed, 0 failed, 0 skipped on four GH200s, including real multi-process KV transfers, 2P/2D peer failure, cancellation, prefix reuse and speculative execution. NIXL 1.4.1 with UCX_MODULES=cuda,cma and UCX_TLS=self,sm,cuda_copy,cuda_ipc,tcp; this qualifies same-node correctness, not cross-node RDMA performance.

  • Initial NIXL preflights failed at native interpreter shutdown. GDB identified system/bundled libnl-route destructor interposition. Restricting UCX module loading removed the conflict, followed by two independent clean create/register/deregister/exit checks before the four-GPU regression was admitted.

GSM8K-87 against the original #128 43aa99e: 60/87 → 59/87, with 3 paired losses / 2 gains, delta −1.1494 pp. This fails the predeclared 1 pp regression gate. Tokens are identical on only 62/87 questions. Strict-parse failures are 20 → 19; length-limited responses are 0 → 1 (ID 182, wrong in both arms). Lost IDs: 203/473/881; gained IDs: 819/871. Independent CPU rescoring verified every response, question/prompt hash and four-loop exit trace.

The protocol uses the fixed 87 source IDs, 3-shot raw prompts, greedy BF16, four loops, max_new_tokens 1024 and max_length 2048; dataset revision 740312add88f781978c0658806c59bc2815b9866. Protocol SHA256: c475328ed0eb1d20e825a02541c7a3d8899e4686fecb91748e9c6900f79fa1f0. These are new measurements with the current package set, not the historical 62/87 result. The old reference resumed a verified immutable 60-example prefix after a 900-second timeout; all 87 candidate examples ran together. The resumed reference is used only for accuracy, not speed. This selected regression cohort does not qualify full-dataset or adaptive quality.

A second full reference run uses current #128 8a19a9a1d298cdc35765882c29c0fddf2e17664d, merged with the same main 21f59c9928f73426dfce0bf9a7fe537758b88ed3 as this PR. It also scores 60/87 and is token-identical to the old reference on 87/87 questions. Against this matching base, the candidate still scores 59/87, with the same 3 losses / 2 gains, 62/87 token-identical and −1.1494 pp: FAIL under the unchanged 1 pp gate. The observed loss therefore persists with the matching main revision. Full-model attribution between the fused decode change and the requested prefill-selection change remains open; the finite decode bitwise checks do not establish full-model token equivalence.

GSM8K attribution (one GH200, 9874287). Same frozen 87 questions, prompts, budgets and scorer; only the model path and package versions (torch 2.14.1) were rebased. Triton, batch 1, eager, four loops, greedy.

Arm Tree Prefill Eager merge Correct
R #128 8a19a9a packed (global Triton capability) separate 60/87
F 9874287 + global packed capability restored packed fused 60/87
P #128 6b02350 (per-instance opt-in) per-row separate 59/87
C 9874287 (two runs) per-row fused 59/87, 59/87
M main 21f59c9 per-row main's decode 59/87

R/F and P/C are token-identical on 87/87, and the two C runs are identical, so the fused eager merge changes no tokens. Every pair that switches the prefill path is identical on 62/87, with the same 3 losses / 2 gains as above, so the earlier regression is the packed-prefill arithmetic that #128 enabled globally for Triton. #128 6b02350 now keeps that opt-in. Against the corrected #128 this PR changes no tokens (0 losses). Against main the #128 + #130 stack scores 59/87 vs 59/87 (1 loss, 1 gain, 0.0 pp). R and C match the earlier runs above token for token.

Both full comparisons, including the failed accuracy verdicts, are retained. This is not an accuracy-qualified or fully green result. Raw outputs, source/control manifests and failure logs were archived with per-member SHA256 verification and independently read back; models and datasets remain staged for reproduction.

Reproduction for the core GPU suite:

python -m pytest --run-gpu -q -rA \
  -k 'not flash_attn and not test_packed_metadata_supports_query_lengths_and_per_request_depths' \
  tests/test_{attention,attention_64bit,speculative,async_speculative,cuda_graph,kv_cache,split_prefill,mixed_execution,huginn,input_preparation,backend_selection}.py

The new validation is a correctness/accuracy follow-up. The historical speed/profile numbers below have not been requalified at the merged/fixed head. At 50e8f7e, _EAGER_STATE was grow-only per (device, stream); 9874287 bounds the number of streams that keep a workspace (see above).

Historical measurements — deed3e9 only

The following original report is retained for provenance. Its H100 performance, profiles and accuracy results belong to deed3e9 and the original package environment, not to f17578b or 50e8f7e.

Original pre-merge report

Purpose

Stacked on #128: only the last commit is new. Review it after #128, or against
#128's head 43aa99e.

Small-batch eager decode is host-bound (GPU idle 74–80% in the #128 profile),
and Ouro makes 96 attention calls per token, so every allocation, launch and
kernel argument on that path counts.

Change, in vllm_rlt/kernels/triton_attention.py:

  • One launch in eager mode: the last split of each (row, head) to finish
    merges all parts in the same launch. An acq_rel atomic ticket per (row, head)
    makes the earlier parts visible to it. The merge reads the parts in index
    order, so which program merges never changes the result.
  • Reused buffers: the partial maxima, normalizers and accumulators share
    one FP32 buffer (two fewer kernel arguments than three buffers), kept per
    (device, stream) and grow-only. Calls on one stream are ordered, so they
    never overlap. The key is the raw current-stream handle, as Triton's launcher
    reads it, so no Stream object is built per call.
  • No PyTorch in-place op: the merging program resets its ticket counter, so
    a call needs no memset. The engine decodes under torch.inference_mode, so
    the cached buffers can be inference tensors, which an in-place op outside
    inference mode rejects (a first version of this change failed
    test_speculative.py that way; see the test result).
  • CUDA graph capture: as in [Kernel] Triton packed paged prefill and split-KV decode attention #128, the split kernel plus the separate merge
    kernel, with the (single) buffer in the graph pool. A replay has no launch
    cost, and the plain split kernel is the faster one on the GPU.

Decode numerics in this historical experiment were bitwise identical to #128; this statement does not apply to packed prefill.

Test Plan

vllm-rlt Version / Commit: this PR deed3e9 (tree identical to the tested commit), #128 43aa99e, main 9d601ce. H100 64 GB, CUDA Triton backend, BF16, Ouro-1.4B (HF ByteDance/Ouro-1.4B @ 574fa66c).

  1. Lint/format: ruff check, ruff format --check (ruff 0.16.8).
  2. CPU regression: OMP_NUM_THREADS=1 python -m pytest -q.
  3. GPU regression on this branch: pytest --run-gpu over test_attention.py, test_attention_64bit.py, test_speculative.py, test_async_speculative.py, test_cuda_graph.py, test_kv_cache.py, test_pd.py (FA4 cases excluded: not installed). New tests:
    • test_split_kv_graph_and_eager_merges_are_bitwise_equal (FP32, BF16): eager vs captured bitwise; a smaller call then a larger one (workspace shrink and grow); two streams at once.
    • test_split_kv_eager_workspace_crosses_inference_mode: first call under torch.inference_mode, then outside, then a smaller batch, then inside again; bitwise equal, ticket counters back at zero.
  4. Bitwise: [Kernel] Triton packed paged prefill and split-KV decode attention #128 vs this PR on 36 seeded decode cases (this PR eager, so the fused path).
  5. Microbenchmark, paged_attention called directly: host µs per call in a host-bound enqueue loop, GPU µs per call by CUDA-graph replay.
  6. Eager speed (sync, no graphs), ABBA order on one GPU, CPUs bound to its NUMA node: prompts 512/2K × rows 1/4, 128 output tokens, trace exit at full depth.
  7. Graph speed (async + CUDA graphs, the PD decode worker's configuration), [Kernel] Triton packed paged prefill and split-KV decode attention #128 vs this PR in ABBA order: prompts 2K/8K × rows 1/16.
  8. Eager profiles (untimed) of one 512-token request, 64 tokens: host time per attention call inside the engine.
  9. GSM8K-87 native (benchmarks/gsm8k.py run --backend native, then compare) vs [Kernel] Triton packed paged prefill and split-KV decode attention #128 and main.

Runs: run 1 = main, #128 and the first version of this change (three buffers, Stream-object key), order A1 B1 H1 H2 B2 A2. Run 2 = main, the first version and this PR, order A1 H1 T1 T2 H2 A2, plus #128 vs this PR under graphs. Same harness, node type and controls in both.

Test Result

Lint/format: ruff check passes. ruff format --check reports the same 5 files as on main (none touched here); no new findings.

CPU regression: 596 passed, 193 skipped (GPU-only), 0 failed.

GPU regression: 167 passed, 0 failed, 0 errors, 0 skipped (H100, --run-gpu, FA4 cases excluded), including both new tests.

The first version of the fused path (with counters.zero_() before each launch) failed 2 of 166 there: test_speculative.py::test_reused_hidden_logits_and_all_kv_match_serial_oracle[*-cuda-triton], RuntimeError: Inplace update to inference tensor outside InferenceMode. The serial oracle reused, outside inference mode, buffers the engine had allocated inside it. The counter reset in the kernel removes the in-place op; test_split_kv_eager_workspace_crosses_inference_mode covers it.

Bitwise vs #128: 36/36 seeded decode cases equal, max |Δ| 0.

Eager speed (TPOT, mean of the two repeats per arm):

Prompt × rows main first version this PR first version / main this PR / main #128 / main (run 1)
512 × 1 45.17 ms 47.49 ms 46.03 ms 1.051 1.019 1.086
512 × 4 46.01 ms 48.59 ms 46.94 ms 1.056 1.020 1.089
2K × 1 45.00 ms 47.61 ms 46.18 ms 1.058 1.026 1.092
2K × 4 57.24 ms 50.25 ms 48.83 ms 0.878 0.853 0.912

Host time per attention call (run 2):

Variant Microbenchmark, hot loop, 1×512 Inside the engine, eager profile, mean (under the profiler)
main 20.8 µs —
#128 (run 1) 42.2 µs —
first version 28.0 µs 81.8 µs
this PR 23.6 µs 69.8 µs

Inside the engine the first version's call cost 12 µs more than this PR's, against 4.4 µs in the hot loop: the launcher cost of each extra argument and of the per-call Stream object is larger in the engine.

GPU time per call (CUDA-graph replay, run 2): 1 row × 512 7.1 µs (#128: 6.9), 1 × 8K 57.5 µs, 16 × 2K 191.3 µs, 64 × 2K 741.1 µs; main 29.3 / 669.8 / 228.0 / 708.6 µs. The packed buffer's offsets cost about 0.2 µs at the smallest shape and nothing measurable elsewhere.

Graph speed (TPOT, async + CUDA graphs, run 2):

Prompt × rows #128 this PR this PR / #128
2K × 1 19.79 ms 19.80 ms 1.000
2K × 16 46.66 ms 46.64 ms 1.000
8K × 1 23.46 ms 23.52 ms 1.002
8K × 16 53.18 ms 53.21 ms 1.001

Under CUDA graphs the captured path is #128's two kernels, so graph-mode speed is unchanged.

GSM8K-87 (benchmarks/gsm8k.py run --backend native, eager, then compare; same packages as #128):

Memory: the eager workspace is now persistent per (device, stream) that runs eager decode: rows × heads × 16 × (head_dim + 2) × 4 bytes plus 4 bytes per (row, head), at the largest batch seen. For Ouro-1.4B (16 heads, head_dim 128) that is about 130 KiB per row, or about 8.3 MiB for 64 rows. Captured graphs keep their own buffer in the graph pool, as in #128.

Not done / limitations:

  • Eager small-batch decode is still about 2–2.6% slower per token than main (1.019–1.026), within the repeat spread of this measurement. The launch has 32 arguments against main's 27, and Triton's per-launch cost grows with them.
  • FA4 is not installed on the test machines, so its path was not exercised (unchanged by this PR).

🤖 Generated with Claude Code

Review fixes and current validation used Codex assistance.

0z5a and others added 2 commits October 15, 2026 21:19
Eager decode at small batch is host-bound. Each split-KV call allocated
three partial-state buffers and launched a second (merge) kernel, and inside
the engine Triton's per-launch Python cost grows with the argument count;
Ouro makes 96 attention calls per token.

In eager mode the last split of each (row, head) to finish now merges all
parts in the same launch. An acq_rel atomic ticket per (row, head) makes the
earlier parts visible, and the merge reads the parts in index order, so the
result does not depend on which program merges. The partial maxima,
normalizers and accumulators share one FP32 buffer, reused per (device,
stream) and grow-only, keyed by the raw current-stream handle. The merging
program resets its ticket, so a call needs no memset and PyTorch never
updates the reused buffers in place: the engine decodes under
torch.inference_mode, so they can be inference tensors.

Under CUDA graph capture the separate merge kernel stays: a replay has no
launch cost, and the plain split kernel is faster on the GPU. Both paths are
bitwise identical to the two-kernel split-KV decode.

Tests: eager vs captured bitwise equality (workspace shrink and grow, two
streams at once), and calls crossing torch.inference_mode with the ticket
counters back at zero.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Preserve the split-KV branch's packed prefill and exact-decode verification
capabilities alongside upstream GPU-owned context-length support.

Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>

@hsliuustc0106 hsliuustc0106 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maintainer self-review. Verdict: INC — no defect proven in the code, but the evidence is not on the reviewed head, so it cannot merge yet.

Reviewed head f17578bca05934fa2b3d4cb8b841d730d8d08794 (a merge commit of deed3e9 + 21f59c9). The GPU and GSM8K numbers in the description were produced on deed3e9 when the base was 9d601ce. The tested tree and the head differ in six PR-relevant files (vllm_rlt/attention/backend.py, worker/speculative.py, worker/async_speculative.py, tests/test_backend_selection.py, tests/test_speculative.py, tests/test_async_speculative.py), so the old-SHA measurements cannot validate this head (review rule: no old evidence for a newer revision).

Reach of the capability flag is wider than the description: attention/backend.py:50 enables Triton packed_prefill=True globally. Consumers include worker/model_runner.py:821 (split prefill, atomic_prefill=False + existing Triton settings) and model_runner.py:609 (requires_boundary_kv CODA boundary attention); after metadata.bind, core/kv_cache_manager.py:729-739 switches to the new packed kernel in triton_attention.py:647+, which is not bitwise-identical to the decode kernel (the PR itself notes this in tests/test_speculative.py:44-50). The author's GPU list does not include tests/test_split_prefill.py (which has @pytest.mark.gpu + attention_backend="triton", e.g. :163-174), test_mixed_execution.py, or the Huginn boundary paths — exactly the engine paths touched.

Static review found no arithmetic defect in split/merge (triton_attention.py:366-423), eager buffer reuse (:463-477) or the graph-capture dual kernel (:486-537). Residual unproven risks: the merge kernel may JIT on first use inside graph capture (tests warm up outside capture only, tests/test_attention.py:95-113), and torch._C._cuda_getCurrentRawStream is a private API (:465).

What was run (CPU, no CUDA): full pytest -q -m "not gpu" → 781 passed, 34 skipped, 330 deselected; targeted test_attention*/test_backend_selection/test_speculative → 102 passed; ruff check clean, ruff format --check 111 files. Not run: GPU, CUDA Graph capture, GSM8K, FA4.

To unblock (either path):

  1. Re-run the GPU evidence on f17578b and add the missing coverage (test_split_prefill.py, test_mixed_execution.py, Huginn + Triton boundary, prefill_uva=True engine cases); or
  2. Make packed_prefill opt-in via the explicit prefill_uva flag instead of turning it on globally, keeping the default path unchanged.

Also note the author-reported eager small-batch regression (1.9–2.6% slower than main, H100) must stay in the description as a diagnostic measurement.

@hsliuustc0106 hsliuustc0106 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Holding review until GPU evidence is reproduced at the current head

The change itself looks reasonable on inspection, and the CPU side is green here (pytest -q: 781 passed / 364 skipped at f17578b; ruff check passes). But we cannot accept the GPU/accuracy evidence as validation of this head, for two concrete reasons.

1. The evidence SHA is not the head SHA

  • Head f17578bca05934fa2b3d4cb8b841d730d8d08794 is a merge commit with parents deed3e9 and 21f59c9.
  • The test plan says the run was at deed3e9, whose base was 9d601ce. Merging current main (21f59c9, including the #118 stage/loop refactor) changed the tree substantially: git diff --stat deed3e9 f17578b reports 118 files, +15445/−5439.
  • So the GPU tests, GSM8K-87 comparison, and the "bitwise identical to #128" claim were all produced against a tree that is no longer the PR head. Please rerun them at f17578b and say so explicitly in the test result section.

2. Triton packed_prefill is enabled globally, not only on the prefill_uva path

vllm_rlt/attention/backend.py:50 sets packed_prefill=True for the Triton capabilities. Beyond the prefill_uva gate, this capability is also consumed by:

  • vllm_rlt/worker/model_runner.py:~820 — packed=batch.has_chunks and capabilities.packed_prefill, i.e. split/chunked prefill on Triton without prefill_uva;
  • vllm_rlt/worker/model_runner.py:~609 — the requires_boundary_kv (Huginn) CODA boundary attention path, also gated only by the capability.

That means this PR changes kernel selection for existing Triton configurations, and your GPU regression list (test_attention, test_attention_64bit, test_speculative, test_async_speculative, test_cuda_graph, test_kv_cache, test_pd) does not include tests/test_split_prefill.py, tests/test_mixed_execution.py, or Huginn with Triton — exactly the paths touched. Your own test comment (tests/test_speculative.py:48-51) already records that the packed kernel is not bitwise equal to the decode kernel, so packed numerics need their own validation.

Requested before we can evaluate merge

  1. Rerun the GPU regressions at f17578b, adding at least tests/test_split_prefill.py, tests/test_mixed_execution.py, Huginn with the Triton backend (requires_boundary_kv / CODA path), and one engine-level prefill_uva=True case; restate the packed-vs-per-row numerics at the merged head.
  2. Or restrict the capability change: keep Triton's packed_prefill false by default and declare it only when prefill_uva is explicitly requested, so existing split-prefill and Huginn configurations keep their current kernel selection.
  3. Update the test-plan section with the head SHA the numbers belong to.

Two smaller notes from the review environment, not blockers: _EAGER_STATE grows per (device, stream) without bound, and the split/merge layout assumes both kernels share grid dimension 0 (states = tl.num_programs(0) * NUM_HEADS * SPLITS), which is worth a comment or assertion. Also note ruff format --check is version-sensitive (the repository does not pin ruff), so state the ruff version alongside the result.

We will re-review as soon as the evidence is attached to the current head; nothing here asks for a code rewrite.

Preserve split-prefill, mixed execution and Huginn kernel selection when prefill_uva is disabled. Cover per-selection capabilities and the real Triton UVA engine path.

Co-authored-by: Codex <noreply@openai.com>
Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
@0z5a

0z5a commented Oct 21, 2026

Copy link
Copy Markdown
Contributor Author

@hsliuustc0106 Following your review, I pushed 50e8f7e, based on f17578b. The PR description now separates validation of this exact tree from the historical deed3e9 measurements. The new GSM8K comparison fails the predeclared accuracy gate, as detailed below.

  • Capability scope: Triton now advertises packed_prefill=False by default. An explicit packed-prefill requirement from prefill_uva enables it only for that selected backend instance. Added tests cover explicit/auto selection, absence of capability leakage into later selections, and a real Triton prefill_uva=True engine case with metadata growth/reuse. The original merged head reproduced four mixed-execution GPU failures; all four pass with this restriction.
  • Current-tree GPU coverage: Arrhenius ARM64 GH200, torch 2.14.0+cu130, Triton 3.8.0. The core suite passes 338 tests, including split prefill, mixed execution, Huginn/Triton CODA, input preparation, speculative/async speculative, CUDA graphs, KV and attention. 31 FA4 cases were deselected because FA4 is unavailable; no selected case skipped. The separate four-GPU PD suite passes 37/37, covering real multi-process KV transfer and 2P/2D failure handling. This is same-node correctness coverage.
  • Numerics: decode versus [Kernel] Triton packed paged prefill and split-KV decode attention #128 passes 36/36 bitwise comparisons, max absolute difference 0 across FP32/FP16/BF16, head dimensions 32/64/128 and four seeds. Packed BF16 prefill passes the declared rtol=atol=0.02 against both per-row decode and a dense FP32 oracle on three shapes, 298 query rows each; maximum absolute differences are 0.0078125 / 0.015625 / 0.015625. Packed prefill is not bitwise equivalent, and packed_matches_decode=False remains in effect.
  • CPU/lint: the exact head has 783 passed / 365 skipped; Ruff 0.16.8 check and formatting of the three changed files pass. GitHub CI also passes lint/format and 803 CPU tests / 365 skipped, on integration merge 91f46f7 (50e8f7e into main 4a35034); that is separate from the exact-head GPU evidence.

GSM8K-87: #128 60/87 → this head 59/87, with 3 paired losses / 2 gains, −1.1494 percentage points, exceeding the frozen 1 pp loss budget. A second complete #128 reference at 8a19a9a, merged with the same main 21f59c9 as this head, also scores 60/87 and is token-identical to the old 43aa99e reference on 87/87 questions. The candidate is token-identical to either reference on 62/87. Thus the observed loss persists with the matching main revision. Strict-parse failures are 20 → 19; length-limited outputs are 0 → 1. All three runs' raw responses were independently rescored and checked for matching IDs, prompt hashes and four-loop depth. The fixed protocol is greedy BF16 Ouro-1.4B, four loops, 3-shot raw prompts, max_new_tokens 1024, max_length 2048; model/data pins and protocol hash are in the updated description.

The accuracy regression remains open; the finite decode comparisons do not establish full-model token equivalence. Full-model attribution between fused decode and the prefill-selection change has not been isolated. The old small-batch eager regression and other performance numbers remain explicitly scoped to deed3e9; ABBA/profiling has not been rerun for this head. The grow-only workspace limitation remains documented. Complete success/failure outputs were archived and checksum-verified; the accuracy threshold was not relaxed.

Keep at most 16 per-stream eager partial-state buffers and evict the oldest
on the miss path only, so a hit is still one dict lookup. Read the raw stream
handle the way Triton's launcher does, with the same public fallback, and
note at the launch site that the merge grid's dimension 0 must match the
split grid's. Kernel sources are unchanged.

Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@0z5a

0z5a commented Oct 22, 2026 •

Copy link
Copy Markdown
Contributor Author

@hsliuustc0106 I pushed 9874287 for the non-blocking notes in your reviews. The PR description has a new section for this head.

  • _EAGER_STATE growth: at most 16 (device, stream) entries keep a workspace, and a new stream evicts the oldest one. Eviction runs only on the miss/grow path, so a cache hit is still a single dict lookup. Evicting while that stream's kernels are still running is safe, because the caching allocator reuses a freed block only on the stream that allocated it. A new CPU test covers eviction order, hit reuse and growth.
  • Private stream API: the handle is now read the same way Triton's own launcher reads it, torch._C._cuda_getCurrentRawStream with a fallback to torch.cuda.current_stream(index).cuda_stream.
  • Grid dimension 0: the launch site now has a comment saying the merge grid's dimension 0 must match the split grid's.

Both @triton.jit kernel bodies are unchanged from 50e8f7e.

Validation of this exact tree on one GH200: the core GPU suite passed 339/339 in 6 of 7 full runs. The first run aborted inside a CUDA graph capture in test_mixed_execution while the garbage collector was running. 50e8f7e passed 6 of 6 runs in the same environment. Decode is still 36/36 bitwise equal to #128.

I traced the abort to an existing bug that is also on main, not to this change. torch.cuda.graph no longer runs gc.collect() before capture. Reference cycles left by an earlier engine (graphs, streams, pinned buffers) can therefore be freed by an automatic collection in the middle of a later capture. Forcing one collection inside every capture gives the same 13 failures on 50e8f7e and on this head, and collecting before capture removes all of them. That fix is in a separate PR against main, #141, with a GPU regression test that aborts on main and passes with the fix.

I did not rerun the PD suite, GSM8K or the performance measurements at this head. The GSM8K regression is still open.

@hsliuustc0106 hsliuustc0106 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Verdict: hold for accuracy qualification; no additional proven implementation defect. Reviewed 987428740c1024e40ca2bc28a6ba8e711c00a0f2, including the changes since the maintainer review at f17578b, and compared this follow-up with #128's 8a19a9a tree. Both PRs have base/merge base 21f59c9 against current main; their separate merge commits mean #130 is not a literal descendant of #128's merge commit. The remote heads were unchanged before posting.

The fused eager merge addresses the documented launch/allocation overhead while retaining the separate merge under CUDA graphs. The compatibility alternative from the earlier review is now implemented: backend.py:117–123 enables packed prefill only on the selected instance when explicitly required. The default shared capability stays false, packed_matches_decode stays false for Triton, and tests cover explicit/auto selection without capability leakage plus a real prefill_uva=True engine case.

I found no additional proven defect in the fused merge, ticket-counter reset, stream ownership, workspace eviction, graph-owned buffers or speculative gate. The earlier workspace-bound, raw-stream fallback and shared-grid documentation notes are addressed. The prerequisite's global capability finding is tracked separately in the #128 review; it is resolved in this follow-up.

Remaining acceptance gap: the paired accuracy gate still fails. At 50e8f7e, the matching-main #128 reference scored 60/87, versus 59/87 for the candidate, with 3 paired losses / 2 gains and −1.1494 percentage points, exceeding the frozen 1 pp loss budget. The reported comparison remains unresolved, and GSM8K was not rerun at 9874287. Attribution between fused decode and the prefill-selection change is still open. The 36/36 bitwise decode checks and packed-prefill tolerances do not establish full-model accuracy qualification. Resolve/isolate this numerical change and provide matched validation under the declared accuracy contract before treating the follow-up as qualified.

Existing evidence checked: CPU CI 39776318498 passed lint/format and 804 tests / 365 skipped, but actually checked integration commit e8c1880 (9874287 merged into 4a35034). The contributor reports six complete 339-pass core-suite runs at the exact current head and one capture abort; these are contributor device results, not locally reproduced validation. Four-GPU PD validation belongs to 50e8f7e, and historical performance measurements belong to deed3e9; neither was rerun at this head. FA4 remains unqualified.

Posting-time dependency update: the capture-abort fix #141 is now integrated in upstream main at 81a540d. I verified its parents (ff05fe1, 45565cc) and the added pre-capture gc.collect(), despite PR metadata still reporting #141 open. That fix is absent from the reviewed 9874287 snapshot; its merge does not retroactively validate this head or resolve the separate accuracy gap.

This was a source and existing-evidence review using the repository skill. No PR code or GPU jobs were run locally, and no source fixes were made.

0z5a added a commit to 0z5a/vllm-rlt that referenced this pull request Oct 24, 2026
The shared Triton capabilities no longer declare packed_prefill, so split
prefill, mixed execution and Huginn boundary attention keep per-row prefill.
Only a backend selected with required packed_prefill (prefill_uva) gets it,
on that instance. Same change as InferMatrixLab#130's 50e8f7e, applied to this branch.

Co-authored-by: Codex <noreply@openai.com>
Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
@0z5a

0z5a commented Oct 24, 2026

Copy link
Copy Markdown
Contributor Author

@hsliuustc0106 Thanks for the review. I isolated the GSM8K-87 change. It comes entirely from the prefill selection, not from the fused eager merge. With the corrected prerequisite, this PR changes no tokens.

I ran a 2×2 on one GH200 with the frozen 87-question protocol. Records, prompts, budgets and scorer are unchanged; only the model path and the package versions (torch 2.14.1) were rebased. The engine is the same as before: Triton, batch 1, eager, four loops, greedy. Each pair of arms differs in one factor only:

Arm Tree Prefill Eager merge Correct
R #128 8a19a9a packed (global Triton capability) separate 60/87
F this PR 9874287 + global packed capability restored packed fused 60/87
P #128 6b02350 (per-instance opt-in) per-row separate 59/87
C this PR 9874287 per-row fused 59/87
C (repeat) this PR 9874287 per-row fused 59/87
M main 21f59c9 per-row (main's decode) 59/87
  • Fused merge: R vs F and P vs C are token-identical on 87/87 questions. The two C runs are also 87/87 identical, so the runs are deterministic.
  • Prefill selection: every pair that changes the prefill path (R vs P, F vs C, R vs C) is token-identical on only 62/87 questions. Each of these pairs has the same 3 losses (203, 473, 881) and 2 gains (819, 871). Six of the 25 differing questions diverge at the first generated token, which comes from prefill.
  • Reproducibility: these results match the earlier runs on a different GH200 token for token. R matches the earlier [Kernel] Triton packed paged prefill and split-KV decode attention #128 reference on 87/87, and C matches the earlier candidate on 87/87.
  • Main: main itself scores 59/87. Against main, this PR (C) has 1 loss (473) and 1 gain (871), 0.0 pp, with 65/87 questions token-identical; the remaining differences come from [Kernel] Triton packed paged prefill and split-KV decode attention #128's split-KV decode arithmetic. The packed-prefill arm R differs from main by +1 (2 gains, 1 loss).

So the 60 → 59 difference in the earlier comparison was the packed-prefill arithmetic that #128 enabled globally for Triton. That is the behavior your #128 review flagged as P1. 6b02350 now restricts it to explicit prefill_uva selection, and main also uses per-row prefill for Triton. Measured against the corrected #128, this PR changes no tokens, so it passes the 1 pp gate with 0 losses. Against main, the whole #128 + #130 stack is within the gate at 0.0 pp. The earlier comparison's verdict stays on record: it compared against the packed-prefill reference.

I have not rerun the PD suite or the performance measurements at 9874287. The raw samples for all runs are archived.

bjf-frz commented Oct 24, 2026

Copy link
Copy Markdown
Collaborator

Please isolate the remaining full-model output differences by comparing #130 against the latest #128 with identical prefill kernel selection, model settings, and runtime environment. The existing comparison changed both the decode implementation and prefill selection, so it does not establish which change caused the divergence.

The 36 bitwise-equal decode cases are useful, but they do not explain why only 62/87 full-model generations match and the declared accuracy gate fails (60/87 → 59/87, with 3 losses and 2 gains). Please trace the first divergence under matched settings, then rerun the paired accuracy check and eager performance comparison on the final commit, including the #141 capture fix. Keep the historical performance results separate from validation of that final tree.

Picks up InferMatrixLab#128 (now merged) and the pre-capture gc.collect() from InferMatrixLab#141.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
@0z5a

0z5a commented Oct 25, 2026

Copy link
Copy Markdown
Contributor Author

@bjf-frz Thanks. I merged main (which now contains #128 and the #141 capture fix) into this branch as e4cb75f. Against main c6730cc, the diff is only vllm_rlt/kernels/triton_attention.py and tests/test_attention.py. Everything below was run on that final tree.

Matched comparison and first divergence. My previous comment has the 2×2 that varies one factor at a time. With the prefill selection held fixed, this PR's decode change is token-identical on 87/87 questions, both with packed prefill (60 vs 60) and with per-row prefill (59 vs 59). The 62/87 matches and the 3 losses / 2 gains appear only in pairs that change the prefill selection. That is the packed-prefill arithmetic #128 originally enabled globally for Triton; 6b02350 made it an explicit opt-in.

Paired accuracy on the final tree. Frozen GSM8K-87 protocol (same records, prompts, budgets and scorer; Ouro-1.4B, Triton, batch 1, eager, 4 loops, greedy), same GH200, same environment (torch 2.14.1, Triton 3.8.0), back to back:

Tree Correct Token-identical to main
main c6730cc 59/87 —
this PR e4cb75f 59/87 87/87

Every generated token matches, so there is no first divergence to trace. main c6730cc is also token-identical (87/87) to #128's head 6b02350, so the #141 fix does not change the numerics. The gate passes with 0 losses.

GPU tests on the final tree (attention, attention-64bit, speculative, async speculative, CUDA graphs, KV cache, split prefill, mixed execution, Huginn, input preparation, backend selection, FlashAttention; FA3 3.0.1, FA4 4.0.0b30): 395 passed, 3 skipped. The skips are the two FA4-only packed-prefill cases and one FA4 fp32-operand case. There was no capture abort, and no test copy was patched. The CPU suite passes (806 passed, 35 skipped).

Eager performance on the final tree. ABBA order (main, PR, PR, main, main, PR, PR, main) on one GH200 GPU, 72 CPU cores, performance governor, the same bench for both trees (Ouro-1.4B bf16, decode 128, 5 requests at concurrency 1 and 2 rounds of 4 at concurrency 4). TPOT medians in ms, ratio = PR / main:

Prompt × concurrency main (4 runs) PR (4 runs) Ratio of medians Round 1 Round 2
512 × 1 85.0 / 93.6 / 90.8 / 94.6 89.1 / 78.4 / 92.4 / 81.2 0.924 0.938 0.937
512 × 4 88.4 / 93.8 / 95.1 / 93.8 91.0 / 80.1 / 91.7 / 85.5 0.940 0.939 0.938
2048 × 1 84.8 / 91.5 / 90.9 / 94.8 92.8 / 79.1 / 92.6 / 81.9 0.956 0.975 0.939
2048 × 4 94.4 / 99.2 / 101.1 / 102.1 98.3 / 89.0 / 98.1 / 89.8 0.938 0.968 0.925

The PR is not slower in any cell or round. The run-to-run spread within one tree is about 15–18%, though (the node was shared with other jobs, and two of the four PR runs sit near 80 ms while the other two sit near 92 ms). That spread is larger than the effect, so I read this as no eager regression on the final tree, not as a measured speedup. I don't claim a number from it.

Historical results, kept separate. The H100 numbers in the PR description (this PR 1.019–1.026× and #128 1.086–1.092× of the then-main) belong to deed3e9 and the old environment, and stay under Historical measurements there; they are not validation of e4cb75f. For context only: main now contains #128, and the historical PR / #128 ratio, 0.936–0.940, has the same size as the ratios above. I did not rerun the four-GPU PD suite on the final tree.

Raw samples, logs and JSON for all runs above are archived.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants