[AR-Diffusion] Allocate K and V separately so the compiled path stops cloning the pool

[linzhenpl07]

Summary

The AR-Diffusion paged-write custom op declares both key_pool and value_pool as mutated:

@torch.library.custom_op(
    "vllm_omni::ar_diffusion_paged_write_attn",
    mutates_args=("key_pool", "value_pool"),
)

They were the two halves of one allocation:

kv = torch.empty(2, num_blocks, block_size, num_kv_heads, head_dim, ...)
k_pools.append(kv[0].reshape(...))
v_pools.append(kv[1].reshape(...))

so they alias a single storage. Inductor's reinplace pass can then re-inplace only the first of the two mutated arguments:

For node auto_functionalized_v2, attempted to reinplace range(0, 2).
We were unable to reinplace [1]

The second hits the "mutated arg aliases a graph input" case in _inductor/fx_passes/reinplace.py, so auto_functionalized_v2's clone of the entire V pool survives into the compiled graph:

buf13 = empty_strided_cuda((545587200, ), (1, ), torch.bfloat16)   # 1.02 GiB
# Topologically Sorted Source Nodes: [ar_diffusion_paged_write_attn]
triton_poi_fused_1.run(arg10_1, buf13, 545587200, stream=stream0)

and in the FX graph:

%as_strided_default = aten.as_strided(%arg10_1, [545587200], [1], 0)
%clone_default      = aten.clone(%as_strided_default)

That is one full pool copied per compiled region per denoising step. At 480x768 with 40 layers and 4 DMD steps that is roughly 163 GiB of copies per committed chunk.

Fix

Allocate K and V as separate tensors. kv_pools[layer][0] and [1] still address K and V, so key_cache() / value_cache() and every other caller are unchanged.

Measured

One RTX PRO 6000 Blackwell (sm_120), LingBot-World 2.0 at 480x768 bf16, single session, steady state over chunks 6-9 — the sliding window (sink=9 + window=9 frames, 3 frames per block) fills at chunk 6, and latency is flat within 0.4% after that:

	chunk latency	RTF
eager	4.049 s	5.40
compiled, before this change	crashed	—
compiled, after	3.668 s	4.89

1.10x, and the compiled path runs at all.

Before the change, the clone kernel faulted during inductor's Triton config selection:

File ".../torch/_inductor/runtime/triton_heuristics.py", line 1176, in autotune_to_one_config
torch.AcceleratorError: CUDA error: an illegal memory access was encountered

confirmed under CUDA_LAUNCH_BLOCKING=1, so it is the clone kernel itself rather than an async report from an earlier one. On other architectures the copy presumably completes and is simply slow, which is likely why this went unnoticed.

The plateau still occurs at chunk 6 after the change, so compilation introduces nothing that grows with session length.

Ruled out along the way

Recorded because each of these looks plausible and none of them is the cause — all four act on a different layer than the one that decides the clone:

hypothesis	result
allocator holding cached blocks from the fp32 -> bf16 load	torch.cuda.empty_cache() released 0.00 GiB
dynamic=True producing multiple Triton configs	dynamic=False still crashes
torch._dynamo.mark_static_address() on the pools	clone still generated
memory pressure during autotune	12 GiB more headroom, still crashes
int32 byte-offset overflow past 2 GiB	pool reduced to 1.53 GiB/layer, still crashes

The clone is a post-grad FX pass decision; none of the above can influence it. TORCH_LOGS="+torch._inductor.fx_passes.reinplace" names the failing argument directly and is what located it.

Test

tests/diffusion/ar_diffusion/ — 83 pass, ruff check and format clean.

• test_paged_pool_layout_exposes_flat_slot_views updated for the new layout, keeping its real invariant: the flat slot views alias their block-shaped cache, so a slot write is visible through the layout the attention kernel reads.
• test_key_and_value_caches_do_not_share_storage added, pinning what this fix establishes — K and V have distinct storages per layer, the flat views point at the right one, and a negative control that filling V leaves K untouched.

Note

The (2, num_blocks, ...) shape was described as following FlashAttention's block-table layout. The kernel takes k and v as separate arguments (flash_attn_varlen_func(q=..., k=key_cache, v=value_cache, block_table=...)), so the combined allocation was a convenience rather than a requirement; each cache keeps its own (num_blocks, block_size, num_kv_heads, head_dim) shape. If there is a reason to keep one allocation that I have missed, the alternative would be to declare only one of the two as mutated and write V through a different path — but that seems worse.

[chatgpt-codex-connector]

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.
Credits must be used to enable repository wide code reviews.

[vllm-omni-review-bot]

This PR appears to belong to: docs/design/module/cache_management.md.

Module owners: @Isotr0py @princepride @SamitHuang

@linzhenpl07, please review your own changes and leave a short self-review comment describing what you checked. PRs without author self-review may not be assigned a reviewer.

Please take a look when you have a chance. If you would like an automated review, mention @vllm-omni-review-bot in a comment.