Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
erictsai-nv wants to merge 1 commit into
base: main
Choose a base branch
from
+68 −7
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions docs/source/features/kvcache.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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); validated across text, multimodal, and disaggregated workloads |

Separately, Gemma4 hybrid attention and sparse-attention models are routed to
V2 unconditionally: their per-layer buffer layouts cannot be represented by V1's
Expand Down
24 changes: 24 additions & 0 deletions tensorrt_llm/_torch/models/modeling_llama.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1134,6 +1134,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"

@fredricz-20070104 fredricz-20070104 Sep 1, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[MINOR] V2 preference on LlamaForCausalLM is inherited by every subclass and every checkpoint mapping to this architecture

get_preferred_kv_cache_manager_version is a classmethod on LlamaForCausalLM, so it applies to (a) every checkpoint whose config declares architectures: ["LlamaForCausalLM"] — far wider than Llama 3.x, including third-party fine-tunes and derivative models — and (b) every Python subclass of LlamaForCausalLM in the registry, which inherits the hook without opting in. The validation matrix quoted in the description covers Llama3 and Llama4 only, so any derived architecture with a different KV layout silently changes manager on this merge and the only signal would be a post-merge regression.

Two concrete asks: (1) enumerate/confirm which registered classes inherit this hook and that none of them need V1, or push the preference down to the concrete classes you validated; (2) the title is [None] — for a default flip on the most widely used architecture in the repo, a ticket id makes a post-merge revert traceable. Note pretrained_config is accepted but unused, so the hook has no room to opt a variant out today.


