Repository navigation
Conversation
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
left a comment
There was a problem hiding this comment.
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):
- Re-run the GPU evidence on
f17578band add the missing coverage (test_split_prefill.py,test_mixed_execution.py, Huginn + Triton boundary,prefill_uva=Trueengine cases); or - Make
packed_prefillopt-in via the explicitprefill_uvaflag 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
left a comment
There was a problem hiding this comment.
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
f17578bca05934fa2b3d4cb8b841d730d8d08794is a merge commit with parentsdeed3e9and21f59c9. - The test plan says the run was at
deed3e9, whose base was9d601ce. Merging currentmain(21f59c9, including the #118 stage/loop refactor) changed the tree substantially:git diff --stat deed3e9 f17578breports 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
f17578band 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 withoutprefill_uva;vllm_rlt/worker/model_runner.py:~609— therequires_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
- Rerun the GPU regressions at
f17578b, adding at leasttests/test_split_prefill.py,tests/test_mixed_execution.py, Huginn with the Triton backend (requires_boundary_kv/ CODA path), and one engine-levelprefill_uva=Truecase; restate the packed-vs-per-row numerics at the merged head. - Or restrict the capability change: keep Triton's
packed_prefillfalse by default and declare it only whenprefill_uvais explicitly requested, so existing split-prefill and Huginn configurations keep their current kernel selection. - 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>
|
@hsliuustc0106 Following your review, I pushed 50e8f7e, based on
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 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 |
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>
|
@hsliuustc0106 I pushed 9874287 for the non-blocking notes in your reviews. The PR description has a new section for this head.
Both 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 I traced the abort to an existing bug that is also on I did not rerun the PD suite, GSM8K or the performance measurements at this head. The GSM8K regression is still open. |
hsliuustc0106
left a comment
There was a problem hiding this comment.
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.
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>
|
@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:
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. I have not rerun the PD suite or the performance measurements at |
|
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>
|
@bjf-frz Thanks. I merged 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; 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:
Every generated token matches, so there is no first divergence to trace. 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:
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- Raw samples, logs and JSON for all runs above are archived. |
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
mainat 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_uvaopt-in. At merged headf17578b, 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=Trueengine case covers metadata growth and reuse.Final tree
e4cb75f— 2026-10-25main(c6730cc, with #128 and the #141 capture fix) merged into this branch. Diff againstmain:vllm_rlt/kernels/triton_attention.pyandtests/test_attention.pyonly. 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.main59/87, this PR 59/87, token-identical on 87/87, so there is no first divergence.mainis also 87/87 identical to [Kernel] Triton packed paged prefill and split-KV decode attention #1286b02350.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.Review follow-up at
9874287— 2026-10-22This addresses the non-blocking notes in the review of f17578b and the self-review. Both
@triton.jitkernel bodies are unchanged from50e8f7e; only host code, comments and one test change.(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.torch._C._cuda_getCurrentRawStream, falling back totorch.cuda.current_stream(index).cuda_streamif that import fails).PARTSfromtl.num_programs(0).Validation of the exact
9874287tree on one GH200 (SM90), torch 2.14.0+cu130, Triton 3.8.0, transformers 4.55.0, Python 3.12:test_mixed_executionwhile the Python garbage collector was running.50e8f7ein the same environment passed 338/338 in 6 of 6 runs.mainas well:torch.cuda.graphno longer collects garbage before capture (torch.compiler.config.force_cudagraph_gcdefaults toFalse), 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 on50e8f7eand9874287(12cudaErrorStreamCaptureInvalidated, 1 invalid argument). This PR only shifts when the collector happens to run. The fix,gc.collect()before capture, is in a separate PR againstmain, [Bugfix] Collect garbage before CUDA graph capture #141, with a GPU regression test.rtol=atol=0.02, max absolute differences 0.0078125 / 0.015625 / 0.015625, not bitwise equal.Validation of
50e8f7e— 2026-10-21Source:
50e8f7e8b598322a289d63a06f4dfe15b81a22f6, based onf17578bca05934fa2b3d4cb8b841d730d8d08794. 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: #12843aa99eb808d66b545d5f6bc5c592d8041fb3183; additional matched-main GSM8K reference: current #1288a19a9a1d298cdc35765882c29c0fddf2e17664d.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(50e8f7emerged into main4a350344027960a48384626aa646f1765fda6d4f); 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.02against 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 keepspacked_matches_decode=False.Lint:
ruff checkpasses with ruff 0.16.8; all three changed files passruff 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,cmaandUCX_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-routedestructor 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 main21f59c9928f73426dfce0bf9a7fe537758b88ed3as 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.8a19a9a9874287+ global packed capability restored6b02350(per-instance opt-in)9874287(two runs)main21f59c9R/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
6b02350now keeps that opt-in. Against the corrected #128 this PR changes no tokens (0 losses). Againstmainthe #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}.pyThe 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_STATEwas grow-only per(device, stream);9874287bounds 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
deed3e9and the original package environment, not tof17578bor50e8f7e.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: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.
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
Streamobject is built per call.a call needs no memset. The engine decodes under
torch.inference_mode, sothe cached buffers can be inference tensors, which an in-place op outside
inference mode rejects (a first version of this change failed
test_speculative.pythat way; see the test result).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), #12843aa99e,main9d601ce. H100 64 GB, CUDA Triton backend, BF16, Ouro-1.4B (HFByteDance/Ouro-1.4B@574fa66c).ruff check,ruff format --check(ruff 0.16.8).OMP_NUM_THREADS=1 python -m pytest -q.pytest --run-gpuovertest_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 undertorch.inference_mode, then outside, then a smaller batch, then inside again; bitwise equal, ticket counters back at zero.paged_attentioncalled directly: host µs per call in a host-bound enqueue loop, GPU µs per call by CUDA-graph replay.benchmarks/gsm8k.py run --backend native, thencompare) vs [Kernel] Triton packed paged prefill and split-KV decode attention #128 andmain.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 checkpasses.ruff format --checkreports the same 5 files as onmain(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_modecovers it.Bitwise vs #128: 36/36 seeded decode cases equal, max |Δ| 0.
Eager speed (TPOT, mean of the two repeats per arm):
mainmain([Kernel] Triton packed paged prefill and split-KV decode attention #128: 9%).Host time per attention call (run 2):
mainInside 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
Streamobject 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;
main29.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):
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, thencompare; same packages as #128):main: 62 vs 60 correct, 0 lost and 2 gained (the same as [Kernel] Triton packed paged prefill and split-KV decode attention #128; changed BF16 arithmetic, not claimed as an accuracy gain).mainand 466 s for [Kernel] Triton packed paged prefill and split-KV decode attention #128 (those two from an earlier job on the same setup).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:
main(1.019–1.026), within the repeat spread of this measurement. The launch has 32 arguments againstmain's 27, and Triton's per-launch cost grows with them.🤖 Generated with Claude Code
Review fixes and current validation used Codex assistance.