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/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index f6e2ce17eabc..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" ) @@ -1774,6 +1793,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 +1841,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 " 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()