From 5e135b71fca94f4ac20a354da6bb5491d19e9de4 Mon Sep 17 00:00:00 2001 From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> Date: Thu, 10 Sep 2026 01:40:30 -0700 Subject: [PATCH 1/5] [None][feat] Enable KVCacheManagerV2 by default for Llama and Llama4 Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> --- docs/source/features/kvcache.md | 1 + tensorrt_llm/_torch/models/modeling_llama.py | 38 ++++++++++++++++++++ 2 files changed, 39 insertions(+) diff --git a/docs/source/features/kvcache.md b/docs/source/features/kvcache.md index d77b9fdebaac..66e7b30f141d 100644 --- a/docs/source/features/kvcache.md +++ b/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 diff --git a/tensorrt_llm/_torch/models/modeling_llama.py b/tensorrt_llm/_torch/models/modeling_llama.py index 18e1951e70f2..46df832ac489 100644 --- a/tensorrt_llm/_torch/models/modeling_llama.py +++ b/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], From bbcd4bb15f76121099eaaed6e2d0f29d3faeb682 Mon Sep 17 00:00:00 2001 From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> Date: Thu, 10 Sep 2026 22:48:47 -0700 Subject: [PATCH 2/5] [None][test] Validate KV iteration stats in manager-specific views Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> --- .../kv_cache/test_kv_cache_iteration_stats.py | 30 +++++++++++++++---- 1 file changed, 24 insertions(+), 6 deletions(-) diff --git a/tests/integration/defs/kv_cache/test_kv_cache_iteration_stats.py b/tests/integration/defs/kv_cache/test_kv_cache_iteration_stats.py index 6319aa85976e..4456fa02cea8 100644 --- a/tests/integration/defs/kv_cache/test_kv_cache_iteration_stats.py +++ b/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,37 @@ 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 manager-specific 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. + is_v2 = "kvCacheIterationStatsByPoolGroup" in s + expected_fields = NON_SECONDARY_FIELDS if is_v2 else set(ALL_FIELDS) for ws, v in ki.items(): - for field in ALL_FIELDS: - if field not in v: - missing_fields.add(field) + missing_fields = expected_fields - v.keys() + assert not missing_fields, ( + f"Missing kvCacheIterationStats fields for window {ws}: " + f"{sorted(missing_fields)}" + ) + + if is_v2: + 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)}" # --------------------------------------------------------------------------- From 1b9e258b592f972db03adb197d933d4ebb52818a Mon Sep 17 00:00:00 2001 From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> Date: Thu, 10 Sep 2026 22:59:34 -0700 Subject: [PATCH 3/5] [None][test] Assume V2 in KV iteration stats completeness checks Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> --- .../kv_cache/test_kv_cache_iteration_stats.py | 19 ++++++++----------- 1 file changed, 8 insertions(+), 11 deletions(-) diff --git a/tests/integration/defs/kv_cache/test_kv_cache_iteration_stats.py b/tests/integration/defs/kv_cache/test_kv_cache_iteration_stats.py index 4456fa02cea8..cb3826957661 100644 --- a/tests/integration/defs/kv_cache/test_kv_cache_iteration_stats.py +++ b/tests/integration/defs/kv_cache/test_kv_cache_iteration_stats.py @@ -390,7 +390,7 @@ 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 fields in their manager-specific views.""" + """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)) @@ -402,22 +402,19 @@ def test_field_completeness(self, llm_instance, all_collected, request): if ki: entries_with_kv += 1 # V2 reports secondary gauges by cold pool group, not by window. - is_v2 = "kvCacheIterationStatsByPoolGroup" in s - expected_fields = NON_SECONDARY_FIELDS if is_v2 else set(ALL_FIELDS) for ws, v in ki.items(): - missing_fields = expected_fields - v.keys() + missing_fields = NON_SECONDARY_FIELDS - v.keys() assert not missing_fields, ( f"Missing kvCacheIterationStats fields for window {ws}: " f"{sorted(missing_fields)}" ) - if is_v2: - 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)}" - ) + 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" From bea2f00a4c2d9f6d31e41bf0bcfab76a8352ae1e Mon Sep 17 00:00:00 2001 From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> Date: Sun, 13 Sep 2026 22:23:29 -0700 Subject: [PATCH 4/5] [None][test] Cover Llama V2 model preferences Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> --- tests/unittest/llmapi/test_llm_args.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index e9c9925dfb37..0cd17b13658a 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/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", From b0d509842a05033b09fa2b127f7f0eba13963115 Mon Sep 17 00:00:00 2001 From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> Date: Sun, 13 Sep 2026 22:59:16 -0700 Subject: [PATCH 5/5] [None][test] Cover Llama4 NIXL runtime resolution Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> --- tests/unittest/llmapi/test_llm_args.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index 0cd17b13658a..4a8788c50932 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -772,7 +772,7 @@ def test_registered_models_prefer_v2(self) -> None: 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 @@ -785,6 +785,7 @@ def test_registered_models_keep_v2_on_nixl(self): get_registered_model_class architectures = ( + "Llama4ForConditionalGeneration", "DeepseekV3ForCausalLM", "DeepseekV32ForCausalLM", "GlmMoeDsaForCausalLM",