From 503f1af77c450909263be400c57c479a4dc5570d Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Tue, 15 Sep 2026 03:50:02 +0000 Subject: [PATCH 1/3] [None][chore] Port DFlash weight-load checks, DSpark norm_dim and disagg tests from the Rubin branch Merge-back of the portable subset of the Rubin runtime work. Most of the originally scoped feature (KV cache manager v2 uGPU localization, the DSpark MLA drafter, the disagg benchmark fill gate) is NOT included: its production half is absent from the source branch as well as from main, so porting the corresponding tests would have added red tests rather than coverage. See the PR description for the per-file disposition. Production: * modeling_dflash: fail loudly when a draft checkpoint omits backbone weights the drafter does not share with the target. `_assert_backbone_complete` walks the constructed module tree rather than a hand-kept list, and `WRAPPER_OWNED_WEIGHTS` gates the wrapper-owned `fc` / `hidden_norm` that a module walk cannot see. Previously `allow_partial_loading=True`, needed for the target-shared `embed_tokens` / `lm_head`, let a truncated checkpoint load clean and leave parameters at `torch.empty`. Also restructures the attention-backend init so every op handle starts at None and adds the `_uses_worker_attention_backend` hook for drafters that bring their own block decode; the FA4 backend added on main is preserved. * modeling_utils: hoist the fused-module component table out of `_load_weights_impl` into `FUSED_MODULE_COMPONENTS` so the check above reads the same table the loader uses instead of keeping a third copy. No behavior change. * modeling_speculative: bound an external drafter's position tables with the target's `max_seq_len`. Without it the drafter falls back to the checkpoint's advertised `max_position_embeddings` and allocates a table sized for a context the runtime never serves. * speculative/dflash: clamp the running context length before deriving the block-decode query positions. `_ctx_len` is clamped to `_max_ctx` only after the step's accepted tokens are folded in, so a request already at the ceiling indexed `num_accepted` positions past the end of the sequence. * dspark_rmsnorm_rope: add an optional `norm_dim` so the fused RMSNorm/RoPE kernel can normalize a prefix of the row and pass the remaining rope lanes through raw, as DeepSeek-style MLA needs. Defaults to the whole row and is bit-identical to the previous behavior there. `norm_dim` is now also forwarded to the support predicate, which the source branch omitted. * kv_cache_manager_v2/_block_radix_tree: document why the stale-tail prune requires every life cycle to be pageless, mirroring the C++ implementation. Tests: * New `test_block_radix_tree_stale_prune` covers that predicate, including the negative control that dead tails still get pruned. * New `test_disaggregated_multinode` is a two-node DSpark disagg accuracy harness. It skips unless run under a 2-node, 1-task-per-node Slurm allocation; the exact srun invocation is in the file header. It is deliberately left out of the test lists, as no current stage uses that layout. * `test_disagg_index_mapper_early_release` gains a case asserting that `release_index_slot` detaches every page-index view before the slot is reused. The import is adjusted for the `pyexecutor.kv_cache` package move. * `test_cache_reuse_adapter` gains a case for the DSpark gen-init draft reserve being removed before the SWA trim. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- .../dspark_rmsnorm_rope_custom_op.py | 19 +- .../blackwell/dspark_rmsnorm_rope.py | 36 +- tensorrt_llm/_torch/models/modeling_dflash.py | 125 +++++- .../_torch/models/modeling_speculative.py | 11 + tensorrt_llm/_torch/models/modeling_utils.py | 16 +- tensorrt_llm/_torch/speculative/dflash.py | 12 +- .../kv_cache_manager_v2/_block_radix_tree.py | 5 + .../accuracy/test_disaggregated_multinode.py | 395 ++++++++++++++++++ .../test_disagg_index_mapper_early_release.py | 42 +- .../disaggregated/test_cache_reuse_adapter.py | 40 +- .../test_block_radix_tree_stale_prune.py | 151 +++++++ 11 files changed, 819 insertions(+), 33 deletions(-) create mode 100644 tests/integration/defs/accuracy/test_disaggregated_multinode.py create mode 100644 tests/unittest/kv_cache_manager_v2_tests/test_block_radix_tree_stale_prune.py diff --git a/tensorrt_llm/_torch/custom_ops/dspark_rmsnorm_rope_custom_op.py b/tensorrt_llm/_torch/custom_ops/dspark_rmsnorm_rope_custom_op.py index 7f341c4999d8..98ed0d4e3b75 100644 --- a/tensorrt_llm/_torch/custom_ops/dspark_rmsnorm_rope_custom_op.py +++ b/tensorrt_llm/_torch/custom_ops/dspark_rmsnorm_rope_custom_op.py @@ -41,6 +41,7 @@ def is_fused_dspark_rmsnorm_rope_supported( freqs: torch.Tensor, num_heads: int, rope_dim: int, + norm_dim: int | None = None, ) -> bool: """Return whether tensors satisfy the production fused-op contract.""" if _get_dspark_arch_str() is None or not all(t.is_cuda for t in (x, weight, freqs)): @@ -51,10 +52,15 @@ def is_fused_dspark_rmsnorm_rope_supported( return False if x.ndim < 2 or x.shape[-1] % 32 != 0: return False - if weight.shape != (x.shape[-1],): - return False if rope_dim < 0 or rope_dim > x.shape[-1] or rope_dim % 2 != 0: return False + # norm_dim defaults to the whole row; the only other supported value is the + # nope prefix, where the weight spans just that prefix. + effective_norm_dim = x.shape[-1] if norm_dim is None else norm_dim + if effective_norm_dim not in (x.shape[-1], x.shape[-1] - rope_dim): + return False + if effective_norm_dim % 32 != 0 or weight.shape != (effective_norm_dim,): + return False if (x.shape[-1] - rope_dim) % 32 != 0 or (rope_dim // 2) % 32 != 0: return False rows = x.numel() // x.shape[-1] @@ -144,6 +150,7 @@ def _compile_fused_dspark_rmsnorm_rope( apply_weight: bool, apply_rmsnorm: bool, inverse_rope: bool, + norm_dim: int, ): rows = cute.sym_int() freq_rows = cute.sym_int() @@ -151,7 +158,7 @@ def _compile_fused_dspark_rmsnorm_rope( cutlass.BFloat16, (rows, hidden_dim), stride_order=(1, 0) ) weight_fake = cute.runtime.make_fake_compact_tensor( - cutlass.BFloat16, (hidden_dim,), stride_order=(0,) + cutlass.BFloat16, (norm_dim,), stride_order=(0,) ) freqs_fake = cute.runtime.make_fake_compact_tensor( cutlass.Float32, @@ -170,6 +177,7 @@ def _compile_fused_dspark_rmsnorm_rope( apply_weight, apply_rmsnorm, inverse_rope, + norm_dim=norm_dim, ) return cute.compile( kernel, @@ -296,9 +304,10 @@ def cute_dsl_dspark_rmsnorm_rope( apply_weight: bool, apply_rmsnorm: bool, inverse_rope: bool, + norm_dim: int | None = None, ) -> torch.Tensor: """Apply fused RMSNorm and adjacent-pair RoPE to contiguous BF16 rows.""" - if not is_fused_dspark_rmsnorm_rope_supported(x, weight, freqs, num_heads, rope_dim): + if not is_fused_dspark_rmsnorm_rope_supported(x, weight, freqs, num_heads, rope_dim, norm_dim): raise ValueError( "cute_dsl_dspark_rmsnorm_rope requires contiguous BF16 tensors on " "an SM100 or SM103 GPU with a valid FP32 frequency view; " @@ -316,6 +325,7 @@ def cute_dsl_dspark_rmsnorm_rope( apply_weight, apply_rmsnorm, inverse_rope, + x.shape[-1] if norm_dim is None else norm_dim, ) compiled(x_flat, weight, freqs, output) return output.view(original_shape) @@ -332,6 +342,7 @@ def _( apply_weight: bool, apply_rmsnorm: bool, inverse_rope: bool, + norm_dim: int | None = None, ) -> torch.Tensor: return torch.empty_like(x) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dspark_rmsnorm_rope.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dspark_rmsnorm_rope.py index 71e3171543af..f14c51aa3650 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dspark_rmsnorm_rope.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dspark_rmsnorm_rope.py @@ -26,6 +26,7 @@ def __init__( apply_weight: bool, apply_rmsnorm: bool, inverse_rope: bool, + norm_dim: int | None = None, ): if hidden_dim % self.num_threads != 0: raise ValueError( @@ -44,6 +45,22 @@ def __init__( self.apply_weight = apply_weight self.apply_rmsnorm = apply_rmsnorm self.inverse_rope = inverse_rope + # Which prefix the RMS is taken over, and how far the norm scale and the + # weight reach. DSpark normalizes the whole row (the default, bit-identical + # to before this knob existed); DeepSeek-style MLA normalizes only the + # kv_lora_rank latent and leaves k_pe raw, which is nope_dim here. + self.norm_dim = hidden_dim if norm_dim is None else norm_dim + if self.norm_dim not in (hidden_dim, self.nope_dim): + raise ValueError( + f"norm_dim must be hidden_dim ({hidden_dim}) or nope_dim " + f"({self.nope_dim}); got {self.norm_dim}" + ) + self.norm_covers_rope = self.norm_dim == hidden_dim + if self.norm_dim % self.num_threads != 0: + raise ValueError( + f"norm_dim must be divisible by {self.num_threads}; got {self.norm_dim}" + ) + self.norm_elements_per_thread = self.norm_dim // self.num_threads if self.nope_dim % self.num_threads != 0: raise ValueError( f"nope_dim must be divisible by {self.num_threads}; got {self.nope_dim}" @@ -85,12 +102,12 @@ def kernel( inverse_rms = cutlass.Float32(1.0) if cutlass.const_expr(self.apply_rmsnorm): sum_sq = cutlass.Float32(0.0) - for item in cutlass.range_constexpr(self.elements_per_thread): + for item in cutlass.range_constexpr(self.norm_elements_per_thread): dim = tidx + item * self.num_threads value = cutlass.Float32(x[row, dim]) sum_sq += value * value sum_sq = cute.arch.warp_reduction_sum(sum_sq) - inverse_rms = cute.math.rsqrt(sum_sq / self.hidden_dim + self.eps) + inverse_rms = cute.math.rsqrt(sum_sq / self.norm_dim + self.eps) for item in cutlass.range_constexpr(self.nope_elements_per_thread): dim = tidx + item * self.num_threads @@ -105,11 +122,16 @@ def kernel( pair = tidx + item * self.num_threads real_dim = self.nope_dim + pair * 2 imag_dim = real_dim + 1 - real = cutlass.Float32(x[row, real_dim]) * inverse_rms - imag = cutlass.Float32(x[row, imag_dim]) * inverse_rms - if cutlass.const_expr(self.apply_weight): - real *= cutlass.Float32(weight[real_dim]) - imag *= cutlass.Float32(weight[imag_dim]) + real = cutlass.Float32(x[row, real_dim]) + imag = cutlass.Float32(x[row, imag_dim]) + # Outside norm_dim the rope lanes are passed through raw: no RMS + # scale, no weight. weight is only norm_dim long in that case. + if cutlass.const_expr(self.norm_covers_rope): + real *= inverse_rms + imag *= inverse_rms + if cutlass.const_expr(self.apply_weight): + real *= cutlass.Float32(weight[real_dim]) + imag *= cutlass.Float32(weight[imag_dim]) cos = cutlass.Float32(freqs[freq_row, pair, 0]) sin = cutlass.Float32(freqs[freq_row, pair, 1]) if cutlass.const_expr(self.inverse_rope): diff --git a/tensorrt_llm/_torch/models/modeling_dflash.py b/tensorrt_llm/_torch/models/modeling_dflash.py index 45a31bccce6a..3f68ab592ecb 100644 --- a/tensorrt_llm/_torch/models/modeling_dflash.py +++ b/tensorrt_llm/_torch/models/modeling_dflash.py @@ -26,7 +26,7 @@ get_dflash_trtllm_gen_ops, ) from ..speculative.interface import SpeculativeDecodingMode -from .modeling_utils import get_model_architecture, register_draft_model +from .modeling_utils import FUSED_MODULE_COMPONENTS, get_model_architecture, register_draft_model def dspark_layer_window_size( @@ -323,6 +323,12 @@ class DFlashForCausalLM(nn.Module): Reference: https://arxiv.org/pdf/2602.06036 """ + # Whether ``dflash_attention_backend`` drives this drafter's block decode. + # Subclasses that bring their own attention set this False: neither backend + # can express every drafter shape, the ops behind them are optional + # dependencies, and the worker's per-backend shape checks do not apply. + _uses_worker_attention_backend = True + def __init__(self, draft_config, *, dflash_attention_backend: str = "VANILLA"): """Build the draft model, resolving its architecture from the draft config (falling back to a model_type-derived name when the checkpoint uses a @@ -374,23 +380,32 @@ def __init__(self, draft_config, *, dflash_attention_backend: str = "VANILLA"): "block_size", getattr(pretrained_config, "block_size", None) ) self.dflash_attention_backend = dflash_attention_backend + if self.dflash_attention_backend not in ("VANILLA", "TRTLLM", "FA4"): + raise ValueError( + "DFlash attention backend must be VANILLA, TRTLLM or FA4, got " + f"{self.dflash_attention_backend!r}." + ) # Each backend loads only its own ops; the rest stay None so the # shared paged prologue can read them unconditionally. + self._dflash_flash_attention = None self._dflash_trtllm_gen_ops = None self._dflash_fa4_fwd = None self._dflash_paged_append = None - if self.dflash_attention_backend == "VANILLA": + if not self._uses_worker_attention_backend: + # Still validated above so a typo fails here rather than silently, + # but no op set is loaded: this drafter calls none of them. + logger.info_once( + f"{type(self).__name__} brings its own block decode; " + f"attention_backend={self.dflash_attention_backend!r} is not used.", + key=f"dflash_own_attention_{type(self).__name__}", + ) + elif self.dflash_attention_backend == "VANILLA": self._dflash_flash_attention = get_dflash_flash_attention() elif self.dflash_attention_backend == "TRTLLM": self._dflash_trtllm_gen_ops = get_dflash_trtllm_gen_ops() - elif self.dflash_attention_backend == "FA4": + else: self._dflash_fa4_fwd = get_dflash_fa4_fwd() self._dflash_paged_append = get_dflash_paged_append() - else: - raise ValueError( - "DFlash attention backend must be VANILLA, TRTLLM or FA4, got " - f"{self.dflash_attention_backend!r}." - ) self._dflash_trtllm_gen_workspace = None self._dflash_trtllm_gen_counters = None self.register_buffer("_dflash_batch_indices", None, persistent=False) @@ -727,6 +742,21 @@ def load_weights(self, weights: Dict, weight_mapper=None, **kwargs): else: remapped[key] = value + # Wrapper-owned and built FROM the checkpoint, so they are never in + # draft_model_full's module tree and _assert_backbone_complete cannot + # see them: a module conjured from data cannot be reported missing by + # walking modules. One tuple drives both the extraction below and this + # check, so the two cannot drift. Without fc the drafter has no capture + # projection, has_target_features stays False, _ctx_len never advances + # and it drafts from an empty context forever (the `hasattr` guards on + # that path are degradation, not a supported mode). + wrapper_missing = [k for k in self.WRAPPER_OWNED_WEIGHTS if k not in remapped] + if wrapper_missing: + raise ValueError( + f"{type(self).__name__}: checkpoint is missing {wrapper_missing}, " + "which this wrapper owns and builds from the checkpoint." + ) + # Load DFlash-specific weights directly if "fc.weight" in remapped: self.fc = nn.Linear( @@ -751,9 +781,10 @@ def load_weights(self, weights: Dict, weight_mapper=None, **kwargs): self.hidden_norm.weight.data.copy_(remapped["hidden_norm.weight"]) del remapped["hidden_norm.weight"] - # Load remaining weights into the draft model. - # DFlash checkpoints don't include embed_tokens or lm_head, so allow partial loading - # since those modules won't find matching weights. + # allow_partial_loading is what lets the shared modules below be absent. + # It is also why a truncated checkpoint loads clean, so gate it on the + # subclass's own declaration of what may legitimately be missing. + self._assert_backbone_complete(remapped, weight_mapper) self.draft_model_full.load_weights( weights=remapped, weight_mapper=weight_mapper, allow_partial_loading=True ) @@ -805,6 +836,78 @@ def take(key: str) -> torch.Tensor: self.candidate_selector = selector return {k: v for k, v in weights.items() if k not in consumed} + #: Tensors the wrapper itself owns: not in draft_model_full, built from the + #: checkpoint in load_weights, and required. + WRAPPER_OWNED_WEIGHTS = ("fc.weight", "hidden_norm.weight") + + #: Parameter-name prefixes this drafter takes from the target instead of its + #: own checkpoint. Everything else in ``draft_model_full`` must be provided. + #: GQA DFlash checkpoints ship neither embedding nor head; an MLA drafter + #: ships its own embedding and overrides this. + WEIGHTS_SHARED_WITH_TARGET = ("embed_tokens", "lm_head") + + def _assert_backbone_complete(self, weights: Dict, weight_mapper=None) -> None: + """Fail on a checkpoint missing weights the drafter does not share. + + The truth is the constructed module tree, not a hand-kept list that + rots as the backbone changes. + + Checked at module granularity. For a PLAIN module that is the hole the + flag cannot close: the loader skips one whose subtree filters to + nothing (modeling_utils.py `if module_weights:`) whatever the flag + says. For a FUSED module `allow_partial_loading=False` would catch it + (linear.py asserts all three shards) -- but the flag has to stay True + for the target-shared modules, so the check covers that case here + instead, and requires every component rather than any. + + Missing parameters INSIDE a present component stay tolerated: a + checkpoint with all three weights but only `q_proj.bias` takes the same + per-shard copy and leaves the rest at `torch.empty`. Module-granular + checking cannot see that and does not pretend to. + """ + provided = set(weights) + + # Whichever fusion table the load below will actually use, not a third + # copy: with a mapper modeling_utils dispatches to _load_weights_impl_v2 + # and the mapper's own table applies; without one it falls back to + # _load_weights_impl, whose table is FUSED_MODULE_COMPONENTS. An empty + # `mapping` means init_model_and_config has not run, so the mapper has + # no table to offer yet and the constant is still the right answer. + fusion = dict(getattr(weight_mapper, "mapping", None) or FUSED_MODULE_COMPONENTS) + + def _has(prefix: str) -> bool: + return any(k == prefix or k.startswith(prefix + ".") for k in provided) + + def _supplied(module_name: str) -> bool: + # A fused module is named once here and stored unfused in the + # checkpoint. ALL components must be present, not any: the fused + # load path is happy with a subset under allow_partial_loading + # (linear.py load_weights_fused_qkv_helper) and leaves the absent + # shards at torch.empty -- uninitialised device memory, not zeros. + if _has(module_name): + return True + for fused, parts in fusion.items(): + if fused in module_name: + return all(_has(module_name.replace(fused, p)) for p in parts) + return False + + missing = sorted( + { + name.rsplit(".", 1)[0] + for name, _ in self.draft_model_full.named_parameters() + if not any(part in name for part in self.WEIGHTS_SHARED_WITH_TARGET) + and not _supplied(name.rsplit(".", 1)[0]) + } + ) + if missing: + raise ValueError( + f"{type(self).__name__}: checkpoint provides no weights for " + f"{missing[:8]}{' ...' if len(missing) > 8 else ''}. These are " + f"not in WEIGHTS_SHARED_WITH_TARGET " + f"({', '.join(self.WEIGHTS_SHARED_WITH_TARGET)}), so loading " + "would leave them randomly initialized." + ) + def load_weights_from_target_model(self, target_model: torch.nn.Module) -> None: """Share embed_tokens and lm_head from the target model.""" self.draft_model_full.model.embed_tokens = target_model.model.embed_tokens diff --git a/tensorrt_llm/_torch/models/modeling_speculative.py b/tensorrt_llm/_torch/models/modeling_speculative.py index 5af73e0add7c..54b0d987dc8c 100755 --- a/tensorrt_llm/_torch/models/modeling_speculative.py +++ b/tensorrt_llm/_torch/models/modeling_speculative.py @@ -1405,6 +1405,17 @@ def external_drafter_config_kwargs(model_config, spec_config) -> dict: spec_config=None, # Avoid recursive spec-dec max_num_tokens=model_config.max_num_tokens, moe_max_num_tokens=model_config.moe_max_num_tokens, + # Bounds the drafter's position tables. Without it the field stays + # None and they fall back to the checkpoint's advertised + # max_position_embeddings -- 1,048,576 for K3, a ~256 MiB complex64 + # table per rank for a context the runtime bounds far below that. + # + # The user's value, NOT the engine's. py_executor_creator raises + # model_engine_max_seq_len past this and never writes it back, so a + # drafter that indexes absolute positions must read the raised value at + # runtime rather than have this line predict it -- reproducing that + # arithmetic here is what let the two drift apart in the first place. + max_seq_len=model_config.max_seq_len, ) # Only the embedded DSpark draft shares the target's EPLB namespace (its # stages are target decoder blocks registered into the target's balancer). diff --git a/tensorrt_llm/_torch/models/modeling_utils.py b/tensorrt_llm/_torch/models/modeling_utils.py index 0c16cb0972f3..9f19293deae8 100755 --- a/tensorrt_llm/_torch/models/modeling_utils.py +++ b/tensorrt_llm/_torch/models/modeling_utils.py @@ -1512,6 +1512,17 @@ def run_concurrently(func, raise +#: Modules stored unfused in a checkpoint and fused in the module tree, mapped +#: to the components they are built from. This is the table the mapper-less +#: load path below uses; HfWeightMapper.map_weights installs the same pairs for +#: the mapper path. Anything reasoning about what a checkpoint must provide has +#: to read one of the two rather than keep a third copy. +FUSED_MODULE_COMPONENTS = { + 'qkv_proj': ['q_proj', 'k_proj', 'v_proj'], + 'gate_up_proj': ['gate_proj', 'up_proj'], +} + + def _load_weights_impl(model: Union[nn.Module, DecoderModelForCausalLM], weights: Dict, skip_modules: List[str] = [], @@ -1538,10 +1549,7 @@ def _load_weights_impl(model: Union[nn.Module, DecoderModelForCausalLM], model.config, 'num_key_value_heads' ) and model.config.num_key_value_heads is not None else model.config.num_attention_heads - params_map = { - 'qkv_proj': ['q_proj', 'k_proj', 'v_proj'], - 'gate_up_proj': ['gate_proj', 'up_proj'] - } + params_map = dict(FUSED_MODULE_COMPONENTS) device_id = local_mpi_rank() def load_single_module(name, module): diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index 939e7366d327..95f0025499ef 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -1422,11 +1422,13 @@ def prepare_1st_drafter_inputs( j_block = torch.arange(query_tokens_per_req, dtype=torch.long, device="cuda") offsets_kp1 = torch.arange(K_plus_1, dtype=torch.long, device="cuda") - query_position_ids = ( - ctx_len_gen.unsqueeze(1) - + gen_num_accepted.long().unsqueeze(1) - + j_block.unsqueeze(0) - ) + # _ctx_len is clamped to _max_ctx only AFTER this step's accepted + # tokens are folded in (see the update below), so the running length + # used here has to be clamped on its own -- otherwise a request that + # already sits at the ceiling indexes num_accepted positions past + # any position the sequence can legitimately reach. + ctx_len_now = (ctx_len_gen + gen_num_accepted.long()).clamp_(max=self._max_ctx) + query_position_ids = ctx_len_now.unsqueeze(1) + j_block.unsqueeze(0) ctx_position_ids = ctx_len_gen.unsqueeze(1) + offsets_kp1.unsqueeze(0) # Go through embed_tokens.forward (NOT .weight[...]) so TP-sharded diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py index 5ccf8e687c74..1624aa628bd2 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py @@ -593,6 +593,11 @@ def clear_stale_blocks_after_page_unlink( # check if consecutive available blocks is sufficient for window_size. (TRTLLM-8802) # But for simplicity, we leave it for now. curr = start + # Only detach blocks with no live page in ANY life cycle: a childless tip + # that lost this life cycle's page may still hold live pages of other life + # cycles; detaching it would orphan the committed chain of an in-flight + # sequence. Mirrors Block::clearStaleBlocksAfterPageUnlink in + # cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp. while ( ( isinstance(curr, Block) diff --git a/tests/integration/defs/accuracy/test_disaggregated_multinode.py b/tests/integration/defs/accuracy/test_disaggregated_multinode.py new file mode 100644 index 000000000000..873c23f1bd39 --- /dev/null +++ b/tests/integration/defs/accuracy/test_disaggregated_multinode.py @@ -0,0 +1,395 @@ +# 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 contextlib +import os +import re +import time +from concurrent.futures import Future +from pathlib import Path +from typing import Iterator, Optional + +import openai +import pytest +import requests + +from tensorrt_llm.llmapi import CompletionOutput, RequestOutput, SamplingParams +from tensorrt_llm.llmapi.llm_args import DSparkDecodingConfig, LlmArgs +from tensorrt_llm.llmapi.tokenizer import DeepseekV4Tokenizer +from tests.unittest.llmapi.apps.openai_server import RemoteDisaggOpenAIServer, RemoteOpenAIServer + +from ..conftest import llm_models_root, skip_pre_blackwell +from .accuracy_core import GSM8K, LlmapiAccuracyTestHarness +from .test_disaggregated_serving import DuckLLM, MyThreadPoolExecutor, Result, run_accuracy_test + +# Run from the repository root without installing the repository into the image: +# +# TRTLLM_TEST_ROOT=/lustre/fsw/coreai_comparch_trtllm/lizhiz/tllm_dsv4 +# TRTLLM_IMAGE_ROOT=/lustre/share/coreai_comparch_trtllm/lizhiz/rubin_dsv4 +# TRTLLM_TEST_IMAGE="${TRTLLM_IMAGE_ROOT}/b8bc42c3ae/trtllm.sqsh" +# TRTLLM_TEST_LOG="${TRTLLM_TEST_ROOT}/build_images/b8bc42c3ae/test_logs" +# mkdir -p "${TRTLLM_TEST_LOG}" +# srun \ +# --job-name=coreai_comparch_aarwlt-dsv4.disagg-acc \ +# --partition=batch-xdr \ +# --account=coreai_comparch_aarwlt \ +# --nodes=2 \ +# --ntasks=2 \ +# --ntasks-per-node=1 \ +# --time=04:00:00 \ +# --output="${TRTLLM_TEST_LOG}/dsv4_disagg_%j.log" \ +# --error="${TRTLLM_TEST_LOG}/dsv4_disagg_%j.log" \ +# --container-image="${TRTLLM_TEST_IMAGE}" \ +# --container-mounts=/lustre:/lustre \ +# --container-workdir="${TRTLLM_TEST_ROOT}" \ +# bash -lc '\ +# export LLM_MODELS_ROOT=/lustre/fsw/coreai_comparch_trtllm/common/llm-models; \ +# export HF_HOME=/lustre/fsw/coreai_comparch_trtllm/lizhiz/hf_cache; \ +# export TLLM_LOG_LEVEL=INFO; export PYTHONPATH=; \ +# unset TRTLLM_MOE_A2A_DISABLE_CFT_COUNTED_WRITES; \ +# TEST_FILE=tests/integration/defs/accuracy/test_disaggregated_multinode.py; \ +# python3 -m pytest -q -s \ +# "${TEST_FILE}::TestDeepSeekV4ProDSparkMultinode::test_gsm8k_1p1d_dep4"' +# +# Slurm node rank 0 hosts the context worker and disaggregated frontend. Node +# rank 1 hosts the generation worker. The test waits for both workers before +# starting the frontend because DSpark makes generation startup substantially +# slower than context startup. + + +def _expand_slurm_nodelist(nodelist: str) -> list[str]: + if not nodelist: + return [] + + groups = [] + group_chars = [] + bracket_depth = 0 + for char in nodelist: + if char == "[": + bracket_depth += 1 + elif char == "]": + bracket_depth -= 1 + + if char == "," and bracket_depth == 0: + groups.append("".join(group_chars)) + group_chars = [] + else: + group_chars.append(char) + groups.append("".join(group_chars)) + + nodes = [] + for group in groups: + match = re.fullmatch(r"(.+?)\[(.+)]", group) + if match is None: + nodes.append(group) + continue + + prefix, suffixes = match.groups() + for suffix in suffixes.split(","): + range_match = re.fullmatch(r"(\d+)-(\d+)", suffix) + if range_match is None: + nodes.append(f"{prefix}{suffix}") + continue + + start_text, end_text = range_match.groups() + width = len(start_text) + nodes.extend( + f"{prefix}{value:0{width}d}" for value in range(int(start_text), int(end_text) + 1) + ) + return nodes + + +def _wait_for_endpoint_ready( + url: str, + timeout: int, + interval: int = 3, +) -> None: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + try: + response = requests.get(url, timeout=10) + if response.status_code == 200: + return + except requests.RequestException: + pass + time.sleep(interval) + raise TimeoutError(f"Endpoint {url} was not ready within {timeout} seconds") + + +def _wait_for_endpoint_down( + url: str, + timeout: int, + interval: int = 1, +) -> None: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + try: + requests.get(url, timeout=10) + except requests.RequestException: + return + time.sleep(interval) + raise TimeoutError(f"Endpoint {url} remained up after {timeout} seconds") + + +NODE_RANK = int(os.environ.get("SLURM_NODEID", 0)) +NODE_LIST = _expand_slurm_nodelist(os.environ.get("SLURM_NODELIST", "")) +SLURM_NTASKS_PER_NODE = int(os.environ.get("SLURM_NTASKS_PER_NODE", 1)) + +CTX_SERVER_PORT = 8001 +GEN_SERVER_PORT = 8002 +DISAGG_SERVER_PORT = 8000 +SERVER_START_TIMEOUT = 7200 + +MODEL_NAME = "deepseek-ai/DeepSeek-V4-Pro" + + +MODEL_PATH = str(Path(llm_models_root()) / "DeepSeek-V4-Pro-DSpark") + +EXTRA_EVALUATOR_KWARGS = { + "apply_chat_template": True, + "system_prompt": ( + "Solve the problem carefully. End your response with a final line " + "exactly in the form #### , using the simplest numeric form " + "without units or trailing zeros." + ), +} + + +def _require_two_node_allocation() -> None: + if len(NODE_LIST) != 2: + pytest.skip("This test requires exactly two Slurm nodes") + if SLURM_NTASKS_PER_NODE != 1: + pytest.skip("This test requires one pytest task per Slurm node") + + +def _is_context_node() -> bool: + return NODE_RANK == 0 + + +def _is_generation_node() -> bool: + return NODE_RANK == 1 + + +def _worker_env() -> dict[str, str]: + env = { + key: value + for key, value in os.environ.items() + if not key.startswith(("OMPI_", "PMIX_", "PMI_", "SLURM_")) + and key + not in { + "MASTER_ADDR", + "MASTER_PORT", + "UCX_TLS", + "UCX_NET_DEVICES", + } + } + # Pyxis starts these worker processes as root; preserve the explicit + # Open MPI container override for mpi4py dynamic worker spawn. + env.update( + { + "OMPI_ALLOW_RUN_AS_ROOT": "1", + "OMPI_ALLOW_RUN_AS_ROOT_CONFIRM": "1", + "TLLM_INDEXER_MQA_LOGITS_ELEM_BUDGET": "1073741824", + # Avoid auto-selecting the NVL72 RDMA VF, whose container-visible + # IPv6 address cannot be bound by UCX on this cluster. + "UCX_NET_DEVICES": os.environ.get("UCX_NET_DEVICES", "eth0"), + "UCX_TLS": "tcp,self,sm,cuda_copy,cuda_ipc", + } + ) + return env + + +def _worker_config(enable_dspark: bool) -> dict[str, object]: + config: dict[str, object] = { + "attn_backend": "TRTLLM", + "tensor_parallel_size": 4, + "moe_expert_parallel_size": 4, + "enable_attention_dp": True, + "moe_config": { + # DeepSeek-V4 routed experts use NVFP4 in this test. Use the + # MegaMoE CuTe DSL path for the NVFP4 checkpoint. + "backend": "MEGAMOE_CUTEDSL", + }, + # Keep the DEP4 worker envelope aligned with the aggregate guard. A + # batch-128 DSpark graph leaves too little headroom for the post-capture + # 4096-token warmup on rank 0. + "max_batch_size": 64, + "cuda_graph_config": {"max_batch_size": 64}, + "max_seq_len": 4096, + "max_num_tokens": 4096, + "kv_cache_config": { + "enable_block_reuse": False, + "free_gpu_memory_fraction": 0.5, + }, + "enable_chunked_prefill": False, + "disable_overlap_scheduler": True, + "enable_iter_perf_stats": True, + "print_iter_log": True, + "custom_tokenizer": "deepseek_v4", + "cache_transceiver_config": { + "backend": "NIXL", + "transceiver_runtime": "PYTHON", + "max_tokens_in_buffer": 4096, + }, + } + if enable_dspark: + config["speculative_config"] = { + "decoding_type": "DSpark", + "max_draft_len": 5, + "speculative_model": MODEL_PATH, + } + return config + + +@pytest.fixture(scope="module") +def worker() -> Iterator[Optional[RemoteOpenAIServer]]: + _require_two_node_allocation() + if _is_context_node(): + port = CTX_SERVER_PORT + elif _is_generation_node(): + port = GEN_SERVER_PORT + else: + yield None + return + + with RemoteOpenAIServer( + MODEL_PATH, + port=port, + cli_args=["--tp_size", "4", "--pp_size", "1"], + host="0.0.0.0", + env=_worker_env(), + llmapi_launch=False, + rank=0, + extra_config=_worker_config(enable_dspark=_is_generation_node()), + ) as server: + yield server + + +@pytest.fixture(scope="module") +def disagg_server( + worker: Optional[RemoteOpenAIServer], +) -> Iterator[Optional[RemoteDisaggOpenAIServer]]: + del worker + if _is_context_node(): + _wait_for_endpoint_ready( + f"http://{NODE_LIST[1]}:{GEN_SERVER_PORT}/health", + timeout=SERVER_START_TIMEOUT, + ) + with RemoteDisaggOpenAIServer( + ctx_servers=[f"{NODE_LIST[0]}:{CTX_SERVER_PORT}"], + gen_servers=[f"{NODE_LIST[1]}:{GEN_SERVER_PORT}"], + port=DISAGG_SERVER_PORT, + llmapi_launch=False, + env=_worker_env(), + ) as server: + yield server + else: + yield None + + +@contextlib.contextmanager +def _accuracy_llm( + server: RemoteDisaggOpenAIServer, +) -> Iterator[DuckLLM]: + client = openai.OpenAI( + api_key=RemoteOpenAIServer.DUMMY_API_KEY, + base_url=server.url_for("v1"), + timeout=1_800_000, + ) + args = LlmArgs(model=MODEL_PATH) + args.quant_config.quant_algo = "FP8_BLOCK_SCALES" + args.speculative_config = DSparkDecodingConfig( + max_draft_len=5, + speculative_model=MODEL_PATH, + ) + tokenizer = DeepseekV4Tokenizer.from_pretrained(MODEL_PATH) + + with MyThreadPoolExecutor(max_workers=128) as thread_pool: + + def send_request( + prompt: str, + sampling_params: Optional[SamplingParams], + streaming: bool, + ) -> RequestOutput: + if sampling_params is None: + sampling_params = SamplingParams() + response = client.completions.create( + model=MODEL_PATH, + prompt=prompt, + stream=streaming, + max_tokens=sampling_params.max_tokens, + n=sampling_params.n, + temperature=( + sampling_params.temperature if sampling_params.top_p is not None else 0 + ), + top_p=sampling_params.top_p, + stop=sampling_params.stop, + seed=sampling_params.seed, + ) + result = Result( + id=0, + sampling_params=sampling_params, + outputs=[ + CompletionOutput(text=choice.text, index=index) + for index, choice in enumerate(response.choices) + ], + ) + output = RequestOutput._from_generation_result(result, prompt=prompt) + setattr(output, "result", result.result) + return output + + def generate_async( + prompt: str, + sampling_params: Optional[SamplingParams] = None, + streaming: bool = False, + ) -> Future[RequestOutput]: + future = thread_pool.submit(send_request, prompt, sampling_params, streaming) + thread_pool.futures.append(future) + return future + + yield DuckLLM(args, tokenizer, generate_async) + + +@pytest.mark.timeout(14400) +@pytest.mark.skip_less_device_memory(140000) +@skip_pre_blackwell +class TestDeepSeekV4ProDSparkMultinode(LlmapiAccuracyTestHarness): + def test_gsm8k_1p1d_dep4( + self, + disagg_server: Optional[RemoteDisaggOpenAIServer], + monkeypatch: pytest.MonkeyPatch, + ) -> None: + health_url = f"http://{NODE_LIST[0]}:{DISAGG_SERVER_PORT}/health/" + if _is_context_node(): + assert disagg_server is not None + monkeypatch.setenv("INTEGRATION_TEST", "0") + with _accuracy_llm(disagg_server) as llm: + run_accuracy_test( + llm, + MODEL_NAME, + ["GSM8K"], + extra_evaluator_kwargs={GSM8K: EXTRA_EVALUATOR_KWARGS}, + ) + disagg_server.terminate() + elif _is_generation_node(): + _wait_for_endpoint_ready( + health_url, + timeout=SERVER_START_TIMEOUT, + ) + _wait_for_endpoint_down( + health_url, + timeout=14400, + ) + else: + raise AssertionError(f"Unexpected Slurm node rank {NODE_RANK}") diff --git a/tests/unittest/_torch/disaggregation/test_disagg_index_mapper_early_release.py b/tests/unittest/_torch/disaggregation/test_disagg_index_mapper_early_release.py index 7807ef566f48..0ba43fa2ab81 100644 --- a/tests/unittest/_torch/disaggregation/test_disagg_index_mapper_early_release.py +++ b/tests/unittest/_torch/disaggregation/test_disagg_index_mapper_early_release.py @@ -19,7 +19,7 @@ iterations of requests accumulate in-flight. """ -from unittest.mock import MagicMock +from unittest.mock import MagicMock, call import pytest import torch @@ -258,6 +258,46 @@ def test_gather_k_block_offsets_matches_beam_zero_index_select(self): assert torch.count_nonzero(destination[:, :, 1] != -1) == 0 assert torch.count_nonzero(destination[:, 4, 0] != -1) == 0 + def test_release_detaches_page_index_views_before_slot_reuse(self): + from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 + from tensorrt_llm.bindings.internal.batch_manager.kv_cache_manager_v2_utils import ( + IndexMapper, + ) + + manager = object.__new__(KVCacheManagerV2) + manager.max_beam_width = 2 + manager.num_pools = 3 + index_mapper = IndexMapper(max_batch_size=1, max_beam_width=2) + index_mapper.add_new_sequence(1) + events = [] + manager.index_mapper = MagicMock(wraps=index_mapper) + manager.index_mapper.remove_sequence.side_effect = lambda request_id: ( + events.append(("remove", request_id)), + index_mapper.remove_sequence(request_id), + )[1] + manager._early_freed_index_requests = set() + kv_cache = MagicMock() + kv_cache.set_base_page_index_buf.side_effect = ( + lambda beam_idx, pool_idx, value: events.append(("detach", beam_idx, pool_idx, value)) + ) + manager.kv_cache_map = {1: kv_cache} + + manager.release_index_slot(1) + + expected_calls = [ + call(beam_idx, pool_idx, None) + for beam_idx in range(manager.max_beam_width) + for pool_idx in range(manager.num_pools) + ] + assert kv_cache.set_base_page_index_buf.call_args_list == expected_calls + assert events == [ + ("detach", beam_idx, pool_idx, None) + for beam_idx in range(manager.max_beam_width) + for pool_idx in range(manager.num_pools) + ] + [("remove", 1)] + assert not _has_sequence(index_mapper, 1) + assert manager._early_freed_index_requests == {1} + class TestFreeResourcesDoubleReleaseSafety: """Test that free_resources handles already-released IndexMapper slots.""" diff --git a/tests/unittest/disaggregated/test_cache_reuse_adapter.py b/tests/unittest/disaggregated/test_cache_reuse_adapter.py index c661cc4314bb..bdcbc4d212f1 100644 --- a/tests/unittest/disaggregated/test_cache_reuse_adapter.py +++ b/tests/unittest/disaggregated/test_cache_reuse_adapter.py @@ -194,6 +194,7 @@ def _build_transceiver_for_kv_slice( sliding_window_size=None, cached_tokens: int = 0, is_generation_only: bool = False, + num_draft_tokens: int = 0, ): """Stub a KvCacheTransceiverV2 so _create_chunk runs without dist setup. @@ -207,7 +208,13 @@ def _build_transceiver_for_kv_slice( kv_head_num_per_rank=1, sliding_window_size=sliding_window_size, ) - total_blocks = (prompt_len + num_extra_kv_tokens + tokens_per_block - 1) // tokens_per_block + total_blocks = ( + prompt_len + + (num_draft_tokens if is_generation_only else 0) + + num_extra_kv_tokens + + tokens_per_block + - 1 + ) // tokens_per_block if block_ids is None: block_ids = np.arange(total_blocks, dtype=np.int64) else: @@ -353,6 +360,37 @@ def test_dspark_disagg_boundary_keeps_only_initialized_swa(self, prompt_len): block_ids[:-1], ) + def test_dspark_gen_init_draft_reserve_is_removed_before_swa_trim(self): + """Remove gen-init draft reserve before selecting the valid SWA suffix.""" + prompt_len = 1144 + tokens_per_block = 128 + num_extra_kv_tokens = 4 + num_draft_tokens = 5 + total_prompt_blocks = (prompt_len + tokens_per_block - 1) // tokens_per_block + sliding_window_size = 133 + stale_end = max( + 0, + (prompt_len + 1 - sliding_window_size) // tokens_per_block, + ) + valid_prompt_blocks = total_prompt_blocks - stale_end + block_ids = np.arange(200, 200 + valid_prompt_blocks + 1, dtype=np.int64) + transceiver, req = _build_transceiver_for_kv_slice( + num_extra_kv_tokens=num_extra_kv_tokens, + num_draft_tokens=num_draft_tokens, + prompt_len=prompt_len, + tokens_per_block=tokens_per_block, + block_ids=block_ids, + sliding_window_size=sliding_window_size, + is_generation_only=True, + ) + + kv_slice = transceiver._create_chunk(req) + + np.testing.assert_array_equal( + kv_slice.block_ids_per_layer_groups[0], + block_ids[:-1], + ) + # --------------------------------------------------------------------------- # CacheReuseAdapter.get_cached_token_count_per_layer_group: SWA clamp. diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_block_radix_tree_stale_prune.py b/tests/unittest/kv_cache_manager_v2_tests/test_block_radix_tree_stale_prune.py new file mode 100644 index 000000000000..6517e7e1b760 --- /dev/null +++ b/tests/unittest/kv_cache_manager_v2_tests/test_block_radix_tree_stale_prune.py @@ -0,0 +1,151 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. +"""Pure unit tests for the stale-block tail prune in ``clear_stale_blocks_after_page_unlink``. + +Regression guard for a multi-life-cycle tree: unlinking one life cycle's page +from a childless tip must not detach the block while another life cycle still +holds a live page there. Hybrid models (Kimi K3: MLA attention + KDA/SSM) are +the ones that hit this, because they are the ones with more than one life cycle. + +The C++ mirror was fixed upstream in PR #17323 +(``Block::clearStaleBlocksAfterPageUnlink`` in +``cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp``); +this covers the Python implementation, which is the one selected by +``TLLM_KV_CACHE_MANAGER_V2_BACKEND=python``. +""" + +import unittest +from collections.abc import Iterator +from importlib.util import find_spec +from typing import TYPE_CHECKING, cast + +if not TYPE_CHECKING and find_spec("kv_cache_manager_v2") is not None: + from kv_cache_manager_v2 import TokenId + from kv_cache_manager_v2._block_radix_tree import Block, BlockRadixTree, ReuseScope + from kv_cache_manager_v2._life_cycle_registry import ( + AttnLifeCycle, + LifeCycleId, + LifeCycleRegistry, + ) +else: + from tensorrt_llm.runtime.kv_cache_manager_v2 import TokenId + from tensorrt_llm.runtime.kv_cache_manager_v2._block_radix_tree import ( + Block, + BlockRadixTree, + ReuseScope, + ) + from tensorrt_llm.runtime.kv_cache_manager_v2._life_cycle_registry import ( + AttnLifeCycle, + LifeCycleId, + LifeCycleRegistry, + ) + + +class _TwoLifeCycles: + """Minimal ``LifeCycleRegistry`` stand-in: only ``size`` is reached here. + + ``Block.storage`` is sized from ``num_life_cycles``, which is all the prune + predicate needs. Two life cycles is the smallest tree that can express + "this block lost life cycle 0's page but still holds life cycle 1's". + """ + + size = LifeCycleId(2) + + @property + def ssm_life_cycle_id(self) -> None: + return None + + def attention_life_cycles(self) -> Iterator[tuple[object, object]]: + return iter(()) + + +class _DeadPageRef: + """Stand-in for ``rawref.ref[CommittedPage]``: occupied slot, dead referent. + + The predicate under test only compares the slot against ``None``. Returning + ``None`` from ``__call__`` keeps ``Block._release_pages`` on its + already-collected path, so teardown never dereferences a fake page. + """ + + def __call__(self) -> None: + return None + + +# Windowed, no sink blocks: keeps ``clear_stale_blocks_after_page_unlink`` off +# the ``remove_subtree`` branch (that branch fires for full attention or sink +# blocks and would drop the subtree regardless of the tail-prune predicate), +# so the test isolates the prune loop. +_WINDOWED_ATTN = AttnLifeCycle(window_size=64, num_sink_blocks=0) + +_LC_UNLINKED = LifeCycleId(0) +_LC_OTHER = LifeCycleId(1) + + +class TestStaleTailPrune(unittest.TestCase): + def _build_chain(self) -> "tuple[BlockRadixTree, object, Block, Block]": + """Root -> first -> tip, two life cycles, tokens_per_block=2.""" + tree = BlockRadixTree(cast(LifeCycleRegistry, _TwoLifeCycles()), tokens_per_block=2) + root = tree.add_or_get_existing(ReuseScope()) + first = Block([TokenId(1), TokenId(2)], root) + tip = Block([TokenId(3), TokenId(4)], first) + self.assertEqual(len(tip.storage), 2) + return tree, root, first, tip + + def test_tip_with_live_page_in_another_life_cycle_is_kept(self) -> None: + tree, root, first, tip = self._build_chain() + # Life cycle 0's page was just unlinked; life cycle 1 still holds one. + tip.storage[_LC_UNLINKED] = None + tip.storage[_LC_OTHER] = cast(object, _DeadPageRef()) + + Block.clear_stale_blocks_after_page_unlink(tip, _LC_UNLINKED, _WINDOWED_ATTN) + + # Detaching the tip here would orphan the committed chain of an + # in-flight sequence that is still using life cycle 1. + self.assertIn(tip.key, first.next) + self.assertIs(first.next[tip.key], tip) + self.assertIsNotNone(tip._prev()) + self.assertIn(first.key, root.next) + + def test_tip_with_no_live_page_anywhere_is_pruned(self) -> None: + # Negative control: the fix must not stop the prune it is supposed to + # allow, otherwise dead tails accumulate forever. + tree, root, first, tip = self._build_chain() + tip.storage[_LC_UNLINKED] = None + tip.storage[_LC_OTHER] = None + + Block.clear_stale_blocks_after_page_unlink(tip, _LC_UNLINKED, _WINDOWED_ATTN) + + # tip is detached, and the walk continues up through `first`, which is + # now itself a childless tip with no pages; the emptied root is then + # dropped from the tree. + self.assertNotIn(tip.key, first.next) + self.assertNotIn(first.key, root.next) + self.assertEqual(tree.next, {}) + + def test_tip_keeps_when_only_the_unlinked_life_cycle_is_empty(self) -> None: + # Same as the first case but with the roles of the two life cycles + # swapped, so the test cannot pass by hard-coding an index. + tree, root, first, tip = self._build_chain() + tip.storage[_LC_OTHER] = None + tip.storage[_LC_UNLINKED] = cast(object, _DeadPageRef()) + + Block.clear_stale_blocks_after_page_unlink(tip, _LC_OTHER, _WINDOWED_ATTN) + + self.assertIn(tip.key, first.next) + self.assertIs(first.next[tip.key], tip) + + +if __name__ == "__main__": + unittest.main() From 48fc2fab3db60e5657f88143c4c4cf5cde290695 Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Tue, 15 Sep 2026 07:16:41 +0000 Subject: [PATCH 2/3] [None][fix] keep the external drafter's KV pool at the drafter's own dtype An external drafter has its own KV cache, sized and allocated separately from the target's. It was inheriting the target's quantization config, so an fp8-KV target silently forced the drafter's pool to fp8 as well -- charging the budget at a dtype the drafter never asked for and does not necessarily support. Neutralize the target's KV quant config when building the draft cache manager for an external drafter, and fall the draft KV dtype back to "auto" when the inherited value is "fp8". Ported from the Rubin branch, where this landed as 33624323a6 and was then lost when a later replay commit (2ba428f26f) re-took pyexecutor/_util.py from main. Restored on the internal branch as 8626bd6577; content survival of the original measured 0%. NOT RUN: no GPU here. Needs an external-drafter plus fp8-KV-target speculative decoding run to validate. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/_util.py | 38 +++++++++++++++++++++++++ 1 file changed, 38 insertions(+) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index f6e2ce17eabc..ade58cfe0b0e 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -1774,6 +1774,29 @@ def _create_one_model_draft_kv_cache_manager( # Get the effective draft config (explicit draft_config if available, # otherwise fall back to target model config for MTP). effective_draft_config = self._get_effective_draft_config() + if (self._speculative_config.spec_dec_mode.is_external_drafter() + and effective_draft_config.quant_config is not None + and effective_draft_config.quant_config.quant_mode. + has_fp8_kv_cache()): + # The args-level kv_cache_config.dtype sync stamps the TARGET's + # fp8 KV algo onto every loaded model, including the external + # drafter. The drafter stores and reads its pool in its weights + # dtype (DFlash validates a bf16 pool and otherwise falls back to + # the max_seq_len-dense private arena, which OOMs at long + # context), so the pool dtype must follow the drafter, not the + # target. + logger.info( + "External drafter KV pool keeps the drafter dtype; dropping " + "the fp8 KV quant algo inherited from the target.") + neutral_quant = copy.copy(effective_draft_config.quant_config) + neutral_quant.kv_cache_quant_algo = None + # QuantConfig.quant_mode is a cached_property; the copy carries + # the already-computed cache, so drop it for the mutation to take. + neutral_quant.__dict__.pop("quant_mode", None) + effective_draft_config = copy.copy(effective_draft_config) + effective_draft_config._frozen = False + effective_draft_config.quant_config = neutral_quant + effective_draft_config._frozen = True kv_cache_config = (kv_cache_config_override if kv_cache_config_override is not None else self._kv_cache_config) @@ -1799,6 +1822,21 @@ def _create_one_model_draft_kv_cache_manager( f"from {draft_kv_config.pool_ratio} to [1.0] for its single " "layer group.") draft_kv_config.pool_ratio = [1.0] + draft_quant = effective_draft_config.quant_config + draft_has_fp8_kv = bool( + draft_quant is not None + and draft_quant.layer_quant_mode.has_fp8_kv_cache()) + if draft_kv_config.dtype == "fp8" and not draft_has_fp8_kv: + # kv_cache_config.dtype describes the TARGET pool. An external + # drafter without fp8-KV quantization stores and reads its pool in + # its own dtype (DFlash validates the pool as bf16 and otherwise + # falls back to the max_seq_len-dense private arena, which OOMs at + # long context). Resolve the draft pool from the drafter's quant + # config instead. + logger.info( + "Separate one-model draft KV cache keeps the drafter dtype " + "(kv_cache_config.dtype=fp8 applies to the target pool only).") + draft_kv_config.dtype = "auto" if uses_vswa_kv_cache_layout(draft_kv_config.max_attention_window): logger.info( f"Derived draft KV cache max_attention_window for separate " From b4934158a4463733ac911e85d7ab5d2d7f0e9fdd Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Fri, 18 Sep 2026 01:10:53 +0000 Subject: [PATCH 3/3] [None][fix] restore the V2 draft-KV max_tokens derivation dropped in the rebase On user/lizhiz/rubin-advance, configure_kv_cache_capacity re-derives max_tokens from the final estimated budget once estimation finishes whenever the V2 KV cache manager is in use. The rebase onto main dropped that block; main never had it either. Without it, a one-model speculative-decoding draft KV cache that shares the config reads the same max_gpu_total_bytes as the target. That cap does not scale with the draft's much smaller per-token footprint, which depends on num_local_layers, so both managers claim the whole budget and the draft OOMs. Deriving max_tokens restores V1's behaviour: V2's quota becomes min(max_gpu_total_bytes, max_tokens * bytes_per_token), which picks the layer-scaled draft budget when the draft manager reads the shared config. An explicit user-provided max_tokens still wins. Restored as it stands on rubin-advance, using this branch's existing self._is_kv_cache_manager_v2 attribute rather than recomputing the issubclass check locally. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/_util.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index ade58cfe0b0e..193aa58af9f3 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -1566,6 +1566,25 @@ def configure_kv_cache_capacity(self, f"max_gpu_total_bytes={self._max_gpu_total_bytes_in / (GB):.2f} GiB is provided. New max memory is {kv_cache_max_memory / (GB):.2f} GiB" ) + if self._is_kv_cache_manager_v2: + # KVCacheManagerV2 normally uses max_gpu_total_bytes alone, so we'd + # just restore the user-provided max_tokens here. However, when a + # one-model speculative-decoding draft KV cache shares the same + # config, max_gpu_total_bytes doesn't scale with the draft's much + # smaller per-token byte footprint (which depends on + # num_local_layers), so both managers read the same cap and the + # draft double-claims the budget -> OOM. Mirror V1's behaviour: + # derive max_tokens from the final estimated memory so V2's quota + # = min(max_gpu_total_bytes, max_tokens * bytes_per_token) + # naturally picks the layer-scaled draft budget when the draft + # manager reads the shared config. Respect an explicit + # user-provided max_tokens if present. + if self._max_kv_tokens_in is not None: + self._kv_cache_config.max_tokens = self._max_kv_tokens_in + else: + self._kv_cache_config.max_tokens = (self._get_kv_size_per_token( + ).tokens_for_budget(kv_cache_max_memory)) + logger.info( f"Estimated max memory in KV cache : {kv_cache_max_memory / (GB):.2f} GiB" )