[0z5a]

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. They chose to
submit the two kernels as one PR; the router is a separate, later PR.

This is a performance change, not a bugfix, so there is no failure evidence to link.

Purpose

Without FlashAttention 4, both attention paths of the Triton backend were
bounded by one program per (row, head) walking a whole context serially.

Prefill (commit 2ab1ba3). Problem: a prompt ran the single-query kernel
once per prompt row, so every key was read once per query row and prefill
traffic grew quadratically. Change:

• add a varlen paged prefill kernel with the FA4 packed interface: cu_seqlens_q, one block table per sequence, bottom-right causal alignment;
• tile 64 queries × 64 keys per program with online softmax;
• gather K/V through the block table, with GQA and non-power-of-two head sizes;
• use exact IEEE products for FP32 inputs.

The Triton backend now declares packed_prefill. Speculative verification
compares against serial decoding token for token, so it needs the packed kernel
to match the decode kernel bitwise. Only FA4 guarantees that, so a new
capability packed_matches_decode (FA4 only) gates it, and speculative
verification keeps per-row attention on Triton.

Decode (commit 43aa99e). Problem: the longest row in a decode batch set
every attention call's latency, about 0.3 ms per call at 10K tokens, and Ouro
makes 96 calls per token. Change:

• split each row's context into NUM_SPLITS = 16 equal, tile-aligned parts, one program per (row, head, split) producing an FP32 partial softmax state, and merge them in a second kernel;
• keep the partition a function of the row's own length only, so a row's result is batch-invariant, which speculative verification relies on;
• keep the grid static under CUDA graphs.

Files changed:

• vllm_rlt/kernels/triton_attention.py: both kernels;
• vllm_rlt/kernels/paged_attention.py: dispatch;
• vllm_rlt/attention/backend.py: capabilities;
• vllm_rlt/worker/speculative.py: the gate;
• tests: tests/test_attention.py, tests/test_backend_selection.py, tests/test_speculative.py.

FA4, the torch backend and the decode path selection are unchanged.

Test Plan

vllm-rlt Version / Commit: B = 43aa99e (this PR), A = main 9d601ce. GPU runs used a 4× H100 64 GB node, the CUDA Triton backend, BF16 weights/activations/KV, and 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 alone: pytest --run-gpu over test_attention.py, test_attention_64bit.py, test_speculative.py, test_async_speculative.py, test_kv_cache.py, test_cuda_graph.py, test_pd.py (FA4 cases excluded: not installed). New tests:
   • test_triton_paged_prefill_matches_per_query_reference: out-of-order pages, GQA, head dims 64/96/128.
   • test_split_kv_decode_matches_reference_and_is_batch_invariant, FP32 and BF16. It covers lengths 0/1/31/32/33/511/512/513/4097/10000, the torch reference, zero output for empty rows, and bitwise equality of a row alone vs in a batch.
4. Speed A/B, A = main and B = this PR, ABBA order (A1, B1, B2, A2) on one GPU:
   • controls: CPUs bound to that GPU's NUMA node; engine as in the PD decode worker (async scheduling, CUDA graphs, Triton, BF16); trace exit at full depth (4 loops every token, equal work in A and B); no prefix caching; fresh random prompts;
   • cells: prompts 2K/8K × concurrency 1/4/16, 128 output tokens each;
   • metrics: TTFT, TPOT and end-to-end latency per request, peak allocated/reserved memory, KV pool size.
5. Single-request profiling (separate diagnostic runs, not timed): one warmed request, 2K prompt and 64 tokens, under torch.profiler, for A and for B. Shown as an A-above-B timeline with CPU step, CUDA runtime and GPU kernel tracks, prefill and decode separated, plus a zoom on one mid-decode engine step on the same time scale.
6. Real-checkpoint numerics: first-token logits and 64-token greedy continuations on real-text prompts, packed vs per-row and B vs a torch (non-Triton) reference.
7. End-to-end PD replay: 2P2D Ouro-1.4B on one node, closed-loop AgentX session traces. Earlier stacked A/B; see the notes in the results.

