[shipiyouniao]

Summary

Part of #490 (packed KV storage). #492 has merged as 055f05569a8a23f930a4de3ff0204a60ecd82142; this PR is now one packed-storage commit directly on that main revision. It no longer carries the profiling change, whose affected contract starts at vLLM 0.27+. #500 and #494 are separate children of this storage step.

vLLM 0.28 CUDA attention backends expose a four-dimensional (blocks, heads, tokens, K+V) cache instead of the five-dimensional layout with a separate K/V axis. The old adapter rejects that shape during graph-profiler temporary-cache initialization.

This change allocates one combined K/V backing pool per layer and builds views using the backend's stride order. It handles NHD and HND, contiguous and non-contiguous layer storage, and differing logical/kernel block sizes. EngineCore and worker use the same physical block geometry. Unsupported layouts and padded page geometries are rejected before an incompatible view is used; padded-page support remains tracked by #426.

Capability discovery includes HND for packed K/V, while legacy split K/V remains NHD-only. An unimplemented backend stride-order method means logical order, matching vLLM's contract. The rebase preserves main's cross-layer-sharing, heterogeneous-view and lifecycle changes.

The packed engine geometry is CUDA-scoped. ROCm retains its existing split-KV path because ROCM_ATTN still uses a five-dimensional cache on 0.28. Packed ROCm backends remain explicitly unsupported rather than assigning incompatible engine/worker geometry; this PR does not claim AMD runtime validation.

The native vLLM 0.28 Triton scale-cache implementation already respects tensor strides and storage offsets. The older five-dimensional KVCached override is therefore not installed on that version.

Storage and View Construction

The worker's _build_packed_kv_views() constructs a view, not a copy or a second K/V allocation. Shape alone is insufficient: NHD and HND expose the same four axes but use different inner strides.

if kv_layout == "NHD":
    inner_strides = (width, heads * width, 1)
elif kv_layout == "HND":
    inner_strides = (kernel_block_size * width, width, 1)

The allocation path selects unified storage for packed K/V and passes the matching buffer count to the native allocator. A logical block split into smaller kernel blocks changes the view's block count and stride, not the scheduler's block-ID contract.

Before

sequenceDiagram
    participant R as V1 runner or graph profiler
    participant B as Attention backend
    participant A as KVCached adapter
    R->>B: Get cache shape
    B-->>R: Four-dimensional packed K/V
    R->>A: Allocate cache
    A->>A: Apply legacy explicit-K/V-axis validation
    A-->>R: Unsupported shape

Loading

After

sequenceDiagram
    participant R as V1 runner or graph profiler
    participant B as Attention backend
    participant A as KVCached adapter
    participant F as Virtual backing storage
    R->>B: Get shape and layout
    B-->>R: Packed dimensions and NHD or HND
    R->>A: Allocate using logical and kernel block sizes
    A->>A: Validate supported geometry
    A->>F: Create unified K/V backing
    A->>A: Build per-layer strided views
    A-->>R: Backend-compatible views
    Note over A,F: Physical pages remain demand-mapped

Loading

For contiguous layer storage, the block stride includes all layers and each view has its own layer offset. For non-contiguous storage, each layer has separate backing. Padding and unsupported layouts fail validation; the diagram is not a claim that every four-dimensional backend format is supported.

Validation

Tested commit: 5b74b3a0f6b12943a0c8fe4b38b9579d3568e782.

• Full pre-commit, GPU-free C++ tests, and MyPy on Python 3.9-3.13 passed.
• CPU suite: 496 passed on each Python version from 3.9 through 3.13.
• Packed KV and heterogeneous-view regression tests with the compiled extension and vLLM 0.28 installed: 62 passed (including backend-contract tests).
• Tesla T4, PyTorch 2.13.0+cu129, vLLM 0.28.0, Qwen2.5-0.5B-Instruct, FP16, TP=1, Triton attention, V1 explicitly selected. CUDA Graphs and prefix caching were enabled; asynchronous scheduling was disabled for this scoped layout comparison.

Serving mode	Startup	Repeated functional prompts	Concurrent requests	Generated output tokens	Final health
Native vLLM	Passed	Stable	16/16	1,024/1,024	HTTP 200
KVCached, non-contiguous	Passed	Matched native	16/16	1,024/1,024	HTTP 200
KVCached, contiguous	Passed	Matched native	16/16	1,024/1,024	HTTP 200

The functional prompts covered a fixed reply, simple arithmetic, a factual question, and Chinese text. Full responses and token counts were checked, not just HTTP status.

A second native run also completed all 16 requests. Concurrent long outputs matched the first native run in 15/16 cases for both KVCached layouts and for the second native run itself. This is not a bitwise-equivalence claim for concurrent generation; the fixed functional responses remained identical across all runs.

No startup/serving CUDA error, illegal access or EngineCore failure was observed. Shutdown is not signed off by this layout test: the native and patched runs logged forced EngineCore termination and resource-tracker warnings, and one patched run logged an output-handler EngineDeadError after shutdown had begun. These are retained in the evidence rather than reported as a completely clean lifecycle run. The GPU returned to idle and the test container did not restart.

All eight real CUDA VMM probes passed, covering both layer layouts, NHD/HND, and logical/kernel block-size ratios 1 and 2. Each probe maps 600 blocks across physical-page boundaries, writes and checks distinct K/V values for three layers, releases the pages, forces a physical-headroom rejection using an actual CUDA allocation, releases that pressure, and verifies all 600 blocks can be allocated and used again. After the external CUDA allocation is freed, the probe allows bounded retries for main's 100 ms physical-capacity cache; every case recovered after five 20 ms retries. The original immediate-recovery assertion no longer matches that merged cache contract. Its initial failure is retained separately and is not counted as a passing run.

The native C++ source tree, build definition and Python VMM binding are identical to the already-built validation binary. That binary was reused and its import verified against this checkout. The Python adapter, CPU matrix and real-device probes above were rerun on the exact head. A dedicated local Codex review completed with no remaining actionable findings.

Scope

This is a V1 packed-storage step, not complete vLLM 0.28 support. HND view correctness is tested independently; full-model serving here covers Triton NHD only. FlashAttention/FlashInfer serving, ROCm packed storage, quantized/padded layouts, V2, hybrid prefix-cache lifecycle, and distributed failure ordering still need their own validation. These single-GPU results do not establish TP/MPS correctness or a performance comparison.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[RixinLiu]

Added FlashAttention 2 coverage on A100. On vLLM 0.28 at 5b74b3a, both KVCached layouts matched native on all 48 complete outputs per layout. Setup: Qwen2.5-0.5B-Instruct BF16, CUDA graphs and APC enabled.

I also reproduced a startup failure on vLLM 0.27: the backend returns packed KV, but the >=0.28.0 gate selects split geometry. Fixed in 054179b by moving the boundary to 0.26 and updating the version tests. After the fix, both layouts on 0.27 also matched native on all 48 complete outputs. The CPU suite passed 501 tests; the focused shape/packed-view suite passed 56.