Forward-port the tkv/turbo-attn vLLM seam onto upstream v0.31.0

[arbi-dev]

Forward-ports the turbo-attn vLLM seam from upstream v0.28.0 to v0.31.0 (db9527a46, the latest stable release; vllm/vllm-openai:v0.31.0 is built from exactly this commit: image label org.opencontainers.image.revision=db9527a46873…, index digest sha256:c1c9f6fd5c10…).

Shape

Same method as #28: the whole carry (git diff v0.28.0 origin/main at 969c125e9) squashed into one commit on v0.28.0 and cherry-picked onto v0.31.0 as a 3-way merge against the real base. Result: v0.31.0 + 1 carry commit.

Carry on fork main (nothing else has merged since #28; no other forward-port branch exists):

commit	what	folded in
b0a14f2eb	#28, the v0.28.0 forward-port	yes
76a3e6d72	fix(kv-seam): hash_block_size on the per-group pool view (post-port)	yes, then dropped with the per-group pools (below)
6d4c75611 / merge 969c125e9	#29 auto-default --attention-backend TURBO_ATTN for plugin kv dtypes (post-port)	yes, kept

49 files / +4160 → 18 files / +1289 (12 of them under vllm/). Most of the shrink is upstream's v0.29–v0.31 KV-cache layout refactor (vllm#51718 and follow-ups): one backing allocation for every layer, per-layer [B, H, N, C] views placed by the engine from the spec's geometry (num_heads, state_content_size_bytes), and AttentionBackend.customize_spec applied to every layer's spec by both model runners. get_kv_cache_shape is no longer called at all. The seam's allocator patches were written against the code that refactor deleted, so they were re-expressed on it or dropped, and the spec-class hook moved to upstream's customize_spec.

This needs the matching turbo-attn plugin change: arbicity/turbo-attn deps/engine-forks-v0.31.0-v0.5.21 @ a2c4bd7d6 (customize_spec, layouts, [B,1,N,C]→3-D views, padded-slot review off on placed views). A v0.31 image without it will not serve tkv.

Conflicts (18 files) and resolutions

file	conflict	resolution
config/cache.py	upstream added layout helpers + device_memory_utilization around our registry and str field	upstream's code verbatim; registry + cache_dtype: str re-applied on top; skip-layers-dtype field dropped (below)
engine/arg_utils.py	upstream's XPU batch-invariant default sits where #29's auto-default was	both, upstream's first; AttentionBackendEnum now imported at module level so our local import went
envs.py	upstream widened VLLM_KV_CACHE_LAYOUT to the new layouts	upstream's; our two env vars dropped (below)
platforms/cuda.py	upstream added TRITON_FLASHINFER/TRITON_FLASH_ATTN mm-prefix candidates	upstream's lists; TURBO_ATTN append re-applied
model_executor/warmup/kernel_warmup.py	import block churn	upstream verbatim; our two warmups dropped (below)
v1/worker/gpu_worker.py	upstream added set_torch_threads_for_runtime() at the end of load_model	hooks run after it (serving thread count); CuTeDSL-cache init and post-capture spec warmup removed
v1/attention/selector.py	upstream dropped get_required_kv_cache_layout + set_kv_cache_layout (layout is now negotiated by the engine core) and the module logger	MLA wrapping re-applied; the layout block stays gone; logger re-added for the wrap log line
model_executor/layers/attention/mla_attention.py	upstream's spec now adds is_index_group_leader and TMA-row block_stride_alignment after construction	upstream's; our get_kv_cache_spec_class("mla") short-circuit dropped (it skipped both). The plugin's MLA wrapper now packs through customize_spec. _get_gather_op merged cleanly onto upstream's new CPU branch
v1/kv_cache_interface.py	KVCacheConfig grew transfer/hisparse fields where per_group_num_blocks was	upstream's; only aggregated_layer_count re-applied
v1/core/kv_cache_utils.py (9 hunks)	upstream rewrote planning: single backing allocation, KVCacheTensor(layers, layer_stride, block_stride, offset), layout-aware strides	upstream's file; fused-page support re-applied as get_tensor_slots in _get_kv_cache_bytes_per_block and the tensor emission loop
v1/core/kv_cache_coordinator.py, kv_cache_manager.py (2+5)	per-group BlockPools vs upstream's shared pool	upstream's (dropped, below)
v1/core/single_type_kv_cache_manager.py	registry lookup rewritten (roles)	upstream's (get_manager_class obsolete, below)
v1/worker/gpu_model_runner.py, v1/worker/gpu/attn_utils.py	_reshape_kv_cache* deleted upstream (moved into allocate_kv_cache/create_kv_cache_views)	upstream's; both of our hunks lived in the deleted functions
v1/worker/utils.py	KVBlockZeroer rewritten for kernel-block ratios	upstream's
tests/v1/worker/test_gpu_worker.py	upstream added memory-accounting tests in the same tail	upstream's file + our hook tests appended
tests/kernels/attention/test_flash_attn.py		upstream's (FA2 patch dropped)

Auto-merged files were re-read against v0.31.0 too; attention.py, backend.py, platforms/interface.py, llm_base_proposer.py changed under the review below.

Per-hunk review against v0.31.0

file / hunk	what it does	verdict	reason
config/cache.py registry, str field, validator	plugin --kv-cache-dtype values	still needed	upstream has no plugin dtype registry; a Literal field makes plugin dtypes unconstructible
engine/arg_utils.py choices/type	--kv-cache-dtype lists plugin dtypes	still needed	same
engine/arg_utils.py TURBO_ATTN auto-default (#29)	named backend for plugin dtypes	still needed	hooks dispatch before attention groups exist
config/attention.py -→_	--attention-backend turbo-attn	still needed	documented spelling in turbo-attn
backends/registry.py TURBO_ATTN	dedicated slot	still needed, fixed	was = None like CUSTOM, so CUSTOM was an Enum alias of TURBO_ATTN; now the plugin class path (distinct value, still is_overridden()-gated)
platforms/cuda.py append	TURBO_ATTN as auto candidate once registered	still needed	covers selection that bypasses EngineArgs (draft models); also replaces the proposer hunk below
v1/attention/selector.py dtype check	validate_cache_dtype instead of get_args(CacheDType)	still needed
v1/attention/selector.py MLA wrapping	wrap a stock MLA backend	still needed	no upstream equivalent
backend.py on_model_loaded, on_draft_model_loaded, on_kv_cache_initialized, adjust_kv_budget, wraps_mla_backend, resolve_user_selected_backend	lifecycle hooks	still needed	all consumed by turbo-attn; no upstream worker hooks
gpu_worker.py dispatch	fail-loud, every backend in use	still needed
mla_attention.py _get_gather_op	packed-latent gather in chunked prefill	still needed	consumed by TkvMLAImpl; no upstream hook
attention.py _tq_layer_idx	per-layer bit widths at impl init	still needed
kv_cache_interface.py + kv_cache_utils.py aggregated_layer_count	fused (composite) pages	conflicting → re-expressed	v0.31's planner sums one page per layer; fused layers now share one page slot and alias one region
backend.py/attention.py/mla_attention.py/platforms/interface.py get_kv_cache_spec_class	backend-declared spec class	obsolete → moved to plugin	upstream AttentionBackend.customize_spec, applied by both runners and by _align_hybrid_block_size
kv_cache_interface.py + single_type_kv_cache_manager.py get_manager_class	spec picks its manager	obsolete	KVCacheSpecRegistry.get_manager_class walks the MRO; TKV specs subclass Full/SlidingWindow specs
backend.py get_supported_kv_cache_dtypes	dynamic dtype list	obsolete	the wrapper sets the supported_kv_cache_dtypes ClassVar (plugin)
backend.py uses_per_group_kv_pool, get_gather_op (backend-level), on_kv_manager_created		obsolete	no consumer in turbo-attn
kv_cache_coordinator.py, kv_cache_manager.py, kv_cache_utils.py per-group BlockPools, _MultiGroupPoolView (+76a3e6d72), sampler reserve, padded-page block fill, drain hook	O(1) mamba pool split	conflicting → dropped	v0.31 overlays every group on one pool by design (KVCacheConfig.num_blocks_of, single allocate_kv_cache buffer); re-expressing split pools means re-doing upstream's allocator. Hybrid models take upstream's alignment, which sizes the TKV page through customize_spec
worker/utils.py zeroer clamp	per-group block counts	obsolete	only existed for split pools
gpu_model_runner.py/gpu/attn_utils.py kv_cache_spec passthrough, backend-managed dtype coercion	get_kv_cache_shape args	obsolete	upstream no longer calls get_kv_cache_shape; views come from the spec (vllm#51718)
spec_decode_warmup.py + gpu_worker call	warm EAGLE prep / rejection expand kernels	obsolete	upstream registers the same four kernels with its JIT warmup registry: aed894c19 (vllm#56323)
turbo_attn_warmup.py, utils/cutedsl_cache.py + calls	dummy-run prefill warmup, persistent CuTeDSL cache	obsolete (plugin)	turbo-attn warms every served prefill kernel in on_kv_cache_initialized (tkv.engine.prewarm) and arms its own CuTeDSL cache in tkv/kernels/cute_dsl_cache.py
rotary_embedding/common.py registry	survive FA4's namespace flash_attn	obsolete	upstream imports under suppress(ModuleNotFoundError): 1f60771c7 (vllm#42679)
llm_base_proposer.py draft inherits backend	draft finds TURBO_ATTN	obsolete	the CUDA candidate append already makes it selectable
--kv-cache-dtype-skip-layers-dtype, VLLM_KV_CACHE_SKIP_LAYERS_DTYPE, VLLM_SAMPLER_RESERVE_MIB		dropped	unused by turbo-attn; skip-dtype also conflicts with upstream's skip_page_size_padded (assumes auto)
_custom_ops.py Triton fused fp8 GEMM		dropped	the gemm package is not in the image; dead
FA2 varlen split-K (cmake/apply_patches.cmake, patch, vllm_flash_attn.cmake, flash_attn_interface.py, test)		dropped	never reached the image (overlay ships no compiled _vllm_fa2_C; turbo-attn excluded the .py), and the FA pin moved f3e1a4f→9cd61de
notify-turbo-attn.yml		dropped	failing since 2026-08-03; turbo-attn polls the fork instead (poll-engine-forks.yml, vllm-project#741)
tests/evals/gsm8k startup waits		dropped	unrelated to the seam

Tests: kept and rewritten for v0.31: test_kv_seam_invariants.py (now the plugin dtype registry), test_aggregated_layer_count.py (one aliased page per fused set; equal-byte capacity), test_mla_wrapper_selection.py, test_gpu_worker.py hook tests, test_arg_utils.py (#29). Dropped with their code: test_kv_drain_hook.py, test_per_group_blockpool.py, test_tkv_hybrid_block_align.py, test_cutedsl_cache.py.

PROTOCOL.md is rewritten for this seam (it also fixes #28's drift: _get_gather_op returns the stock op, not None, and serves chunked-context prefill, not decode).

Known gaps

• Hybrid capacity: dropping the per-group pools puts hybrid (mamba/GDN + attention) models on upstream's shared pool, so their reported KV capacity can be lower than on the v0.28 fork. Correctness is upstream's path.
• Smart-mix composite on padded pages: if a hybrid model's attention page is ever padded (upstream pads it only when the mamba page sets the maximum and the ratio is not integral), the composite region view refuses the non-contiguous page instead of serving.
• get_kv_cache_spec_class no longer exists, so a plugin that implemented only it (not customize_spec) silently gets stock specs. turbo-attn implements both.
• test_plugin_kv_cache_dtype_auto_selects_turbo_attn_backend needs a device for create_engine_config; it errors in a CPU-only container exactly like upstream's own test_arg_utils.py cases do.

Verification

CPU only, in vllm/vllm-openai:v0.31.0 with this branch's 12 vllm/ files overlaid (as the image build does), on 10.1.0.202, no GPU visible:

• seam tests: 37 passed, 9 skipped (need the CUDA allocator), 1 env error (above).
• upstream tests/v1/core/test_kv_cache_utils.py, test_prefix_caching.py, test_single_type_kv_cache_manager.py, tests/engine/test_arg_utils.py, tests/v1/worker/test_gpu_worker.py: no new failures vs stock v0.31.0; the remaining failures are device-inference errors that fail identically on stock when run in isolation.
• turbo-attn tests/vllm_plugin, tests/integrations/vllm, tests/test_vllm_*.py with the plugin branch: 2098 passed, 260 skipped; 2 failures are stale on turbo-attn main (_backend_resolution._MODEL_LOADED gone; logger test assumes no vLLM), 7 errors are the GPU e2e fixtures.
• compileall clean, no conflict markers.

Not yet runtime-tested on a GPU: that is the engine-image build and GPU validation in turbo-attn on the new base. Merge after that is green.

🤖 Generated with Claude Code