Test Result

Lint/format: ruff check passes. ruff format --check reports the same 5 files on main itself under ruff 0.16.8; no new findings.

CPU regression: 43aa99e 596 passed, 190 skipped (GPU-only), 0 failed. main 9d601ce on the same host: 595 passed, 185 skipped.

GPU regression, this branch alone: JUnit, H100, --run-gpu, FA4 cases excluded:

• this PR head 43aa99e: 188 passed, 0 failed, 0 errors, 0 skipped;
• the packed-prefill commit 2ab1ba3 alone: 186 passed, 0 failed;
• main 9d601ce on the same node: 182 passed, 0 failed.

The extra tests on the branch are the new kernel tests. No failure is attributable to either commit.

Speed A/B (main vs this PR, ABBA, async + CUDA graphs, fixed 4 loops, 128 output tokens): median per cell; A1/A2 and B1/B2 are the two repeats of each arm.

Prompt × rows	Metric	A1	A2	B1	B2	B/A
2K × 1	TTFT	727.4 ms	727.8 ms	120.0 ms	118.2 ms	0.164
2K × 1	TPOT	34.3 ms	34.3 ms	19.4 ms	19.4 ms	0.566
2K × 4	TTFT	2.79 s	2.79 s	0.476 s	0.470 s	0.170
2K × 4	TPOT	43.0 ms	43.0 ms	25.1 ms	25.1 ms	0.585
2K × 16	TTFT	7.17 s	7.15 s	1.20 s	1.19 s	0.167
2K × 16	TPOT	84.4 ms	84.3 ms	46.8 ms	46.7 ms	0.554
8K × 1	TTFT	10.19 s	10.19 s	0.561 s	0.560 s	0.055
8K × 1	TPOT	82.5 ms	82.5 ms	23.1 ms	23.1 ms	0.280
8K × 4	TTFT	40.0 s	40.0 s	2.22 s	2.21 s	0.055
8K × 4	TPOT	126.0 ms	126.0 ms	38.3 ms	38.3 ms	0.304
8K × 16	TTFT	135.5 s	135.5 s	13.6 s	13.7 s	0.101
8K × 16	TPOT	162.1 ms	162.1 ms	53.3 ms	53.3 ms	0.329

• Repeats of the same arm differ by at most 1.6%, and every request in all four arms completed: 58 per arm.
• Peak memory is allocated 49.10 GiB for A and 48.87 GiB for B; reserved 49.42 / 49.28 GiB.
• With gpu_memory_utilization=0.85, the KV pool is 46.03 GiB for A and 45.80 GiB for B. B profiles 0.23 GiB more transient memory: the split-KV workspace and the packed-prefill tiles.

Single-request profiling: one warmed request, 2K prompt and 64 tokens, async + CUDA graphs. Each arm ran separately under torch.profiler (Kineto/CUPTI); these runs were not part of the timed A/B.

Arm	Phase	Wall	GPU busy	GPU idle	Kernels	Attention	GEMM	Other
A (main)	prefill	739.5 ms	699.2 ms	40.3 ms	5837	614.1 ms	30.1 ms	55.1 ms
B (this PR)	prefill	153.9 ms	97.5 ms	56.4 ms	5448	13.7 ms	29.3 ms	54.4 ms
A (main)	decode	2259.9 ms	2026.0 ms	233.8 ms	349446	1039.9 ms	524.9 ms	464.3 ms
B (this PR)	decode	1365.3 ms	1111.1 ms	254.2 ms	355883	119.9 ms	529.2 ms	465.2 ms

