fix(FIX-KV-GROUP-LAYER-COUNT): resolve the placeholder KV group names, so the budget bounds the allocation again (#1963, #1966)

docs/FEATURES.md

@@ -386,7 +386,7 @@ CPU elementwise GEMM (f32/f16/bf16) runs AVX2 and AVX-512 tiers on x86 where the
 | ROCm | W0 community-verified on 5 gfx archs; classic-dense and GDN-hybrid e2e run all-native; correctness gaps remain | 49 registered ops including the GDN state/conv/postconv/recurrence set, MoE combine/gate, and keep-quant expert GEMM; APU managed-allocation branch remains unverified. [ROCm guide](ROCM.md) |
 | XPU, TPU | Not started | CUDA, CPU, Metal and Vulkan are the built backends |
 | Custom logits processors on CUDA | Open, not root-caused | Segfaults in a CUDA build, 232/232 green on CPU |
-| Memory budgeting (`ROAD-V1-MEM`, #83) | M1+M2 landed (absolute bytes) | `--kv-cache-memory` sizes the KV pool from an absolute byte budget (ABI v16, group-aware divisor); `--num-blocks` overrides; `--gpu-memory-utilization` needs the M3 profile run (dgx-gated). See `specs/kv-sizing.md` |
+| Memory budgeting (`ROAD-V1-MEM`, #83) | M1+M2 landed (absolute bytes) | `--kv-cache-memory` sizes the KV pool from an absolute byte budget (ABI v16, per-layer divisor since #1963 — the group-aware one counted placeholder names and overshot by the layer count); `--num-blocks` overrides; `--gpu-memory-utilization` needs the M3 profile run (dgx-gated). See `specs/kv-sizing.md` |
 | Gemma4 MoE ROCm FP8 + SharedK-WMMA | Partial | Dual-GPU FP8 resident experts, SharedK-WMMA prefill (RDNA4); decode-graph and forward extract deferred. Env `VT_GEMMA4_*`/`VT_ATTN_*`, seam `test_gemma4_rocm_fp8_seams`. [spec](../.agents/specs/gemma4-rocm-fp8-moe.md) |
 
 ## How to read this page

docs/USAGE.md

@@ -208,6 +208,24 @@ Note that `--kv-cache-memory` is what turns the halved block into twice the
 pool. Without it the server falls back to a fixed block count, and `fp8` then
 halves the KV bytes for the same context instead.
 
+**`--kv-cache-memory` now bounds the whole pool, and it did not before.** The
+value is an absolute budget for the paged KV cache, and the engine sizes the
+block count so that everything it allocates fits inside it — which is what vLLM
+means by the flag. Until #1963 the divisor counted one layer per KV group while
+the engine allocated a buffer per layer, so the same number bought as many times
+the memory as the model has attention layers: 8.5 GiB of buffers for
+`--kv-cache-memory 1073741824` on the 27B. If you tuned this flag against the
+old behaviour, the same value now gives a shorter served context; raise it, and
+the auto-fit line on stderr tells you what it settled on.
+
+`--num-blocks` is unaffected: it names a per-layer block count and always did,
+so a launch line that sizes the pool that way means exactly what it meant
+before. Only the byte budget converts differently. The recurrent-state clamp
+(#1983) reads the resolved block count, so at a fixed `--kv-cache-memory` it
+now seats fewer concurrent sequences than it did — it is being told the pool's
+true size for the first time, and the `INFO recurrent-state budget:` line names
+what it compared.
+
 **It costs you the fast attention kernels, and we have not measured the net.**
 An fp8 KV cache is read by the tiled prefill and block decode kernels only.
 FA-2 prefill, all three FA-2 decode topologies, the WMMA ladder and the

include/vllm/v1/kv_cache_interface.h

@@ -78,6 +78,7 @@
 #include <memory>
 #include <optional>
 #include <string>
+#include <string_view>
 #include <vector>
 
 #include "vllm/v1/kv_cache_dtype.h"
@@ -559,6 +560,66 @@ struct KVCacheConfig {
 // spec's own `page_size_bytes()` throws (deferred quantized-KV math).
 int64_t KVBytesPerBlock(const KVCacheConfig& config);
 
+// The model layer index encoded in an upstream-style KV layer name, or nullopt
+// when the name carries no layer identity.
+//
+// Upstream's `KVCacheGroupSpec.layer_names` holds real module paths
+// ("model.layers.5.self_attn.attn", "backbone.layers.12.mixer"); this returns
+// the integer of the `.layers.<N>.` segment. It deliberately returns nullopt
+// for a PLACEHOLDER group name — "fa", "gdn", "mla", "kda", "fa_draft",
+// "encoder" — because such a name names no layer at all, and that is the
+// discriminator `ResolveKVCacheGroupLayerNames` and the runner's
+// `GroupLayerMask` both key on.
+//
+// Lived in `gpu/runner.cpp`'s anonymous namespace until
+// FIX-KV-GROUP-LAYER-COUNT needed the same parse in the sizing path. One
+// function rather than two copies, because a second derivation of one rule is
+// the thing that can disagree.
+std::optional<int64_t> KVCacheLayerIndexOfName(std::string_view name);
+
+// FIX-KV-GROUP-LAYER-COUNT (#1963, #1966). Replace PLACEHOLDER group names with
+// the real per-layer names the runner's own allocation classification implies,
+// so that every consumer weighting a group by `layer_names.size()` weights it
+// by the layers the runner will actually allocate for.
+//
+// WHY THIS EXISTS. `layer_names` is upstream's per-layer name list
+// (`kv_cache_utils.py:1208-1210` appends every layer sharing a spec object), and
+// upstream bounds its allocation with `max(len(group.layer_names) ...)`
+// (`:1399`) over that list, dividing the budget by the same count it multiplies
+// the allocation by (`:1005-1008`, `:1409-1416`). Thirty-three of our
+// thirty-four registries publish a single placeholder string instead, so
+// `KVBytesPerBlock` divided a byte budget by ONE layer's page while the runner
+// allocated one buffer PER layer — 8.5 GiB for a 1 GiB budget on the 27B — and
+// the #371 recurrent-state OOM guard read 0.90 GiB against a 43.40 GiB
+// allocation.
+//
+// THE CLASSIFICATION IS THE RUNNER'S OWN, not a second derivation of the
+// model's shape:
+//   - a layer is recurrent iff the config has a Mamba group AND `layer_types`
+//     is non-empty AND `layer_types[l] == "linear_attention"` — the predicate at
+//     `gpu/runner.cpp`'s `is_gdn` fallback;
+//   - the TARGET attention group is the FIRST non-eagle attention-kind group,
+//     which is the runner's own first-wins selection, and it covers every
+//     non-recurrent layer;
+//   - a SECOND attention group is the speculative draft layer: exactly ONE
+//     layer, at index `num_hidden_layers` (upstream's MTP head index,
+//     `qwen3_5_mtp.py:105-112`), because the runner allocates exactly one draft
+//     buffer and breaks;
+//   - a THIRD or later attention group gets an EMPTY list, because the runner
+//     allocates no buffer for it at all. Zero is the honest count. No registry
+//     emits one.
+//
+// A REGISTRY THAT ALREADY PUBLISHES REAL NAMES IS NEVER OVERWRITTEN. If any
+// group carries a name `KVCacheLayerIndexOfName` resolves, this returns with the
+// config untouched. `NemotronHForCausalLM` is that case and it knows more than
+// the fallback can — its `layer_types` is empty and its MoE blocks cache
+// nothing, which is exactly what #810 fixed.
+//
+// Idempotent: a second call finds real names and returns.
+void ResolveKVCacheGroupLayerNames(KVCacheConfig& config,
+                                   int64_t num_hidden_layers,
+                                   const std::vector<std::string>& layer_types);
+
 // KV-FP8 W3 — HALF-SIZED KV BLOCKS. Rewrite every ATTENTION spec in `config` to
 // the resolved KV storage dtype, then hand the fp8 interpretation and the
 // per-tensor scales to the same specs.

include/vllm/v1/worker/gpu/runner.h

@@ -355,6 +355,21 @@ class GPUModelRunner final : public ModelRunnerBase {
   // value the old HF-config arithmetic could not.
   int64_t fa_page_size_bytes() const { return fa_page_size_bytes_; }
 
+  // FIX-KV-GROUP-LAYER-COUNT (#1963, #1966). What `initialize_kv_cache`
+  // ALLOCATED, summed over every buffer it created, so a gate can compare the
+  // sizing arithmetic against the allocation rather than against a second copy
+  // of the same formula. Both are 0 before `initialize_kv_cache` runs.
+  //
+  //   ...paged_bytes() — the block-scaled half: one buffer per full-attention
+  //     layer plus the speculative draft layer's. This is the half a
+  //     `--kv-cache-memory` budget is supposed to bound, and
+  //     `KVBytesPerBlock(cfg) * cfg.num_blocks` is supposed to equal it.
+  //   ...allocated_bytes() — that plus the recurrent (GDN/Mamba) conv and SSM
+  //     state, which is sized per sequence slot and not per block, and which
+  //     `recurrent_state_bytes(cfg, max_num_reqs)` is supposed to equal.
+  int64_t kv_cache_allocated_paged_bytes() const;
+  int64_t kv_cache_allocated_bytes() const;
+
   // #810: the per-layer KV class `initialize_kv_cache` RESOLVED, index == model
   // layer index, one entry per hidden layer. `kNone` is a layer that no KV
   // cache group named and that therefore caches nothing — NemotronH's 23
@@ -469,9 +484,15 @@ class GPUModelRunner final : public ModelRunnerBase {
       return backend_resident_ ? backend_data_ : host_data_.data();
     }
 
+    // The byte size this buffer was constructed with — what the allocation
+    // COST, not what a formula predicts it cost. `kv_cache_allocated_bytes()`
+    // below sums these (FIX-KV-GROUP-LAYER-COUNT, #1963).
+    size_t bytes() const { return bytes_; }
+
    private:
     vt::Device device_;
     bool backend_resident_ = false;
+    size_t bytes_ = 0;
     void* backend_data_ = nullptr;
     std::vector<uint8_t> host_data_;
   };

src/vllm/entrypoints/model_loader.cpp

@@ -43,6 +43,7 @@
 #include "vllm/platforms/interface.h"  // CurrentPlatform() — SelectQueue
 #include "vllm/v1/core/hybrid_kv_budget.h"
 #include "vllm/v1/core/kv_cache_utils.h"  // check_enough_kv_cache_memory (M4)
+#include "vllm/v1/kv_cache_interface.h"  // FIX-KV-GROUP-LAYER-COUNT resolver
 #include "vllm/v1/structured_output/backend_native.h"  // MakeNativeBackendFactory
 #include "vllm/v1/structured_output/jump_forward.h"     // JumpForwardEnabled (SW3)
 #include "vt/dtype.h"
@@ -1439,14 +1440,43 @@ std::optional<vllm::SpeculativeConfig> LoadedEngine::ResolveSpecConfig(
 vllm::v1::KVCacheConfig LoadedEngine::MakeKVCacheMaybeSpec(
     const LoadedModel& model, const HfConfig& config, int block_size,
     int num_blocks, const std::optional<vllm::SpeculativeConfig>& spec) {
+  vllm::v1::KVCacheConfig kv;
   if (spec.has_value()) {
     // Speculation is Qwen3.5/3.6-only at this pin (both gate checkpoints); build
     // the widened spec KV directly (extra GDN k+1 state slots + widened conv row
     // + the `fa_draft` full-attn group). MakeQwen3_5KVCacheSpec(num_spec>0).
-    return vllm::MakeQwen3_5KVCacheSpec(config, block_size, num_blocks,
-                                        spec->ResolvedNumSpeculativeTokens());
+    kv = vllm::MakeQwen3_5KVCacheSpec(config, block_size, num_blocks,
+                                      spec->ResolvedNumSpeculativeTokens());
+  } else {
+    kv = ModelRegistry::MakeKVCache(model, config, block_size, num_blocks);
   }
-  return ModelRegistry::MakeKVCache(model, config, block_size, num_blocks);
+  // FIX-KV-GROUP-LAYER-COUNT (#1963, #1966). THE single funnel: both branches
+  // above return through here, so one call reaches every architecture and both
+  // the probe and the resized config MakeKVCacheResolved builds.
+  //
+  // Thirty-three of the thirty-four registries publish ONE placeholder name per
+  // KV group, and `KVBytesPerBlock` / `recurrent_state_bytes` read
+  // `layer_names.size()` as the layer count. Without this line a
+  // `--kv-cache-memory` budget is divided by ONE layer's page and then
+  // multiplied by every layer when the runner allocates: 1 GiB in, 8.5 GiB
+  // allocated on the 27B, and the #371 recurrent-state OOM guard reads 0.90 GiB
+  // against a 43.40 GiB allocation. Upstream cannot have that bug because the
+  // count that divides the budget and the count that sizes the allocation are
+  // the same expression over the same list (`kv_cache_utils.py:1399`,
+  // `:1005-1008`, `:1409-1416`).
+  //
+  // ORDERING, against `ResolveMaxNumSeqs` (#1983, which landed first): that
+  // resolver reads `kv_cfg_`, which is `MakeKVCacheResolved`'s result, so it
+  // always sees names this call has already resolved. Its seat count is
+  // `num_blocks`-linear, and `num_blocks` is the one input of its arithmetic
+  // this change moves — which is the point: upstream's `num_blocks` is
+  // per-layer (`kv_cache_utils.py:1008` divides by `num_layers`), and that is
+  // the meaning its unification against one attention page assumes. Before this
+  // line the byte-budget path handed it a count inflated by the layer count, so
+  // its clamp was too permissive. The two fixes agree; they do not fight.
+  vllm::v1::ResolveKVCacheGroupLayerNames(kv, config.num_hidden_layers,
+                                          config.layer_types);
+  return kv;
 }
 
 int LoadedEngine::ResolveNumBlocks(const EngineParams& params,

src/vllm/v1/kv_cache_interface.cpp

@@ -266,6 +266,104 @@ int64_t KVBytesPerBlock(const KVCacheConfig& config) {
   return bytes;
 }
 
+// See kv_cache_interface.h for the argument. Body moved verbatim from
+// `gpu/runner.cpp`'s `LayerIndexOfName`.
+std::optional<int64_t> KVCacheLayerIndexOfName(std::string_view name) {
+  constexpr std::string_view kSep = ".layers.";
+  const size_t at = name.find(kSep);
+  if (at == std::string_view::npos) return std::nullopt;
+  size_t i = at + kSep.size();
+  const size_t start = i;
+  int64_t value = 0;
+  while (i < name.size() && name[i] >= '0' && name[i] <= '9') {
+    value = value * 10 + (name[i] - '0');
+    if (value > (1 << 20)) return std::nullopt;  // not a layer index
+    ++i;
+  }
+  if (i == start) return std::nullopt;                         // ".layers.mixer"
+  if (i < name.size() && name[i] != '.') return std::nullopt;   // ".layers.5x"
+  return value;
+}
+
+void ResolveKVCacheGroupLayerNames(
+    KVCacheConfig& config, int64_t num_hidden_layers,
+    const std::vector<std::string>& layer_types) {
+  // Nothing to name. A caller with no layer count cannot be given one here, and
+  // inventing one is how the placeholder count became a layer count in the
+  // first place.
+  if (num_hidden_layers <= 0) return;
+
+  // A registry that already publishes REAL names knows more than the fallback
+  // classification below: NemotronH's `layer_types` is EMPTY and its MoE blocks
+  // register no attention module at all, so overwriting it would re-introduce
+  // exactly the 52-against-6 mis-classification #810 removed. One resolvable
+  // name anywhere is the whole config's answer, which also makes this
+  // idempotent.
+  for (const KVCacheGroupSpec& group : config.kv_cache_groups) {
+    for (const std::string& name : group.layer_names) {
+      if (KVCacheLayerIndexOfName(name).has_value()) return;
+    }
+  }
+
+  bool has_mamba_group = false;
+  for (const KVCacheGroupSpec& group : config.kv_cache_groups) {
+    if (group.kv_cache_spec != nullptr &&
+        group.kv_cache_spec->kind() == KVCacheSpecKind::kMamba) {
+      has_mamba_group = true;
+    }
+  }
+
+  // The runner's own `is_gdn` fallback predicate, and nothing else. A dense
+  // model (no Mamba group) and a hybrid whose config does not spell
+  // `layer_types` both classify every layer as attention here, which is what the
+  // runner does with them too.
+  std::vector<std::string> recurrent;
+  std::vector<std::string> attention;
+  for (int64_t l = 0; l < num_hidden_layers; ++l) {
+    const size_t idx = static_cast<size_t>(l);
+    const bool is_gdn = has_mamba_group && idx < layer_types.size() &&
+                        layer_types[idx] == "linear_attention";
+    if (is_gdn) {
+      recurrent.push_back("model.layers." + std::to_string(l) + ".linear_attn");
+    } else {
+      attention.push_back("model.layers." + std::to_string(l) +
+                          ".self_attn.attn");
+    }
+  }
+
+  bool target_named = false;
+  bool draft_named = false;
+  for (KVCacheGroupSpec& group : config.kv_cache_groups) {
+    if (group.kv_cache_spec == nullptr) continue;
+    const KVCacheSpecKind kind = group.kv_cache_spec->kind();
+    if (kind == KVCacheSpecKind::kMamba) {
+      group.layer_names = recurrent;
+      continue;
+    }
+    if (dynamic_cast<const AttentionSpec*>(group.kv_cache_spec.get()) ==
+        nullptr) {
+      continue;  // not an attention group and not recurrent: leave it alone.
+    }
+    if (!group.is_eagle_group && !target_named) {
+      group.layer_names = attention;
+      target_named = true;
+    } else if (!draft_named) {
+      // The speculative draft layer. Upstream registers the MTP head as one
+      // extra decoder layer at index num_hidden_layers
+      // (`qwen3_5_mtp.py:105-112`), and the runner allocates exactly one buffer
+      // for it before breaking out of its search.
+      group.layer_names = {"model.layers." +
+                           std::to_string(num_hidden_layers) +
+                           ".self_attn.attn"};
+      draft_named = true;
+    } else {
+      // The runner allocates NO buffer for a further attention group, so the
+      // honest weight is zero. Unreachable for every registry shipping today.
+      group.layer_names.clear();
+    }
+  }
+}
+
 namespace {
 
 // The ONE arithmetic statement W3 makes about block sizing, written where it can