feat(BACKEND-ROCM): select the attention backend in the runner

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` | `<dir>/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) |

examples/bench/bench_core.h

@@ -555,8 +555,13 @@ inline BenchResult RunBench(const BenchConfig& cfg) {
           std::max(max_prompt, static_cast<int>(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);

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

include/vllm/v1/attention/registry.h

@@ -85,6 +85,21 @@ std::unique_ptr<AttentionBackend> 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",

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<PagedKvCache>& 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<std::string>& attn_backend_names() const {
+    return attn_backend_names_;
+  }
   const std::vector<GdnStateCache>& 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<std::unique_ptr<CacheBuffer>> conv_buf_;
   std::vector<PagedKvCache> attn_kv_;
   std::vector<GdnStateCache> 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<std::string> attn_backend_names_;
 
   // ── KV-EXTERNAL-CACHE (LMCache) worker-side store/load ──────────────────────
   // Non-owning; null (default) = inert. See set_kv_connector.

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

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") {

src/vllm/v1/attention/registry.cpp

@@ -120,4 +120,34 @@ std::unique_ptr<AttentionBackend> 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<int64_t> 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<int64_t> 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

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<FaDims> 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<char> mla_layer_mask;
   layer_kv_class_.assign(static_cast<size_t>(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<size_t>(num_blocks_) * static_cast<size_t>(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<size_t>(l)]->kind()
+          : kv_cache_config
+                .kv_cache_groups[static_cast<size_t>(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<size_t>(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<int>(queue_.device.type),
+                     static_cast<long long>(num_blocks_),
+                     static_cast<long long>(fa_block_size),
+                     static_cast<long long>(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<int>(queue_.device.type),
+                     static_cast<long long>(num_blocks_),
+                     static_cast<long long>(fa_block_size),
+                     static_cast<long long>(fa_dims[i].num_kv_heads),
+                     static_cast<long long>(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);
   }
 

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;