• Prefill: A runs the per-row decode kernel over every prompt row, 614 ms of attention. B's packed kernel takes 13.7 ms. GEMM and other kernels are unchanged, so the 586 ms TTFT saving is all attention.
• Decode: attention drops from 1040 ms to 120 ms, which is 109 ms of split parts plus 11 ms of merge. GEMM and the small kernels are unchanged. The 895 ms decode saving accounts for 97% of the attention saving (920 ms).
• What runs: B's trace shows _paged_prefill_kernel during prefill and _paged_attention_kernel + _paged_attention_merge_kernel during decode. A's trace shows only _paged_attention_kernel in both phases. So the measured speedup comes from the new paths.
• In both arms the GPU sits idle 234–254 ms during decode, from host-side scheduling between graph replays. The PR does not change that idle time.
• Traces (Chrome trace format, open in Perfetto): A (main), B (this PR).

Numerics (Ouro-1.4B BF16, fixed 4 loops):

Comparison	First-token logits	64-token greedy
packed vs per-row prefill (16 prompts, 256–8K)	top-1 16/16, top-5 overlap 5/5, max KL 1.9e-3, max |Δlogit| 0.48	8/16 identical; others diverge at tokens 10–57
this PR vs torch reference (12 prompts, 256–2K)	top-1 12/12, max KL 2.2e-3	token agreement 0.781, 7/12 identical (the per-row decode kernel: 0.695, 6/12)

• Both kernels are batch-invariant: a row computed alone equals the same row in a batch, bitwise.

Earlier stacked A/B (each commit vs its parent on a local tree that also carries the PD routing changes of a separate PR; same node and run per comparison):

Change	Measure	Before	After
packed prefill	single-GPU prefill tok/s, 1K / 4K / 8K / 16K prompt	4.7K / 1.5K / 0.8K / 0.41K	11.9K / 15.7K / 14.3K / 11.9K (2.5×–29×)
packed prefill	PD 2P2D n64, busy-window out tok/s	117.4	301.8 / 297.2
split-KV decode	decode TPOT, async + graphs, 2K × 1 / 4 / 16 rows	35.4 / 37.4 / 44.7 ms	19.7 / 24.7 / 41.2 ms
split-KV decode	decode TPOT, 8K × 1 / 4 rows	83.5 / 86.7 ms	23.3 / 37.2 ms
split-KV decode	PD 2P2D kv routing, n96 / n128 out tok/s	236.5 / 246.3	359.9 / 370.7

• The PD replays completed every request in both arms: 1001 at n64, 1503 at n96, 1994 at n128.
• In the PD replays, 90–96% of requests are token-identical between arms, about the same as a same-kernel repeat (93–98%). This is consistent with a changed BF16 reduction order.

