[cecilsd]

Required human contributor answers

1. What outcome are you trying to achieve?

Reduce decode latency for looped models using the Triton backend with ordinary CUDA Graphs. The contributor delegated the specific optimization choice and requested real-weight validation on the other supported models.

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

GPU time spent on the per-layer K/V scatter, which repeats in every executed loop. On Ouro-1.4B at four loops, the ordinary graph path executes 192 indexed writes per generated token.

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

Reuse the existing graph KV writer for ordinary unmasked Triton graphs. It writes K and V in one kernel, bringing Ouro's 192 writes down to 96 for the same cache addresses and bytes. The graph still reads the current block and offset tensors on every replay.

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

Operators serving Ouro, Dory, Nanbeige, and Huginn through Triton CUDA Graphs need lower decode token latency with the same generated tokens and cache contents.

PR necessity and bugfix evidence

A PR is needed to change the writer selected by ordinary Triton graph capture. This is a performance change.

Purpose

_DeviceCache currently calls PyTorch indexed assignment once for keys and once for values in ordinary unmasked captures. The existing graph writer already supports both tensors and strided inputs. When no dynamic active-row count is present, every row in the ordinary capture may write. Masked writes retain the existing dispatch and the other attention backends keep their current path.

ExactLoop's observation that per-layer work is multiplied by recurrence motivates targeting this path. ExactLoop preprint

Test Plan

vllm-rlt Version / Commit: a87f6c8 versus ea17646.

• pre-commit run --all-files on Linux.
• pytest --run-gpu tests/test_cuda_graph.py tests/test_dory_gpu.py -k 'not flash_attn', pytest --run-gpu tests/test_huginn.py -k graph and pytest --run-gpu tests/test_async_speculative.py -k 'graph and not flash' on RTX 5090.
• The new regression checks FP32, FP16 and BF16 with strided source tensors. Replay uses changed page addresses and values, compares the entire K/V pools to indexed writes, and verifies the ordinary Triton path avoids indexed assignment.
• BF16 real-checkpoint A/B, one RTX 5090 per run, PyTorch 2.11.0+cu128 and Triton 3.6.0. Replay the current main writer and the production writer with the same loop count, inputs, KV capacity and graph configuration within each pair. Use three interleaved unprofiled repeats at each measured concurrency, with 128 deterministic token IDs per prompt. Count output tokens and token gaps after all measured requests have emitted their first token. Check full generated token ID sequences in greedy and fixed-seed sampled cases.
• Checkpoints pinned to Ouro-1.4B 574fa66 (4 loops), Dory 4B BF16 ee517dd (2 loops), Nanbeige4.2-3B b82e54b (2 loops), and Huginn-0125 bb6621b (32 loops). The benchmark uses 16-token KV blocks and fixes their count per checkpoint at 368, 192, 192, and 1296 respectively. Greedy cases generate 48 tokens per request, or 24 for Huginn. Sampled cases generate 16, with temperature=0.7, top_k=24, top_p=0.9 and seeds starting at 1243.

Test Result

• The regression reproduced the two indexed writes on main in all three dtypes. The patch passed the full Linux pre-commit checks, 145 related CUDA Graph and Dory cases, 3 Huginn CUDA Graph cases, and 12 speculative Graph cases. One speculative case was skipped.
• Independent ops-only and full profiles of one warmed synchronous Ouro decode token both saw 192 indexed writes become 96 Triton writes. Their summed GPU kernel time went from approximately 485 µs to 70 µs. This is kernel time from a diagnostic capture. The latency and throughput table uses unprofiled runs.

Real-checkpoint results

Cells show the median of three runs. Percentages are the median of the three paired relative changes.

Checkpoint	Scheduling	Requests	Decode token/s	p95 ITL (ms)	TTFT (ms)	Peak allocated (MiB)
Ouro	Sync Graph	1	63.44 → 65.26 (+2.67%)	15.91 → 15.43 (-3.21%)	47.42 → 47.68	4007 → 3982
Ouro	Sync Graph	8	352.08 → 361.02 (+2.52%)	22.94 → 22.28 (-2.95%)	121.93 → 121.79	4146 → 4122
Ouro	Async Graph	1	68.59 → 70.75 (+3.15%)	14.67 → 14.22 (-3.08%)	48.32 → 48.10	4104 → 4120
Ouro	Async Graph	8	386.94 → 396.55 (+2.56%)	20.80 → 20.32 (-2.29%)	122.91 → 122.54	4243 → 4259
Dory	Sync Graph	1	53.41 → 54.29 (+1.47%)	18.98 → 18.49 (-3.02%)	62.79 → 60.40	8349 → 8324
Dory	Sync Graph	8	321.06 → 323.58 (+0.78%)	24.99 → 24.87 (-0.70%)	152.06 → 152.00	8626 → 8601
Dory	Async Graph	1	56.97 → 57.81 (+1.46%)	17.61 → 17.40 (-1.39%)	64.41 → 63.37	8670 → 8695
Dory	Async Graph	8	343.67 → 348.48 (+1.40%)	23.42 → 23.07 (-1.42%)	155.55 → 155.49	8947 → 8972
Nanbeige	Sync Graph	1	80.25 → 81.46 (+1.44%)	12.48 → 12.32 (-1.27%)	24.52 → 24.51	8485 → 8461
Nanbeige	Sync Graph	8	511.47 → 520.58 (+1.59%)	15.67 → 15.52 (-1.62%)	128.98 → 128.91	8868 → 8844
Nanbeige	Async Graph	1	85.30 → 86.71 (+1.69%)	11.77 → 11.58 (-1.64%)	25.29 → 25.62	8783 → 8799
Nanbeige	Async Graph	8	558.81 → 568.30 (+1.70%)	14.38 → 14.15 (-1.62%)	130.45 → 129.95	9166 → 9182
Huginn	Sync Graph	1	11.16 → 11.24 (+0.92%)	90.31 → 89.69 (-1.27%)	180.62 → 180.36	10366 → 10333
Huginn	Sync Graph	4	40.14 → 40.64 (+1.21%)	100.55 → 99.32 (-1.04%)	421.13 → 420.26	10610 → 10578

Dory was also tested with 1,056 input tokens, crossing its 1,024-token sliding window. These runs used 560 KV blocks and generated 32 greedy or 16 sampled tokens per request.

Checkpoint	Scheduling	Requests	Decode token/s	p95 ITL (ms)	TTFT (ms)	Peak allocated (MiB)
Dory	Sync Graph	1	47.65 → 48.20 (+1.17%)	21.11 → 20.78 (-1.46%)	504.17 → 506.77	9177 → 9153
Dory	Sync Graph	4	167.23 → 169.16 (+1.16%)	24.17 → 23.71 (-1.87%)	1104.58 → 1103.13	9298 → 9273
Dory	Async Graph	1	50.19 → 50.91 (+1.41%)	20.05 → 19.77 (-1.41%)	507.49 → 511.22	9794 → 9819
Dory	Async Graph	4	176.71 → 179.09 (+1.38%)	22.77 → 22.47 (-1.35%)	1108.51 → 1108.93	9915 → 9940

• All 276 matched request pairs produced identical token IDs, including 69 fixed-seed sampled pairs. All runs reached their expected loop depth and length limit, released their KV blocks, and completed without graph capture or eager fallback in the measured window.
• Huginn was tested with synchronous graph decode, its supported scheduling mode. The optional FlashAttention test configurations were excluded from GPU validation.