[None][feat] Enable KVCacheManagerV2 by default for Llama and Llama4

docs/source/features/kvcache.md

@@ -89,6 +89,7 @@ Models that select the V2 manager by default:
 | DeepSeek-V4 | Sparse attention attaches auxiliary per-layer buffers |
 | GPT-OSS | Sliding window on every other layer (VSWA), so the sliding-window and full-attention pools are sized independently |
 | Gemma3 / Gemma4 (text and multimodal) | Alternating sliding-window and full-attention layers (VSWA); same independent pool sizing |
+| Llama / Llama4 | Uniform KV pool layout; chunked attention does not partition the pools |
 
 Separately, Gemma4 hybrid attention and sparse-attention models are routed to
 V2 unconditionally: their per-layer buffer layouts cannot be represented by V1's

tensorrt_llm/_torch/models/modeling_llama.py

@@ -1,3 +1,17 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 import copy
 import os
 from typing import Any, Dict, List, Literal, Optional, Tuple, Union
@@ -1134,6 +1148,14 @@ def forward(
 @register_auto_model("LlamaForCausalLM")
 class LlamaForCausalLM(SpecDecOneEngineForCausalLM[LlamaModel, LlamaConfig]):
 
+    @classmethod
+    def get_preferred_kv_cache_manager_version(
+        cls,
+        pretrained_config: Any = None,
+    ) -> Literal["V2"]:
+        """Prefer KV cache manager V2 for Llama."""
+        return "V2"
+
     @classmethod
     def get_preferred_transceiver_runtime(
         cls,
@@ -1505,6 +1527,22 @@ def call_with_text_prompt(
 class Llama4ForConditionalGeneration(SpecDecOneEngineForCausalLM[Llama4Model,
                                                                  Llama4Config]):
 
+    @classmethod
+    def get_preferred_kv_cache_manager_version(
+        cls,
+        pretrained_config: Any = None,
+    ) -> Literal["V2"]:
+        """Prefer KV cache manager V2 for Llama4."""
+        return "V2"
+
+    @classmethod
+    def get_preferred_transceiver_runtime(
+        cls,
+        pretrained_config: Any = None,
+    ) -> Optional[Literal["CPP", "PYTHON"]]:
+        """Prefer the Python transceiver for Llama4 NIXL disaggregated serving."""
+        return "PYTHON"
+
     def __init__(
         self,
         model_config: ModelConfig[Llama4Config],

tests/integration/defs/kv_cache/test_kv_cache_iteration_stats.py

@@ -73,6 +73,13 @@
     "iterIntraDeviceCopyBytes",
 ]
 
+SECONDARY_FIELDS = {
+    "secondaryMaxNumBlocks",
+    "secondaryFreeNumBlocks",
+    "secondaryUsedNumBlocks",
+}
+NON_SECONDARY_FIELDS = set(ALL_FIELDS) - SECONDARY_FIELDS
+
 TEST_NAMES = {
     1: "Cold start",
     2: "Partial block reuse",
@@ -383,26 +390,34 @@ def test_rapid_fire(self, llm_instance, all_collected, request):
         assert total_alloc > 0, "iterAllocTotalBlocks = 0 across all entries"
 
     def test_field_completeness(self, llm_instance, all_collected, request):
-        """Field completeness — verify all 18 fields present across all collected stats."""
+        """Field completeness — verify fields in their V2 window and cold-pool views."""
         # If running standalone (no prior tests), generate some traffic
         if not all_collected:
             llm_instance.generate(["Hello world"], SamplingParams(max_tokens=16))
             collect_stats(llm_instance, all_collected)
 
         entries_with_kv = 0
-        missing_fields = set()
         for s in all_collected:
             ki = s.get("kvCacheIterationStats")
             if ki:
                 entries_with_kv += 1
+                # V2 reports secondary gauges by cold pool group, not by window.
                 for ws, v in ki.items():
-                    for field in ALL_FIELDS:
-                        if field not in v:
-                            missing_fields.add(field)
+                    missing_fields = NON_SECONDARY_FIELDS - v.keys()
+                    assert not missing_fields, (
+                        f"Missing kvCacheIterationStats fields for window {ws}: "
+                        f"{sorted(missing_fields)}"
+                    )
+
+                for group, v in s.get("kvCacheIterationStatsByColdPoolGroup", {}).items():
+                    missing_fields = SECONDARY_FIELDS - v.keys()
+                    assert not missing_fields, (
+                        f"Missing kvCacheIterationStatsByColdPoolGroup fields for group {group}: "
+                        f"{sorted(missing_fields)}"
+                    )
 
         print(f"  Entries with kvCacheIterationStats: {entries_with_kv}/{len(all_collected)}")
         assert entries_with_kv > 0, "no entries contain kvCacheIterationStats"
-        assert len(missing_fields) == 0, f"Missing fields: {sorted(missing_fields)}"
 
 
 # ---------------------------------------------------------------------------

tests/unittest/llmapi/test_llm_args.py

@@ -737,11 +737,13 @@ def test_explicit_value_overrides_model_preference(self, user_setting,
 
         assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is user_setting
 
-    def test_registered_models_prefer_v2(self):
+    def test_registered_models_prefer_v2(self) -> None:
         from tensorrt_llm._torch.models.modeling_utils import \
             get_registered_model_class
 
         architectures = (
+            "LlamaForCausalLM",
+            "Llama4ForConditionalGeneration",
             "DeepseekV3ForCausalLM",
             "DeepseekV32ForCausalLM",
             "GlmMoeDsaForCausalLM",
@@ -770,7 +772,7 @@ def test_registered_models_prefer_v2(self):
             assert model_cls is not None
             assert model_cls.get_preferred_kv_cache_manager_version() == "V2"
 
-    def test_registered_models_keep_v2_on_nixl(self):
+    def test_registered_models_keep_v2_on_nixl(self) -> None:
         """Models preferring V2 and the Python transceiver keep V2 on NIXL.
 
         Both sentinels start at 'auto'; production resolves the transceiver
@@ -783,6 +785,7 @@ def test_registered_models_keep_v2_on_nixl(self):
             get_registered_model_class
 
         architectures = (
+            "Llama4ForConditionalGeneration",
             "DeepseekV3ForCausalLM",
             "DeepseekV32ForCausalLM",
             "GlmMoeDsaForCausalLM",