GSM8K-87, native A/B (the repository's benchmarks/gsm8k.py: run --backend native on each checkout, then compare):

• Settings: BF16; 4 loops; greedy; batch 1; max 1024 new tokens; the gsm8k-87 question IDs.
• Packages: lm-eval 0.4.9.2, transformers 4.55.0, datasets 3.6.0, tokenizers 0.21.4, triton 3.7.1, torch 2.12.1.
• The fixture's HF baseline pins torch 2.13.0, so its 59/87 reference is not applied here. This is a paired main-vs-PR comparison.

Arm	Correct	Unparseable	Length-limited	Tokens/question
A (main), run 1	60/87	19	0	105
A (main), run 2	60/87	19	0	105
B (this PR), run 1	62/87	19	0	103
B (this PR), run 2	62/87	19	0	103

• B vs A: 0 lost and 2 gained (delta +2.3 pp, paired stderr 1.6 pp). The harness gate (≥ −1 pp) passes. Answer disagreements: 3 questions (871, 881, 1030).
• The +2 is within noise and comes from changed arithmetic, so it is not claimed as an accuracy gain. The claim is no loss.
• Repeats of each arm are identical token for token. So the 14/87 questions whose token sequences differ between A and B (median first divergence at token 29) differ because of the kernels' changed BF16 arithmetic, not run-to-run noise.
• Timing: this harness runs sync, without CUDA graphs, at batch 1 with ~500-token contexts. There B's generation took 466 s vs 434 s for A (+7%), reproducibly. See the eager-mode section below.

Eager mode (sync, no CUDA graphs; this is the GSM8K harness configuration). Separate ABBA A/B, same controls as above, 128 output tokens:

Prompt × rows	TTFT B/A	TPOT B/A
512 × 1	0.773	1.087
512 × 4	0.458	1.090
2K × 1	0.168	1.098
2K × 4	0.168	0.915

• Regression: eager decode at small batch is about 9% slower per token, roughly +4 ms/token at 1 row.

• Eager profile (512-token request):

  • decode is host-bound, with the GPU idle 74–80% of the time;
  • this PR cuts GPU attention time 5× (285 → 56 ms over 64 tokens);
  • each attention call does more host work, and Ouro makes 96 calls per token.

• Microbenchmark: paged_attention called directly, host µs per call measured in a host-bound enqueue loop.

  Variant	Host µs per call
  main	20.9
  this PR	42
  this PR, buffers preallocated	28.5

  In this loop the extra ~21 µs per call is about 14 µs of per-call partial-state allocations plus about 7.5 µs for the second (merge) launch. At 96 calls per token that is only about 2 ms of the ~4 ms per-token gap, because a hot loop understates the cost inside the engine (next item).

• In-engine attribution (corrected after posting; the first version of this description left about 2 ms per token unattributed). From the eager profile above, the host time of each decode-attention call on the CPU timeline: from the end of the preceding operator to the start of the next operator after the call's own allocations and launches.

  Arm	Host time per attention call, mean (under the profiler)
  main	63.7 µs
  this PR	103.2 µs

  • +39.5 µs per call × 16 calls per engine step ≈ 632 µs, against a measured step-time difference of 598 µs. So the eager gap is inside the attention call; the rest of the step did not get slower.
  • The merge launch is about 23 µs of it: a 19.8 µs Python gap in Triton's launcher, plus 3.3 µs in the driver's launch call. The other ~16 µs are the three extra partial-state allocations and the Python around them.
  • Without the profiler, the ~4 ms per-token gap over 96 calls is also about 40 µs per call, about twice the hot loop's 21 µs. Under the profiler the merge launch alone costs about 3× the loop's 7.5 µs. Triton's per-launch Python cost grows with the number of kernel arguments and is higher inside the engine than in a hot loop.

• GPU time per call (CUDA-graph replay), main vs this PR:

  Shape	main	this PR
  1 row × 512	29.1 µs	6.9 µs
  1 row × 8K	674 µs	57 µs
  16 rows × 2K	228 µs	191 µs
  64 rows × 2K	704 µs	741 µs

  At 64 uniform 2K rows the split adds about 5% GPU time. Real decode batches with mixed lengths are dominated by the long rows.

• Follow-up: a fix is being validated for a follow-up PR. In eager mode it merges in the same launch and reuses the partial-state buffers per stream; under graph capture it keeps the separate merge.

  • Its eager path, the same-launch merge, was measured on its own: 22.0 µs host per call with preallocated buffers, against 20.9 µs for main.
  • Its outputs were bitwise identical to this PR in 36/36 seeded cases, and GSM8K-87 matched token for token.
  • A first attempt that only reused a workspace did not help, because its own per-call lookup cost about as much as the allocations.

• Where it inverts: with 4 rows at 2K the GPU work dominates and the PR is 8.5% faster.

• Under CUDA graphs, the PD decode worker's configuration, launches are replayed and the PR is 1.7–3.6× faster per token (table above).

• TTFT is faster in every eager cell (packed prefill).

Memory:

• The prefill kernel adds no persistent buffers.
• The decode kernel allocates an FP32 workspace per call of rows × heads × 16 × (head_dim + 2) × 4 bytes. For Ouro-1.4B (16 heads, head_dim 128) that is about 130 KiB per row, or about 8.3 MiB for 64 rows. It is transient in eager mode and lives in the CUDA-graph pool when captured.
• Peak memory per arm: see the speed A/B above. With the same utilization setting, B leaves a 0.5% smaller KV pool.

Not done / limitations:

• The eager small-batch decode regression above (about 9% per token at 1 row) is open in this PR; its fix is the follow-up described above.
• NUM_SPLITS = 16 is fixed, not tuned per shape. At 16 rows × 2K the decode gain is small (1.08×), because the GPU is already mostly busy there.
• FA4 is not installed on the test machines, so its path was not exercised (unchanged by this PR).

🤖 Generated with Claude Code

[hsliuustc0106]

Verdict: request changes before #128 merges independently. Reviewed head 8a19a9a1d298cdc35765882c29c0fddf2e17664d, against base/merge base 21f59c9928f73426dfce0bf9a7fe537758b88ed3; the remote head was rechecked before posting.

The tiled prefill and split-KV decode target documented serial attention bottlenecks. Retaining per-row prefill by default, with an explicit packed opt-in, is the simpler compatibility alternative.

[P1] Make packed prefill an explicit opt-in — vllm_rlt/attention/backend.py:50. Setting packed_prefill=True in the shared Triton capabilities changes kernel selection for existing split/mixed execution and Huginn boundary callers even when prefill_uva=False. Those callers consume the capability directly, so this arithmetic change reaches beyond the explicitly requested packed-prefill feature. The identical capability and prefill implementation at #130's f17578b produced four mixed-execution GPU token-oracle failures in the contributor's follow-up validation; all four passed after the capability was restricted to an explicitly selected instance.

Move #130's per-instance opt-in correction into this prerequisite before merging #128 alone. That correction is present in #130's 50e8f7e and current 9874287, but absent from this reviewed #128 head; it should not be counted as already resolved here.

Scope: reviewed the full frozen diff, attention dispatch and metadata callers, causal/page/depth handling, split/merge reductions, speculative routing, tests, descriptions and existing evidence using the repository review skill. No PR code or GPU jobs were run locally for this review, and the reported GPU failures were not independently reproduced here. #128's historical GPU, accuracy and speed results at 43aa99e do not by themselves qualify the merged 8a19a9a tree. No source fixes were made.

[0z5a]

@hsliuustc0106 Thanks. I pushed 6b02350 on top of 8a19a9a. It moves #130's per-instance opt-in into this PR:

• The shared Triton capabilities no longer declare packed_prefill. Split prefill, mixed execution and Huginn boundary attention therefore keep per-row prefill, whatever prefill_uva is set to.
• create_backend enables packed_prefill only on a Triton instance selected with required_capabilities.packed_prefill (the prefill_uva=True path). packed_matches_decode stays false.
• New tests: test_triton_packed_prefill_is_per_selection_opt_in (explicit and auto selection, no capability leak into a later default selection) and a prefill_uva=True Triton case in test_input_preparation.

The diff is the same as #130's 50e8f7e for these three files. #130's other changes are not part of it.

Validation at 6b02350:

• The CPU suite passes (783 passed, 34 skipped) and Ruff check passes.
• One GH200 (torch 2.14.1, Triton 3.8.0, FA3 3.0.1, FA4 4.0.0b30) ran the attention, attention-64bit, split-prefill, mixed-execution, Huginn, input-preparation, backend-selection, CUDA-graph, speculative, KV-cache and FlashAttention tests.
  • With main's pre-capture gc.collect() from [Bugfix] Collect garbage before CUDA graph capture #141 added to a copy of the tree: 357 passed, 3 skipped. The skips are the two FA4-only packed-prefill cases and one FA4 fp32-operand case.
  • Without it, the run aborted inside a CUDA graph capture while garbage was being collected. An earlier run on this tree failed one speculative graph case with an invalidated capture. Both are the intermittent issue [Bugfix] Collect garbage before CUDA graph capture #141 fixes; this branch predates it.

I did not rerun GSM8K or any performance measurement at this head.

[bjf-frz]

lgtm