diff --git a/docs/USAGE.md b/docs/USAGE.md
index f3c859e12..cce59b785 100644
--- a/docs/USAGE.md
+++ b/docs/USAGE.md
@@ -2217,7 +2217,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` | `0` (auto, resolves to `256`) | KV block count, and vLLM's `num_gpu_blocks_override`. It wins over every other sizing knob. `0` means auto, which uses `--kv-cache-memory` when that is set and otherwise falls back to `256` blocks |
| `--kv-cache-memory BYTES` | `0` (unset) | Absolute KV-pool size in bytes, vLLM's `kv_cache_memory_bytes`. The block count is this budget divided by the model's own bytes per block, summed across its KV groups, so it is correct on MLA and heterogeneous-KV architectures too. It ignores `--gpu-memory-utilization`, as vLLM does. A budget smaller than one KV block is refused at startup |
| `--gpu-memory-utilization F` | `0.92` | **Accepted, and it does not size anything yet.** See [What `--gpu-memory-utilization` does not do yet](#what---gpu-memory-utilization-does-not-do-yet) |
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.h b/include/vllm.h
index a41ebdf0f..416c8da6a 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/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 ccae4d24a..1d992a0a5 100644
--- a/include/vllm/v1/worker/gpu/runner.h
+++ b/include/vllm/v1/worker/gpu/runner.h
@@ -224,6 +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 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.
@@ -817,6 +831,9 @@ class GPUModelRunner final : public ModelRunnerBase {
std::vector> conv_buf_;
std::vector attn_kv_;
std::vector gdn_state_;
+ // 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/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
diff --git a/src/vllm/entrypoints/openai/server_main.cpp b/src/vllm/entrypoints/openai/server_main.cpp
index 291e6bb21..148801106 100644
--- a/src/vllm/entrypoints/openai/server_main.cpp
+++ b/src/vllm/entrypoints/openai/server_main.cpp
@@ -441,6 +441,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 be0fcba86..4ac0cebf5 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)
@@ -502,6 +504,22 @@ GPUModelRunner::CacheBuffer::~CacheBuffer() {
void GPUModelRunner::initialize_kv_cache(const KVCacheConfig& kv_cache_config) {
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.
@@ -689,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");
@@ -809,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;
@@ -887,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:
@@ -904,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();
@@ -912,6 +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 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') {
+ 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 f6f21ed5f..fd39f21f0 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;
@@ -437,6 +438,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 882323a81..2a2fb5165 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,14 @@ 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, 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);
const PagedKvCache& kv = runner.attn_kv()[0];
@@ -597,7 +606,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 +642,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());