@classmethod
def get_preferred_transceiver_runtime(
cls,
Expand Down Expand Up @@ -1505,6 +1513,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"

@fredricz-20070104 fredricz-20070104 Sep 1, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[MAJOR] Llama4 prefers the Python transceiver unconditionally, including when the KV manager resolves to V1

get_preferred_transceiver_runtime returns "PYTHON" with no dependency on whether V2 was actually selected — it does not inspect pretrained_config and has no branch. The PR rationale is "so disaggregated serving over NIXL retains V2", but the hook cannot express that condition.

Concrete trigger: a Llama4 NIXL disagg run where the user sets kv_cache_config.use_kv_cache_manager_v2=False, or where the resolver auto-demotes to V1 (e.g. the two-model speculative-decoding route). transceiver_runtime is still "auto", so the model hook wins and the run comes up as V1 manager + Python transceiver — a combination the Validation section never mentions (it lists Llama4 V2 on B200: quickstart, TP2/PP2, chunked prefill, one-model Eagle3). Before this PR that same config used the C++ transceiver, so this is a silent default change for a path with no stated coverage.

Suggested fix: make the preference conditional on V2 actually being in effect, or return None so the existing default applies when V2 is not chosen:

@classmethod
def get_preferred_transceiver_runtime(
    cls,
    pretrained_config: Any = None,
) -> Optional[Literal["CPP", "PYTHON"]]:
    # Only meaningful together with KVCacheManagerV2; leave CPP as-is for V1.
    return "PYTHON" if cls.get_preferred_kv_cache_manager_version(
        pretrained_config) == "V2" else None

If the resolver already couples the two decisions (not visible in this diff), please say so explicitly, and state whether any Llama4 disagg run exercised V1 + Python transceiver.


def __init__(
self,
model_config: ModelConfig[Llama4Config],
Expand Down
46 changes: 39 additions & 7 deletions tests/integration/defs/kv_cache/test_kv_cache_iteration_stats.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,10 +49,18 @@
"primaryMaxNumBlocks",
"primaryFreeNumBlocks",
"primaryUsedNumBlocks",
"primaryEvictableNumBlocks",
"primaryPeakFreeNumBlocks",
"primaryPeakUsedNumBlocks",
"primaryPeakEvictableNumBlocks",
# Instantaneous gauges — secondary (host) pool
"secondaryMaxNumBlocks",
"secondaryFreeNumBlocks",
"secondaryUsedNumBlocks",
"secondaryEvictableNumBlocks",
"secondaryPeakFreeNumBlocks",
"secondaryPeakUsedNumBlocks",
"secondaryPeakEvictableNumBlocks",
# Per-iteration deltas — context phase
"iterAllocTotalBlocks",
"iterAllocNewBlocks",
Expand All @@ -71,8 +79,22 @@
# Intra-device (GPU → GPU) block copies
"iterIntraDeviceCopyBlocks",
"iterIntraDeviceCopyBytes",
# Host blocks dropped instead of being copied to another cold tier
"iterHostDroppedBlocks",
"iterHostDroppedBytes",
]

SECONDARY_FIELDS = {
"secondaryMaxNumBlocks",
"secondaryFreeNumBlocks",
"secondaryUsedNumBlocks",
"secondaryEvictableNumBlocks",
"secondaryPeakFreeNumBlocks",
"secondaryPeakUsedNumBlocks",
"secondaryPeakEvictableNumBlocks",
}
NON_SECONDARY_FIELDS = set(ALL_FIELDS) - SECONDARY_FIELDS

TEST_NAMES = {
1: "Cold start",
2: "Partial block reuse",
Expand Down Expand Up @@ -383,26 +405,36 @@ 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
for ws, v in ki.items():
for field in ALL_FIELDS:
if field not in v:
missing_fields.add(field)
is_v2 = "kvCacheIterationStatsByPoolGroup" in s

@fredricz-20070104 fredricz-20070104 Sep 1, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[MINOR] Manager version inferred from a payload key; pool-group view contents never validated

is_v2 = "kvCacheIterationStatsByPoolGroup" in s infers the manager version from the presence of a serialized key rather than from the resolved config. Two consequences: (1) if that key is renamed or becomes conditional, a V2 run is silently graded against set(ALL_FIELDS) and fails with a misleading "Missing kvCacheIterationStats fields: secondary*" message instead of a clear version-detection error; (2) the pool-group view is used only as a flag — its own field contents are never asserted anywhere in this test, so a V2 regression that empties kvCacheIterationStatsByPoolGroup down to a bare dict still passes.

Suggested fix: derive the expectation from the LLM's resolved kv_cache_config.use_kv_cache_manager_v2 (available via llm_instance), and add a field-completeness check over s["kvCacheIterationStatsByPoolGroup"].values() the same way the cold view is checked.

expected_fields = NON_SECONDARY_FIELDS if is_v2 else set(ALL_FIELDS)

@fredricz-20070104 fredricz-20070104 Sep 1, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[MAJOR] ALL_FIELDS gains 10 newly-required fields with no serializer change in this PR

expected_fields is set(ALL_FIELDS) on the V1 arm and NON_SECONDARY_FIELDS on the V2 arm, and ALL_FIELDS just grew by primaryEvictableNumBlocks, primaryPeakFreeNumBlocks, primaryPeakUsedNumBlocks, primaryPeakEvictableNumBlocks, secondaryEvictable*, secondaryPeak*, iterHostDroppedBlocks, iterHostDroppedBytes (lines 52-55, 60-63, 83-84). The old docstring said "all 18 fields"; the list is now 26.

No serializer or stats-producer change is present in this PR, so the new names are asserted purely on the assumption that the current stats payload already emits them. Failure mode is concrete: if the V1 (or the V2 hot) view does not emit primaryPeak* or iterHostDropped*, test_field_completeness raises on the first entry with kvCacheIterationStats — and because the assert now lives inside the loop instead of being aggregated at the end, the message names only that one entry's gap rather than the full missing set as before.

Please confirm these field names match what the serializer emits today on both arms (and, if they were added by an earlier PR, reference it), or scope the newly-added names to the manager version that produces them.

for v in ki.values():
missing_fields = expected_fields - v.keys()
assert not missing_fields, (
f"Missing kvCacheIterationStats fields: {sorted(missing_fields)}"
)

if is_v2 and "kvCacheIterationStatsByColdPoolGroup" in s:

@fredricz-20070104 fredricz-20070104 Sep 1, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[MINOR] Secondary/cold-pool fields are never verified on the V2 arm if the cold view is absent

On the V2 arm, SECONDARY_FIELDS are removed from the hot-view expectation and only checked inside if is_v2 and "kvCacheIterationStatsByColdPoolGroup" in s:. If the fixture never produces host/cold-tier offloading — which is the likely case for a short generate(["Hello world"], max_tokens=16) standalone run — that branch never executes and the whole secondary-pool field set goes unasserted, so the test cannot fail even if V2 stopped reporting cold-pool stats entirely. With the Llama default now flipping to V2, that is exactly the arm this test is supposed to guard.

Suggested fix: either configure the fixture so a cold pool group is guaranteed (non-zero host cache / forced offload) and then require the key unconditionally, or assert explicitly that the key is expected-absent for this configuration so the skip is visible rather than silent:

cold = s.get("kvCacheIterationStatsByColdPoolGroup")
assert cold is not None, "V2 entry has no kvCacheIterationStatsByColdPoolGroup"

cold_stats = s["kvCacheIterationStatsByColdPoolGroup"]
for v in cold_stats.values():
missing_fields = SECONDARY_FIELDS - v.keys()
assert not missing_fields, (
"Missing kvCacheIterationStatsByColdPoolGroup fields: "
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)}"


# ---------------------------------------------------------------------------
Expand Down
4 changes: 4 additions & 0 deletions tests/unittest/llmapi/test_llm_args.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -711,6 +711,8 @@ def test_registered_models_prefer_v2(self):
"Gemma4ForCausalLM",
"Gemma4ForConditionalGeneration",
"Gemma4UnifiedForConditionalGeneration",
"LlamaForCausalLM",
"Llama4ForConditionalGeneration",
)
for architecture in architectures:
model_cls = get_registered_model_class(architecture)
Expand Down Expand Up @@ -751,6 +753,8 @@ def test_registered_models_keep_v2_on_nixl(self):
"Gemma4ForCausalLM",
"Gemma4ForConditionalGeneration",
"Gemma4UnifiedForConditionalGeneration",
"LlamaForCausalLM",
"Llama4ForConditionalGeneration",
)
for architecture in architectures:
model_cls = get_registered_model_class(architecture)
Expand Down
Loading