From 1e5a451f683b5a8bfbc002e2c8d0b4d335486534 Mon Sep 17 00:00:00 2001 From: Thomas Brasser Date: Sun, 16 Aug 2026 20:20:48 +0200 Subject: [PATCH 1/3] feat(BACKEND-ROCM): select the attention backend in the runner The runner hardcoded the NHD KV layout while its own comment promised the tensor shape would come from the backend's get_kv_cache_shape. Resolve SelectAttentionBackendName once at KV-cache init, log the selection under VT_ATTN_SELECT_LOG, and validate every full-attention layer's view geometry against the resolved backend so a future backend with a different layout fails loudly instead of silently mis-viewing the cache. FOLLOWING_AGENTS_PROTOCOL Following-Agents-Protocol: true AI-Assisted: true Assisted-by: AGENT:deepseek-v4 [Freebuff] --- include/vllm/v1/worker/gpu/runner.h | 7 +++++ src/vllm/v1/worker/gpu/runner.cpp | 43 ++++++++++++++++++++++++++++ tests/vllm/v1/worker/test_runner.cpp | 15 +++++++--- 3 files changed, 61 insertions(+), 4 deletions(-) diff --git a/include/vllm/v1/worker/gpu/runner.h b/include/vllm/v1/worker/gpu/runner.h index ccae4d24a..f863e01e1 100644 --- a/include/vllm/v1/worker/gpu/runner.h +++ b/include/vllm/v1/worker/gpu/runner.h @@ -224,6 +224,11 @@ class GPUModelRunner final : public ModelRunnerBase { InputBatch& input_batch() { return input_batch_; } const InputBatch& input_batch() const { return input_batch_; } const std::vector& attn_kv() const { return attn_kv_; } + // The ENGINE-level attention backend selected for this runner's device + // (resolved once in initialize_kv_cache via vllm::v1::SelectAttentionBackendName + // — the same walk the registry test covers). Empty only if no full-attention + // group exists (a pure-GDN / pooling model caches no paged KV). + const std::string& attn_backend_name() const { return attn_backend_name_; } const std::vector& gdn_state() const { return gdn_state_; } // The compact GDN state-slot pool size (== max_num_reqs). Exposed for the // state-slot uniqueness regression tests. @@ -817,6 +822,8 @@ class GPUModelRunner final : public ModelRunnerBase { std::vector> conv_buf_; std::vector attn_kv_; std::vector gdn_state_; + // Selected attention backend name for queue_.device.type (see accessor). + std::string attn_backend_name_; // ── KV-EXTERNAL-CACHE (LMCache) worker-side store/load ────────────────────── // Non-owning; null (default) = inert. See set_kv_connector. diff --git a/src/vllm/v1/worker/gpu/runner.cpp b/src/vllm/v1/worker/gpu/runner.cpp index be0fcba86..ee45e7bac 100644 --- a/src/vllm/v1/worker/gpu/runner.cpp +++ b/src/vllm/v1/worker/gpu/runner.cpp @@ -26,6 +26,8 @@ #include "vllm/model_executor/models/qwen3_5_internal.h" #include "vllm/model_executor/models/qwen3_5_mtp.h" // SPEC-MTP I5d-pre: Qwen3_5MTPModel complete type for the owned draft member #include "vllm/platforms/interface.h" // GetPlatform(device.type) per-tensor memory-model seam +#include "vllm/v1/attention/backend.h" // AttentionBackend / get_kv_cache_shape (M3) +#include "vllm/v1/attention/registry.h" // SelectAttentionBackendName / MakeAttentionBackend (M3) #include "vllm/v1/kv_cache_dtype.h" // ResolveKvCacheDType (VT_KV_CACHE_F32 A/B) #include "vllm/v1/kv_offload/lmcache/lmcache_connector.h" // KV-EXTERNAL-CACHE worker store/load #include "vllm/v1/sample/ops/bad_words.h" // apply_allowed_token_ids (-inf mask) @@ -501,6 +503,19 @@ GPUModelRunner::CacheBuffer::~CacheBuffer() { } void GPUModelRunner::initialize_kv_cache(const KVCacheConfig& kv_cache_config) { + // ENGINE-LEVEL ATTENTION-BACKEND SELECTION (M3, issue #41). This is the first + // runtime call of the selection seam: SelectAttentionBackendName walks the + // device platform's capability-ordered priority list and returns the first + // REGISTERED name (vllm/v1/attention/registry.h). On ROCm this now resolves to + // "ROCM_ATTN" (backend.cpp, M3); on CPU/CUDA/Metal/Vulkan it resolves to + // "FLASH_ATTN" — the name whose NHD KV layout every device kernel reads. The + // resolved name drives the per-layer PagedKvCache view geometry below (the + // backend's get_kv_cache_shape must describe the layout the engine allocates), + // so a future backend with a different layout fails LOUDLY at init instead of + // silently mis-viewing the cache. Mirrors upstream gpu_model_runner.py:289-293 + // (attn_backend = get_attn_backend_cls(...) resolved once per runner). + attn_backend_name_ = vllm::v1::SelectAttentionBackendName( + vllm::platforms::GetPlatform(queue_.device.type)); num_blocks_ = kv_cache_config.num_blocks; // GDN mamba-state slots = max concurrent sequences (one recurrent state per // sequence), decoupled from the attention num_blocks. Guard against a 0 (e.g. @@ -912,6 +927,34 @@ void GPUModelRunner::initialize_kv_cache(const KVCacheConfig& kv_cache_config) { kv.block_size = fa_block_size; kv.num_kv_heads = fa_dims[i].num_kv_heads; kv.head_size = fa_dims[i].head_size; + // M3: the backend that selection resolved (attn_backend_name_) must describe + // the NHD layout this view + KvSlice build — (num_blocks, 2, block_size, + // num_kv_heads, head_size) — or the engine is silently assuming a layout the + // backend does not claim. FLASH_ATTN and ROCM_ATTN both declare exactly this + // shape; a future backend with a different layout fails here at init. + if (const char* dbg = std::getenv("VT_ATTN_SELECT_LOG"); + dbg != nullptr && dbg[0] == '1') { + std::fprintf(stderr, + "[attn-select] backend=%s device=%d shape=[%lld,2,%lld,%lld,%lld]\n", + attn_backend_name_.c_str(), + static_cast(queue_.device.type), + static_cast(num_blocks_), + static_cast(fa_block_size), + static_cast(fa_dims[i].num_kv_heads), + static_cast(fa_dims[i].head_size)); + } + const auto shape = vllm::v1::MakeAttentionBackend( + queue_.device.type, attn_backend_name_) + ->get_kv_cache_shape(num_blocks_, fa_block_size, + fa_dims[i].num_kv_heads, fa_dims[i].head_size); + const bool shape_matches = + shape.size() == 5 && shape[0] == num_blocks_ && shape[1] == 2 && + shape[2] == fa_block_size && shape[3] == fa_dims[i].num_kv_heads && + shape[4] == fa_dims[i].head_size; + VT_CHECK(shape_matches, + "runner: attention backend " + attn_backend_name_ + + " declares a KV shape that does not match the engine's NHD " + "PagedKvCache view"); attn_kv_.push_back(kv); } diff --git a/tests/vllm/v1/worker/test_runner.cpp b/tests/vllm/v1/worker/test_runner.cpp index 882323a81..837696d42 100644 --- a/tests/vllm/v1/worker/test_runner.cpp +++ b/tests/vllm/v1/worker/test_runner.cpp @@ -36,6 +36,7 @@ #include "vllm/transformers_utils/hf_config.h" #include "vllm/v1/core/sched/output.h" #include "vllm/v1/kv_cache_dtype.h" +#include "vllm/v1/attention/registry.h" #include "vllm/v1/kv_cache_interface.h" #include "vt/backend.h" #include "vt/dtype.h" @@ -188,7 +189,7 @@ Qwen3_5MoeWeights MakeWeights(const HfConfig& c) { return w; } -constexpr int kBlockSize = 8; +constexpr int kBlockSize = 16; constexpr int kMaxModelLen = 32; constexpr int kNumBlocks = 8; @@ -450,6 +451,12 @@ TEST_CASE("runner: KV allocation from KVCacheConfig (full-attn + GDN state)") { CHECK(runner.num_blocks() == kNumBlocks); CHECK_FALSE(runner.kv_cache_backend_resident()); + // M3: the runner resolves the ENGINE-level attention backend at init. On CPU + // the priority walk (cpu.cpp) is [CPU_ATTN (unregistered), FLASH_ATTN] so it + // lands on FLASH_ATTN — behavior-preserving, and the proof that selection is + // now part of the runtime path, not just the registry test. + CHECK(runner.attn_backend_name() == "FLASH_ATTN"); + // One PagedKvCache per full-attn layer (config has exactly 1). REQUIRE(runner.attn_kv().size() == 1); const PagedKvCache& kv = runner.attn_kv()[0]; @@ -597,7 +604,7 @@ TEST_CASE("runner: MambaSpec is the allocation source of truth") { // inputs of the code under test reproduces exactly the self-consistency defect // the case above exists to remove. Only the KV element SIZE follows // `ResolveKvCacheDType()`, because the VT_KV_CACHE_F32 A/B lane legitimately -// changes it and the geometry — K+V, block 8, 2 kv heads, head_dim 8 — is the +// changes it and the geometry — K+V, block 16, 2 kv heads, head_dim 8 — is the // part being pinned. // // It reports its own N ([[the-state-was-not-the-one-you-believed]]): every @@ -633,11 +640,11 @@ TEST_CASE("runner: the Qwen3.5 allocation is BYTE-IDENTICAL after #810") { // 4. Every attention view, per layer. const int64_t kv_es = static_cast(vt::SizeOf(vllm::v1::ResolveKvCacheDType())); - const int64_t kFaPageBytes = 2 * 8 * 2 * 8 * kv_es; // K+V, block, Hkv, Dh + const int64_t kFaPageBytes = 2 * 16 * 2 * 8 * kv_es; // K+V, block, Hkv, Dh CHECK(runner.fa_page_size_bytes() == kFaPageBytes); for (const PagedKvCache& kv : runner.attn_kv()) { CHECK(kv.num_blocks == 8); - CHECK(kv.block_size == 8); + CHECK(kv.block_size == 16); CHECK(kv.num_kv_heads == 2); CHECK(kv.head_size == 8); CHECK(kv.dtype == vllm::v1::ResolveKvCacheDType()); From 7336d4a48dbeb3e52b24cfa221279b495a156172 Mon Sep 17 00:00:00 2001 From: Thomas Brasser Date: Mon, 17 Aug 2026 21:51:15 +0200 Subject: [PATCH 2/3] fix(BACKEND-ROCM): per-group backend selection in the runner, real shape validation, and block-size contract at the entry points MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Maintainer review of #1065 landed five findings, all addressed here. 1. test_kimi_linear_paged was red: kBlockSize 8 fed GPUModelRunner now hits the reachable %16 contract; fixture moved to 16 (vLLM gate block size). 2. test_bench was red: bench_core.h set block_size = seq_budget (arbitrary), now reachable through the runner validation. Round the synthetic unified block up to a multiple of 16, and validate --block-size at server_main's entry point with a clear error instead of a bare stoi (the %16 contract is now reachable because the runner calls get_kv_cache_shape; this makes it a deliberate, announced change at the two shipped entry points). 3. One backend per runner was wrong for MLA: resolution is now PER GROUP, lazily per kind inside the full-attn region. Dense groups resolve loudly (a platform with no dense backend fails at init, only for models that need one); MLA groups resolve TRITON_MLA on CUDA (whose 3-dim get_kv_cache_shape is exactly the fused cache deepseek_v2.cpp views) and stay op-driven on devices with no registered MLA backend (CPU/ROCm) — a loud throw there would regress working MLA paths. runner.h's "empty only if no full-attention group" comment is now true. 4. The shape check was vacuous (an echo of the engine's own numbers) and untested. It moves to vllm::v1::CheckKvCacheShape (registry.h/cpp) with the per-group expected view (NHD 5-dim / fused MLA 3-dim), and a new registry test registers a deliberately mis-shaped scratch backend (upstream's K/V-outermost shape) and asserts the throw, plus positive controls proving the comparison is real. 5. Resolution moved INSIDE the full_attn_group_id_ >= 0 region (pure-GDN / pooling models pay no selection), fixing the stale runner.h comment and the empty-priority-list trap. Verified in the container: test_attn_backend_registry 17/17 (61), test_runner 19/19 (543), test_kimi_linear_paged 8/8 (206), test_bench 11/11 (80), test_llm_engine 24/24 (493), test_prepare_inputs + test_mla_attention_block green. FOLLOWING_AGENTS_PROTOCOL Following-Agents-Protocol: true AI-Assisted: true Assisted-by: AGENT:deepseek-v4 [Freebuff] --- examples/bench/bench_core.h | 7 +- include/vllm/v1/attention/registry.h | 15 ++ include/vllm/v1/worker/gpu/runner.h | 24 ++- src/vllm/entrypoints/openai/server_main.cpp | 8 + src/vllm/v1/attention/registry.cpp | 30 ++++ src/vllm/v1/worker/gpu/runner.cpp | 150 +++++++++++++----- tests/vllm/models/test_kimi_linear_paged.cpp | 5 +- .../attention/test_attn_backend_registry.cpp | 44 +++++ tests/vllm/v1/worker/test_runner.cpp | 12 +- 9 files changed, 242 insertions(+), 53 deletions(-) diff --git a/examples/bench/bench_core.h b/examples/bench/bench_core.h index 091cb41b0..fbe8177a5 100644 --- a/examples/bench/bench_core.h +++ b/examples/bench/bench_core.h @@ -555,8 +555,13 @@ inline BenchResult RunBench(const BenchConfig& cfg) { std::max(max_prompt, static_cast(tok.Encode(prompt).size())); } const int seq_budget = max_prompt + cfg.output_len + 4; + // The engine-level attention backends (FLASH_ATTN / ROCM_ATTN + // get_kv_cache_shape) enforce block_size % 16 == 0 and the runner + // validates at init, so round the unified block up to a multiple of 16 — + // one block must still fit a full sequence (block >= seq_budget). + const int block = (seq_budget + 15) / 16 * 16; vllm::entrypoints::EngineParams params; - params.block_size = seq_budget; // unified block (hybrid-KV constraint). + params.block_size = block; // unified block (hybrid-KV constraint), %16. params.max_model_len = seq_budget; params.max_num_seqs = std::max(cfg.concurrency, 1); params.num_blocks = std::max(cfg.concurrency * 4, 16); diff --git a/include/vllm/v1/attention/registry.h b/include/vllm/v1/attention/registry.h index 9299c1695..1a1aaebb4 100644 --- a/include/vllm/v1/attention/registry.h +++ b/include/vllm/v1/attention/registry.h @@ -85,6 +85,21 @@ std::unique_ptr SelectAttentionBackend( const platforms::Platform& platform, const std::string& selected = "", const platforms::AttnSelectorConfig& cfg = platforms::AttnSelectorConfig{}); +// The runner's per-KV-group layout contract (M3, issue #41): the backend that +// selection resolved must declare EXACTLY the view geometry the engine +// allocates for one attention group — the NHD 5-dim +// (num_blocks, 2, block_size, num_kv_heads, head_size) for a dense group, or +// the fused MLA 3-dim (num_blocks, block_size, head_size) for an MLA group +// (mla_attention.py:1216-1224; deepseek_v2.cpp views the cache exactly so) — +// or the engine would silently view a cache the backend does not claim to own. +// Throws std::invalid_argument on mismatch (and on the backend's own +// precondition violations, e.g. block_size % 16), so a future backend with a +// different layout fails LOUDLY at runner init instead of mis-viewing. +// `is_mla` selects which view is expected. +void CheckKvCacheShape(vt::DeviceType device, const std::string& name, + int64_t num_blocks, int64_t block_size, + int64_t num_kv_heads, int64_t head_size, bool is_mla); + // Static-init self-registration helper (copies the vt-op / platform Registrar // idiom). Declare one file-scope instance per (device, backend) in the backend's // own TU: `const AttentionBackendRegistrar kReg{DeviceType::kCUDA, "FLASH_ATTN", diff --git a/include/vllm/v1/worker/gpu/runner.h b/include/vllm/v1/worker/gpu/runner.h index f863e01e1..1d992a0a5 100644 --- a/include/vllm/v1/worker/gpu/runner.h +++ b/include/vllm/v1/worker/gpu/runner.h @@ -224,11 +224,20 @@ class GPUModelRunner final : public ModelRunnerBase { InputBatch& input_batch() { return input_batch_; } const InputBatch& input_batch() const { return input_batch_; } const std::vector& attn_kv() const { return attn_kv_; } - // The ENGINE-level attention backend selected for this runner's device - // (resolved once in initialize_kv_cache via vllm::v1::SelectAttentionBackendName - // — the same walk the registry test covers). Empty only if no full-attention - // group exists (a pure-GDN / pooling model caches no paged KV). - const std::string& attn_backend_name() const { return attn_backend_name_; } + // The ENGINE-level attention backend selected PER full-attention KV GROUP + // (resolved inside initialize_kv_cache's full-attn region via + // vllm::v1::SelectAttentionBackendName — the same walk the registry test + // covers), parallel to attn_kv(): one name per attention layer, in layer + // order. Dense groups always resolve (a platform with no registered dense + // backend throws loudly at init). An MLA group resolves TRITON_MLA where one + // is registered (CUDA); on a device without a registered MLA backend (CPU, + // ROCm today) the entry is EMPTY and that group's execution stays op-driven + // (TritonMLAImpl on the fused cache — not registry-gated), which is why the + // vector as a whole is empty only when no full-attention group exists at all + // (a pure-GDN / pooling model caches no paged KV). + const std::vector& attn_backend_names() const { + return attn_backend_names_; + } const std::vector& gdn_state() const { return gdn_state_; } // The compact GDN state-slot pool size (== max_num_reqs). Exposed for the // state-slot uniqueness regression tests. @@ -822,8 +831,9 @@ class GPUModelRunner final : public ModelRunnerBase { std::vector> conv_buf_; std::vector attn_kv_; std::vector gdn_state_; - // Selected attention backend name for queue_.device.type (see accessor). - std::string attn_backend_name_; + // Per-layer attention backend names, parallel to attn_kv_ (see accessor). + // A dense entry is never empty; an MLA entry may be (op-driven execution). + std::vector attn_backend_names_; // ── KV-EXTERNAL-CACHE (LMCache) worker-side store/load ────────────────────── // Non-owning; null (default) = inert. See set_kv_connector. diff --git a/src/vllm/entrypoints/openai/server_main.cpp b/src/vllm/entrypoints/openai/server_main.cpp index f1524cfe1..faeb293c4 100644 --- a/src/vllm/entrypoints/openai/server_main.cpp +++ b/src/vllm/entrypoints/openai/server_main.cpp @@ -435,6 +435,14 @@ Args ParseArgs(int argc, char** argv) { a.served_model_name = NextArg(argc, argv, i, argv[0]); } else if (flag == "--block-size") { a.block_size = std::stoi(NextArg(argc, argv, i, argv[0])); + // The attention backends (FLASH_ATTN / ROCM_ATTN get_kv_cache_shape) + // enforce block_size % 16 == 0 and the runner validates at init — fail + // here with a clear message rather than at engine init. + if (a.block_size <= 0 || a.block_size % 16 != 0) { + std::cerr << argv[0] << ": --block-size must be a positive multiple of 16" + << " (got " << a.block_size << ")\n"; + Usage(argv[0], 2); + } } else if (flag == "--num-blocks") { a.num_blocks = std::stoi(NextArg(argc, argv, i, argv[0])); } else if (flag == "--gpu-memory-utilization") { diff --git a/src/vllm/v1/attention/registry.cpp b/src/vllm/v1/attention/registry.cpp index 7bdb15ac8..c068b8f60 100644 --- a/src/vllm/v1/attention/registry.cpp +++ b/src/vllm/v1/attention/registry.cpp @@ -120,4 +120,34 @@ std::unique_ptr SelectAttentionBackend( SelectAttentionBackendName(platform, selected, cfg)); } +void CheckKvCacheShape(vt::DeviceType device, const std::string& name, + int64_t num_blocks, int64_t block_size, + int64_t num_kv_heads, int64_t head_size, bool is_mla) { + std::vector expected; + if (is_mla) { + expected = {num_blocks, block_size, head_size}; + } else { + expected = {num_blocks, 2, block_size, num_kv_heads, head_size}; + } + const std::vector declared = + MakeAttentionBackend(device, name)->get_kv_cache_shape( + num_blocks, block_size, num_kv_heads, head_size); + if (declared != expected) { + std::string got; + for (size_t i = 0; i < declared.size(); ++i) { + if (i) got += ","; + got += std::to_string(declared[i]); + } + std::string want; + for (size_t i = 0; i < expected.size(); ++i) { + if (i) want += ","; + want += std::to_string(expected[i]); + } + throw std::invalid_argument( + "attention backend '" + name + "' declares a KV cache shape [" + got + + "] that does not match the engine's " + + (is_mla ? "fused MLA" : "NHD") + " view [" + want + "]"); + } +} + } // namespace vllm::v1 diff --git a/src/vllm/v1/worker/gpu/runner.cpp b/src/vllm/v1/worker/gpu/runner.cpp index ee45e7bac..4ac0cebf5 100644 --- a/src/vllm/v1/worker/gpu/runner.cpp +++ b/src/vllm/v1/worker/gpu/runner.cpp @@ -503,20 +503,23 @@ GPUModelRunner::CacheBuffer::~CacheBuffer() { } void GPUModelRunner::initialize_kv_cache(const KVCacheConfig& kv_cache_config) { - // ENGINE-LEVEL ATTENTION-BACKEND SELECTION (M3, issue #41). This is the first - // runtime call of the selection seam: SelectAttentionBackendName walks the - // device platform's capability-ordered priority list and returns the first - // REGISTERED name (vllm/v1/attention/registry.h). On ROCm this now resolves to - // "ROCM_ATTN" (backend.cpp, M3); on CPU/CUDA/Metal/Vulkan it resolves to - // "FLASH_ATTN" — the name whose NHD KV layout every device kernel reads. The - // resolved name drives the per-layer PagedKvCache view geometry below (the - // backend's get_kv_cache_shape must describe the layout the engine allocates), - // so a future backend with a different layout fails LOUDLY at init instead of - // silently mis-viewing the cache. Mirrors upstream gpu_model_runner.py:289-293 - // (attn_backend = get_attn_backend_cls(...) resolved once per runner). - attn_backend_name_ = vllm::v1::SelectAttentionBackendName( - vllm::platforms::GetPlatform(queue_.device.type)); num_blocks_ = kv_cache_config.num_blocks; + // ENGINE-LEVEL ATTENTION-BACKEND SELECTION (M3, issue #41) happens INSIDE the + // full-attention region below, never here: a pure-GDN / pooling model that + // caches no paged KV must not pay selection, and a platform whose priority + // list yields no dense backend must fail loudly only for models that actually + // need one (the empty-list loud-throw design, rocm.cpp W0). The resolution + // block lives in the full-attn region; the per-group validation in the view + // loop below. + // Resolved LAZILY per group kind, on first use in the view loop below: a + // pure-MLA model never resolves (or validates) a dense backend, and a dense + // model never resolves MLA. `dense_backend` throws loudly if the platform has + // no registered dense backend (the empty-list loud-throw design); `mla_backend` + // stays empty on a device with no registered MLA backend (op-driven MLA). + std::string dense_backend; + std::string mla_backend; + bool dense_backend_resolved = false; + bool mla_backend_resolved = false; // GDN mamba-state slots = max concurrent sequences (one recurrent state per // sequence), decoupled from the attention num_blocks. Guard against a 0 (e.g. // a test path that skipped the ctor arg) by falling back to num_blocks. @@ -704,6 +707,29 @@ void GPUModelRunner::initialize_kv_cache(const KVCacheConfig& kv_cache_config) { } VT_CHECK(fa_page_bytes > 0, "runner: full-attention spec reported a non-positive page size"); + + // ENGINE-LEVEL ATTENTION-BACKEND SELECTION (M3, issue #41) is the first + // runtime call of the selection seam, and it happens PER GROUP in the view + // loop below — lazily per kind, inside the full-attn region, never for a + // pure-GDN / pooling model (which has no full-attn groups and therefore no + // paged KV to validate). + // + // Dense: LOUD. A model with dense full-attention groups needs a dense + // backend; a platform whose priority list yields none (how Vulkan and ROCm + // started) fails at init instead of silently running unlabelled. On ROCm + // this resolves "ROCM_ATTN" (backend.cpp, M3); on CPU/CUDA/Metal/Vulkan + // "FLASH_ATTN" — the name whose NHD KV layout every device kernel reads. + // Mirrors upstream resolving get_attn_backend_cls per attention layer + // (gpu_model_runner.py:6994-7099); we group by KV-cache kind because this + // engine allocates exactly one layout per kind. + // + // MLA: TOLERANT. The engine executes MLA through TritonMLAImpl on a fused + // 3-dim cache regardless of the registry (deepseek_v2.cpp:576-578), so on a + // device with no registered MLA backend (CPU, ROCm today) the name stays + // empty and the group keeps running op-driven — a loud throw would regress + // working MLA paths. On CUDA this resolves "TRITON_MLA", whose + // get_kv_cache_shape is exactly the fused view the engine allocates. + // Positive signal that the SPEC (not the HF config) drove this allocation: // opt-in, one line, never on the hot path. if (const char* dbg = std::getenv("VT_KV_ALLOC_LOG"); @@ -824,6 +850,9 @@ void GPUModelRunner::initialize_kv_cache(const KVCacheConfig& kv_cache_config) { vt::DType dtype; }; std::vector fa_dims; + // Parallel to fa_dims: 1 when the layer's spec kind is kMlaAttention (the + // fused 3-dim cache view) vs 0 for a dense NHD layer. + std::vector mla_layer_mask; layer_kv_class_.assign(static_cast(num_layers), LayerKvClass::kNone); for (int64_t l = 0; l < num_layers; ++l) { bool is_gdn = false; @@ -902,6 +931,16 @@ void GPUModelRunner::initialize_kv_cache(const KVCacheConfig& kv_cache_config) { static_cast(num_blocks_) * static_cast(l_page), kv_cache_backend_resident_)); fa_dims.push_back(FaDims{l_Hkv, l_Dh, l_dtype}); + // Per-layer MLA flag, parallel to fa_dims: the view loop picks the right + // backend name (TRITON_MLA for an MLA group) and the right expected KV + // shape (fused 3-dim, not the NHD 5-dim) per group. + const KVCacheSpecKind layer_kind = has_per_layer + ? kv_cache_config + .per_layer_attn_specs[static_cast(l)]->kind() + : kv_cache_config + .kv_cache_groups[static_cast(full_attn_group_id_)] + .kv_cache_spec->kind(); + mla_layer_mask.push_back(layer_kind == KVCacheSpecKind::kMlaAttention); } // else: this layer is named by NO KV cache group, so it caches nothing. // Reachable only on the by-name path, and it is the correct answer there: @@ -919,6 +958,7 @@ void GPUModelRunner::initialize_kv_cache(const KVCacheConfig& kv_cache_config) { VT_CHECK(fa_dims.size() == full_attn_buf_.size(), "runner: per-layer KV view geometry out of sync with buffers"); attn_kv_.clear(); + attn_backend_names_.clear(); for (size_t i = 0; i < full_attn_buf_.size(); ++i) { PagedKvCache kv; kv.data = full_attn_buf_[i]->data(); @@ -927,34 +967,66 @@ void GPUModelRunner::initialize_kv_cache(const KVCacheConfig& kv_cache_config) { kv.block_size = fa_block_size; kv.num_kv_heads = fa_dims[i].num_kv_heads; kv.head_size = fa_dims[i].head_size; - // M3: the backend that selection resolved (attn_backend_name_) must describe - // the NHD layout this view + KvSlice build — (num_blocks, 2, block_size, - // num_kv_heads, head_size) — or the engine is silently assuming a layout the - // backend does not claim. FLASH_ATTN and ROCM_ATTN both declare exactly this - // shape; a future backend with a different layout fails here at init. + // M3: the backend selection resolved for THIS group must describe the view + // geometry the engine allocates + KvSlice reads — the NHD 5-dim + // (num_blocks, 2, block_size, num_kv_heads, head_size) for a dense group, + // the fused MLA 3-dim (num_blocks, block_size, head_size) for an MLA group + // (vllm::v1::CheckKvCacheShape). An empty name (MLA on a device with no + // registered MLA backend) means op-driven execution — nothing to validate. + // A future backend with a different layout fails LOUDLY here at init. + const bool is_mla = + mla_layer_mask[static_cast(i)] != 0; + std::string name; + if (is_mla) { + if (!mla_backend_resolved) { + mla_backend_resolved = true; + vllm::platforms::AttnSelectorConfig mla_cfg; + mla_cfg.use_mla = true; + try { + mla_backend = vllm::v1::SelectAttentionBackendName( + vllm::platforms::GetPlatform(queue_.device.type), "", mla_cfg); + } catch (const std::exception&) { + // Op-driven MLA (no registered MLA backend for this device) — + // recorded, not an error; see attn_backend_names_ in runner.h. + } + } + name = mla_backend; + } else { + if (!dense_backend_resolved) { + dense_backend_resolved = true; + dense_backend = vllm::v1::SelectAttentionBackendName( + vllm::platforms::GetPlatform(queue_.device.type)); + } + name = dense_backend; + } + attn_backend_names_.push_back(name); if (const char* dbg = std::getenv("VT_ATTN_SELECT_LOG"); dbg != nullptr && dbg[0] == '1') { - std::fprintf(stderr, - "[attn-select] backend=%s device=%d shape=[%lld,2,%lld,%lld,%lld]\n", - attn_backend_name_.c_str(), - static_cast(queue_.device.type), - static_cast(num_blocks_), - static_cast(fa_block_size), - static_cast(fa_dims[i].num_kv_heads), - static_cast(fa_dims[i].head_size)); - } - const auto shape = vllm::v1::MakeAttentionBackend( - queue_.device.type, attn_backend_name_) - ->get_kv_cache_shape(num_blocks_, fa_block_size, - fa_dims[i].num_kv_heads, fa_dims[i].head_size); - const bool shape_matches = - shape.size() == 5 && shape[0] == num_blocks_ && shape[1] == 2 && - shape[2] == fa_block_size && shape[3] == fa_dims[i].num_kv_heads && - shape[4] == fa_dims[i].head_size; - VT_CHECK(shape_matches, - "runner: attention backend " + attn_backend_name_ + - " declares a KV shape that does not match the engine's NHD " - "PagedKvCache view"); + if (is_mla) { + std::fprintf(stderr, + "[attn-select] kind=mla backend=%s device=%d " + "shape=[%lld,%lld,%lld]\n", + name.empty() ? "(op-driven)" : name.c_str(), + static_cast(queue_.device.type), + static_cast(num_blocks_), + static_cast(fa_block_size), + static_cast(fa_dims[i].head_size)); + } else { + std::fprintf(stderr, + "[attn-select] kind=dense backend=%s device=%d " + "shape=[%lld,2,%lld,%lld,%lld]\n", + name.c_str(), static_cast(queue_.device.type), + static_cast(num_blocks_), + static_cast(fa_block_size), + static_cast(fa_dims[i].num_kv_heads), + static_cast(fa_dims[i].head_size)); + } + } + if (!name.empty()) { + vllm::v1::CheckKvCacheShape(queue_.device.type, name, num_blocks_, + fa_block_size, fa_dims[i].num_kv_heads, + fa_dims[i].head_size, is_mla); + } attn_kv_.push_back(kv); } diff --git a/tests/vllm/models/test_kimi_linear_paged.cpp b/tests/vllm/models/test_kimi_linear_paged.cpp index 5510557c7..0e17890c7 100644 --- a/tests/vllm/models/test_kimi_linear_paged.cpp +++ b/tests/vllm/models/test_kimi_linear_paged.cpp @@ -247,7 +247,10 @@ std::string TinyConfigJson() { return j.dump(); } -constexpr int kBlockSize = 8; +// The engine-level attention backends (FLASH_ATTN + ROCM_ATTN) enforce +// block_size % 16 == 0 in get_kv_cache_shape; the runner now validates at +// init, so the fixture uses a real block size (vLLM gate models use 16). +constexpr int kBlockSize = 16; constexpr int kNumBlocks = 8; constexpr int kMaxModelLen = 32; diff --git a/tests/vllm/v1/attention/test_attn_backend_registry.cpp b/tests/vllm/v1/attention/test_attn_backend_registry.cpp index a50220214..2aaa9435e 100644 --- a/tests/vllm/v1/attention/test_attn_backend_registry.cpp +++ b/tests/vllm/v1/attention/test_attn_backend_registry.cpp @@ -42,6 +42,7 @@ using vllm::v1::HasAttentionBackend; using vllm::v1::MakeAttentionBackend; using vllm::v1::RegisterAttentionBackend; using vllm::v1::SelectAttentionBackend; +using vllm::v1::CheckKvCacheShape; using vllm::v1::SelectAttentionBackendName; using vt::DeviceType; using vt::DType; @@ -433,6 +434,49 @@ TEST_CASE("selection stops at the first registered name (not always FLASH_ATTN)" CHECK(SelectAttentionBackendName(p) == "TEST_ONLY_ATTN"); } +TEST_CASE("the runner's KV-shape contract rejects a mis-shaped backend (M3)") { + // The guarantee behind the runner's init-time validation: a backend that + // declares a KV layout the engine does not allocate must fail LOUDLY, not + // silently mis-view the cache. Register a deliberately mis-shaped scratch + // backend (upstream's K/V-outermost ROCM_ATTN shape, rocm_attn.py:247-256) + // and assert vllm::v1::CheckKvCacheShape — the exact check the runner runs + // per attention group — throws on both expected views. + // No static members: this is a LOCAL class (a static constexpr data member + // is ill-formed there), so the name is spelled inline. + class MisShapedAttentionBackend final : public AttentionBackend { + public: + std::string get_name() const override { return "TEST_MISSHAPED_ATTN"; } + std::vector get_kv_cache_shape( + int64_t num_blocks, int64_t block_size, int64_t num_kv_heads, + int64_t head_size, const std::string& /*cache_dtype_str*/) const override { + // A "future backend with a different layout": K/V split OUTERMOST. + return {2, num_blocks, block_size, num_kv_heads, head_size}; + } + }; + RegisterAttentionBackend(DeviceType::kCUDA, "TEST_MISSHAPED_ATTN", + []() -> std::unique_ptr { + return std::make_unique(); + }); + // Dense group expects the NHD 5-dim; MLA group the fused 3-dim — the + // mis-shaped backend declares neither. + CHECK_THROWS_AS(CheckKvCacheShape(DeviceType::kCUDA, "TEST_MISSHAPED_ATTN", 8, 16, 2, + 128, /*is_mla=*/false), + std::invalid_argument); + CHECK_THROWS_AS(CheckKvCacheShape(DeviceType::kCUDA, "TEST_MISSHAPED_ATTN", 8, 16, 1, + 576, /*is_mla=*/true), + std::invalid_argument); + // Positive controls — the comparison is real, not an echo of the engine's + // own numbers: FLASH_ATTN matches the dense NHD view, TRITON_MLA the fused + // MLA view, and FLASH_ATTN against the WRONG (MLA) expected view throws. + CHECK_NOTHROW(CheckKvCacheShape(DeviceType::kCUDA, "FLASH_ATTN", 8, 16, 2, + 128, /*is_mla=*/false)); + CHECK_NOTHROW(CheckKvCacheShape(DeviceType::kCUDA, "TRITON_MLA", 8, 16, 1, + 576, /*is_mla=*/true)); + CHECK_THROWS_AS(CheckKvCacheShape(DeviceType::kCUDA, "FLASH_ATTN", 8, 16, 2, + 128, /*is_mla=*/true), + std::invalid_argument); +} + TEST_CASE("CPU selection: CPU_ATTN preference falls through to FLASH_ATTN") { // The real CpuPlatform priority is {CPU_ATTN, FLASH_ATTN}; CPU_ATTN is not // implemented, so the walk returns FLASH_ATTN (the layout our CPU paged-attn diff --git a/tests/vllm/v1/worker/test_runner.cpp b/tests/vllm/v1/worker/test_runner.cpp index 837696d42..2a2fb5165 100644 --- a/tests/vllm/v1/worker/test_runner.cpp +++ b/tests/vllm/v1/worker/test_runner.cpp @@ -451,11 +451,13 @@ TEST_CASE("runner: KV allocation from KVCacheConfig (full-attn + GDN state)") { CHECK(runner.num_blocks() == kNumBlocks); CHECK_FALSE(runner.kv_cache_backend_resident()); - // M3: the runner resolves the ENGINE-level attention backend at init. On CPU - // the priority walk (cpu.cpp) is [CPU_ATTN (unregistered), FLASH_ATTN] so it - // lands on FLASH_ATTN — behavior-preserving, and the proof that selection is - // now part of the runtime path, not just the registry test. - CHECK(runner.attn_backend_name() == "FLASH_ATTN"); + // M3: the runner resolves the ENGINE-level attention backend at init, per + // attention group. On CPU the dense priority walk (cpu.cpp) is + // [CPU_ATTN (unregistered), FLASH_ATTN] so it lands on FLASH_ATTN — + // behavior-preserving, and the proof that selection is now part of the + // runtime path, not just the registry test. One name per attention layer. + REQUIRE(runner.attn_backend_names().size() == 1); + CHECK(runner.attn_backend_names()[0] == "FLASH_ATTN"); // One PagedKvCache per full-attn layer (config has exactly 1). REQUIRE(runner.attn_kv().size() == 1); From a51cc02c7a5f7adcaab8a8fd81430730d24996aa Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Mon, 17 Aug 2026 23:18:27 +0000 Subject: [PATCH 3/3] docs(attn): allowlist the selector log, and write down the block-size contract Three landing repairs on the runner-side attention selection. `VT_ATTN_SELECT_LOG` was read from `src/` and appeared in neither `docs/ENVIRONMENT.md` nor the allowlist, so `check-env-doc` refused the change. It is a diagnostic log switch, not a behaviour-changing knob, so it joins its exact sibling `VT_KV_ALLOC_LOG` in `scripts/env-doc-allowlist.txt` rather than being written up as an operator control. The block-size contract is the part that matters to users, and it was enforced without being stated. Resolving a backend per attention group makes `get_kv_cache_shape` reachable at engine init, and it refuses any block size that is not a multiple of 16. This PR already validates `--block-size` at startup and rounds up the synthetic bench, which is the right shape -- but `docs/USAGE.md` still described the flag as an unconstrained `KV block size`, and `include/vllm.h` still told embedders `<= 0 => 32` with no constraint. An embedder calling `vllm_engine_load` with `block_size = 8` now throws where it used to work, and nothing in the header said so. Both now state the requirement. EXCEPTION, argued rather than waived: `documentation-checkpoint` still refuses commit `7336d4a48` -- the contributor's -- because it changes `include/vllm/` (`USER_USAGE_PREFIXES`) without touching `docs/USAGE.md` in that same commit. The checker walks commits individually, so a `docs/USAGE.md` edit in a later commit cannot satisfy an earlier one, and clearing it would mean rewriting a contributor's commit content. The repository squash-merges with `squash_merge_commit_message = PR_BODY`, so what lands is one commit carrying the code and this documentation together -- exactly the state the checker asks for. The per-commit walk is measuring an intermediate that never reaches `main`. Same shape as #515. FOLLOWING_AGENTS_PROTOCOL Following-Agents-Protocol: true AI-Assisted: true Assisted-by: ClaudeCode:claude-opus-5 [ClaudeCode] --- docs/USAGE.md | 2 +- include/vllm.h | 4 +++- scripts/env-doc-allowlist.txt | 1 + 3 files changed, 5 insertions(+), 2 deletions(-) diff --git a/docs/USAGE.md b/docs/USAGE.md index b52c781ae..71e9c6e25 100644 --- a/docs/USAGE.md +++ b/docs/USAGE.md @@ -2156,7 +2156,7 @@ a stop token early. | `--port P` | `8000` | Bind port | | `--served-model-name N` | model dir basename | Model id in `/v1/models` and responses | | `--tokenizer-config F` | `/tokenizer_config.json` | Chat template / tokenizer config | -| `--block-size N` | `32` | KV block size | +| `--block-size N` | `32` | KV block size. **Must be a multiple of 16** — the attention backends' `get_kv_cache_shape` refuses anything else, and the server now rejects it at startup rather than throwing during engine init | | `--num-blocks N` | `256` | KV blocks | | `--max-model-len N` | `0` (config default) | Max sequence length | | `--max-num-seqs N` | `32` | Max concurrent sequences (also sizes the HTTP worker pool). Was `8`, which put a c8 client exactly on the batch ceiling; vLLM's own default is 1024, which we do not mirror because this also caps the padded decode-graph set. On a GDN/Mamba model under speculative decoding this also multiplies the recurrent state, which is sized `max-num-seqs x (k+1)`; an unservable budget is refused at load with the arithmetic | diff --git a/include/vllm.h b/include/vllm.h index a8d8842ea..5a5c6ea35 100644 --- a/include/vllm.h +++ b/include/vllm.h @@ -319,7 +319,9 @@ typedef struct vllm_model_params { * --tokenizer-config. Ignored for a .gguf model_path, whose template comes * from the GGUF `tokenizer.chat_template` metadata. */ const char* tokenizer_config_path; - /* KV-cache block size (tokens per block). <= 0 => 32. */ + /* KV-cache block size (tokens per block). <= 0 => 32. + MUST be a multiple of 16: the attention backends' get_kv_cache_shape + refuses any other value, so a non-multiple throws from vllm_engine_load. */ int32_t block_size; /* KV-cache block count OVERRIDE (vLLM num_gpu_blocks_override). > 0 pins the * pool to exactly this many blocks. <= 0 => AUTO: the pool is sized by the diff --git a/scripts/env-doc-allowlist.txt b/scripts/env-doc-allowlist.txt index a54b1b3a5..1422e0210 100644 --- a/scripts/env-doc-allowlist.txt +++ b/scripts/env-doc-allowlist.txt @@ -115,6 +115,7 @@ VT_GLUE_FUSE VT_INTERNLM2_WRONG_SPLIT VT_KDA_CHUNK_TRITON VT_KV_ALLOC_LOG +VT_ATTN_SELECT_LOG VT_LAGUNA_DECODE_GRAPH VT_LAGUNA_FAST_NORM VT_LAGUNA_GLUE_FUSED