Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h`:
- Around line 940-951: getSizePerHead() can return -1 for recurrent-state pools
created by WindowBlockManager::WindowBlockManager(), which then propagates
through getSizePerHeadPerWindow()/getSizePerHeadForWindow() and the nanobind
API; change getSizePerHead() to guard against negative sizePerHead by checking
mPools.front().sizePerHead and returning 0 when it is less than 0 (keep the
existing behavior of returning 0 when mPools.empty()), so callers never observe
a negative head dimension.

In `@examples/auto_deploy/model_registry/configs/gemma4_moe.yaml`:
- Line 13: Remove the duplicate max_batch_size entry nested inside the
cuda_graph_config block so the file only defines max_batch_size once at the top
level; specifically edit the gemma4_moe.yaml to delete the nested
"max_batch_size" key under "cuda_graph_config" and leave only the top-level
"max_batch_size: 512", keeping "cuda_graph_config" containing just the
"batch_sizes" array to match the other configs.

In `@tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py`:
- Around line 1326-1356: When cache_loc_per_group is provided but may be
partial, missing groups silently reuse prior staged values; fix by normalizing
the per-group maps before staging: for each group_idx in
1..self.num_window_groups-1, ensure cache_loc_per_group, cu_num_pages_per_group,
extra_page_per_seq_per_group, and kv_page_offset_per_group contain an entry (use
cache_loc, cu_num_pages, extra_page_per_seq, and kv_page_offset or a
per-sequence [0]*self.num_sequences for kv_page_offset as the group-0 defaults),
then proceed to call self._stage_arg for each suffix as currently done;
alternatively raise a clear error if any map is missing an entry — implement
this normalization where the current "if cache_loc_per_group is not None:" loop
runs so subsequent calls to self._stage_arg always see explicit values for every
group.

In
`@tensorrt_llm/_torch/auto_deploy/custom_ops/attention/triton_paged_attention.py`:
- Around line 1163-1174: The current max_cached = page_counts *
kv_cache.shape[3] overcounts tokens by treating every page as full; instead
compute the actual addressable tokens per sample as full_pages * page_size +
last_page_len (and zero when page_counts==0). Replace the max_cached line with
logic using page_counts, kv_cache.shape[3] (page_size) and the per-sample
last-page lengths (e.g. cu_last_page_lens): full_pages = torch.clamp(page_counts
- 1, min=0); max_cached = full_pages * page_size + torch.where(page_counts > 0,
cu_last_page_lens, torch.zeros_like(cu_last_page_lens)); then continue to use
cache_len_capped and seq_len_with_cache as before.

In `@tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py`:
- Around line 887-923: The loop currently sets num_active =
len(all_indices[front_removed:]) so the "extra" slot is never extracted; instead
compute the intended active block count first (e.g., active_count = determine
the number of active pages for this window — often window_size or a configured
per-group active size), then slice all_indices into active_pages =
all_indices[front_removed : front_removed + active_count] and extra_page =
all_indices[front_removed + active_count] if that index exists else -1; replace
uses of active_indices/num_active with active_pages/len(active_pages) and set
extra_page_per_seq (and extra_page_per_seq_per_group) from extra_page, updating
cache_loc, cache_loc_per_group, cu_num_pages, cu_num_pages_per_group,
kv_page_offset and kv_page_offset_per_group accordingly (refer to
variables/all_indices, front_removed, cache_loc, extra_page_per_seq,
cache_loc_per_group, extra_page_per_seq_per_group, cu_num_pages,
cu_num_pages_per_group in ad_executor.py).

In `@tensorrt_llm/_torch/auto_deploy/shim/interface.py`:
- Around line 121-131: This file is missing the required NVIDIA/Apache header;
add the repository-standard copyright/license header (with the latest
modification year) to the top of
tensorrt_llm/_torch/auto_deploy/shim/interface.py, before any imports or code,
following the same format used in other source files (NVIDIA copyright line +
Apache 2.0 license block and SPDX identifier), so the module (containing symbols
like KVCacheManager, MambaHybridCacheManager, self._kv_cache_manager,
self._kv_group_windows) complies with the project's header policy.
- Around line 916-924: The transform-time _kv_group_windows can become stale
because KVCacheManager may clamp windows during construction; after calling
self._assign_kv_cache_views(...) you must refresh the cached group windows from
the manager so downstream callers (e.g., ad_executor.py using get_cache_indices
and get_num_front_blocks_removed) see the finalized per-window sizes—update
self._kv_group_windows from the KVCacheManager's authoritative windows (via
self._kv_cache_manager.windows or the manager API that exposes per-window sizes)
immediately after _assign_kv_cache_views returns, keeping _kv_cache_config_tuned
and block_offset_multiplier logic unchanged.
- Around line 814-823: The branch that returns self.info.max_seq_len when
total_max_tokens is None and SWA is present incorrectly overrides an explicit
VSWA budget stored on self._kv_cache_config_original; update the logic in the
function (around the total_max_tokens /
self._has_swa_window(size_per_head_per_window) check) to first detect and
preserve an existing user-set max_tokens on self._kv_cache_config_original
(e.g., check getattr(self._kv_cache_config_original, "max_tokens", None) or
similar) and return that value when present instead of collapsing to
self.info.max_seq_len, so that _prepare_kv_cache_config and config.max_tokens
keep the intended budget.

In `@tensorrt_llm/_torch/auto_deploy/transform/library/kvcache.py`:
- Line 343: The loop in kvcache.py uses enumerate but never uses the index
variable (for idx, (attn_node, qkv, cache_in_nodes, constants, group_idx) in
enumerate(layer_infos)), triggering Ruff B007; fix by either removing enumerate
and iterating directly over layer_infos (for attn_node, qkv, cache_in_nodes,
constants, group_idx in layer_infos) or rename the index to _idx if you want to
keep enumerate, updating the loop header in the same location.
- Around line 319-350: Pass 2 only replaces per-layer standard metadata
(layer_meta_nodes_std) but leaves group-sensitive inputs in meta_nodes_extra
untouched, so multi-group backends like TrtllmAttention (which populate
meta_nodes_extra via prepare_trtllm_metadata) still reference group-0 nodes;
update the loop over layer_infos to also build a per-group copy of
meta_nodes_extra (e.g., layer_meta_nodes_extra = list(meta_nodes_extra)) for
group_idx > 0 and, using the same std_arg_names / vswa_group_nodes mapping you
use for layer_meta_nodes_std, replace any entries in layer_meta_nodes_extra
whose argument names match keys in vswa_group_nodes[group_idx] with the
corresponding vswa_group_nodes[group_idx][arg_name]; ensure downstream uses the
per-group layer_meta_nodes_extra where prepare_trtllm_metadata or
TrtllmAttention expects group-specific inputs so VSWA routing is correct.

In
`@tests/unittest/auto_deploy/singlegpu/transformations/library/test_kv_cache.py`:
- Around line 764-769: CachedSequenceInterface is being instantiated without the
required max_num_tokens argument; update each creation of
CachedSequenceInterface (the instances inside this test file) to pass
max_num_tokens=default_max_num_tokens(max_seq_len, batch_size) so the
constructor matches its signature—specifically add that parameter to the calls
that currently set max_seq_len, max_batch_size, device, and kv_cache_config.

---

Outside diff comments:
In `@cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp`:
- Around line 653-669: Block IDs are currently treated as global but
WindowBlockManager instances hosted under BlockManager (mWindowBlockManagers)
make IDs unique only per window; update the pin/unpin contract so pinned handles
are window-qualified (e.g., pair<windowSize, blockId> or a per-window map) and
propagate this everywhere: change BlockManager::storeBlocksForReuse() to record
window-qualified IDs when pinBlocks==true, change
BlockManager::unpinBlocksById() to group incoming handles by window and call the
appropriate WindowBlockManager unpin path (dispatch into each WindowBlockManager
in mWindowBlockManagers rather than always using the first), and update any
callers and types that assume flat IDs to accept and pass the new
window-qualified pin handle.

In `@tensorrt_llm/_torch/pyexecutor/resource_manager.py`:
- Around line 554-561: The adjusted attention-window sizes returned by
_validate_and_adjust_attention_windows may rewrite
self.max_attention_window_vec, but the per-window override dicts
(size_per_head_per_window and dtype_per_window) are still using the original
keys; update those dicts after calling _validate_and_adjust_attention_windows to
remap or rebuild their entries to the adjusted window sizes (i.e., iterate the
original override entries and replace keys with the corresponding values from
the new self.max_attention_window_vec or reconstruct the dicts from the adjusted
vector) so lookups in the C++ ctor and get_buffers() use the clamped keys rather
than falling back to scalar defaults.
- Around line 1412-1429: The sizing loop currently uses the fallback
self.head_dim and dtype for all windows; update it to use the
per-window/per-layer head-dim and dtype maps when computing cache_size_per_token
and cache_size_bytes_per_token so mixed-shape VSWA pools are sized correctly:
inside calculate_cache_size_per_token (and where you call get_size_in_bytes and
the NVFP4 scaling branch) look up the actual head_dim and dtype for the current
window/layers (e.g., use self.head_dim_map or self.head_dim_per_layer and
self.dtype_map/self.dtype_per_layer) instead of self.head_dim/self.dtype,
compute per-layer/per-window contributions accordingly, and then sum them to
produce required_mem_bytes_per_seq and to feed
KVCacheManager.calculate_scaling_factor_size_bytes so blocks_per_window reflects
the real C++ layout.

In
`@tests/unittest/auto_deploy/singlegpu/transformations/library/test_kv_cache.py`:
- Around line 1-2: Add the NVIDIA copyright + Apache-2.0 header at the top of
the file, remove the legacy typing imports, and replace usages of typing.List
and typing.Optional with Python 3.10+ builtins/generics: delete "from typing
import List, Optional", keep "from types import SimpleNamespace", change any
"List[T]" to "list[T]" and any "Optional[T]" to "T | None", and update the
function/method return annotation referenced around line 49 to use the new
"list[...]" or "| None" form as appropriate (adjust imports accordingly so no
typing.List/Optional remain).

---

Nitpick comments:
In `@cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp`:
- Around line 8698-8720: The test unnecessarily constructs a full BlockManager
and calls BlockManager::allocatePools() for a pure GenerationRequest bookkeeping
test; instead create the needed windowSizeToMetadata/BlocksPerWindow data
directly (or via a small helper) and construct a GenerationRequest with that
metadata so no GPU pools are allocated. Replace the BlockManager instantiation
and blockManager.allocatePools() calls with construction of the BlocksPerWindow
/ maxAttentionWindowVec used in the BlockManager call, then pass that metadata
into GenerationRequest (or a test helper) so only
removeFrontBlock/clearCacheBlocks/getNumFrontBlocksRemoved logic on
GenerationRequest is exercised without invoking BlockManager::allocatePools or
creating a CudaStream.

In `@tests/integration/defs/accuracy/test_llm_api_autodeploy.py`:
- Around line 1104-1155: Add a perf-test entry for this KV-cache/CUDA-graph
path: create or update an l0_perf.yml entry and a matching qa llm_perf_*.yml
entry that reference the new/changed scenario exercised by test_bf16 (the
AutoDeployLLM invocation using model_factory="Gemma4ForConditionalGeneration",
compile_backend="torch-cudagraph", attn_backend="triton_paged", kv_cache_config
and cuda_graph_config), so scheduled and pre-merge perf regressions are caught;
ensure the perf test captures relevant config tags (kv-cache, cuda-graph,
compile_backend) and points to the appropriate benchmark task(s) (e.g.,
MMLU/GSM8K) used in the test.
- Around line 1100-1102: The test hard-codes MODEL_PATH to
llm_models_root()/gemma/... which bypasses the shared resolver; replace the
manual path with a call to hf_id_to_local_model_dir(MODEL_NAME) so the test uses
the central resolver (keep MODEL_NAME and EXTRA_EVALUATOR_KWARGS unchanged),
i.e. set MODEL_PATH = hf_id_to_local_model_dir(MODEL_NAME) and import or
reference hf_id_to_local_model_dir where used.

In
`@tests/unittest/auto_deploy/singlegpu/transformations/library/test_kv_cache.py`:
- Around line 897-898: Move the two local "import math" statements out of the
test functions and place a single "import math" at the module level with the
other imports; locate the inline imports by searching for "import math" inside
the test functions in this test_kv_cache module and remove them there so all
tests use the top-level module import instead.