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], 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..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 @@ -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)}" # --------------------------------------------------------------------------- diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index e9c9925dfb37..4a8788c50932 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", @@ -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",