From 594bf14e8a9d4e724b19561b0a10763f46f90ac1 Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Mon, 14 Sep 2026 15:32:59 -0700 Subject: [PATCH 1/8] [https://nvbugs/6661846][fix] Initialize FP8 storage in MSA layout tests (#19078) Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- 3rdparty/patches/msa_strided_paged_kv.patch | 21 ++++++++----------- tests/integration/test_lists/waives.txt | 1 - .../attention/sparse/msa/test_msa_backend.py | 14 ++++++++----- 3 files changed, 18 insertions(+), 18 deletions(-) diff --git a/3rdparty/patches/msa_strided_paged_kv.patch b/3rdparty/patches/msa_strided_paged_kv.patch index 136be361a4ae..adee1b3dd0c5 100644 --- a/3rdparty/patches/msa_strided_paged_kv.patch +++ b/3rdparty/patches/msa_strided_paged_kv.patch @@ -362,7 +362,7 @@ diff --git a/python/fmha_sm100/cute/test_sparse_atten.py b/python/fmha_sm100/cut index 21c777e..b5f078b 100644 --- a/python/fmha_sm100/cute/test_sparse_atten.py +++ b/python/fmha_sm100/cute/test_sparse_atten.py -@@ -61,6 +61,81 @@ DECODE_DIM = 128 +@@ -61,6 +61,78 @@ DECODE_DIM = 128 DECODE_KV_TOKEN_SWEEP = tuple(2**exp for exp in range(3, 21)) @@ -387,14 +387,11 @@ index 21c777e..b5f078b 100644 + +def test_prepare_paged_hnd_input_materializes_unpacked_tokens() -> None: + pages, heads, page_size, head_dim = 5, 2, 128, 128 -+ storage = torch.empty( -+ pages, -+ heads, -+ page_size * 2, -+ head_dim, -+ dtype=torch.float8_e4m3fn, ++ storage = torch.arange( ++ pages * heads * page_size * 2 * head_dim, ++ dtype=torch.int32, + device="cuda", -+ ) ++ ).remainder(16).to(torch.float8_e4m3fn).reshape(pages, heads, page_size * 2, head_dim) + view = storage[:, :, ::2, :] + + prepared = sparse_interface._prepare_paged_hnd_input(view, page_size) @@ -406,11 +403,11 @@ index 21c777e..b5f078b 100644 +def test_prepare_paged_hnd_input_materializes_unaligned_outer_stride() -> None: + pages, heads, page_size, head_dim = 5, 2, 128, 128 + outer_stride = heads * page_size * head_dim + 1 -+ storage = torch.empty( ++ storage = torch.arange( + pages * outer_stride, -+ dtype=torch.float8_e4m3fn, ++ dtype=torch.int32, + device="cuda", -+ ) ++ ).remainder(16).to(torch.float8_e4m3fn) + view = storage.as_strided( + (pages, heads, page_size, head_dim), + (outer_stride, page_size * head_dim, head_dim, 1), @@ -444,7 +441,7 @@ index 21c777e..b5f078b 100644 @contextmanager def _nvtx_range(message: str): torch.cuda.nvtx.range_push(message) -@@ -1786,6 +1861,74 @@ def test_sparse_page_atten( +@@ -1786,6 +1858,74 @@ def test_sparse_page_atten( _assert_forward_close(out, out_ref, out_pt.float(), lse, lse_ref) diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 5e1748453451..35305a169afa 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -297,7 +297,6 @@ test_e2e.py::test_ptp_quickstart_advanced[Nemotron-Nano-9B-v2-nvfp4-NVIDIA-Nemot test_e2e.py::test_ptp_quickstart_bert[TRTLLM-BertForSequenceClassification-bert/bert-base-uncased-yelp-polarity] SKIP (https://nvbugs/6605819) test_e2e.py::test_ptp_quickstart_bert[VANILLA-BertForSequenceClassification-bert/bert-base-uncased-yelp-polarity] SKIP (bug pending, tracked in PR 17414) unittest/_torch/attention/sparse/dsa/test_req_idx_per_token.py::test_on_update_kv_lens_rebuilds_stale_map SKIP (https://nvbugs/6574939) -unittest/_torch/attention/sparse/msa/test_msa_backend.py::test_msa_paged_hnd_input_materializes_unaligned_outer_stride SKIP (https://nvbugs/6661846) unittest/_torch/attention/test_attention_backends.py::test_attention_backend[exaone_moe_gqa_swa128-ctx-bf16-HND-p32-v1] SKIP (https://nvbugs/6668773) unittest/_torch/attention/test_attention_backends.py::test_attention_backend[qwen2_0_5b_gqa_hd64-ctx-bf16-HND-p32-v1] SKIP (https://nvbugs/6641268) unittest/_torch/modeling/test_gemma4_e2e_dummy.py::test_e2e_text_31b_dummy SKIP (https://nvbugs/6607482) diff --git a/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py b/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py index 210db3093c41..9af80597eed9 100644 --- a/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py +++ b/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py @@ -403,10 +403,14 @@ def test_msa_paged_hnd_input_materializes_unaligned_outer_stride() -> None: pages, heads, page_size, head_dim = 5, 2, 128, 128 outer_stride = heads * page_size * head_dim + 1 - storage = torch.empty( - pages * outer_stride, - dtype=torch.float8_e4m3fn, - device="cuda", + storage = ( + torch.arange( + pages * outer_stride, + dtype=torch.int32, + device="cuda", + ) + .remainder(16) + .to(torch.float8_e4m3fn) ) view = storage.as_strided( (pages, heads, page_size, head_dim), @@ -417,7 +421,7 @@ def test_msa_paged_hnd_input_materializes_unaligned_outer_stride() -> None: assert prepared.is_contiguous() assert prepared.data_ptr() != view.data_ptr() - torch.testing.assert_close(prepared, view) + torch.testing.assert_close(prepared, view, rtol=0, atol=0) def test_per_token_valid_blocks_multi_token_decode(): From 2ebe4e65193c4232e6a326ee7dbf14770f6c4e74 Mon Sep 17 00:00:00 2001 From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> Date: Tue, 15 Sep 2026 07:04:41 +0800 Subject: [PATCH 2/8] [https://nvbugs/6732123][fix] Correct V2 KV cache quota estimation (#18988) Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> --- examples/layer_wise_benchmarks/run.py | 7 + tensorrt_llm/_torch/pyexecutor/_util.py | 51 +++-- .../kv_cache/kv_cache_manager_v2.py | 203 +++++++++++------- .../tools/layer_wise_benchmarks/runner.py | 9 +- .../defs/accuracy/test_llm_api_pytorch.py | 22 +- .../test_deepseek_v4_cache_manager.py | 13 +- .../kv_cache/test_kv_cache_budget_split.py | 107 ++++++++- .../kv_cache/test_kv_cache_estimation.py | 92 +++----- .../kv_cache/test_kv_cache_manager_v2.py | 73 +++++++ .../kv_cache/test_mamba_cache_manager.py | 3 + .../tools/test_layer_wise_benchmarks.py | 2 + 11 files changed, 392 insertions(+), 190 deletions(-) diff --git a/examples/layer_wise_benchmarks/run.py b/examples/layer_wise_benchmarks/run.py index 496e43f8eed1..eb712bd66164 100644 --- a/examples/layer_wise_benchmarks/run.py +++ b/examples/layer_wise_benchmarks/run.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + import argparse import itertools import json @@ -41,6 +44,9 @@ def comma_separated_floats(s): parser.add_argument("--max-batch-size", type=int) parser.add_argument("--tokens-per-block", type=int) parser.add_argument("--max-seq-len", type=int) +parser.add_argument( + "--use-kv-cache-manager-v2", action=argparse.BooleanOptionalAction, default="auto" +) group = parser.add_mutually_exclusive_group() group.add_argument("--enable-attention-dp", action="store_true", dest="enable_attention_dp") group.add_argument("--no-enable-attention-dp", action="store_false", dest="enable_attention_dp") @@ -212,6 +218,7 @@ def comma_separated_floats(s): enable_swa_scratch_reuse=args.enable_swa_scratch_reuse, spec_config=spec_config, vision_config=args.vision_config, + use_kv_cache_manager_v2=args.use_kv_cache_manager_v2, ) attn_workspace = torch.empty((0,), device="cuda", dtype=torch.int8) logger.info("Layer-wise benchmarks: Create KV cache manager ... Done") diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 78d1b6d921aa..cfcb6e849212 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -801,7 +801,7 @@ def _per_manager_cache_cost(self, tokens_per_block=self._tokens_per_block, max_seq_len=self._max_seq_len, max_batch_size=self._max_batch_size, - max_num_tokens=self._max_num_tokens if is_draft else 0, + max_num_tokens=self._max_num_tokens, kv_cache_config=kv_cache_config, spec_config=self._speculative_config, is_draft=is_draft, @@ -836,14 +836,29 @@ def _get_kv_size_per_token(self, use_separate_draft_kv_cache=use_separate_draft_kv_cache) if self._is_encoder_decoder(): total += CacheCost.from_raw(self._get_cross_kv_size_per_token()) + draft_cost = self._get_draft_cache_cost( + kv_cache_config, + use_separate_draft_kv_cache=use_separate_draft_kv_cache, + ) + if draft_cost is not None: + total += draft_cost + return total + + def _get_draft_cache_cost( + self, + kv_cache_config: KvCacheConfig, + *, + use_separate_draft_kv_cache: bool, + ) -> Optional[CacheCost]: + """Return the draft manager's standalone cache cost, if it has one.""" if self._draft_model_engine is not None: draft_model_config = self._draft_model_engine.model.model_config draft_kv_cache_manager_cls = self._get_model_kv_cache_manager_cls( self._draft_model_engine, kv_cache_config) - total += self._per_manager_cache_cost(draft_kv_cache_manager_cls, - draft_model_config, - kv_cache_config) - elif use_separate_draft_kv_cache: + return self._per_manager_cache_cost(draft_kv_cache_manager_cls, + draft_model_config, + kv_cache_config) + if use_separate_draft_kv_cache: # One-model draft with separate KV cache layout. # Pass num_layers explicitly since the HF config may report a # different layer count than what is actually used at runtime @@ -861,20 +876,19 @@ def _get_kv_size_per_token(self, effective_draft_config, draft_kv_cache_config, is_disagg=self._is_disagg) - total += self._per_manager_cache_cost( - draft_kv_cache_manager_cls, - effective_draft_config, - draft_kv_cache_config, - is_draft=True) + return self._per_manager_cache_cost(draft_kv_cache_manager_cls, + effective_draft_config, + draft_kv_cache_config, + is_draft=True) elif self._mapping.is_last_pp_rank(): # EAGLE3/MTP: draft layers only on last PP rank - total += self._per_manager_cache_cost( + return self._per_manager_cache_cost( self._kv_cache_manager_cls, effective_draft_config, draft_kv_cache_config, num_layers=self._get_num_draft_layers(), is_draft=True) - return total + return None def _cal_max_memory(self, peak_memory, total_gpu_memory, fraction, allocated_bytes: int) -> int: @@ -1732,7 +1746,6 @@ def _get_target_and_draft_cache_costs( """Per-manager KV cache costs for target and draft layers.""" target_kv_cache_config = (kv_cache_config if kv_cache_config is not None else self._kv_cache_config) - total_kv = self._get_kv_size_per_token(target_kv_cache_config) use_separate_draft_kv_cache = ( self._should_create_separate_draft_kv_cache()) target_kv = self._per_manager_cache_cost( @@ -1740,8 +1753,14 @@ def _get_target_and_draft_cache_costs( self._model_engine.model.model_config, target_kv_cache_config, use_separate_draft_kv_cache=use_separate_draft_kv_cache) - draft_kv = CacheCost(slope=total_kv.slope - target_kv.slope, - intercept=total_kv.intercept - target_kv.intercept) + # Estimate the draft component directly so its independently modelled + # affine intercept is preserved exactly. + draft_kv = self._get_draft_cache_cost( + target_kv_cache_config, + use_separate_draft_kv_cache=use_separate_draft_kv_cache, + ) + if draft_kv is None: + return None costs = (target_kv, draft_kv) if any(cost.slope < 0 or cost.intercept < 0 or ( cost.slope == 0 and cost.intercept == 0) for cost in costs): @@ -1843,7 +1862,7 @@ def _split_kv_cache_budget_for_draft( raise ValueError( f"KV cache GPU budget ({total_budget / GB:.2f} GiB) is " f"insufficient after the combined fixed cost " - f"({intercept_total / GB:.2f} GiB, e.g. mamba SSM state) " + f"({intercept_total / GB:.2f} GiB, e.g. SWA or mamba state) " f"for target+draft. Increase free_gpu_memory_fraction or " f"max_gpu_total_bytes, or reduce max_batch_size (the fixed " f"cost scales with batch size).") diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index 354bc144d26a..47655d32cece 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -126,12 +126,16 @@ "iter_partial_reused_blocks", "iter_missed_blocks", ) + KV_CACHE_ITERATION_STATS_POOL_GROUP_FIELDS = tuple( field_name for field_name in KV_CACHE_ITERATION_STATS_DELTA_FIELDS if field_name not in KV_CACHE_ITERATION_STATS_REUSE_FIELDS ) +# Shared by capacity estimation and growth, in addition to speculative tokens. +BASE_GENERATION_TOKEN_COUNT = 1 + # The guard page is a permanent sequence that must never collide with any other # permanent sequence id. The CUDA-graph machinery reserves a whole window below # ``CUDA_GRAPH_DUMMY_REQUEST_ID``: spec-decode capture uses @@ -445,7 +449,7 @@ def _estimate_swa_cache_size( *, context: bool, scratch: bool, - generation_capacity_headroom: Optional[int] = None, + generation_capacity_headroom: int = BASE_GENERATION_TOKEN_COUNT, ) -> tuple[int, int]: tokens_per_block = int(tokens_per_block) size_per_token = 0 @@ -453,15 +457,12 @@ def _estimate_swa_cache_size( scratch_keys = set() for layer_size, window_size in zip(layer_sizes, attention_windows): if window_size is not None and window_size > 0: - if generation_capacity_headroom is None: - window_blocks = math.ceil(window_size / tokens_per_block) - else: - # Match DFlash's retained boundary page and capacity reserved - # ahead of committed history for the next draft step. - window_blocks = ( - math.ceil((window_size + generation_capacity_headroom - 1) / tokens_per_block) - + 1 - ) + # Match AttnLifeCycle.get_stale_range(): the live interval contains + # window_size + generation_capacity_headroom - 1 tokens. Across all + # page offsets, that interval touches at most the count below. + window_blocks = ( + math.ceil((window_size + generation_capacity_headroom - 2) / tokens_per_block) + 1 + ) window_tokens = window_blocks * tokens_per_block if not context: size_per_request += window_tokens * layer_size @@ -477,16 +478,58 @@ def _estimate_swa_cache_size( return size_per_token, size_per_request -def _get_dflash_generation_kv_capacity_headroom(spec_config) -> Optional[int]: - """DFlash KV capacity reserved ahead of committed history.""" - from tensorrt_llm._torch.speculative.interface import SpeculativeDecodingMode +def _get_generation_kv_capacity(spec_config, *, is_draft: bool) -> tuple[int, int]: + """Return draft-token reserve and total capacity headroom over history.""" + from tensorrt_llm._torch.speculative import get_num_extra_kv_tokens - if spec_config is None or spec_config.spec_dec_mode != SpeculativeDecodingMode.DFLASH: - return None + if spec_config is None: + return 0, BASE_GENERATION_TOKEN_COUNT + reserve = spec_config.max_total_draft_tokens + if ( + is_draft + and getattr(spec_config, "use_dynamic_tree", False) + and getattr(spec_config, "dynamic_tree_max_topK", 0) > 0 + ): + draft_loop_tokens = spec_config.dynamic_tree_max_topK * spec_config.max_draft_len + reserve = max(reserve, draft_loop_tokens) + dynamic_reserve = reserve - spec_config.max_total_draft_tokens + headroom = ( + get_num_extra_kv_tokens(spec_config) + spec_config.tokens_per_gen_step + dynamic_reserve + ) + return reserve, headroom - from tensorrt_llm._torch.speculative import get_num_extra_kv_tokens - return get_num_extra_kv_tokens(spec_config) + spec_config.tokens_per_gen_step +def _estimate_cache_size_components( + layer_sizes: Sequence[int], + attention_windows: Sequence[int | None], + tokens_per_block: int, + *, + scratch: bool, + generation_capacity_headroom: int, +) -> tuple[int, int, int]: + """Return context/generation bytes per token and generation bytes per request. + + Static profiling and runtime quota conversion must charge the same SWA + retention pages and context scratch space. Resume-watermark normalization + is separate from these usable-capacity costs. + """ + full_attn_size = _estimate_full_attn_size_per_token(layer_sizes, attention_windows) + context_swa_size, _ = _estimate_swa_cache_size( + layer_sizes, attention_windows, tokens_per_block, context=True, scratch=scratch + ) + generation_swa_size, generation_swa_per_request = _estimate_swa_cache_size( + layer_sizes, + attention_windows, + tokens_per_block, + context=False, + scratch=False, + generation_capacity_headroom=generation_capacity_headroom, + ) + return ( + full_attn_size + context_swa_size, + full_attn_size + generation_swa_size, + generation_swa_per_request, + ) def _get_single_swa_pool_slot_bytes( @@ -546,6 +589,7 @@ def _get_static_cache_size_layer_components( *, max_seq_len: Optional[int] = None, kv_cache_config: Optional[KvCacheConfig] = None, + is_external_draft: bool = False, ) -> tuple[List[int], List[Optional[int]]]: config = model_config.pretrained_config @@ -609,6 +653,17 @@ def normalize_window_size(window_size: Optional[int]) -> Optional[int]: attention_windows = [ normalize_window_size(window_size) for window_size in local_window_pattern ] + layer_types = getattr(config, "layer_types", None) + if not is_external_draft and isinstance(layer_types, (list, tuple)) and layer_types: + # Estimate native full-attention layers as growing with token capacity, + # even when a cache-window override bounds their runtime retention. + # Charging max_batch_size saturated windows would turn a capacity cap + # into a mandatory reservation. Native SWA and external draft windows + # retain their existing fixed costs; runtime window sizing is unchanged. + attention_windows = [ + None if layer_types[layer_idx % len(layer_types)] == "full_attention" else window + for layer_idx, window in zip(local_layer_ids, attention_windows) + ] return layer_sizes, attention_windows @@ -1177,15 +1232,9 @@ def __init__( if not self._supports_reuse_match_backoff: self.reuse_match_backoff = 0 # Mirror V1's KV reserve sizing (see V1 __init__ for rationale). - self._kv_reserve_draft_tokens = self.max_total_draft_tokens - if ( - self.is_draft - and spec_config is not None - and getattr(spec_config, "use_dynamic_tree", False) - and getattr(spec_config, "dynamic_tree_max_topK", 0) > 0 - ): - draft_loop_tokens = spec_config.dynamic_tree_max_topK * spec_config.max_draft_len - self._kv_reserve_draft_tokens = max(self.max_total_draft_tokens, draft_loop_tokens) + self._kv_reserve_draft_tokens, self._generation_kv_capacity_headroom = ( + _get_generation_kv_capacity(spec_config, is_draft=self.is_draft) + ) self.event_buffer_max_size = kv_cache_config.event_buffer_max_size self.enable_stats = enable_stats @@ -1601,11 +1650,8 @@ def create_cold_page_codec(cache_config: object) -> Optional[object]: self.max_seq_len = int(max_num_tokens) # Pad max_blocks_per_seq to next multiple of 4 (copy_block_offsets kernel). - # Account for max single-sequence capacity = seq_len + extra KV tokens + - # _kv_reserve_draft_tokens (see __init__) + 1 base decode token. - max_seq_capacity = ( - self.max_seq_len + self.num_extra_kv_tokens + self._kv_reserve_draft_tokens + 1 - ) + # Include the same maximum generation lead used by allocation and sizing. + max_seq_capacity = self.max_seq_len + self._generation_kv_capacity_headroom self.max_blocks_per_seq = ( max_seq_capacity + self._ledger_tokens_per_block - 1 ) // self._ledger_tokens_per_block @@ -2204,33 +2250,26 @@ def _get_max_tokens_from_quota(self, quota: int) -> float: def _get_max_tokens_from_quota_impl(self, quota: int) -> float: layer_sizes, attention_windows = self._get_runtime_cache_size_layer_components() - full_attn_size_per_token = _estimate_full_attn_size_per_token( - layer_sizes, attention_windows - ) - context_swa_size_per_token, _ = _estimate_swa_cache_size( + ( + context_size_per_token, + generation_size_per_token, + generation_swa_size_per_request, + ) = _estimate_cache_size_components( layer_sizes, attention_windows, self.tokens_per_block, - context=True, scratch=self.enable_swa_scratch_reuse, - ) - ( - generation_swa_size_per_token, - generation_swa_size_per_request, - ) = _estimate_swa_cache_size( - layer_sizes, attention_windows, self.tokens_per_block, context=False, scratch=False + generation_capacity_headroom=self._generation_kv_capacity_headroom, ) size_per_batch = self.max_batch_size * generation_swa_size_per_request if quota < size_per_batch: return 0 - context_size_per_token = full_attn_size_per_token + context_swa_size_per_token context_limit_quota = self.max_num_tokens * context_size_per_token + size_per_batch if quota <= context_limit_quota: if context_size_per_token <= 0: return float("inf") return (quota - size_per_batch) / context_size_per_token - generation_size_per_token = full_attn_size_per_token + generation_swa_size_per_token if generation_size_per_token <= 0: return float("inf") return self.max_num_tokens + (quota - context_limit_quota) / generation_size_per_token @@ -2247,34 +2286,24 @@ def _get_quota_from_max_tokens(self, max_tokens: int) -> int: def _get_quota_from_max_tokens_impl(self, max_tokens: int) -> int: layer_sizes, attention_windows = self._get_runtime_cache_size_layer_components() - full_attn_size_per_token = _estimate_full_attn_size_per_token( - layer_sizes, attention_windows - ) ( - context_swa_size_per_token, - _, - ) = _estimate_swa_cache_size( + context_size_per_token, + generation_size_per_token, + generation_swa_size_per_request, + ) = _estimate_cache_size_components( layer_sizes, attention_windows, self.tokens_per_block, - context=True, scratch=self.enable_swa_scratch_reuse, - ) - ( - generation_swa_size_per_token, - generation_swa_size_per_request, - ) = _estimate_swa_cache_size( - layer_sizes, attention_windows, self.tokens_per_block, context=False, scratch=False + generation_capacity_headroom=self._generation_kv_capacity_headroom, ) context_tokens = min(max_tokens, self.max_num_tokens) generation_tokens = max_tokens - context_tokens - generation_quota = ( - max_tokens * full_attn_size_per_token - + generation_tokens * generation_swa_size_per_token + return int( + context_tokens * context_size_per_token + + generation_tokens * generation_size_per_token + self.max_batch_size * generation_swa_size_per_request ) - context_extra_quota = context_tokens * context_swa_size_per_token - return int(generation_quota + context_extra_quota) def _get_event_num_blocks_per_cache_level( self, @@ -3181,7 +3210,7 @@ def _required_gen_capacity(self, req: LlmRequest, current_capacity: int) -> int: Grows *current_capacity* by 1 + draft tokens. """ - return current_capacity + 1 + self._generation_draft_slots(req) + return current_capacity + BASE_GENERATION_TOKEN_COUNT + self._generation_draft_slots(req) def _generation_draft_slots(self, req: LlmRequest) -> int: """Physical draft width reserved for one iteration. Dynamic-tree draft pools @@ -4873,6 +4902,11 @@ def get_cache_size_per_token( num_layers=num_layers, max_seq_len=max_seq_len, kv_cache_config=kv_cache_config, + is_external_draft=( + is_draft + and spec_config is not None + and spec_config.spec_dec_mode.is_external_drafter() + ), ) reuse_backoff_enabled = ( kv_cache_config is not None @@ -4889,37 +4923,40 @@ def get_cache_size_per_token( backoff, max_seq_len, ) - full_attn_size_per_token = _estimate_full_attn_size_per_token( - layer_sizes, attention_windows + _, generation_capacity_headroom = _get_generation_kv_capacity( + spec_config, is_draft=is_draft ) - dflash_headroom = _get_dflash_generation_kv_capacity_headroom(spec_config) - is_dflash_draft = is_draft and dflash_headroom is not None - generation_capacity_headroom = dflash_headroom if is_dflash_draft else None - swa_size_per_token, swa_size_per_request = _estimate_swa_cache_size( + ( + context_size_per_token, + cache_size_per_token, + swa_size_per_request, + ) = _estimate_cache_size_components( layer_sizes, attention_windows, tokens_per_block, - context=False, - scratch=False, + scratch=( + kv_cache_config is not None + and kv_cache_config.enable_swa_scratch_reuse + and not is_draft + ), generation_capacity_headroom=generation_capacity_headroom, ) - context_swa_size_per_token = 0 - if is_dflash_draft: - context_swa_size_per_token, _ = _estimate_swa_cache_size( - layer_sizes, - attention_windows, - tokens_per_block, - context=True, - scratch=False, - ) + # The affine slope covers all tokens; context additionally retains SWA + # pages for the current token batch beyond the generation windows. fixed_cost = ( - swa_size_per_request * max_batch_size + context_swa_size_per_token * max_num_tokens + swa_size_per_request * max_batch_size + + (context_size_per_token - cache_size_per_token) * max_num_tokens ) - cache_size_per_token = full_attn_size_per_token + swa_size_per_token bytes_per_slot = _get_single_swa_pool_slot_bytes( layer_sizes, attention_windows, tokens_per_block ) - if is_dflash_draft and bytes_per_slot is not None: + if is_draft and fixed_cost > 0 and bytes_per_slot is not None: + # Without a config, all windows are full attention, so no SWA slot exists. + assert kv_cache_config is not None + # The affine intercept is the configured quota needed to preserve + # the fixed usable capacity at V2's resume watermark. Keep the + # allocator's page rounding inside the manager estimator rather + # than extending the generic CacheCost model with pool geometry. required_slots = math.ceil(fixed_cost / bytes_per_slot) resume_util = float(np.float32(kv_cache_config.max_util_for_resume)) fixed_cost = math.ceil(required_slots / resume_util) * bytes_per_slot diff --git a/tensorrt_llm/tools/layer_wise_benchmarks/runner.py b/tensorrt_llm/tools/layer_wise_benchmarks/runner.py index 23d51a4e019e..f31f17446190 100644 --- a/tensorrt_llm/tools/layer_wise_benchmarks/runner.py +++ b/tensorrt_llm/tools/layer_wise_benchmarks/runner.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + import contextlib import functools import inspect @@ -6,7 +9,7 @@ import weakref from dataclasses import replace from enum import IntEnum -from typing import Optional +from typing import Literal, Optional import torch @@ -887,6 +890,7 @@ def create_kv_cache_manager( enable_swa_scratch_reuse=False, spec_config: Optional[DecodingBaseConfig] = None, vision_config: Optional[str] = None, + use_kv_cache_manager_v2: bool | Literal["auto"] = "auto", ) -> KVCacheManager: # Please refer to `tensorrt_llm/_torch/pyexecutor/py_executor_creator.py` for `tokens_per_block` with Runner.vision_config_ctx(vision_config): @@ -909,6 +913,9 @@ def create_kv_cache_manager( * round_up(max_seq_len + 1, tokens_per_block), enable_block_reuse=False, enable_swa_scratch_reuse=enable_swa_scratch_reuse, + # Every dummy request uses max_seq_len tokens, so this is the actual average. + avg_seq_len=max_seq_len, + use_kv_cache_manager_v2=use_kv_cache_manager_v2, ) kv_cache_manager_cls = get_kv_cache_manager_cls(model_config, kv_cache_config) kv_cache_dtype = { diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 034082a64191..d283757e9495 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -4484,10 +4484,12 @@ def test_w4_1gpu_suspend_resume(self) -> None: f"Question: what was the code word? Answer in one word." for n in needles ] - # Long generation so each suspend/resume round trip spans many decode - # steps and the suspended requests are held across a long active run of - # the request that owns the pool - sampling = SamplingParams(max_tokens=128, temperature=0.0) + # Keep every request resident for the full decode length. An early EOS + # can release enough pages to avoid the suspend/resume path. GPT-OSS + # has multiple EOS IDs, so min_tokens alone does not prevent early stops. + sampling = SamplingParams(max_tokens=128, + ignore_eos=True, + temperature=0.0) def _drain_stats() -> tuple[int, int, int, int, int]: # Single drain: llm.get_stats() consumes the per-iteration records, @@ -4553,11 +4555,8 @@ def _common_prefix_len(a: list[int], b: list[int]) -> int: # Checked against the *measured* pool, never against max_tokens -- # the manager inflates that knob by 1 / max_util_for_resume, so a # config-derived bound is wrong by ~1.43x (see the sizing note on - # kv_cache_config). Under VSWA the pool holds far more tokens per - # byte than a full-attention model, so it takes twelve requests to - # overflow it: observed on B200 at concurrent_peak ~4764 vs - # pool_tokens ~3328, a ~1.4x margin that reliably suspends/resumes a - # couple of in-flight requests. Failing here means the workload + # kv_cache_config). The allocated pool also depends on page-aligned + # window retention. Failing here means the workload # stopped being contended, NOT that the KV path broke. assert concurrent_peak > pool_tokens, ( f"workload precondition failed: the pool ({pool_tokens} tokens) " @@ -4588,6 +4587,7 @@ def _common_prefix_len(a: list[int], b: list[int]) -> int: f"common_prefix_tokens={prefixes}") for i in range(len(prompts)): print(f"[I-10 out {i}] needle={needles[i]} " + f"ref_tokens={len(ref_ids[i])} con_tokens={len(con_ids[i])} " f"ref_recall={needles[i] in ref_txt[i].upper()} " f"con_recall={needles[i] in con_txt[i].upper()} " f"ref={ref_txt[i]!r} con={con_txt[i]!r}") @@ -4599,6 +4599,10 @@ def _common_prefix_len(a: list[int], b: list[int]) -> int: assert all(len(ids) > 0 for ids in con_ids), ( "a request produced no tokens under suspend/resume contention " "(possible V2 scheduler deadlock or illegal-memory-access crash)") + assert all(len(ids) == sampling.max_tokens for ids in con_ids), ( + "fixed-length contention workload ended early: " + f"expected {sampling.max_tokens} tokens per request, " + f"got {[len(ids) for ids in con_ids]}") # (2) The ACTIVE<->SUSPENDED state machine genuinely fired (not mere # queuing): an in-flight request was suspended under pressure and later # recovered. These per-iteration manager counters are the direct signal; diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py index 6db4d0c58a91..45749726466b 100644 --- a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py @@ -36,11 +36,14 @@ from tensorrt_llm._torch.pyexecutor._util import CacheCost from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, LlmRequestState from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests -from tensorrt_llm._torch.speculative.interface import SpeculativeDecodingMode from tensorrt_llm._utils import binding_to_torch_dtype from tensorrt_llm.bindings import DataType, SamplingConfig from tensorrt_llm.bindings.internal.batch_manager import CacheType as CacheTypeCpp -from tensorrt_llm.llmapi.llm_args import DeepSeekV4SparseAttentionConfig, KvCacheConfig +from tensorrt_llm.llmapi.llm_args import ( + DeepSeekV4SparseAttentionConfig, + DraftTargetDecodingConfig, + KvCacheConfig, +) from tensorrt_llm.mapping import Mapping from tensorrt_llm.runtime.kv_cache_manager_v2 import BatchDesc, KVCacheDesc, PageIndexMode from tensorrt_llm.runtime.kv_cache_manager_v2._common import BAD_PAGE_INDEX @@ -1730,11 +1733,7 @@ def test_swa_scratch_reuse_disabled_by_config_for_main_manager(self): cache_manager.shutdown() def test_swa_scratch_reuse_uses_extra_kv_tokens_for_rewind(self): - spec_config = SimpleNamespace( - max_draft_len=7, - max_total_draft_tokens=7, - spec_dec_mode=SpeculativeDecodingMode.DRAFT_TARGET_ONE_MODEL, - ) + spec_config = DraftTargetDecodingConfig(max_draft_len=7, speculative_model="draft") cache_manager, _ = self._create_deepseek_v4_cache_manager( tokens_per_block=self.tokens_per_block, max_batch_size=1, diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_budget_split.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_budget_split.py index e2575f414ea2..5a6220330125 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_budget_split.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_budget_split.py @@ -14,9 +14,11 @@ # limitations under the License. """Tests for KV cache budget splitting between target and draft managers.""" +import math from types import SimpleNamespace from unittest.mock import Mock +import numpy as np import pytest from tensorrt_llm._torch.pyexecutor._util import CacheCost, KvCacheCreator @@ -47,8 +49,8 @@ def _make_creator( """Minimal KvCacheCreator for budget-split helpers. ``*_intercept`` model the affine fixed cost (e.g. mamba SSM state) that a - manager pays per batch regardless of token count. The draft cost is derived - as ``total - target`` for both slope and intercept. + manager pays per batch regardless of token count. The draft mock receives + the component-wise ``total - target`` values directly. """ c = object.__new__(KvCacheCreator) @@ -60,6 +62,7 @@ def _make_creator( ) c._tokens_per_block = 64 c._max_seq_len = 1024 + c._max_num_tokens = 0 c._max_batch_size = 1 c._speculative_config = None c._mapping = Mock() @@ -78,11 +81,92 @@ def _make_creator( ) ) c._should_create_separate_draft_kv_cache = Mock(return_value=True) + c._get_draft_cache_cost = Mock( + return_value=CacheCost( + slope=total_kv_per_token - target_kv_per_token, + intercept=total_kv_intercept - target_kv_intercept, + ) + ) return c class TestSplitGpuBudgetForDraft: + @pytest.mark.parametrize( + "long_window, max_batch_size", [(32768, 2048), (32768, 1), (131072, 1)] + ) + def test_native_full_target_cost_preserves_draft_capacity( + self, long_window, max_batch_size + ) -> None: + """Bounded full layers must not reserve saturated windows or starve the draft.""" + max_seq_len = 131072 + target_layer_bytes = 1024 + draft_layer_bytes = 512 + # Include 25% slack above the full history plus native SWA window cost. + single_request_bytes = ( + 12 * (max_seq_len + 128) * target_layer_bytes + max_seq_len * draft_layer_bytes + ) + total_budget = single_request_bytes * 5 // 4 if max_batch_size == 1 else 10 * GB + c = _make_creator(max_gpu_total_bytes=total_budget) + del c._get_draft_cache_cost + c._kv_cache_config.enable_block_reuse = False + c._kv_cache_config.max_attention_window = [128, long_window] + c._kv_cache_manager_cls = KVCacheManagerV2 + c._max_seq_len = max_seq_len + c._max_batch_size = max_batch_size + c._max_num_tokens = 8192 + c._mapping = Mock(enable_attention_dp=False, tp_size=2) + c._mapping.pp_layers.return_value = list(range(24)) + c._mapping.is_last_pp_rank.return_value = True + c._speculative_config = SimpleNamespace( + spec_dec_mode=SpeculativeDecodingMode.EAGLE3_ONE_MODEL, + max_draft_len=3, + max_total_draft_tokens=3, + tokens_per_gen_step=4, + use_dynamic_tree=False, + _use_shared_kv_cache=False, + ) + target_model_config = SimpleNamespace( + is_encoder_decoder=False, + quant_config=None, + pretrained_config=SimpleNamespace( + num_hidden_layers=24, + hidden_size=2880, + num_attention_heads=64, + num_key_value_heads=8, + head_dim=64, + layer_types=["sliding_attention", "full_attention"] * 12, + ), + get_num_attention_layers=lambda: 24, + ) + draft_model_config = SimpleNamespace( + quant_config=None, + pretrained_config=SimpleNamespace( + num_hidden_layers=1, + hidden_size=2880, + num_attention_heads=64, + num_key_value_heads=4, + head_dim=64, + ), + ) + c._model_engine.model.model_config = target_model_config + c._draft_model_engine = None + c._get_effective_draft_config = Mock(return_value=draft_model_config) + c._get_num_draft_layers = Mock(return_value=1) + + target, draft = c._split_kv_cache_budget_for_draft("max_gpu_total_bytes") + assert draft is not None + assert target.max_gpu_total_bytes > 0 + assert draft.max_gpu_total_bytes > 0 + if max_batch_size == 1: + # A necessary capacity bound, not just an assertion of the split formula. + # Counting native SWA in the slope instead would fail this bound. + assert draft.max_gpu_total_bytes >= max_seq_len * draft_layer_bytes + assert target.max_gpu_total_bytes >= 12 * (max_seq_len + 128) * target_layer_bytes + assert target.max_gpu_total_bytes + draft.max_gpu_total_bytes == total_budget + assert c._kv_cache_config.max_gpu_total_bytes == total_budget + assert c._kv_cache_config.max_attention_window == [128, long_window] + @pytest.mark.parametrize( "mode", [ @@ -136,7 +220,14 @@ def get_cache_size_per_token(model_config, *args, **kwargs): creator._mapping = Mock(enable_attention_dp=False, tp_size=1) creator._mapping.pp_layers.return_value = [0] creator._mapping.is_last_pp_rank.return_value = True - creator._speculative_config = SimpleNamespace(spec_dec_mode=mode) + creator._speculative_config = SimpleNamespace( + spec_dec_mode=mode, + max_draft_len=1, + max_total_draft_tokens=0, + tokens_per_gen_step=1, + use_dynamic_tree=False, + _use_shared_kv_cache=False, + ) creator._model_engine = SimpleNamespace( model=SimpleNamespace(model_config=target_model_config) ) @@ -153,9 +244,13 @@ def get_cache_size_per_token(model_config, *args, **kwargs): ) # The draft layer stores 64 bytes/token in a fixed 512-token window. + # Generation retains one additional 64-token boundary block and the + # 128-token context budget consumes two more blocks. # Leaking the target's 16K window would instead count it as 64 bytes/token. cost = creator._get_kv_size_per_token() - assert cost == CacheCost(slope=10, intercept=512 * 64) + usable_slots = 11 + configured_slots = math.ceil(usable_slots / float(np.float32(0.95))) + assert cost == CacheCost(slope=10, intercept=configured_slots * 64 * 64) assert len(draft_kv_configs) == 1 draft_kv_config = draft_kv_configs[0] assert draft_kv_config.max_attention_window == [512] @@ -269,9 +364,7 @@ def test_returns_none_when_draft_kv_zero(self): def test_fixed_only_draft_uses_manager_estimated_quota(self): total_gpu = 10 * GB - slot_bytes = 327_680 - configured_slots = 2_561 - configured_bytes = configured_slots * slot_bytes + configured_bytes = 2_561 * 327_680 c = _make_creator( max_gpu_total_bytes=total_gpu, total_kv_per_token=80, diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py index 6bad069afaa4..6c819f101fb0 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py @@ -10,7 +10,6 @@ share, not all copies. """ -import math from dataclasses import dataclass from types import SimpleNamespace from unittest.mock import Mock, patch @@ -23,7 +22,6 @@ from tensorrt_llm._torch.pyexecutor._util import CacheCost, KvCacheCreator from tensorrt_llm._torch.pyexecutor.config_utils import get_layer_attention_window from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 -from tensorrt_llm._torch.speculative.interface import SpeculativeDecodingMode from tensorrt_llm.inputs.multimodal import MultimodalParams from tensorrt_llm.llmapi.llm_args import ( KvCacheConfig, @@ -633,7 +631,9 @@ def get_num_attention_layers(self): tokens_per_block=64, max_seq_len=4096, max_batch_size=3, - kv_cache_config=KvCacheConfig(max_attention_window=[2048, 2048, 4096]), + kv_cache_config=KvCacheConfig( + max_attention_window=[2048, 2048, 4096], enable_swa_scratch_reuse=False + ), ) ) scratch_size_per_token = CacheCost.from_raw( @@ -643,76 +643,27 @@ def get_num_attention_layers(self): tokens_per_block=64, max_seq_len=4096, max_batch_size=3, - kv_cache_config=KvCacheConfig(max_attention_window=[2048, 2048, 4096]), - enable_swa_scratch_reuse=True, + kv_cache_config=KvCacheConfig( + max_attention_window=[2048, 2048, 4096], enable_swa_scratch_reuse=True + ), ) ) # Per layer: K+V * kv_heads * head_dim * bf16 bytes = 2 * 2 * 8 * 2. - expected = CacheCost(slope=64, intercept=3 * 2 * 2048 * 64) + expected = CacheCost(slope=64, intercept=3 * 2 * (2048 + 64) * 64) assert no_scratch_size_per_token == expected assert scratch_size_per_token == expected -def test_v2_dflash_draft_cost_covers_context_and_generation_slots(): - class FakeDraftModelConfig: - quant_config = None - pretrained_config = SimpleNamespace( - hidden_size=32, - num_attention_heads=4, - num_key_value_heads=2, - ) - - def get_num_attention_layers(self): - return 1 - - spec_config = SimpleNamespace( - spec_dec_mode=SpeculativeDecodingMode.DFLASH, - max_draft_len=4, - tokens_per_gen_step=5, - ) - mapping = Mock(enable_attention_dp=False, tp_size=1) - mapping.pp_layers.return_value = [0] - tokens_per_block = 32 - max_batch_size = 128 - max_num_tokens = 4096 - - cost = CacheCost.from_raw( - KVCacheManagerV2.get_cache_size_per_token( - FakeDraftModelConfig(), - mapping, - tokens_per_block=tokens_per_block, - max_seq_len=4096, - max_batch_size=max_batch_size, - max_num_tokens=max_num_tokens, - kv_cache_config=KvCacheConfig(max_attention_window=[512]), - spec_config=spec_config, - is_draft=True, - ) - ) - - # One layer stores K+V * 2 KV heads * 8 head dim * BF16 = 64 B/token. - slot_bytes = tokens_per_block * 64 - # Runtime generation capacity can lead history by max_draft_len - 1 - # (get_num_extra_kv_tokens) plus tokens_per_gen_step: 3 + 5 = 8. - generation_blocks_per_request = math.ceil((512 + 8 - 1) / tokens_per_block) + 1 - assert generation_blocks_per_request == 18 - context_slots = max_num_tokens // tokens_per_block - expected_usable_slots = max_batch_size * generation_blocks_per_request + context_slots - assert expected_usable_slots == 2_432 - # float32(0.95) requires 2561 configured slots for 2432 slots to remain - # resumable. The estimator owns this manager-specific quota normalization. - expected_configured_slots = 2_561 - assert cost == CacheCost(slope=0, intercept=expected_configured_slots * slot_bytes) - - -def test_v2_static_cache_size_preserves_window_pattern_phase_across_pp() -> None: +@pytest.mark.parametrize("max_seq_len", [256, 512]) +def test_v2_static_cache_size_preserves_window_pattern_phase_across_pp(max_seq_len) -> None: class FakeModelConfig: quant_config = None pretrained_config = SimpleNamespace( hidden_size=32, num_attention_heads=4, num_key_value_heads=2, + layer_types=["sliding_attention", "full_attention"] * 2 + ["sliding_attention"], ) def get_num_attention_layers(self) -> int: @@ -726,7 +677,7 @@ def get_num_attention_layers(self) -> int: FakeModelConfig(), mapping, tokens_per_block=64, - max_seq_len=256, + max_seq_len=max_seq_len, max_batch_size=1, kv_cache_config=KvCacheConfig(max_attention_window=[128, 256]), ) @@ -759,7 +710,10 @@ class UnsupportedKVCacheManagerV2(KVCacheManagerV2): tokens_per_block=64, max_seq_len=4096, max_batch_size=3, - kv_cache_config=KvCacheConfig(max_attention_window=[64]), + kv_cache_config=KvCacheConfig( + max_attention_window=[64], + max_util_for_resume=1.0, + ), spec_config=spec_config, ) @@ -797,18 +751,20 @@ class UnsupportedKVCacheManagerV2(KVCacheManagerV2): "kv_cache_config": KvCacheConfig( enable_block_reuse=False, max_attention_window=[64], + max_util_for_resume=1.0, ) } ), ) ) - # W=64 occupies one page; one-model draft reuse retains W+D=65 and - # therefore charges two pages in single and separate KVCM layouts. - assert no_draft == CacheCost(slope=0, intercept=3 * 64 * 64) + # W=64 plus the base generation token can retain two boundary pages. + # One-model draft reuse extends retention to W+D=65; its two-token + # generation step can therefore cross into a third page. + assert no_draft == CacheCost(slope=0, intercept=3 * 128 * 64) assert unsupported == no_draft assert block_reuse_disabled == no_draft - assert single_kvcm == target == draft == CacheCost(slope=0, intercept=3 * 128 * 64) + assert single_kvcm == target == draft == CacheCost(slope=0, intercept=3 * 192 * 64) def test_creator_uses_v2_affine_cache_cost(): @@ -822,6 +778,7 @@ def get_cache_size_per_token(model_config, mapping, **kwargs): creator._tokens_per_block = 64 creator._max_seq_len = 1024 creator._max_batch_size = 3 + creator._max_num_tokens = 1024 creator._kv_cache_config = KvCacheConfig() creator._speculative_config = None @@ -839,6 +796,7 @@ def test_v2_quota_from_max_tokens_models_context_swa_scratch(): manager.tokens_per_block = 64 manager.max_batch_size = 4 manager.max_num_tokens = 1000 + manager._generation_kv_capacity_headroom = 1 manager.get_layer_bytes_per_token = lambda local_layer_idx, data_role: [10, 10, 20][ local_layer_idx ] @@ -847,12 +805,12 @@ def test_v2_quota_from_max_tokens_models_context_swa_scratch(): manager.enable_swa_scratch_reuse = False no_scratch_quota = manager._get_quota_from_max_tokens(max_tokens) - assert no_scratch_quota == (max_tokens * 20 + manager.max_num_tokens * 20 + 4 * 2 * 128 * 10) + assert no_scratch_quota == (max_tokens * 20 + manager.max_num_tokens * 20 + 4 * 2 * 192 * 10) assert manager._get_max_tokens_from_quota(no_scratch_quota) == max_tokens manager.enable_swa_scratch_reuse = True scratch_quota = manager._get_quota_from_max_tokens(max_tokens) - assert scratch_quota == (max_tokens * 20 + manager.max_num_tokens * 10 + 4 * 2 * 128 * 10) + assert scratch_quota == (max_tokens * 20 + manager.max_num_tokens * 10 + 4 * 2 * 192 * 10) assert manager._get_max_tokens_from_quota(scratch_quota) == max_tokens diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2.py index 56697c5bd689..0ec5aaba4be4 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2.py @@ -41,9 +41,11 @@ from tensorrt_llm.conversation_params import ConversationParams from tensorrt_llm.llmapi.llm_args import ( BlockReuseConfig, + DFlashDecodingConfig, Eagle3DecodingConfig, KvCacheConfig, MTPDecodingConfig, + PARDDecodingConfig, ) from tensorrt_llm.mapping import Mapping from tensorrt_llm.runtime.kv_cache_manager_v2 import ( @@ -1349,6 +1351,77 @@ def set_prepopulated_prompt_len(self, length: int, tokens_per_block: int) -> Non self.context_current_position = length +@pytest.mark.parametrize("config_cls", [DFlashDecodingConfig, PARDDecodingConfig]) +def test_external_draft_estimated_quota_supports_allocation_and_resume( + config_cls: type[DFlashDecodingConfig] | type[PARDDecodingConfig], +) -> None: + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + init_cuda_once() + spec = config_cls(max_draft_len=4, speculative_model="draft") + model_config = SimpleNamespace( + quant_config=None, + pretrained_config=SimpleNamespace( + hidden_size=32, + num_attention_heads=4, + num_key_value_heads=2, + layer_types=["full_attention"], + ), + get_num_attention_layers=lambda: 1, + ) + batch_size, context_tokens, window = 64, 512, 504 + config = KvCacheConfig(enable_block_reuse=False, max_attention_window=[window]) + slope, fixed = KVCacheManagerV2.get_cache_size_per_token( + model_config, + Mapping(), + tokens_per_block=32, + max_seq_len=4096, + max_batch_size=batch_size, + max_num_tokens=context_tokens, + kv_cache_config=config, + spec_config=spec, + is_draft=True, + ) + config.max_gpu_total_bytes = slope * context_tokens + fixed + manager = KVCacheManagerV2( + config, + CacheType.SELF, + num_layers=1, + num_kv_heads=2, + head_dim=8, + tokens_per_block=32, + max_seq_len=4096, + max_batch_size=batch_size, + max_num_tokens=context_tokens, + mapping=Mapping(), + dtype=DataType.HALF, + spec_config=spec, + is_draft=True, + ) + caches = [] + stream = torch.cuda.current_stream().cuda_stream + try: + for _ in range(batch_size): + cache = manager.impl.create_kv_cache() + caches.append(cache) + assert cache.resume(stream) + # One context request runs alongside a full generation batch minus one. + assert caches[0].resize(context_tokens + spec.max_draft_len - 1) + # Sweep a complete page so the workload exercises unaligned retention. + for history in range(1024, 1056): + capacity = history + spec.max_draft_len - 1 + spec.tokens_per_gen_step + for cache in caches[1:]: + assert cache.resize(capacity, history) + # A quota that admits allocation must also allow requests to resume. + for cache in caches: + cache.suspend() + assert cache.resume(stream) + finally: + for cache in caches: + cache.close() + manager.shutdown() + + @pytest.fixture def max_num_turns() -> int: return 1 diff --git a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py index 078e528225ae..e292e1c2224d 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py @@ -1797,6 +1797,7 @@ def test_v2_hybrid_warns_when_avg_seq_len_is_missing(monkeypatch): def test_v2_hybrid_rejects_quota_below_live_state_floor(): mgr = object.__new__(MambaHybridCacheManagerV2) + mgr._generation_kv_capacity_headroom = 1 mgr._has_cp_helix = False mgr.max_batch_size = 2 mgr.mapping = Mapping(world_size=1, rank=0, tp_size=1, pp_size=1) @@ -1827,6 +1828,7 @@ def test_v2_hybrid_rejects_quota_below_live_state_floor(): def test_v2_hybrid_pure_mamba_rank_does_not_reserve_attention_page(): mgr = object.__new__(MambaHybridCacheManagerV2) + mgr._generation_kv_capacity_headroom = 1 mgr._has_cp_helix = False mgr.max_batch_size = 2 mgr.mapping = Mapping(world_size=1, rank=0, tp_size=1, pp_size=1) @@ -2054,6 +2056,7 @@ def test_expect_snapshot_points_binding_round_trip(): def test_v2_hybrid_pool_ratio_controls_allocated_memory(): def allocated_memory(pool_ratio): mgr = object.__new__(MambaHybridCacheManagerV2) + mgr._generation_kv_capacity_headroom = 1 mgr._has_cp_helix = False mgr.kv_cache_type = CacheTypeCpp.SELF mgr.head_dim_per_layer = [64, 64] diff --git a/tests/unittest/tools/test_layer_wise_benchmarks.py b/tests/unittest/tools/test_layer_wise_benchmarks.py index d01478a5e3e9..e75eb1a3d07c 100644 --- a/tests/unittest/tools/test_layer_wise_benchmarks.py +++ b/tests/unittest/tools/test_layer_wise_benchmarks.py @@ -277,6 +277,7 @@ def test_nemotron_gen_dep(llm_root, world_size): model_root / "NVIDIA-Nemotron-3-Nano-30B-A3B-BF16", "--layer-indices=4,5,6", "--mamba-ssm-cache-dtype=float16", + "--use-kv-cache-manager-v2", ], cwd=llm_root / "examples" / "layer_wise_benchmarks", env={ @@ -308,6 +309,7 @@ def test_qwen3_next_gen_tep(llm_root, world_size): "--no-enable-attention-dp", "--mamba-ssm-cache-dtype=float16", "--moe-backend=TRTLLM", + "--use-kv-cache-manager-v2", ], cwd=llm_root / "examples" / "layer_wise_benchmarks", env={ From 64fbd92b1071d2df5222c97a83ef5932bb1308f7 Mon Sep 17 00:00:00 2001 From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> Date: Tue, 15 Sep 2026 09:49:56 +0800 Subject: [PATCH 3/8] [None][feat] support disaggregated serving for Qwen3.8-Flash-Next (#18921) Signed-off-by: Bo Deng Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> --- docs/source/_static/config_db.json | 16 + ...-guide-for-qwen3.8-flash-next-on-trtllm.md | 409 +++++++++++++++++ docs/source/deployment-guide/index.rst | 1 + docs/source/models/supported-models.md | 2 +- examples/configs/curated/lookup.yaml | 17 + .../qwen3.8-flash-next-disagg-ctx.yaml | 18 + .../qwen3.8-flash-next-disagg-gen.yaml | 19 + .../configs/curated/qwen3.8-flash-next.yaml | 19 + scripts/generate_config_table.py | 4 + .../backends/sparse/qsa/cache_manager.py | 13 +- .../native/mixers/attention/peer.py | 38 +- .../disaggregation/native/mixers/ssm/peer.py | 208 ++++++--- .../_torch/disaggregation/native/peer.py | 41 +- .../_torch/disaggregation/native/transfer.py | 5 + .../disaggregation/resource/kv_extractor.py | 414 +++++++++--------- .../_torch/disaggregation/resource/page.py | 86 +++- .../_torch/disaggregation/resource/utils.py | 47 +- .../_torch/disaggregation/transceiver.py | 13 +- .../kv_cache/kv_cache_manager_v2.py | 13 +- .../kv_cache/mamba_cache_manager.py | 39 ++ .../accuracy/test_disaggregated_serving.py | 118 +++++ .../defs/accuracy/test_llm_api_pytorch.py | 55 ++- .../test_lists/qa/llm_function_core.txt | 4 + .../test_lists/test-db/l0_gb300.yml | 1 + .../test-db/l0_gb300_multi_gpus.yml | 2 + .../kv_cache/test_mamba_cache_manager.py | 9 +- .../_torch/modeling/test_qwen4_exp_support.py | 28 +- tests/unittest/disaggregated/test_bounce.py | 4 +- .../unittest/disaggregated/test_extractor.py | 248 +++++++++-- .../disaggregated/test_kda_mamba_transfer.py | 12 +- .../disaggregated/test_mamba_transfer.py | 229 +++++++++- tests/unittest/disaggregated/test_peer.py | 112 ++++- 32 files changed, 1840 insertions(+), 404 deletions(-) create mode 100644 docs/source/deployment-guide/deployment-guide-for-qwen3.8-flash-next-on-trtllm.md create mode 100644 examples/configs/curated/qwen3.8-flash-next-disagg-ctx.yaml create mode 100644 examples/configs/curated/qwen3.8-flash-next-disagg-gen.yaml create mode 100644 examples/configs/curated/qwen3.8-flash-next.yaml diff --git a/docs/source/_static/config_db.json b/docs/source/_static/config_db.json index feb0d3ec0383..8f9f4ebec38e 100644 --- a/docs/source/_static/config_db.json +++ b/docs/source/_static/config_db.json @@ -84,6 +84,18 @@ "model_url": "https://huggingface.co/Qwen/Qwen3.8-2.4T-A95B-FP8", "scenario": "Max Throughput (Static EPLB, MTP3)" }, + { + "command": "trtllm-serve Qwen/Qwen3.8-Flash-Next-FP8 --config ${TRTLLM_DIR}/examples/configs/curated/qwen3.8-flash-next.yaml", + "config_filename": "qwen3.8-flash-next.yaml", + "config_github_url": "https://github.com/NVIDIA/TensorRT-LLM/blob/main/examples/configs/curated/qwen3.8-flash-next.yaml", + "config_path": "examples/configs/curated/qwen3.8-flash-next.yaml", + "config_raw_url": "https://raw.githubusercontent.com/NVIDIA/TensorRT-LLM/main/examples/configs/curated/qwen3.8-flash-next.yaml", + "gpu_compatibility": "B200, B300, GB200, GB300", + "model": "Qwen/Qwen3.8-Flash-Next-FP8", + "model_display_name": "Qwen3.8-Flash-Next (FP8)", + "model_url": "https://huggingface.co/Qwen/Qwen3.8-Flash-Next-FP8", + "scenario": "General deployment (MTP3)" + }, { "command": "trtllm-serve nvidia/Qwen3.5-397B-A17B-NVFP4 --config ${TRTLLM_DIR}/examples/configs/curated/qwen3.5.yaml", "config_filename": "qwen3.5.yaml", @@ -4278,6 +4290,10 @@ "display_name": "Qwen3.8-2.4T-A95B (FP8)", "url": "https://huggingface.co/Qwen/Qwen3.8-2.4T-A95B-FP8" }, + "Qwen/Qwen3.8-Flash-Next-FP8": { + "display_name": "Qwen3.8-Flash-Next (FP8)", + "url": "https://huggingface.co/Qwen/Qwen3.8-Flash-Next-FP8" + }, "deepseek-ai/DeepSeek-R1-0528": { "display_name": "DeepSeek-R1", "url": "https://huggingface.co/deepseek-ai/DeepSeek-R1-0528" diff --git a/docs/source/deployment-guide/deployment-guide-for-qwen3.8-flash-next-on-trtllm.md b/docs/source/deployment-guide/deployment-guide-for-qwen3.8-flash-next-on-trtllm.md new file mode 100644 index 000000000000..22d0d2b6da2d --- /dev/null +++ b/docs/source/deployment-guide/deployment-guide-for-qwen3.8-flash-next-on-trtllm.md @@ -0,0 +1,409 @@ +# Deployment Guide for Qwen3.8-Flash-Next on TensorRT LLM - Blackwell Hardware + +## Introduction + +This guide describes how to serve Qwen3.8-Flash-Next with the TensorRT LLM PyTorch backend. It covers the BF16, block-scaled FP8, and NVFP4 checkpoints, aggregated and disaggregated serving, MTP speculative decoding, and image-language input. + +The Hugging Face checkpoint is a composite vision-language model registered as `Qwen4ExpForConditionalGeneration`. Language-only serving flattens it to the `qwen4_exp_text` decoder, which TensorRT LLM registers as `Qwen4ExpForCausalLM`. The same decoder serves both modes; the vision tower is loaded only when the server accepts image input. + +The decoder is a hybrid model rather than a plain MoE transformer. It interleaves Gated DeltaNet (GDN) linear-attention layers with QSA sparse full-attention layers, routes tokens through 512 experts with top-10 routing plus a shared expert, mixes several residual streams with Hyper-Connections, and adds a PLE n-gram side path. One recurrent MTP module ships with the checkpoint. + +## Support Status + +The BF16, block-FP8, and NVFP4 checkpoints support the same feature set: + +| Capability | Status | +|---|---| +| Aggregated and disaggregated serving | Supported | +| MTP speculative decoding with draft length 3 | Supported | +| Image-language input | Supported | +| Tensor, expert, and attention data parallelism (TP, EP, ADP) | Supported | +| Chunked prefill | Supported | +| FP8 KV cache | Supported | +| Prefix caching (block reuse) | Supported | +| KV cache manager V2 | Required, and selected by default | + +The KV-cache dtype and the GDN recurrent-state dtype are selected independently of the model weight precision, through `kv_cache_config.dtype` and `kv_cache_config.mamba_ssm_cache_dtype`. + +## Prerequisites + +* GPU: NVIDIA Blackwell architecture. The BF16 checkpoint also runs on Hopper. +* OS: Linux +* Drivers: CUDA Driver 575 or later +* Docker with NVIDIA Container Toolkit installed +* Python3 and python3-pip (optional, for accuracy evaluation only) + +## Models + +* [Qwen/Qwen3.8-Flash-Next](https://huggingface.co/Qwen/Qwen3.8-Flash-Next) (base, BF16) +* [Qwen/Qwen3.8-Flash-Next-FP8](https://huggingface.co/Qwen/Qwen3.8-Flash-Next-FP8) (block-scaled FP8 routed experts) +* [nvidia/Qwen3.8-Flash-Next-NVFP4](https://huggingface.co/nvidia/Qwen3.8-Flash-Next-NVFP4) (mixed-precision NVFP4) + +The NVFP4 checkpoint is a mixed-precision export: its routed experts are NVFP4, while the PLE n-gram table and the MTP experts are FP8. TensorRT LLM therefore reports `MIXED_PRECISION` rather than `NVFP4` for it. The per-layer quantization metadata is part of the checkpoint contract, so do not override it with a single global quantization algorithm. + +The FP8 checkpoint is pre-quantized. It is not produced by enabling an FP8 option while loading the BF16 checkpoint. Its routed-expert projections carry FP8 weights with 128x128 block scales, while attention, GDN projections, the shared expert, routers, embeddings, and the LM head keep their checkpoint dtypes. + +## GPU Requirements + +The following table lists the minimum resources for each checkpoint. All three configurations enable PLE host offload, which is what makes the host-memory floor part of the requirement. + +| Checkpoint | Platform | Minimum GPUs | GPU memory per GPU | Host memory | +|---|---|---:|---:|---:| +| BF16 | Hopper or Blackwell | 2 | 142 GB | 128 GiB | +| Block-FP8 | Blackwell | 1 | 145 GB | 96 GiB | +| NVFP4 | Blackwell | 1 | 100 GB | 96 GiB | + +## Deployment Steps + +### Run Docker Container + +Run the docker container using the TensorRT LLM NVIDIA NGC image. + +```shell +docker run --rm -it \ +--ipc=host \ +--gpus all \ +-p 8000:8000 \ +-v ~/.cache:/root/.cache:rw \ +--name tensorrt_llm \ +nvcr.io/nvidia/tensorrt-llm/release:x.y.z \ +/bin/bash +``` + +Note: + +* The command mounts your user `.cache` directory to save the downloaded model checkpoints which are saved to `~/.cache/huggingface/hub/` by default. This prevents having to redownload the weights each time you rerun the container. If the `~/.cache` directory doesn't exist please create it using `$ mkdir ~/.cache`. +* You can mount additional directories and paths using the `-v :` flag if needed, such as mounting the downloaded weight paths. +* The command also maps port `8000` from the container to your host so you can access the LLM API endpoint from your host. +* See the for all the available containers. The containers published in the main branch weekly have `rcN` suffix, while the monthly release with QA tests has no `rcN` suffix. Use the `rc` release to get the latest model and feature support. + +If you want to use latest main branch, you can choose to build from source to install TensorRT LLM, the steps refer to [https://nvidia.github.io/TensorRT-LLM/latest/installation/build-from-source.html](https://nvidia.github.io/TensorRT-LLM/latest/installation/build-from-source.html). + +### PLE Table Host Offload + +The PLE n-gram table is a large weight that is read once per token. Offloading it to pinned host memory frees device memory for the KV cache at the cost of host memory and a host-to-device read. Enable it in the environment of every worker before starting the server: + +```shell +export TRTLLM_QWEN4_EXP_PLE_HOST_OFFLOAD=1 +``` + +This is an environment-only option; there is no YAML field for it. It moves the table only. The per-request PLE short-convolution state and n-gram context stay in the cache manager on the device, and remain part of every disaggregated handoff. + +### Recommended Performance Settings + +We maintain YAML configuration files with recommended performance settings in the [`examples/configs`](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/configs) directory. These config files are present in the TensorRT LLM container at the path `/app/tensorrt_llm/examples/configs`. You can use these out-of-the-box, or adjust them to your specific use case. + +Set the TensorRT LLM directory and select one of the configuration files listed below: + +```shell +TRTLLM_DIR=/app/tensorrt_llm # change as needed to match your environment +EXTRA_LLM_API_FILE=${TRTLLM_DIR}/examples/configs/curated/qwen3.8-flash-next.yaml +``` + +| Deployment | Configuration file | +|---|---| +| Aggregated serving | [`qwen3.8-flash-next.yaml`](../../../examples/configs/curated/qwen3.8-flash-next.yaml) | +| Disaggregated context worker | [`qwen3.8-flash-next-disagg-ctx.yaml`](../../../examples/configs/curated/qwen3.8-flash-next-disagg-ctx.yaml) | +| Disaggregated generation worker | [`qwen3.8-flash-next-disagg-gen.yaml`](../../../examples/configs/curated/qwen3.8-flash-next-disagg-gen.yaml) | + +`Qwen4ExpForCausalLM` already selects QSA sparse attention and KV cache manager V2, and already disables block reuse, so these files set only what differs from those defaults. They target the block-FP8 and NVFP4 checkpoints on one GPU; size `max_batch_size`, `max_num_tokens`, and `max_seq_len` for your workload before treating them as a performance baseline. + +To adapt them: + +* **BF16**: set `tensor_parallel_size: 2` and `moe_config.backend: CUTLASS`. +* **Four GPUs**: set `tensor_parallel_size: 4` and `moe_expert_parallel_size: 4`, and leave `moe_tensor_parallel_size` at `1` so each expert keeps its complete intermediate dimension. Add `enable_attention_dp: true` with `enable_lm_head_tp_in_adp: true` for the throughput-oriented layout. +* **No speculative decoding**: remove the `speculative_config` block. + +#### Non-greedy MTP + +The `speculative_config` block in those files is correct for greedy decoding. For requests with a nonzero temperature or top-p, add rejection sampling so the sampled distribution stays correct: + +```yaml +speculative_config: + decoding_type: MTP + max_draft_len: 3 + use_rejection_sampling: true + advanced_sampling_mode: full +``` + +`advanced_sampling_mode: full` keeps per-request top-k and top-p filtering. The runtime detects non-greedy requests automatically. Acceptance depends on the prompt distribution, sampling parameters, checkpoint, and batch shape; a rate measured on one workload is not a model constant. + +#### Prefix caching + +Attention KV blocks alone cannot restore GDN and PLE state, so the model disables block reuse by default. Enable it together with a snapshot policy: + +```yaml +kv_cache_config: + enable_block_reuse: true + enable_partial_reuse: true + copy_on_partial_reuse: true + mamba_state_config: + periodic_snapshot_interval: 256 +``` + +Block reuse stays on only if `mamba_state_config` sets at least one snapshot placement: `periodic_snapshot_interval`, `additional_snapshot_offsets_from_start`, or `additional_snapshot_offsets_from_end`. `enable_branch_snapshot` refines where snapshots land but does not by itself keep reuse enabled. The same settings apply in disaggregated serving; configure them on both workers. A snapshot restores the PLE short-convolution state and n-gram context along with the GDN state, because they share one recurrent page. Check the effective LLM arguments and the per-request `num_reused_blocks` metric rather than assuming the submitted setting took effect. + +### Launch the TensorRT LLM Server + +```shell +trtllm-serve \ + --host 0.0.0.0 --port 8000 \ + --reasoning_parser qwen3_5 \ + --tool_parser qwen3 \ + --config ${EXTRA_LLM_API_FILE} +``` + +The chat template pre-injects a `` block, so reasoning starts at the beginning of the response and the `qwen3_5` reasoning parser applies. This architecture is not in the parser auto-detection table, so pass both parsers explicitly. Thinking is controlled per request through `chat_template_kwargs`, for example `{"chat_template_kwargs": {"enable_thinking": false}}` to answer directly, or `{"chat_template_kwargs": {"enable_thinking": true, "reasoning_effort": "xhigh"}}` for a longer trace. + +### Disaggregated Serving + +Disaggregated serving separates prefill (context) and decode (generation) onto different workers. For this model the transferred request state is larger than an attention KV cache: it also carries the QSA index state, the GDN recurrent state, and the PLE n-gram context and short-convolution state. Only the Python NIXL transceiver can move that state, and this model does not declare a preferred runtime, so both workers must set `cache_transceiver_config.transceiver_runtime` to `PYTHON` explicitly. Leaving it at its `auto` default selects the C++ transceiver, and KV cache manager V2 then rejects the configuration rather than serving wrong results. + +Use [`qwen3.8-flash-next-disagg-ctx.yaml`](../../../examples/configs/curated/qwen3.8-flash-next-disagg-ctx.yaml) and [`qwen3.8-flash-next-disagg-gen.yaml`](../../../examples/configs/curated/qwen3.8-flash-next-disagg-gen.yaml) for the two workers. Both must load the same checkpoint revision, use the same parallel layout, and use the same PLE host-offload setting. Add `speculative_config` to the generation worker only. For long generations, raise `cache_transceiver_config.kv_transfer_timeout_ms` above the expected request duration. + +The orchestrator is launched with a disaggregated config that lists the worker URLs: + +```yaml +hostname: localhost +port: 8000 +backend: pytorch +context_servers: + num_instances: 1 + urls: + - "localhost:8001" +generation_servers: + num_instances: 1 + urls: + - "localhost:8002" +``` + +```bash +# Start each worker first, and wait for both /health endpoints to return 200. +trtllm-serve --host 0.0.0.0 --port 8001 --config ${TRTLLM_DIR}/examples/configs/curated/qwen3.8-flash-next-disagg-ctx.yaml +trtllm-serve --host 0.0.0.0 --port 8002 --config ${TRTLLM_DIR}/examples/configs/curated/qwen3.8-flash-next-disagg-gen.yaml + +trtllm-serve disaggregated -c disagg_config.yaml +``` + +Clients then send OpenAI-compatible requests to the orchestrator on port `8000`. To restrict the worker endpoints to requests coming from the orchestrator, set the same `internal_request_auth_key` value in both worker configs and in the orchestrator config. For the full walkthrough, per-worker GPU placement, and multi-node or SLURM launch, see the [Disaggregated Serving guide](../features/disagg-serving.md). + +### Image-Language Serving + +The same configurations serve image input; the vision tower is loaded unless `disable_mm_encoder: true` is set. For image workloads, size the encoder separately from the decoder: + +```yaml +encoder_max_batch_size: 8 +encoder_max_num_tokens: 65536 +``` + +Image input is supported in aggregated serving with a local encoder. A separate encoder-to-prefill handoff is not part of this guide. + +### Key Configuration Options + +These options control TensorRT LLM behavior and are set in the YAML file passed to `trtllm-serve` with the `--config` argument. + +#### `tensor_parallel_size` + +Sets the tensor-parallel size for attention, GDN, and PLE layers. This should typically match the number of GPUs used by one model instance. + +#### `moe_tensor_parallel_size` + +Sets the tensor-parallel size for routed experts, independently of `tensor_parallel_size`. Keep it at `1` and use expert parallelism instead: the routed-expert intermediate dimension is five 128-element blocks, so a pure MoE tensor-parallel split of 2 or 4 would cut FP8 scale blocks. + +#### `moe_expert_parallel_size` + +Sets the expert-parallel size for MoE layers. Use it together with `moe_tensor_parallel_size: 1` for multi-GPU deployments. + +#### `enable_attention_dp` + +Runs attention and linear-attention layers data-parallel while the MoE layers stay expert-parallel. This is the throughput-oriented layout. Pair it with `enable_lm_head_tp_in_adp: true` to keep the LM head tensor-parallel. + +#### `kv_cache_config.free_gpu_memory_fraction` + +Fraction of free GPU memory reserved for the KV cache and recurrent state after the model is loaded. Reduce it if initialization reports an out-of-memory error. + +#### `kv_cache_config.mamba_ssm_cache_dtype` + +Selects the GDN recurrent-state dtype independently of the attention KV-cache dtype. Supported values are `auto`, `float16`, `bfloat16`, and `float32`. + +#### `kv_cache_config.mamba_state_config.periodic_snapshot_interval` + +Number of tokens between recurrent-state snapshots in the prefix cache. Snapshots at a fixed interval suit a workload whose shared prefixes vary in length; see [Prefix caching](#prefix-caching). + +#### `kv_cache_config.mamba_state_config.additional_snapshot_offsets_from_start`, `..._from_end` + +Snapshot the recurrent state at fixed token offsets measured from the start or the end of each prompt, instead of, or in addition to, a periodic interval. An offset of `0` from the end selects the prompt end, which suits multi-turn workloads that reuse a whole previous turn. Offsets that fall outside a prompt are ignored. + +#### `max_batch_size`, `max_num_tokens`, `max_seq_len` + +Set the maximum number of requests per scheduled batch, the maximum total tokens per scheduled batch, and the maximum length of a single request including generated tokens. + +#### `cuda_graph_config` + +Controls CUDA graph capture and padding. The model defaults to eager execution, so set this explicitly: + +* `enable_padding`: Pads input batches to a captured CUDA graph batch size. +* `max_batch_size`: Largest batch size for which graphs are captured. Set it to match `max_batch_size`. + +#### `moe_config` + +Controls MoE execution: + +* `backend`: Selects the MoE backend. Use `CUTLASS` for BF16 and `TRTLLM` for the FP8 and NVFP4 checkpoints; `DEEPGEMM` is an alternative for block-scaled FP8 on SM100 or later. +* `max_num_tokens`: Limits the tokens processed by one fused MoE invocation before chunking. +* `disable_finalize_fusion`: Uses the separate deterministic finalization path on the CUTLASS BF16 backend. The fused path is numerically valid but not bitwise reproducible, so set this only when exact run-to-run replay is required. + +#### `speculative_config` + +Configures MTP speculative decoding: + +* `decoding_type`: Set to `MTP`. +* `max_draft_len`: Draft length; the checkpoint's recurrent MTP module supports up to `3`. +* `use_rejection_sampling` and `advanced_sampling_mode`: See [Non-greedy MTP](#non-greedy-mtp). + +#### `cache_transceiver_config` + +Configures the disaggregated state transfer. This model requires `backend: NIXL` with `transceiver_runtime: PYTHON`; see [Disaggregated Serving](#disaggregated-serving). + +#### `trust_remote_code` + +Allows Hugging Face to load custom model and tokenizer code from the model repository. Enable it only for trusted model sources. + +See the [`TorchLlmArgs` API reference](https://nvidia.github.io/TensorRT-LLM/llm-api/reference.html#tensorrt_llm.llmapi.TorchLlmArgs) for the complete configuration schema, [KV cache documentation](../features/kvcache.md) for hybrid-state cache settings, and [speculative decoding documentation](../features/speculative-decoding.md) for MTP details. + +## Testing API Endpoint + +### Health Check + +Start a new terminal on the host to test the TensorRT LLM server you just launched. + +```shell +curl -s -o /dev/null -w "Status: %{http_code}\n" "http://localhost:8000/health" +``` + +When the `Status: 200` code is returned, the server is ready for queries. The very first query may take longer due to initialization and compilation. + +### Basic Test + +```shell +curl http://localhost:8000/v1/chat/completions -H "Content-Type: application/json" -d '{ + "model": "", + "messages": [ + { + "role": "user", + "content": "Where is New York?" + } + ], + "max_tokens": 1024, + "top_p": 1.0 +}' -w "\n" +``` + +### Image Request + +An OpenAI-compatible image request uses an `image_url` content part followed by the text instruction. Add more `image_url` parts for a multi-image request and preserve their order: + +```shell +curl http://localhost:8000/v1/chat/completions -H "Content-Type: application/json" -d '{ + "model": "", + "messages": [{ + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"url": "https://example.com/image.jpg"}}, + {"type": "text", "text": "Describe the image briefly."} + ] + }], + "max_tokens": 128 +}' -w "\n" +``` + +The image URL must be reachable from the serving frontend; a supported data URL may be used instead. + +## Running Evaluations to Verify Accuracy (Optional) + +`trtllm-eval` runs the same tasks as the in-tree accuracy suite: + +```shell +export TRTLLM_QWEN4_EXP_PLE_HOST_OFFLOAD=1 +trtllm-eval --model \ + --config config.yaml \ + gsm8k +``` + +The reference scores committed with the model are: + +| Checkpoint | GSM8K | MMLU | MMMU | +|---|--:|--:|--:| +| BF16 | 95.72 | 87.82 | — | +| Block-FP8 with MTP3 | 95.11 | 87.91 | — | +| NVFP4 with MTP3 | 95.57 | 87.57 | 59.22 | + +The GSM8K and MMLU runs use the chat template with `enable_thinking` on and off respectively; see `TestQwen3_8_Flash_Next` in `tests/integration/defs/accuracy/test_llm_api_pytorch.py` for the exact evaluator settings. Small differences of roughly half a point are expected across checkpoint and dependency revisions. + +## Benchmarking Performance + +To benchmark the performance of your TensorRT LLM server you can leverage the built-in `benchmark_serving.py` script. To do this, first create a wrapper `bench.sh` script. + +```shell +cat <<'EOF' > bench.sh +#!/usr/bin/env bash +set -euo pipefail + +MODEL_NAME="" + +concurrency_list="1 2 4 8 16 32 64 128 256" +multi_round=5 +isl=8192 +osl=1024 +result_dir=/tmp/qwen3.8_flash_next_output + +for concurrency in ${concurrency_list}; do + num_prompts=$((concurrency * multi_round)) + python -m tensorrt_llm.serve.scripts.benchmark_serving \ + --model ${MODEL_NAME} \ + --backend openai \ + --dataset-name "random" \ + --random-input-len ${isl} \ + --random-output-len ${osl} \ + --random-prefix-len 0 \ + --random-ids \ + --num-prompts ${num_prompts} \ + --max-concurrency ${concurrency} \ + --ignore-eos \ + --tokenize-on-client \ + --percentile-metrics "ttft,tpot,itl,e2el" +done +EOF +chmod +x bench.sh +``` + +To achieve max throughput, with attention DP on, one needs to sweep up to `concurrency = max_batch_size * num_gpus`. + +If you want to save the results to a file add the following options. + +```shell +--save-result \ +--result-dir "${result_dir}" \ +--result-filename "concurrency_${concurrency}.json" +``` + +For more benchmarking options see [benchmark_serving.py](https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/serve/scripts/benchmark_serving.py) + +Run `bench.sh` to begin a serving benchmark. This will take a long time if you run all the concurrencies mentioned in the above `bench.sh` script. + +```shell +./bench.sh +``` + +Complete initialization, CUDA graph capture, and warmup before measuring; the first requests after startup are not representative of steady state. + +## Troubleshooting Tips + +* If you encounter CUDA out-of-memory errors, try reducing `max_batch_size`, `max_num_tokens`, or `kv_cache_config.free_gpu_memory_fraction`. If the error occurs during CUDA graph capture, also reduce `cuda_graph_config.max_batch_size`. Enabling `TRTLLM_QWEN4_EXP_PLE_HOST_OFFLOAD=1` frees additional device memory at the cost of host memory. +* If weight loading fails on the FP8 or NVFP4 checkpoint after a multi-GPU change, check that `moe_tensor_parallel_size` is `1`. A pure MoE tensor-parallel split of 2 or 4 cuts the routed experts' quantization blocks. +* If block reuse appears to have no effect, check the effective LLM arguments for `enable_block_reuse`. The runtime turns it off when `kv_cache_config.mamba_state_config` configures no snapshot placement, because the recurrent state cannot be restored from attention blocks alone. +* If a disaggregated worker fails to start with a KV cache manager V2 error, check that both workers set `cache_transceiver_config.backend: NIXL` and `transceiver_runtime: PYTHON`. +* If MTP output is empty or incorrect, confirm that the checkpoint contains its MTP weights and that `max_draft_len` is configured identically on every rank. +* If reasoning content is not separated from the answer, confirm that the server was started with `--reasoning_parser qwen3_5`; this architecture is not auto-detected. +* If the container fails to start, verify that the NVIDIA Container Toolkit is properly installed. +* For connection issues, make sure the server port (`8000` in this guide) is not being used by another application. diff --git a/docs/source/deployment-guide/index.rst b/docs/source/deployment-guide/index.rst index 0ac223aa1f3a..aa946acb3a1c 100644 --- a/docs/source/deployment-guide/index.rst +++ b/docs/source/deployment-guide/index.rst @@ -37,6 +37,7 @@ The deployment guides below provide more detailed instructions for serving speci deployment-guide-for-gpt-oss-on-trtllm.md deployment-guide-for-qwen3-on-trtllm.md deployment-guide-for-qwen3.8-qwen3.5-on-trtllm.md + deployment-guide-for-qwen3.8-flash-next-on-trtllm.md deployment-guide-for-kimi-k2-thinking-on-trtllm.md deployment-guide-for-kimi-k3-on-trtllm.md deployment-guide-for-glm-5-on-trtllm.md diff --git a/docs/source/models/supported-models.md b/docs/source/models/supported-models.md index b9588323b1af..150b4bda0bfa 100644 --- a/docs/source/models/supported-models.md +++ b/docs/source/models/supported-models.md @@ -82,7 +82,7 @@ Note: Support for other models may vary. Features marked "N/A" are not applicabl | `Qwen3MoeForCausalLM` | Yes | Yes | Yes | Yes | Yes | EAGLE-3 (Linear, Dynamic) | Yes | Yes | Yes | N/A | Yes | Yes | | `Qwen3NextForCausalLM` [^3] | Yes | Yes | Yes | Untested | Yes | No | Yes | Yes | No | No | Untested | Yes | | `Qwen3_5MoeForCausalLM` | Yes | Yes | Yes | Yes | Yes | MTP | Yes | Untested | Yes | N/A | Untested | Yes | -| `Qwen4ExpForCausalLM` | Yes | Yes | Yes | No | Yes | MTP | Yes | Untested | No | N/A | Yes | Yes | +| `Qwen4ExpForCausalLM` | Yes | Yes | Yes | Yes | Yes | MTP | Yes | Untested | No | N/A | Yes | Yes | | `Llama4ForConditionalGeneration` | Yes | Yes | Yes | Yes | Yes | EAGLE-3 (Linear, Dynamic) | Yes | Yes | Untested | N/A | Yes | Yes | | `GptOssForCausalLM` | Yes | Yes | Yes | Yes | Yes | EAGLE-3 (Linear), DFlash | Yes | Yes | Yes | N/A | Yes | Yes | | `KimiK3ForConditionalGeneration` [^15] [^17] | Yes | Yes | Yes | Yes | Yes | DSpark | Yes | No | Yes | N/A | Yes | Yes | diff --git a/examples/configs/curated/lookup.yaml b/examples/configs/curated/lookup.yaml index 5f2a7cd9536c..da0789bc22e5 100644 --- a/examples/configs/curated/lookup.yaml +++ b/examples/configs/curated/lookup.yaml @@ -34,6 +34,23 @@ config_path: examples/configs/curated/qwen3.8-high-throughput-mtp3.yaml scenario: Max Throughput (Static EPLB, MTP3) gpu_compatibility: "B200, B300, GB200, GB300" +- model: Qwen/Qwen3.8-Flash-Next-FP8 + arch: Qwen4ExpForCausalLM + config_path: examples/configs/curated/qwen3.8-flash-next.yaml + scenario: General deployment (MTP3) + gpu_compatibility: "B200, B300, GB200, GB300" +- model: Qwen/Qwen3.8-Flash-Next-FP8 + arch: Qwen4ExpForCausalLM + config_path: examples/configs/curated/qwen3.8-flash-next-disagg-ctx.yaml + scenario: Disaggregated Context + gpu_compatibility: "B200, B300, GB200, GB300" + disagg: true +- model: Qwen/Qwen3.8-Flash-Next-FP8 + arch: Qwen4ExpForCausalLM + config_path: examples/configs/curated/qwen3.8-flash-next-disagg-gen.yaml + scenario: Disaggregated Generation + gpu_compatibility: "B200, B300, GB200, GB300" + disagg: true - model: nvidia/Qwen3.5-397B-A17B-NVFP4 arch: Qwen3_5MoeForCausalLM config_path: examples/configs/curated/qwen3.5.yaml diff --git a/examples/configs/curated/qwen3.8-flash-next-disagg-ctx.yaml b/examples/configs/curated/qwen3.8-flash-next-disagg-ctx.yaml new file mode 100644 index 000000000000..c1e0ec8f8c46 --- /dev/null +++ b/examples/configs/curated/qwen3.8-flash-next-disagg-ctx.yaml @@ -0,0 +1,18 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# Context (prefill) worker. Pair with qwen3.8-flash-next-disagg-gen.yaml. +# The handoff carries recurrent and PLE state, so both workers must set +# cache_transceiver_config.transceiver_runtime to PYTHON. +trust_remote_code: true +max_batch_size: 16 +disable_overlap_scheduler: true +moe_config: + backend: TRTLLM +kv_cache_config: + free_gpu_memory_fraction: 0.5 + mamba_ssm_cache_dtype: bfloat16 +cache_transceiver_config: + backend: NIXL + transceiver_runtime: PYTHON + max_tokens_in_buffer: 8192 diff --git a/examples/configs/curated/qwen3.8-flash-next-disagg-gen.yaml b/examples/configs/curated/qwen3.8-flash-next-disagg-gen.yaml new file mode 100644 index 000000000000..806edf546aba --- /dev/null +++ b/examples/configs/curated/qwen3.8-flash-next-disagg-gen.yaml @@ -0,0 +1,19 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# Generation (decode) worker. Pair with qwen3.8-flash-next-disagg-ctx.yaml. +# Add speculative_config here, not on the context worker, to run MTP. +trust_remote_code: true +max_batch_size: 16 +cuda_graph_config: + enable_padding: true + max_batch_size: 16 +moe_config: + backend: TRTLLM +kv_cache_config: + free_gpu_memory_fraction: 0.5 + mamba_ssm_cache_dtype: bfloat16 +cache_transceiver_config: + backend: NIXL + transceiver_runtime: PYTHON + max_tokens_in_buffer: 8192 diff --git a/examples/configs/curated/qwen3.8-flash-next.yaml b/examples/configs/curated/qwen3.8-flash-next.yaml new file mode 100644 index 000000000000..7aa43898978a --- /dev/null +++ b/examples/configs/curated/qwen3.8-flash-next.yaml @@ -0,0 +1,19 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# Aggregated serving on one Blackwell GPU, block-FP8 or NVFP4 checkpoint. +# For the BF16 checkpoint use tensor_parallel_size: 2 and moe_config.backend: CUTLASS. +trust_remote_code: true +max_batch_size: 16 +enable_chunked_prefill: true +cuda_graph_config: + enable_padding: true + max_batch_size: 16 +moe_config: + backend: TRTLLM +kv_cache_config: + free_gpu_memory_fraction: 0.5 + mamba_ssm_cache_dtype: bfloat16 +speculative_config: + decoding_type: MTP + max_draft_len: 3 diff --git a/scripts/generate_config_table.py b/scripts/generate_config_table.py index 7539d4b5be81..339428e96e96 100644 --- a/scripts/generate_config_table.py +++ b/scripts/generate_config_table.py @@ -78,6 +78,10 @@ "display_name": "Qwen3.8-2.4T-A95B (FP8)", "url": "https://huggingface.co/Qwen/Qwen3.8-2.4T-A95B-FP8", }, + "Qwen/Qwen3.8-Flash-Next-FP8": { + "display_name": "Qwen3.8-Flash-Next (FP8)", + "url": "https://huggingface.co/Qwen/Qwen3.8-Flash-Next-FP8", + }, "nvidia/Qwen3.5-397B-A17B-NVFP4": { "display_name": "Qwen3.5-397B-A17B (NVFP4)", "url": "https://huggingface.co/nvidia/Qwen3.5-397B-A17B-NVFP4", diff --git a/tensorrt_llm/_torch/attention/backends/sparse/qsa/cache_manager.py b/tensorrt_llm/_torch/attention/backends/sparse/qsa/cache_manager.py index 7029f111cfdf..a70c3aa0aa3b 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/qsa/cache_manager.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/qsa/cache_manager.py @@ -2,10 +2,11 @@ # SPDX-License-Identifier: Apache-2.0 """V2 hybrid cache manager for QSA sparse attention.""" -from typing import TYPE_CHECKING, Optional +from typing import TYPE_CHECKING, Dict, Optional import torch +from tensorrt_llm._torch.disaggregation.resource.page import MapperKind from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import Role from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManagerV2 from tensorrt_llm._utils import TensorWrapper, binding_to_torch_dtype, convert_to_torch_tensor @@ -77,6 +78,16 @@ def __init__( self.qsa_position_layer_id: Optional[int] = None super().__init__(*args, layer_mask=layer_mask, **kwargs) + def get_disagg_role_mapper_kinds(self) -> Dict[DataRole, MapperKind]: + """Index positions are request-wide coordinates, so every rank holds + the same bytes. They carry no head axis, and the INDEXED fallback would + divide their size by the rank's KV-head count, which only cancels out + when both peers shard heads the same way.""" + return { + **super().get_disagg_role_mapper_kinds(), + QSA_INDEX_POSITION: MapperKind.REPLICATED, + } + def _extra_buffers_per_layer( self, *, diff --git a/tensorrt_llm/_torch/disaggregation/native/mixers/attention/peer.py b/tensorrt_llm/_torch/disaggregation/native/mixers/attention/peer.py index 661b6d5cccb6..f141cce93ff4 100644 --- a/tensorrt_llm/_torch/disaggregation/native/mixers/attention/peer.py +++ b/tensorrt_llm/_torch/disaggregation/native/mixers/attention/peer.py @@ -14,6 +14,7 @@ # limitations under the License. from collections.abc import Sequence +from typing import Optional import numpy as np @@ -25,7 +26,8 @@ SpecRegionPair, ) from tensorrt_llm._torch.disaggregation.native.rank_info import RankInfo -from tensorrt_llm._torch.disaggregation.resource.page import MapperKind +from tensorrt_llm._torch.disaggregation.resource.page import CacheKind, KVCachePageTable, MapperKind +from tensorrt_llm._torch.disaggregation.resource.utils import find_replicated_role_mismatch from tensorrt_llm._utils import nvtx_range @@ -481,8 +483,16 @@ class AttentionPolicy: def __init__(self, self_rank_info: RankInfo): self._ri = self_rank_info - def should_send(self, peer_overlap, peer_rank_info) -> bool: - """Attention uses head-duplication routing.""" + def should_send( + self, peer_overlap, peer_rank_info, *, mapper_kind: Optional[MapperKind] = None + ) -> "bool | None": + """Attention uses head-duplication routing. + + REPLICATED views (index-key side caches) hold identical bytes on every + TP rank, so they return None to defer to fan-in election instead. + """ + if mapper_kind == MapperKind.REPLICATED: + return None dup = peer_overlap.duplicate_head_factor if dup <= 1: return True @@ -546,6 +556,24 @@ def _tpb_check(self, local: int, peer: int, peer_ri: RankInfo) -> bool: ) return False + @staticmethod + def validate_peer_compatible( + self_page_table: Optional[KVCachePageTable], + peer_page_table: Optional[KVCachePageTable], + ) -> None: + """Reject a peer whose PAGED groups declare different replicated roles. + + Pool matching drops a view with no counterpart silently, so an + index-key side cache the peer never declares would leave the receiver + holding zeroed state instead of raising. Raises ``ValueError``. + """ + differing = find_replicated_role_mismatch(self_page_table, peer_page_table, CacheKind.PAGED) + if differing: + raise ValueError( + "AttentionPolicy.validate_peer_compatible: replicated roles differ on " + f"overlapping layers: {differing}" + ) + def check_peer_compatible(self, peer_ri: RankInfo) -> bool: a = self._ri.attention b = peer_ri.attention @@ -628,6 +656,8 @@ def build_mapper( peer_buffers_per_layer: int = 1, self_lg=None, peer_lg=None, + self_pv=None, + peer_pv=None, ) -> RegionMapperBase: """Pick the mapper for one view pair. @@ -637,7 +667,7 @@ def build_mapper( The kind only decides the two irreducible semantic differences: - REPLICATED skips head matching entirely (bytes are identical on - every TP rank; fan-in ownership is decided upstream). + every TP rank; fan-in ownership is decided by ``should_send``). - Under head mismatch, HND (INDEXED) slices one contiguous head range per K/V buffer, while NHD must slice inside every token. diff --git a/tensorrt_llm/_torch/disaggregation/native/mixers/ssm/peer.py b/tensorrt_llm/_torch/disaggregation/native/mixers/ssm/peer.py index 76951b754640..5d7c993df49c 100644 --- a/tensorrt_llm/_torch/disaggregation/native/mixers/ssm/peer.py +++ b/tensorrt_llm/_torch/disaggregation/native/mixers/ssm/peer.py @@ -28,6 +28,11 @@ KVCachePageTable, MambaLayerGroup, MapperKind, + PoolView, +) +from tensorrt_llm._torch.disaggregation.resource.utils import ( + find_replicated_role_mismatch, + get_pool_view_global_layer_ids, ) from tensorrt_llm._utils import nvtx_range @@ -343,9 +348,18 @@ class MambaPolicy: def __init__(self, self_rank_info: RankInfo): self._ri = self_rank_info - def should_send(self, peer_overlap, peer_rank_info) -> "bool | None": - """Mamba TP routing: always send (each rank owns unique sharded state). - When mamba_tp == 1 (attention_dp), returns None to signal fan-in election.""" + def should_send( + self, peer_overlap, peer_rank_info, *, mapper_kind: Optional[MapperKind] = None + ) -> "bool | None": + """Mamba TP routing. + + Sharded state (conv/SSM): always send, each rank owns a unique shard. + When mamba_tp == 1 (attention_dp) the state is replicated, so return + None to signal fan-in election. REPLICATED side state holds identical + bytes on every rank regardless of mamba_tp, so it always elects. + """ + if mapper_kind == MapperKind.REPLICATED: + return None mamba_tp, _ = MambaPolicy._mamba_tp(self._ri) if mamba_tp == 1: return None # caller should use fan-in election @@ -364,6 +378,8 @@ def build_mapper( peer_buffers_per_layer: int = 1, self_lg=None, peer_lg=None, + self_pv: Optional[PoolView] = None, + peer_pv: Optional[PoolView] = None, src_layer_off: int = 0, dst_layer_off: int = 0, ) -> RegionMapperBase: @@ -379,11 +395,29 @@ def build_mapper( ``extract_slot``. Under full PP overlap these are 0; under partial PP overlap they identify which slice of the extraction to transfer. - The ``mapper_kind`` discriminates conv (SECTIONED) from ssm (INDEXED). - TP info comes from ``self._ri`` and ``peer_ri``. Per-head / per-section - metadata comes from ``self_lg`` / ``peer_lg`` (MambaLayerGroup). + The ``mapper_kind`` discriminates conv (SECTIONED), ssm (INDEXED) and + replicated side state (REPLICATED). TP info comes from ``self._ri`` + and ``peer_ri``. Per-section / per-head resharding metadata comes + from the views themselves (``PoolView.section_bytes`` / + ``PoolView.bytes_per_head``). """ transfer_layers = len(self_layer_offsets) + + if mapper_kind == MapperKind.REPLICATED: + # Identical bytes on every rank: whole per-layer copy, no TP + # resharding. Fan-in ownership is decided by should_send. + if self_bytes_per_layer != peer_bytes_per_layer: + raise ValueError( + "Replicated state size differs between peers: " + f"local={self_bytes_per_layer}, peer={peer_bytes_per_layer}" + ) + return MambaHeadMatchMapper( + transfer_layers=transfer_layers, + src_layer_off=src_layer_off, + dst_layer_off=dst_layer_off, + block_bytes_per_layer=self_bytes_per_layer, + ) + self_mamba_tp, self_mamba_tp_rank = MambaPolicy._mamba_tp(self._ri) peer_mamba_tp, peer_mamba_tp_rank = MambaPolicy._mamba_tp(peer_ri) tp_match = self_mamba_tp == peer_mamba_tp @@ -396,30 +430,40 @@ def build_mapper( block_bytes_per_layer=self_bytes_per_layer, ) - is_conv = mapper_kind == MapperKind.SECTIONED - if is_conv: + if self_pv is None or peer_pv is None: + raise ValueError("MambaPolicy.build_mapper needs both pool views under a TP mismatch") + + if mapper_kind == MapperKind.SECTIONED: + if self_pv.section_bytes is None or peer_pv.section_bytes is None: + raise ValueError( + f"SECTIONED view {sorted(self_pv.pool_role)} lacks section_bytes on " + "one side; cannot reshard under a TP mismatch" + ) return ConvStateMismatchMapper( transfer_layers=transfer_layers, src_layer_off=src_layer_off, dst_layer_off=dst_layer_off, - self_section_bytes=self_lg.conv_section_bytes, - peer_section_bytes=peer_lg.conv_section_bytes, + self_section_bytes=self_pv.section_bytes, + peer_section_bytes=peer_pv.section_bytes, self_tp_per_dp=self_mamba_tp, peer_tp_per_dp=peer_mamba_tp, self_tp_rank=self_mamba_tp_rank, peer_tp_rank=peer_mamba_tp_rank, ) - # SSM state (INDEXED): head-level granularity - assert self_lg.ssm_bytes_per_head is not None, "ssm_bytes_per_head required for SSM mapper" - assert peer_lg.ssm_bytes_per_head is not None, "ssm_bytes_per_head required for SSM mapper" - self_nheads = self_bytes_per_layer // self_lg.ssm_bytes_per_head - peer_nheads = peer_bytes_per_layer // peer_lg.ssm_bytes_per_head + # INDEXED state: head-level granularity + if self_pv.bytes_per_head is None or peer_pv.bytes_per_head is None: + raise ValueError( + f"INDEXED view {sorted(self_pv.pool_role)} lacks bytes_per_head on one " + "side; cannot reshard under a TP mismatch" + ) + self_nheads = self_bytes_per_layer // self_pv.bytes_per_head + peer_nheads = peer_bytes_per_layer // peer_pv.bytes_per_head return MambaHeadMismatchMapper( transfer_layers=transfer_layers, src_layer_off=src_layer_off, dst_layer_off=dst_layer_off, - bytes_per_head=self_lg.ssm_bytes_per_head, + bytes_per_head=self_pv.bytes_per_head, self_nheads=self_nheads, peer_nheads=peer_nheads, self_tp_per_dp=self_mamba_tp, @@ -483,19 +527,6 @@ def validate_peer_compatible( # invariants below are per-layer-slot quantities, uniform across a # model's recurrent layers, so they apply regardless of which layers # overlap. - if ( - self_mlg.ssm_bytes_per_head is not None - and peer_mlg.ssm_bytes_per_head is not None - and self_mlg.ssm_bytes_per_head != peer_mlg.ssm_bytes_per_head - ): - # TP-invariant: head_dim * d_state * element_size. A mismatch - # means different state shape or SSM cache dtype. - raise ValueError( - "MambaPolicy.validate_peer_compatible: ssm_bytes_per_head differs " - f"(local={self_mlg.ssm_bytes_per_head}, peer={peer_mlg.ssm_bytes_per_head}); " - "check head_dim / d_state / mamba_ssm_cache_dtype" - ) - self_tp, _ = MambaPolicy._mamba_tp(self_ri) peer_tp, _ = MambaPolicy._mamba_tp(peer_ri) @@ -512,37 +543,68 @@ def _check_global(field: str, self_bytes: int, peer_bytes: int) -> None: "only with attention-DP enabled on both sides." ) - # Resolve slot_bytes from pool_views by pool_role (conv_states/ssm_states - # fields were removed; pool_views carry the same geometry). - from tensorrt_llm._torch.disaggregation.resource.page import MAMBA_CONV_ROLE, MAMBA_SSM_ROLE - - def _slot_bytes_by_role(mlg, role): - for pv in mlg.pool_views: - if pv.pool_role == role: - return pv.bytes_per_layer - return None - - self_ssm_slot = _slot_bytes_by_role(self_mlg, MAMBA_SSM_ROLE) - peer_ssm_slot = _slot_bytes_by_role(peer_mlg, MAMBA_SSM_ROLE) - self_conv_slot = _slot_bytes_by_role(self_mlg, MAMBA_CONV_ROLE) - peer_conv_slot = _slot_bytes_by_role(peer_mlg, MAMBA_CONV_ROLE) - - if self_ssm_slot is not None and peer_ssm_slot is not None: - _check_global("ssm slot_bytes", self_ssm_slot, peer_ssm_slot) - if self_conv_slot is not None and peer_conv_slot is not None: - _check_global("conv slot_bytes", self_conv_slot, peer_conv_slot) - - if self_mlg.conv_section_bytes is not None and peer_mlg.conv_section_bytes is not None: - if len(self_mlg.conv_section_bytes) != len(peer_mlg.conv_section_bytes): + # Views pair up by pool_role, exactly as get_pool_mapping matches them. + # A role present on only one side is a peer-declaration mismatch for + # REPLICATED views (checked below); for sharded views it is left to + # pool matching, which skips it. + peer_views = {pv.pool_role: pv for pv in peer_mlg.pool_views} + for self_pv in self_mlg.pool_views: + peer_pv = peer_views.get(self_pv.pool_role) + if peer_pv is None: + continue + role = sorted(self_pv.pool_role) + if self_pv.mapper_kind != peer_pv.mapper_kind: raise ValueError( - "MambaPolicy.validate_peer_compatible: conv section count differs " - f"(local={len(self_mlg.conv_section_bytes)}, " - f"peer={len(peer_mlg.conv_section_bytes)})" + f"MambaPolicy.validate_peer_compatible: mapper kind of role {role} differs " + f"(local={self_pv.mapper_kind.name}, peer={peer_pv.mapper_kind.name})" ) - for i, (s, p) in enumerate( - zip(self_mlg.conv_section_bytes, peer_mlg.conv_section_bytes) + if self_pv.bytes_per_layer is None or peer_pv.bytes_per_layer is None: + continue + + if self_pv.mapper_kind == MapperKind.REPLICATED: + # Identical bytes on every rank: sizes must match exactly. + if self_pv.bytes_per_layer != peer_pv.bytes_per_layer: + raise ValueError( + "MambaPolicy.validate_peer_compatible: replicated state size differs " + f"for role {role} (local={self_pv.bytes_per_layer}, " + f"peer={peer_pv.bytes_per_layer})" + ) + continue + + _check_global(f"{role} slot_bytes", self_pv.bytes_per_layer, peer_pv.bytes_per_layer) + + if self_pv.mapper_kind == MapperKind.SECTIONED: + if self_pv.section_bytes is not None and peer_pv.section_bytes is not None: + if len(self_pv.section_bytes) != len(peer_pv.section_bytes): + raise ValueError( + f"MambaPolicy.validate_peer_compatible: section count of role {role} " + f"differs (local={len(self_pv.section_bytes)}, " + f"peer={len(peer_pv.section_bytes)})" + ) + for i, (s, p) in enumerate(zip(self_pv.section_bytes, peer_pv.section_bytes)): + _check_global(f"{role} section_bytes[{i}]", s, p) + elif ( + self_pv.bytes_per_head is not None + and peer_pv.bytes_per_head is not None + and self_pv.bytes_per_head != peer_pv.bytes_per_head ): - _check_global(f"conv_section_bytes[{i}]", s, p) + # TP-invariant: head_dim * d_state * element_size. A mismatch + # means different state shape or SSM cache dtype. + raise ValueError( + f"MambaPolicy.validate_peer_compatible: bytes_per_head of role {role} " + f"differs (local={self_pv.bytes_per_head}, peer={peer_pv.bytes_per_head}); " + "check head_dim / d_state / mamba_ssm_cache_dtype" + ) + + # Replicated side state (e.g. PLE) is copied whole per layer, so both + # sides must declare the same roles on shared layers. A role missing + # on one side would otherwise be dropped silently by pool matching. + differing = find_replicated_role_mismatch(self_page_table, peer_page_table, CacheKind.STATE) + if differing: + raise ValueError( + "MambaPolicy.validate_peer_compatible: replicated roles differ on " + f"overlapping layers: {differing}" + ) @staticmethod def _mamba_tp(ri: RankInfo) -> Tuple[int, int]: @@ -597,19 +659,19 @@ def mamba_receiver_payload_bytes( if sender_mlg is None or receiver_mlg is None: return 0 - sender_globals = {ll.global_layer_id for ll in sender_mlg.local_layers} - receiver_globals = {ll.global_layer_id for ll in receiver_mlg.local_layers} - overlap = sender_globals & receiver_globals - if not overlap: - return 0 - - from tensorrt_llm._torch.disaggregation.resource.utils import get_physical_pool - - receiver_lg_idx = next( - i for i, lg in enumerate(receiver_page_table.layer_groups) if lg.kind == CacheKind.STATE - ) - per_layer = sum( - get_physical_pool(receiver_page_table, receiver_lg_idx, pv.pool_idx).slot_bytes - for pv in receiver_mlg.pool_views - ) - return len(overlap) * per_layer + sender_views = { + (pool_view.pool_role, pool_view.mapper_kind): pool_view + for pool_view in sender_mlg.pool_views + } + total = 0 + for receiver_view in receiver_mlg.pool_views: + sender_view = sender_views.get((receiver_view.pool_role, receiver_view.mapper_kind)) + if sender_view is None: + continue + sender_layers = set(get_pool_view_global_layer_ids(sender_view, sender_mlg)) + receiver_layers = set(get_pool_view_global_layer_ids(receiver_view, receiver_mlg)) + overlap = sender_layers & receiver_layers + if not overlap: + continue + total += len(overlap) * receiver_view.bytes_per_layer + return total diff --git a/tensorrt_llm/_torch/disaggregation/native/peer.py b/tensorrt_llm/_torch/disaggregation/native/peer.py index ee3a9b81945f..72bd85972e0c 100644 --- a/tensorrt_llm/_torch/disaggregation/native/peer.py +++ b/tensorrt_llm/_torch/disaggregation/native/peer.py @@ -48,7 +48,9 @@ class PeerOverlap: class PeerRegistrar: - # Registry: CacheKind -> PolicyClass. Add new layer types here. + # Registry: CacheKind -> PolicyClass. Add new layer types here. The + # policy owns every mapper kind of its life cycle and picks the mapper + # per view in build_mapper. _POLICY_CLASSES = { CacheKind.PAGED: AttentionPolicy, CacheKind.STATE: MambaPolicy, @@ -158,12 +160,16 @@ def _check_peer_compatible(self, peer_ri: RankInfo) -> bool: # Recurrent-state (Mamba/KDA) layout gate. Raises ValueError with a # field-level diagnostic instead of returning False, so the precise # mismatch reaches the caller of register(). + self_page_table = ( + self._self_ext_cache.page_table if self._self_ext_cache is not None else None + ) MambaPolicy.validate_peer_compatible( self._ri, peer_ri, - self._self_ext_cache.page_table if self._self_ext_cache is not None else None, + self_page_table, peer_ri.page_table, ) + AttentionPolicy.validate_peer_compatible(self_page_table, peer_ri.page_table) self_layers = sum(self._ri.layer_num_per_pp) peer_layers = sum(peer_ri.layer_num_per_pp) @@ -306,8 +312,7 @@ def _peer_l2g(kind: CacheKind) -> Dict[int, int]: def _get_policy(self, kind: CacheKind) -> Union[AttentionPolicy, MambaPolicy]: """Return the policy for a CacheKind (lazily instantiated).""" if kind not in self._policies: - cls = self._POLICY_CLASSES[kind] - self._policies[kind] = cls(self._ri) + self._policies[kind] = self._POLICY_CLASSES[kind](self._ri) return self._policies[kind] def get_kv_map( @@ -393,17 +398,12 @@ def get_kv_map( # the mapper must slice only the overlapping subset. extra_kwargs = {} if self_lg.kind == CacheKind.STATE: - # Position of each overlapping layer in the full sorted local_layer_ids - # used by extract_slot. These are indices into the ptrs array. - self_all_lids = sorted(set(int(e["local_layer_id"]) for e in self_pv.buffer_entries)) - peer_all_lids = sorted(set(int(e["local_layer_id"]) for e in peer_pv.buffer_entries)) - self_lid_to_pos = {lid: i for i, lid in enumerate(self_all_lids)} - peer_lid_to_pos = {lid: i for i, lid in enumerate(peer_all_lids)} - # overlapping_layers are global IDs; map them to local_layer_ids - self_overlap_positions = [self_lid_to_pos[self_g2l[gid]] for gid in overlapping_layers] - peer_overlap_positions = [ - peer_lid_to_pos[peer_g2l[gid]] for gid in overlapping_layers if gid in peer_g2l - ] + # Positions index the ptr array ``extract_slot`` builds, which is + # ordered by view offset. + self_gid_to_pos = {gid: i for i, gid in enumerate(self_global_ids)} + peer_gid_to_pos = {gid: i for i, gid in enumerate(peer_global_ids)} + self_overlap_positions = [self_gid_to_pos[gid] for gid in overlapping_layers] + peer_overlap_positions = [peer_gid_to_pos[gid] for gid in overlapping_layers] # Under contiguous PP partitioning, overlapping layers form a # contiguous block in the ptrs array. extra_kwargs["src_layer_off"] = ( @@ -424,6 +424,8 @@ def get_kv_map( peer_buffers_per_layer=peer_buffers_per_layer, self_lg=self_lg, peer_lg=peer_lg, + self_pv=self_pv, + peer_pv=peer_pv, **extra_kwargs, ) @@ -565,8 +567,8 @@ def should_send_pool( ``pool_idx`` indexes the layer group's ``pool_views`` list (one view per role class; several views may share a physical pool). Each view is kind-homogeneous, so ownership is a single per-view decision: - replicated views use one sender per fan-in group, sharded views - retain head-duplication routing. + replicated views use one sender per fan-in group (the policy returns + None for them), sharded views retain head-duplication routing. For mamba layer groups, ownership is based on mamba's own TP routing (independent of attention's duplicate-head logic): @@ -575,12 +577,9 @@ def should_send_pool( """ layer_group = self._self_ext_cache.page_table.layer_groups[layer_group_id] pool_view = layer_group.pool_views[pool_idx] - if pool_view.mapper_kind == MapperKind.REPLICATED: - return self._owns_tp_fan_in(peer_rank_info) - # Delegate to the policy's should_send. None means fan-in election. policy = self._get_policy(layer_group.kind) - result = policy.should_send(peer_overlap, peer_rank_info) + result = policy.should_send(peer_overlap, peer_rank_info, mapper_kind=pool_view.mapper_kind) return self._owns_tp_fan_in(peer_rank_info) if result is None else result def should_send_aux(self, peer_rank_info: RankInfo) -> bool: diff --git a/tensorrt_llm/_torch/disaggregation/native/transfer.py b/tensorrt_llm/_torch/disaggregation/native/transfer.py index bd8ddb9a0d41..b0fa94029892 100644 --- a/tensorrt_llm/_torch/disaggregation/native/transfer.py +++ b/tensorrt_llm/_torch/disaggregation/native/transfer.py @@ -61,6 +61,7 @@ get_non_empty_aux_indices, ) from tensorrt_llm._torch.disaggregation.native.messenger import ZMQMessenger, decode_message +from tensorrt_llm._torch.disaggregation.native.mixers.attention.peer import AttentionPolicy from tensorrt_llm._torch.disaggregation.native.mixers.ssm.peer import ( MambaPolicy, mamba_receiver_payload_bytes, @@ -2834,6 +2835,10 @@ def _get_sender_info(self, params: DisaggregatedParams) -> RankInfo: self._registrar.self_extractor.page_table, sender_info.page_table, ) + AttentionPolicy.validate_peer_compatible( + self._registrar.self_extractor.page_table, + sender_info.page_table, + ) except ValueError as e: msg = ( f"context peer at '{info_endpoint}' is incompatible: {e} " diff --git a/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py b/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py index 7d7490146a6a..3720238a547d 100644 --- a/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py +++ b/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py @@ -14,10 +14,9 @@ # limitations under the License. from collections import defaultdict -from typing import Dict, List, Optional, Sequence +from typing import Dict, List, Optional, Tuple import numpy as np -import torch from tensorrt_llm._torch.disaggregation.base.region import ( DataLayout, @@ -30,6 +29,7 @@ MAMBA_CONV_ROLE, MAMBA_SSM_ROLE, AttentionLayerGroup, + CacheKind, KVCachePageTable, LayerGroup, LocalLayer, @@ -38,9 +38,11 @@ PhysicalPool, PhysicalPoolGroup, PoolView, + RoleLayout, ) from tensorrt_llm._torch.disaggregation.resource.utils import ( compute_layer_byte_ranges, + get_layer_byte_ranges, get_physical_pool, ) from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import Role @@ -61,9 +63,17 @@ MapperKind.INDEXED, MapperKind.REPLICATED, MapperKind.NHD, + MapperKind.SECTIONED, } ) +# Mapper kinds each life cycle's policy can serve. A SECTIONED attention +# view or an NHD recurrent view has no mapper and is rejected at build time. +_MAPPER_KINDS_BY_CACHE_KIND = { + CacheKind.PAGED: frozenset({MapperKind.INDEXED, MapperKind.REPLICATED, MapperKind.NHD}), + CacheKind.STATE: frozenset({MapperKind.INDEXED, MapperKind.REPLICATED, MapperKind.SECTIONED}), +} + class KVRegionExtractorV1(RegionExtractorBase): """ @@ -80,6 +90,18 @@ def __init__(self, kv_arg): # Assume it is a manager (KVCacheManager / KVCacheManagerV2) self._page_table = build_page_table_from_manager(kv_arg) self._data_layout = DataLayout.HND + # extract_slot runs per request per view, but the page table is fixed + # once built, so resolve each view's layer offsets only once. + self._slot_layouts: Dict[Tuple[int, int], Tuple[np.ndarray, int]] = {} + + def _slot_layout(self, layer_group_id: int, pool_idx: int) -> "Tuple[np.ndarray, int]": + layout = self._slot_layouts.get((layer_group_id, pool_idx)) + if layout is None: + pv = self._page_table.layer_groups[layer_group_id].pool_views[pool_idx] + layer_offsets, bytes_per_layer = get_layer_byte_ranges(pv) + layout = (np.array(sorted(layer_offsets.values()), dtype=np.int64), bytes_per_layer) + self._slot_layouts[(layer_group_id, pool_idx)] = layout + return layout @property def page_table(self) -> KVCachePageTable: @@ -95,24 +117,22 @@ def extract_slot( """Extract per-layer pointers for a single slot (used for mamba state). Returns a SpecRegion with one pointer per layer: - ptr[i] = base + local_layer_id[i] * layer_stride + slot_id * slot_stride + ptr[i] = base + view_offset[i] + slot_id * slot_stride + + View offsets are explicit because a recurrent role may exist on only a + subset of layers or be interleaved with another role in a V2 pool. """ lg = self._page_table.layer_groups[layer_group_id] pv = lg.pool_views[pool_idx] pool = get_physical_pool(self._page_table, layer_group_id, pv.pool_idx) base_ptr = pool.base_address - layer_stride = pool.layer_stride_bytes slot_stride = pool.slot_stride_bytes - assert layer_stride is not None assert slot_stride is not None - local_layer_ids = sorted(set(int(e["local_layer_id"]) for e in pv.buffer_entries)) - ptrs = np.array( - [base_ptr + lid * layer_stride + slot_id * slot_stride for lid in local_layer_ids], - dtype=np.int64, - ) - memory = MemRegionGroup(ptrs=ptrs, bytes_per_region=pool.slot_bytes) + ordered_offsets, bytes_per_layer = self._slot_layout(layer_group_id, pool_idx) + ptrs = ordered_offsets + (base_ptr + slot_id * slot_stride) + memory = MemRegionGroup(ptrs=ptrs, bytes_per_region=bytes_per_layer) return SpecRegion(memory=memory) @nvtx_range("KVRegionExtractorV1.extract") @@ -157,34 +177,54 @@ def extract( # --------------------------------------------------------------------------- -def _build_mamba_pool_views(conv_pool, ssm_pool, local_layers): +def _build_mamba_pool_views( + conv_pool, + ssm_pool, + local_layers, + *, + conv_section_bytes: List[int], + ssm_bytes_per_head: int, +): """Build pool_views for mamba: conv at pool_idx=0, ssm at pool_idx=1. - Conv uses mapper_kind=SECTIONED (section-level granularity for TP split), - SSM uses mapper_kind=INDEXED (head-level granularity). MambaPolicy.build_mapper + Conv uses mapper_kind=SECTIONED with ``section_bytes`` (each section is + TP-sharded independently), SSM uses mapper_kind=INDEXED with + ``bytes_per_head`` (head-level granularity). MambaPolicy.build_mapper dispatches ConvStateMismatchMapper vs MambaHeadMismatchMapper accordingly. """ sorted_lids = [ ll.local_layer_id for ll in sorted(local_layers, key=lambda ll: ll.local_layer_id) ] + conv_layer_stride = int(conv_pool.layer_stride_bytes) + ssm_layer_stride = int(ssm_pool.layer_stride_bytes) return [ PoolView( pool_idx=0, buffer_entries=np.array( - [(lid, 0, conv_pool.slot_bytes) for lid in sorted_lids], dtype=BUFFER_ENTRY_DTYPE + [ + (lid, offset * conv_layer_stride, conv_pool.slot_bytes) + for offset, lid in enumerate(sorted_lids) + ], + dtype=BUFFER_ENTRY_DTYPE, ), pool_role=MAMBA_CONV_ROLE, mapper_kind=MapperKind.SECTIONED, bytes_per_layer=conv_pool.slot_bytes, + section_bytes=[int(x) for x in conv_section_bytes], ), PoolView( pool_idx=1, buffer_entries=np.array( - [(lid, 0, ssm_pool.slot_bytes) for lid in sorted_lids], dtype=BUFFER_ENTRY_DTYPE + [ + (lid, offset * ssm_layer_stride, ssm_pool.slot_bytes) + for offset, lid in enumerate(sorted_lids) + ], + dtype=BUFFER_ENTRY_DTYPE, ), pool_role=MAMBA_SSM_ROLE, mapper_kind=MapperKind.INDEXED, bytes_per_layer=ssm_pool.slot_bytes, + bytes_per_head=int(ssm_bytes_per_head), ), ] @@ -232,95 +272,13 @@ def _build_layer_group_for_mamba( layer_group = MambaLayerGroup( pool_group_idx=pool_group_idx, local_layers=local_layers, - pool_views=_build_mamba_pool_views(conv_pool, ssm_pool, local_layers), - conv_section_bytes=conv_section_bytes, - ssm_bytes_per_head=ssm_bytes_per_head, - ) - return layer_group, pool_group - - -def _slot_stride_bytes(tensor: torch.Tensor) -> int: - return int(tensor.stride(0) * tensor.element_size()) - - -def _build_v2_mamba_state_pool(states: Sequence[torch.Tensor]) -> PhysicalPool: - """Describe affine layer/slot addressing for one V2 Mamba state role.""" - if not states: - raise ValueError("V2 Mamba state pool requires at least one layer") - - first_state = states[0] - base_address = int(first_state.data_ptr()) - num_slots = int(first_state.shape[0]) - slot_bytes = int(first_state[0].numel() * first_state.element_size()) - slot_stride_bytes = _slot_stride_bytes(first_state) - - num_layers = len(states) - if slot_stride_bytes % num_layers != 0: - raise ValueError("V2 Mamba physical slot must divide evenly across layers") - # Each role appears once per layer in its size-class pool. Equal-size SSM - # and convolution states share that pool and are interleaved, so their - # layer stride includes both role payloads. - layer_stride_bytes = slot_stride_bytes // num_layers - - for layer_offset, state in enumerate(states): - state_slot_bytes = int(state[0].numel() * state.element_size()) - if ( - int(state.shape[0]) != num_slots - or state_slot_bytes != slot_bytes - or _slot_stride_bytes(state) != slot_stride_bytes - ): - raise ValueError("V2 Mamba state tensors must share one slot layout per role") - expected_address = base_address + layer_offset * layer_stride_bytes - if int(state.data_ptr()) != expected_address: - raise ValueError("V2 Mamba state tensors must have a uniform layer stride per role") - - return PhysicalPool( - base_address=base_address, - slot_bytes=slot_bytes, - num_slots=num_slots, - slot_stride_bytes=slot_stride_bytes, - layer_stride_bytes=layer_stride_bytes, - ) - - -def _build_layer_group_for_v2_mamba( - manager: MambaHybridCacheManagerV2, pool_group_idx: int -) -> "tuple[MambaLayerGroup, PhysicalPoolGroup]": - local_layers = [ - LocalLayer(local_layer_id=int(lid), global_layer_id=int(gid)) - for gid, lid in sorted(manager.mamba_layer_offsets.items(), key=lambda x: x[1]) - ] - - num_layers = len(local_layers) - expected_offsets = list(range(num_layers)) - if sorted(ll.local_layer_id for ll in local_layers) != expected_offsets: - raise ValueError("V2 Mamba layer offsets must be dense") - if len(manager.all_conv_states) != num_layers or len(manager.all_ssm_states) != num_layers: - raise ValueError("V2 Mamba state tensors must match the layer-offset table") - - first_conv_state = manager.all_conv_states[0] - first_ssm_state = manager.all_ssm_states[0] - conv_pool = _build_v2_mamba_state_pool(manager.all_conv_states) - ssm_pool = _build_v2_mamba_state_pool(manager.all_ssm_states) - if conv_pool.num_slots != ssm_pool.num_slots: - raise ValueError("V2 Mamba convolution and SSM states must have the same number of slots") - - d_conv_m1 = manager.conv_state_shape[1] - conv_elem_size = first_conv_state.element_size() - _, head_dim, d_state = manager.ssm_state_shape - conv_section_bytes = [dim * d_conv_m1 * conv_elem_size for dim in manager.conv_section_dims] - - ssm_elem_size = first_ssm_state.element_size() - ssm_bytes_per_head = head_dim * d_state * ssm_elem_size - - pool_group = PhysicalPoolGroup(pools=[conv_pool, ssm_pool]) - layer_group = MambaLayerGroup( - pool_group_idx=pool_group_idx, - local_layers=local_layers, - pool_views=_build_mamba_pool_views(conv_pool, ssm_pool, local_layers), - conv_section_bytes=conv_section_bytes, - ssm_bytes_per_head=ssm_bytes_per_head, - slot_major_layout=True, + pool_views=_build_mamba_pool_views( + conv_pool, + ssm_pool, + local_layers, + conv_section_bytes=conv_section_bytes, + ssm_bytes_per_head=ssm_bytes_per_head, + ), ) return layer_group, pool_group @@ -329,33 +287,17 @@ def _build_non_kv_layers( manager, layer_groups: List[LayerGroup], pool_groups: List[PhysicalPoolGroup], - *, - has_v2_mamba: bool = False, - v2_mamba_insert_idx: Optional[int] = None, ) -> None: - """Append (or insert) non-KV (recurrent/state) layer groups to the page table. - - Extension point for non-attention layer types. Currently handles Mamba; - add elif branches here for future recurrent/state layer types. + """Append non-KV (recurrent/state) layer groups for a V1 manager. - Args: - has_v2_mamba: If True, the V2 manager owns mamba layers that need - a dedicated pool group appended. - v2_mamba_insert_idx: If set, insert the mamba layer group at this - position (preserving original lifecycle ordering) instead of - appending at the end. + V1 managers have no pool descriptors, so the recurrent layer group is + described from the state tensors. V2 managers describe every life cycle, + recurrent ones included, from ``pool_group_descs`` in + ``_build_page_table_v2``. """ - if isinstance(manager, MambaHybridCacheManagerV2): - if has_v2_mamba and manager.local_num_mamba_layers > 0: - # Append a dedicated pool group (don't mutate the shared V2 entry). - mamba_pg_idx = len(pool_groups) - layer_group, local_pool_group = _build_layer_group_for_v2_mamba(manager, mamba_pg_idx) - pool_groups.append(local_pool_group) - if v2_mamba_insert_idx is not None: - layer_groups.insert(v2_mamba_insert_idx, layer_group) - else: - layer_groups.append(layer_group) - elif isinstance(manager, MambaHybridCacheManager): + if isinstance(manager, MambaHybridCacheManager) and not isinstance( + manager, MambaHybridCacheManagerV2 + ): pool_group_idx = len(pool_groups) layer_group, pool_group = _build_layer_group_for_mamba(manager, pool_group_idx) layer_groups.append(layer_group) @@ -551,6 +493,93 @@ def _compute_global_layer_ids(manager, lg_idx: int) -> List[int]: return [inverse[lid][0] * num_attn_types + inverse[lid][1] for lid in local_layer_ids] +def _build_pool_views_for_variant( + variant, + role_mapper_kinds: Dict, + default_mapper_kind: MapperKind, + layer_group_id: int, + role_layouts: Optional[Dict] = None, +) -> List[PoolView]: + """Bucket one slot-desc variant's coalesced buffers into pool views. + + One view is emitted per ``(physical pool, mapper kind)``; roles sharing a + kind share a view (KEY+VALUE), roles with different kinds in the same + physical pool get separate views (M3 coalesced index-K). ``role_layouts`` + supplies the per-role resharding geometry carried onto the view. + """ + role_layouts = role_layouts or {} + bucket_entries: Dict[tuple, list] = defaultdict(list) + bucket_roles: Dict[tuple, set] = defaultdict(set) + bucket_layouts: Dict[tuple, set] = defaultdict(set) + for pool_idx, coalesced_buffer in enumerate(variant.coalesced_buffers): + single_buffer_size = int(coalesced_buffer.single_buffer_size) + offset = 0 + for buffer_id in coalesced_buffer.buffer_ids: + kind = role_mapper_kinds.get(buffer_id.role, default_mapper_kind) + bucket_key = (pool_idx, kind) + bucket_entries[bucket_key].append((int(buffer_id.layer_id), offset, single_buffer_size)) + bucket_roles[bucket_key].add(str(buffer_id.role)) + bucket_layouts[bucket_key].add(role_layouts.get(buffer_id.role)) + offset += single_buffer_size + + # All ordering below is canonicalization — the page table is + # serialized and matched against peers, so view order (pool, + # then lowest slot offset), entry order (slot offset), and role + # text must not depend on dict/set iteration order. + pool_views = [] + lg_bucket_keys = sorted( + bucket_entries, + key=lambda key: (key[0], min(entry[1] for entry in bucket_entries[key])), + ) + for bucket_key in lg_bucket_keys: + pool_idx, mapper_kind = bucket_key + roles = frozenset(bucket_roles[bucket_key]) + context = ( + f"View(layer_group={layer_group_id}, pool={pool_idx}, " + f"kind={mapper_kind.name}, role={sorted(roles)})" + ) + entries = np.array( + sorted(bucket_entries[bucket_key], key=lambda entry: entry[1]), + dtype=BUFFER_ENTRY_DTYPE, + ) + # Fail fast on invalid geometry and record the uniform + # per-layer region size on the wire. Every kind is + # entries-driven, so the contiguous-layer-region / + # uniform-size invariants apply to all views uniformly. + _, bytes_per_layer = compute_layer_byte_ranges(entries, context=context) + # Roles that share a view must agree on their resharding geometry; + # roles without one contribute None. + layouts = {layout for layout in bucket_layouts[bucket_key] if layout is not None} + if len(layouts) > 1: + raise ValueError(f"{context} mixes roles with different layouts: {sorted(roles)}") + layout = layouts.pop() if layouts else RoleLayout() + if layout.section_bytes is not None and sum(layout.section_bytes) != bytes_per_layer: + raise ValueError( + f"{context} section_bytes {list(layout.section_bytes)} do not sum to " + f"bytes_per_layer={bytes_per_layer}" + ) + if layout.bytes_per_head is not None and bytes_per_layer % layout.bytes_per_head != 0: + raise ValueError( + f"{context} bytes_per_layer={bytes_per_layer} is not a multiple of " + f"bytes_per_head={layout.bytes_per_head}" + ) + pool_views.append( + PoolView( + pool_idx=pool_idx, + buffer_entries=entries, + pool_role=roles, + mapper_kind=mapper_kind, + bytes_per_layer=bytes_per_layer, + section_bytes=( + list(layout.section_bytes) if layout.section_bytes is not None else None + ), + bytes_per_head=layout.bytes_per_head, + ) + ) + + return pool_views + + def _build_page_table_v2(manager) -> KVCachePageTable: """Build a KVCachePageTable from a KVCacheManagerV2. @@ -595,6 +624,10 @@ def _build_page_table_v2(manager) -> KVCachePageTable: f"got it for role {role!s}" ) default_mapper_kind = role_mapper_kinds[Role.ALL] + role_layouts = manager.get_disagg_role_layouts() + for role, layout in role_layouts.items(): + if not isinstance(layout, RoleLayout): + raise ValueError(f"Invalid disaggregation role layout {layout!r} for role {role!s}") def _window_size_for_layer(internal_layer_id: int): if internal_layer_id < len(config.layers): @@ -615,8 +648,6 @@ def _window_size_for_layer(internal_layer_id: int): pool_groups: List[PhysicalPoolGroup] = [] storage_pg_to_list_idx: Dict[int, int] = {} layer_groups_by_id: List[LayerGroup | None] = [None] * len(manager.impl.layer_grouping) - has_v2_mamba: bool = False - v2_mamba_layer_group_ids: set = set() # layer_group_ids handled by _build_non_kv_layers for pg_desc in pool_group_descs: storage_pg_idx = int(pg_desc.pool_group_index) @@ -642,88 +673,58 @@ def _window_size_for_layer(internal_layer_id: int): for variant in pg_desc.slot_desc.variants: layer_group_id = int(variant.layer_group_id) all_internal_layer_ids = list(manager.impl.layer_grouping[layer_group_id]) - if isinstance(manager, MambaHybridCacheManagerV2) and any( - manager._is_local_mamba_layer(int(layer_id)) for layer_id in all_internal_layer_ids - ): - # Record that V2 mamba layers exist; handled by _build_non_kv_layers later. - has_v2_mamba = True - v2_mamba_layer_group_ids.add(layer_group_id) - continue - all_global_layer_ids = _compute_global_layer_ids(manager, layer_group_id) - local_layers = [ LocalLayer(local_layer_id=int(iid), global_layer_id=int(gid)) for iid, gid in zip(all_internal_layer_ids, all_global_layer_ids) ] + # A life cycle is recurrent when its layers are the manager's + # mamba layers. Recurrent and attention layers never share a life + # cycle: their buffer configs differ, so V2 groups them apart. + is_recurrent = isinstance(manager, MambaHybridCacheManagerV2) and any( + manager._is_local_mamba_layer(int(layer_id)) for layer_id in all_internal_layer_ids + ) + if is_recurrent and not all( + manager._is_local_mamba_layer(int(layer_id)) for layer_id in all_internal_layer_ids + ): + raise ValueError( + f"Layer group {layer_group_id} mixes recurrent and attention layers" + ) + cache_kind = CacheKind.STATE if is_recurrent else CacheKind.PAGED + # Bucket buffer entries by (pool, mapper kind). One PoolView is # emitted per bucket and spans every layer of that role class, # so the view count per layer group is bounded by the number of - # role classes — never by the layer count. A physical pool may - # hold several classes (V2 storage coalesces buffers purely by - # size within a layer group, so e.g. MiniMax M3's index-K shares - # the K/V pool when their per-block sizes coincide); each class - # still gets its own view, which keeps peer matching independent - # of that physical coalescing decision. ``pool_role`` stays the - # manager-supplied equivalence label used for peer matching + # role classes — never by the layer count. ``pool_role`` stays + # the manager-supplied equivalence label used for peer matching # without enumerating role names. Buffer offsets within a slot # follow ``buffer_ids`` order: the i-th buffer of a coalesced # buffer lives at ``i * single_buffer_size``. - bucket_entries: Dict[tuple, list] = defaultdict(list) - bucket_roles: Dict[tuple, set] = defaultdict(set) - for pool_idx, coalesced_buffer in enumerate(variant.coalesced_buffers): - single_buffer_size = int(coalesced_buffer.single_buffer_size) - offset = 0 - for buffer_id in coalesced_buffer.buffer_ids: - kind = role_mapper_kinds.get(buffer_id.role, default_mapper_kind) - bucket_key = (pool_idx, kind) - bucket_entries[bucket_key].append( - (int(buffer_id.layer_id), offset, single_buffer_size) - ) - bucket_roles[bucket_key].add(str(buffer_id.role)) - offset += single_buffer_size - - # Emit this layer group's views: one per (pool, mapper-kind - # class of roles). Roles sharing a kind share a view - # (KEY+VALUE); roles with different kinds in the same physical - # pool get separate views (M3 coalesced index-K). - # All ordering below is canonicalization — the page table is - # serialized and matched against peers, so view order (pool, - # then lowest slot offset), entry order (slot offset), and role - # text must not depend on dict/set iteration order. - pool_views = [] - lg_bucket_keys = sorted( - bucket_entries, - key=lambda key: (key[0], min(entry[1] for entry in bucket_entries[key])), + pool_views = _build_pool_views_for_variant( + variant, + role_mapper_kinds, + default_mapper_kind, + layer_group_id, + role_layouts=role_layouts, ) - for bucket_key in lg_bucket_keys: - pool_idx, mapper_kind = bucket_key - roles = frozenset(bucket_roles[bucket_key]) - entries = np.array( - sorted(bucket_entries[bucket_key], key=lambda entry: entry[1]), - dtype=BUFFER_ENTRY_DTYPE, - ) - # Fail fast on invalid geometry and record the uniform - # per-layer region size on the wire. Every kind is - # entries-driven, so the contiguous-layer-region / - # uniform-size invariants apply to all views uniformly. - _, bytes_per_layer = compute_layer_byte_ranges( - entries, - context=( - f"View(layer_group={layer_group_id}, pool={pool_idx}, " - f"kind={mapper_kind.name}, role={sorted(roles)})" - ), - ) - pool_views.append( - PoolView( - pool_idx=pool_idx, - buffer_entries=entries, - pool_role=roles, - mapper_kind=mapper_kind, - bytes_per_layer=bytes_per_layer, + for pool_view in pool_views: + if pool_view.mapper_kind not in _MAPPER_KINDS_BY_CACHE_KIND[cache_kind]: + raise ValueError( + f"Layer group {layer_group_id} ({cache_kind.name}) has no mapper for " + f"{pool_view.mapper_kind.name} view {sorted(pool_view.pool_role)}" ) + + if is_recurrent: + # One slot per request holds every layer of the life cycle; + # per-layer offsets live in the views' buffer entries. + layer_groups_by_id[layer_group_id] = MambaLayerGroup( + pool_group_idx=storage_pg_to_list_idx[storage_pg_idx], + local_layers=local_layers, + pool_views=pool_views, + slot_major_layout=True, ) + continue # Determine layer group metadata. # For managers with virtual layers, internal layer_ids @@ -744,23 +745,12 @@ def _window_size_for_layer(internal_layer_id: int): pool_views=pool_views, ) - # Preserve original lifecycle ordering: mamba groups stay at their - # original layer_group_id positions. _build_non_kv_layers fills them in. + # Layer groups are indexed by layer_group_id (== life cycle id). layer_groups: List[LayerGroup] = [] for layer_group_id, layer_group in enumerate(layer_groups_by_id): - if layer_group is None and layer_group_id not in v2_mamba_layer_group_ids: + if layer_group is None: raise ValueError(f"Missing V2 layer group descriptor for layer group {layer_group_id}") - if layer_group is not None: - layer_groups.append(layer_group) - # For skipped mamba IDs, a placeholder is inserted by _build_non_kv_layers below - - _build_non_kv_layers( - manager, - layer_groups, - pool_groups, - has_v2_mamba=has_v2_mamba, - v2_mamba_insert_idx=min(v2_mamba_layer_group_ids) if v2_mamba_layer_group_ids else None, - ) + layer_groups.append(layer_group) return KVCachePageTable( tokens_per_block=config.tokens_per_block, diff --git a/tensorrt_llm/_torch/disaggregation/resource/page.py b/tensorrt_llm/_torch/disaggregation/resource/page.py index 8798d9b205fd..0d9bb8ce47ca 100644 --- a/tensorrt_llm/_torch/disaggregation/resource/page.py +++ b/tensorrt_llm/_torch/disaggregation/resource/page.py @@ -17,14 +17,16 @@ from dataclasses import dataclass, field from enum import IntEnum -from typing import FrozenSet, List, Optional +from typing import FrozenSet, List, Optional, Tuple import numpy as np BUFFER_ENTRY_DTYPE = np.dtype( [ ("local_layer_id", np.uint32), - ("offset", np.uint32), + # 64-bit: layer-major recurrent-state pools address a layer at + # layer_index * num_slots * slot_bytes, which exceeds 4 GiB. + ("offset", np.uint64), ("size", np.uint32), ] ) @@ -64,9 +66,11 @@ class MapperKind(IntEnum): in ``bytes_per_layer``). View count per layer group is bounded by the number of role classes, never by layer count. - Mamba state pools do not use this enum: Mamba's transfer is dispatched - through :class:`MambaPolicy` which hard-codes the ``is_conv`` switch and - bypasses the attention pool-matching path entirely. + Mamba state pools use SECTIONED for convolution state, INDEXED for SSM + state, and REPLICATED for auxiliary recurrent state. The policy is chosen + by the layer group's :class:`CacheKind` (PAGED -> ``AttentionPolicy``, + STATE -> ``MambaPolicy``); each policy then picks the mapper for a view + from its mapper kind, copying REPLICATED views whole per layer. """ INDEXED = 0 @@ -76,6 +80,21 @@ class MapperKind(IntEnum): SECTIONED = 3 # Sectioned layout: [Sec0|Sec1|...], each section independently TP-sharded +@dataclass(frozen=True) +class RoleLayout: + """Resharding geometry a cache manager declares for one role. + + Carried onto the role's :class:`PoolView` by the page-table builder. + ``section_bytes`` describes a SECTIONED role (one entry per section, in + slot order); ``bytes_per_head`` describes an INDEXED role whose head + count is not derivable from ``RankInfo.attention`` (recurrent SSM state). + Roles with neither need no geometry (REPLICATED, attention K/V). + """ + + section_bytes: Optional[Tuple[int, ...]] = None + bytes_per_head: Optional[int] = None + + @dataclass class PhysicalPool: """Affine view of a physical pool over logical layers and slots. @@ -190,6 +209,14 @@ class PoolView: classes between layers, making the layer stride non-uniform. Set for every kind; ``None`` only in tables serialized by older builders, where consumers re-derive it from the entries. + section_bytes: SECTIONED views only. Byte size of each section of one + layer's region, in slot order (``sum == bytes_per_layer``). Each + section is TP-sharded independently, so a TP mismatch is resolved + per section. + bytes_per_head: INDEXED views whose head count is not derivable from + ``RankInfo.attention`` (recurrent SSM state). ``bytes_per_layer // + bytes_per_head`` is the per-rank head count. Attention K/V leaves + it ``None`` and uses ``kv_heads_per_rank`` instead. """ pool_idx: int @@ -197,6 +224,8 @@ class PoolView: pool_role: FrozenSet[str] = field(default_factory=frozenset) mapper_kind: MapperKind = MapperKind.INDEXED bytes_per_layer: Optional[int] = None + section_bytes: Optional[List[int]] = None + bytes_per_head: Optional[int] = None def to_dict(self) -> dict: return { @@ -207,6 +236,12 @@ def to_dict(self) -> dict: "bytes_per_layer": ( int(self.bytes_per_layer) if self.bytes_per_layer is not None else None ), + "section_bytes": ( + [int(x) for x in self.section_bytes] if self.section_bytes is not None else None + ), + "bytes_per_head": ( + int(self.bytes_per_head) if self.bytes_per_head is not None else None + ), } @staticmethod @@ -225,6 +260,14 @@ def from_dict(data: dict) -> "PoolView": bytes_per_layer=( int(data["bytes_per_layer"]) if data.get("bytes_per_layer") is not None else None ), + section_bytes=( + [int(x) for x in data["section_bytes"]] + if data.get("section_bytes") is not None + else None + ), + bytes_per_head=( + int(data["bytes_per_head"]) if data.get("bytes_per_head") is not None else None + ), ) @@ -298,21 +341,32 @@ def from_dict(cls, data: dict) -> "AttentionLayerGroup": ) -MAMBA_CONV_ROLE = frozenset({"mamba_conv"}) -MAMBA_SSM_ROLE = frozenset({"mamba_ssm"}) +# Native role names of the hybrid Mamba managers (``MambaRole.CONV_STATE`` / +# ``MambaRole.SSM_STATE``), spelled out here because page.py must not import +# the manager. The V1 builder stamps them so V1 and V2 peers match by role. +MAMBA_CONV_ROLE = frozenset({"conv_state"}) +MAMBA_SSM_ROLE = frozenset({"ssm_state"}) @dataclass class MambaLayerGroup(LayerGroup): - """Layer group for Mamba SSM states. + """Layer group for recurrent (STATE life cycle) states. - Pools are accessed the same way as attention: + One slot per request; every view of the group is addressed by the same + slot id. Pools are accessed the same way as attention: pool_groups[pool_group_idx].pools[pool_view.pool_idx] - where pool_idx=0 is conv, pool_idx=1 is ssm. + + Views are matched to a peer by ``pool_role``; how each view reshards + under a TP mismatch is described on the view itself (``mapper_kind`` + plus ``section_bytes`` / ``bytes_per_head``), so a group may carry any + mix of recurrent roles (Mamba2 / GDN conv and SSM state, KDA state, PLE + side state) without the group knowing which is which. + + ``slot_major_layout`` records whether the group's pools pack all layers + of one slot together (V2) or one layer's slots together (V1); it only + affects how the pools are registered with the transfer agent. """ - conv_section_bytes: Optional[List[int]] = None - ssm_bytes_per_head: Optional[int] = None slot_major_layout: bool = False def __post_init__(self) -> None: @@ -324,8 +378,6 @@ def to_dict(self) -> dict: "pool_group_idx": int(self.pool_group_idx), "local_layers": [ll.to_dict() for ll in self.local_layers], "pool_views": [pv.to_dict() for pv in self.pool_views], - "conv_section_bytes": self.conv_section_bytes, - "ssm_bytes_per_head": self.ssm_bytes_per_head, "slot_major_layout": self.slot_major_layout, } @@ -335,12 +387,6 @@ def from_dict(cls, data: dict) -> "MambaLayerGroup": pool_group_idx=int(data["pool_group_idx"]), local_layers=[LocalLayer.from_dict(x) for x in data["local_layers"]], pool_views=[PoolView.from_dict(pv) for pv in data["pool_views"]], - conv_section_bytes=[int(x) for x in data["conv_section_bytes"]] - if data.get("conv_section_bytes") - else None, - ssm_bytes_per_head=int(data["ssm_bytes_per_head"]) - if data.get("ssm_bytes_per_head") - else None, slot_major_layout=bool(data.get("slot_major_layout", False)), ) diff --git a/tensorrt_llm/_torch/disaggregation/resource/utils.py b/tensorrt_llm/_torch/disaggregation/resource/utils.py index ccc0fbc878bd..352ff83c75eb 100644 --- a/tensorrt_llm/_torch/disaggregation/resource/utils.py +++ b/tensorrt_llm/_torch/disaggregation/resource/utils.py @@ -15,7 +15,7 @@ from __future__ import annotations -from typing import Dict, List, Set +from typing import Dict, FrozenSet, List, Optional, Set, Tuple from .page import ( AttentionLayerGroup, @@ -23,6 +23,7 @@ KVCachePageTable, LayerGroup, MambaLayerGroup, + MapperKind, PhysicalPool, PoolView, ) @@ -156,6 +157,50 @@ def get_global_layer_ids(layer_group: AttentionLayerGroup) -> List[int]: return [ll.global_layer_id for ll in layer_group.local_layers] +def get_hosted_global_layer_ids(page_table: KVCachePageTable, kind: CacheKind) -> Set[int]: + """Global layer IDs hosted by the layer groups of *kind* in *page_table*.""" + return { + gid for lg in page_table.layer_groups if lg.kind == kind for gid in get_global_layer_ids(lg) + } + + +def get_replicated_role_layers( + page_table: KVCachePageTable, kind: CacheKind +) -> Set[Tuple[FrozenSet[str], int]]: + """``(pool_role, global_layer_id)`` pairs of every REPLICATED view of *kind*.""" + return { + (pool_view.pool_role, global_layer_id) + for layer_group in page_table.layer_groups + if layer_group.kind == kind + for pool_view in layer_group.pool_views + if pool_view.mapper_kind == MapperKind.REPLICATED + for global_layer_id in get_pool_view_global_layer_ids(pool_view, layer_group) + } + + +def find_replicated_role_mismatch( + self_page_table: Optional[KVCachePageTable], + peer_page_table: Optional[KVCachePageTable], + kind: CacheKind, +) -> List[Tuple[List[str], int]]: + """Replicated roles that only one side declares on a layer both sides host. + + Pool matching drops a view with no counterpart silently, so a role the peer + never declares would leave the receiver holding zeroed state instead of + raising. Returns a sorted list of ``(sorted_role_names, global_layer_id)``; + empty when the two sides agree or either page table is missing. + """ + if self_page_table is None or peer_page_table is None: + return [] + shared = get_hosted_global_layer_ids(self_page_table, kind) & get_hosted_global_layer_ids( + peer_page_table, kind + ) + differing = get_replicated_role_layers(self_page_table, kind) ^ get_replicated_role_layers( + peer_page_table, kind + ) + return sorted((sorted(role), gid) for role, gid in differing if gid in shared) + + def get_layer_group_num_layers(layer_group: AttentionLayerGroup) -> int: """ Number of layers in *layer_group* diff --git a/tensorrt_llm/_torch/disaggregation/transceiver.py b/tensorrt_llm/_torch/disaggregation/transceiver.py index 04d68fae7ce4..fbde326401c0 100644 --- a/tensorrt_llm/_torch/disaggregation/transceiver.py +++ b/tensorrt_llm/_torch/disaggregation/transceiver.py @@ -50,7 +50,10 @@ create_cache_reuse_adapter, ) from tensorrt_llm._torch.disaggregation.resource.page import CacheKind -from tensorrt_llm._torch.disaggregation.resource.utils import get_physical_pool +from tensorrt_llm._torch.disaggregation.resource.utils import ( + get_physical_pool, + get_pool_view_num_layers, +) from tensorrt_llm._torch.distributed.communicator import Distributed from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import ( BlockReusePolicy, @@ -485,9 +488,11 @@ def _slice_num_bytes(self, slice: KVSlice) -> int: for pv in lg.pool_views: pool = get_physical_pool(pt, lg_id, pv.pool_idx) if lg.kind == CacheKind.STATE: - # STATE: n=1 (one slot), but transfer covers all layers. - num_layers = len(lg.local_layers) - total += num_layers * pool.slot_bytes + # STATE: n=1 (one slot), but transfer covers all layers of + # the view. The physical slot may hold several roles, so + # size by the view's per-layer bytes, not the pool's slot. + num_layers = get_pool_view_num_layers(pv) + total += num_layers * pv.bytes_per_layer else: # Attention: n blocks, each slot covers all layers. total += n * pool.slot_bytes diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index 47655d32cece..d7e290843bb1 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -25,7 +25,7 @@ import torch from strenum import StrEnum -from tensorrt_llm._torch.disaggregation.resource.page import MapperKind +from tensorrt_llm._torch.disaggregation.resource.page import MapperKind, RoleLayout from tensorrt_llm._torch.distributed.communicator import Distributed, ReduceOp from tensorrt_llm._torch.utils import maybe_compile from tensorrt_llm._utils import ( @@ -2894,6 +2894,17 @@ def get_disagg_role_mapper_kinds(self) -> dict[DataRole, MapperKind]: """ return {Role.ALL: MapperKind.INDEXED, Role.INDEX_KEY: MapperKind.REPLICATED} + def get_disagg_role_layouts(self) -> dict[DataRole, RoleLayout]: + """Resharding geometry for roles whose mapper kind needs it. + + SECTIONED roles declare ``section_bytes``; INDEXED roles whose head + count is not derivable from the attention topology declare + ``bytes_per_head``. Attention K/V and replicated roles need none, so + the base manager declares nothing. The page-table builder copies the + layout onto the role's ``PoolView``. + """ + return {} + @property def blocks_in_primary_pool(self) -> int: """ diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py index f1284e713296..7287940d5eef 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py @@ -30,6 +30,8 @@ from tensorrt_llm.llmapi.llm_args import DecodingBaseConfig from tensorrt_llm.sampling_params import SamplingParams +from tensorrt_llm._torch.disaggregation.resource.page import (MapperKind, + RoleLayout) from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import ( _RESERVED_REQUEST_IDS, BlockReusePolicy, KVCacheManagerV2, Role) from tensorrt_llm._torch.pyexecutor.kv_cache_stats import \ @@ -3361,6 +3363,43 @@ def ple_layer_cache( return None return conv, ngram + def get_disagg_role_mapper_kinds(self) -> Dict[DataRole, MapperKind]: + """Recurrent roles and how they reshard across TP. + + Convolution state is SECTIONED: its flat per-layer buffer is a + concatenation of sections (Mamba2 ``[x | B | C]``, GDN ``[Q | K | V]``) + that are each TP-sharded independently. SSM state keeps the INDEXED + default (sharded by head). PLE state is computed from replicated + inputs, so every rank holds identical bytes and the transfer copies + whole per-layer regions. + """ + return { + **super().get_disagg_role_mapper_kinds(), + MambaRole.CONV_STATE: + MapperKind.SECTIONED, + MambaRole.PLE_CONV_STATE: + MapperKind.REPLICATED, + MambaRole.PLE_NGRAM_CONTEXT: + MapperKind.REPLICATED, + } + + def get_disagg_role_layouts(self) -> Dict[DataRole, RoleLayout]: + """Per-layer geometry for resharding conv and SSM state.""" + if self.local_num_mamba_layers == 0: + return {} + d_conv_m1 = int(self.conv_state_shape[1]) + conv_elem_size = self.conv_state_dtype.itemsize + _, head_dim, d_state = self.ssm_state_shape + return { + MambaRole.CONV_STATE: + RoleLayout(section_bytes=tuple( + int(dim) * d_conv_m1 * conv_elem_size + for dim in self.conv_section_dims)), + MambaRole.SSM_STATE: + RoleLayout(bytes_per_head=int(head_dim) * int(d_state) * + self.ssm_state_dtype.itemsize), + } + @property def use_gdn_cached_replay_all_layer_commit(self) -> bool: return getattr(self, "_use_gdn_cached_replay_all_layer_commit", False) diff --git a/tests/integration/defs/accuracy/test_disaggregated_serving.py b/tests/integration/defs/accuracy/test_disaggregated_serving.py index cfa98edf5d5f..9760f41f8c28 100644 --- a/tests/integration/defs/accuracy/test_disaggregated_serving.py +++ b/tests/integration/defs/accuracy/test_disaggregated_serving.py @@ -2107,6 +2107,124 @@ def test_auto_dtype(self, use_py_transceiver, mocker): run_accuracy_test(llm, self.MODEL_NAME, ["GSM8K"]) +@pytest.mark.timeout(DEFAULT_TEST_TIMEOUT) +@skip_pre_blackwell +@pytest.mark.skip_less_device(4) +@pytest.mark.skip_less_device_memory(145000) +@pytest.mark.skip_less_host_memory(98304) +class TestQwen3_8_Flash_Next(LlmapiAccuracyTestHarness): + """Block-FP8 Qwen3.8-Flash-Next over the Python NIXL transceiver.""" + + MODEL_NAME = "Qwen/Qwen3.8-Flash-Next" + MODEL_PATH = f"{llm_models_root()}/Qwen3.8-Flash-Next-FP8" + + # Match the aggregate accuracy run: the chat template thinks by default, + # so cap the trace and ask for the bare answer to fit the output budget. + GSM8K_EVALUATOR_KWARGS = dict( + apply_chat_template=True, + fewshot_as_multiturn=True, + system_prompt=("Use at most three short reasoning sentences, then " + "end with `#### NUMBER`. Do not restate the problem."), + chat_template_kwargs=dict(enable_thinking=True, + reasoning_effort="xhigh"), + ) + + @pytest.mark.parametrize( + "snapshot_policy", [None, "interval", "offsets"], + ids=["no_reuse", "prefix_cache", "prefix_cache_offsets"]) + def test_fp8_nixl_python(self, mocker, snapshot_policy): + """Tensor/expert-parallel-2 context to attention-DP2 generation. + + The handoff carries the QSA index state, the Gated DeltaNet state, and + the replicated PLE n-gram and convolution state alongside the KV pages. + + The two sides shard recurrent state differently on purpose. The context + side splits Gated DeltaNet state across both ranks, while attention DP + leaves it unsharded on the generation side (MambaPolicy._mamba_tp + returns 1), so the transfer exercises the sharded-to-replicated path. + PLE state stays replicated on both sides throughout. + + With block reuse the context worker restores part of its recurrent + state from a snapshot instead of computing it, so the transferred state + is only correct if the snapshot carries the PLE roles too. The two + reuse variants place snapshots differently: at a fixed token interval, + or at offsets measured from the prompt start and end. + """ + cache_transceiver_config = { + "backend": "NIXL", + "transceiver_runtime": "PYTHON", + "max_tokens_in_buffer": 8192, + } + kv_cache_config = { + "enable_block_reuse": snapshot_policy is not None, + "mamba_ssm_cache_dtype": "bfloat16", + "free_gpu_memory_fraction": 0.5, + } + # Attention pages alone cannot restore GDN or PLE state, so the runtime + # turns reuse back off unless a snapshot placement is configured. + if snapshot_policy == "interval": + kv_cache_config["mamba_state_config"] = { + "periodic_snapshot_interval": 256 + } + elif snapshot_policy == "offsets": + kv_cache_config["mamba_state_config"] = { + "additional_snapshot_offsets_from_start": [256], + "additional_snapshot_offsets_from_end": [0], + } + common_config = { + "trust_remote_code": True, + "tensor_parallel_size": 2, + "moe_expert_parallel_size": 2, + "max_batch_size": 16, + "cache_transceiver_config": cache_transceiver_config, + "moe_config": { + "backend": "TRTLLM" + }, + "kv_cache_config": kv_cache_config, + } + ctx_server_config = { + **common_config, + "disable_overlap_scheduler": True, + "cuda_graph_config": None, + } + gen_server_config = { + **common_config, + "enable_attention_dp": True, + "enable_lm_head_tp_in_adp": True, + "disable_overlap_scheduler": False, + "cuda_graph_config": { + "max_batch_size": 16, + "enable_padding": True, + }, + } + disaggregated_server_config = { + "hostname": "localhost", + "port": 8000, + "backend": "pytorch", + "context_servers": { + "num_instances": 1, + "urls": ["localhost:8001"] + }, + "generation_servers": { + "num_instances": 1, + "urls": ["localhost:8002"] + } + } + + mocker.patch.object(GSM8K, "MAX_OUTPUT_LEN", 512) + with launch_disaggregated_llm( + disaggregated_server_config, + ctx_server_config, + gen_server_config, + self.MODEL_PATH, + extra_env={"TRTLLM_QWEN4_EXP_PLE_HOST_OFFLOAD": "1"}, + ) as llm: + run_accuracy_test( + llm, + self.MODEL_NAME, ["GSM8K"], + extra_evaluator_kwargs={GSM8K: self.GSM8K_EVALUATOR_KWARGS}) + + @pytest.mark.timeout(DEFAULT_TEST_TIMEOUT) @skip_pre_blackwell @pytest.mark.skip_less_device_memory(80000) diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index d283757e9495..1a86be4c475e 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -6365,11 +6365,16 @@ def _build_llm(self, check_acceptance_length: bool, moe_expert_parallel_size: int = 1, enable_attention_dp: bool = False, - cover_guided_decoding: bool = False) -> LLM: + cover_guided_decoding: bool = False, + enable_block_reuse: bool = False) -> LLM: """Construct the engine shared by both evaluation tasks.""" kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.5, - enable_block_reuse=False, + enable_block_reuse=enable_block_reuse, mamba_ssm_cache_dtype="bfloat16") + if enable_block_reuse: + # Attention pages alone cannot restore GDN or PLE state, so the + # runtime turns reuse back off without a snapshot policy. + kv_cache_config.mamba_state_config.periodic_snapshot_interval = 256 cuda_graph_config = CudaGraphConfig(max_batch_size=self.MAX_BATCH_SIZE, enable_padding=True) mtp_config = (None if max_draft_len is None else MTPDecodingConfig( @@ -6404,7 +6409,8 @@ def _run_evals(self, mocker, moe_expert_parallel_size: int = 1, enable_attention_dp: bool = False, - cover_guided_decoding: bool = False) -> None: + cover_guided_decoding: bool = False, + enable_block_reuse: bool = False) -> None: if not os.path.exists(model_path): pytest.skip(f"Model directory {model_path} does not exist") @@ -6414,16 +6420,16 @@ def _run_evals(self, expected_quant_algo == QuantAlgo.FP8_BLOCK_SCALES and tensor_parallel_size == 4 and moe_backend == "TRTLLM" and max_draft_len == 3 and moe_expert_parallel_size == 4 - and enable_attention_dp) - with self._build_llm( - model_path, - tensor_parallel_size, - moe_backend, - max_draft_len, - check_acceptance_length, - moe_expert_parallel_size=moe_expert_parallel_size, - enable_attention_dp=enable_attention_dp, - cover_guided_decoding=cover_guided_decoding) as llm: + and enable_attention_dp and not enable_block_reuse) + with self._build_llm(model_path, + tensor_parallel_size, + moe_backend, + max_draft_len, + check_acceptance_length, + moe_expert_parallel_size=moe_expert_parallel_size, + enable_attention_dp=enable_attention_dp, + cover_guided_decoding=cover_guided_decoding, + enable_block_reuse=enable_block_reuse) as llm: assert llm.args.quant_config.quant_algo == expected_quant_algo if cover_guided_decoding: assert_guided_decoding_regex(llm) @@ -6477,6 +6483,29 @@ def test_fp8_adp4_mtp3_trtllm_ple_offload(self, moe_expert_parallel_size=4, enable_attention_dp=True) + @skip_pre_blackwell + @pytest.mark.skip_less_device(4) + @pytest.mark.skip_less_device_memory(82000) + @pytest.mark.skip_less_host_memory(98304) + def test_fp8_adp4_mtp3_prefix_cache(self, monkeypatch: pytest.MonkeyPatch, + mocker) -> None: + """test_fp8_adp4_mtp3_trtllm_ple_offload with block reuse enabled. + + A reused prefix restores the Gated DeltaNet and PLE state from a + recurrent-state snapshot instead of recomputing it, so accuracy holds + only if the snapshot carries every role in the recurrent page. + """ + self._run_evals(f"{llm_models_root()}/Qwen3.8-Flash-Next-FP8", + tensor_parallel_size=4, + moe_backend="TRTLLM", + max_draft_len=3, + expected_quant_algo=QuantAlgo.FP8_BLOCK_SCALES, + monkeypatch=monkeypatch, + mocker=mocker, + moe_expert_parallel_size=4, + enable_attention_dp=True, + enable_block_reuse=True) + @skip_pre_blackwell @pytest.mark.skip_less_device(4) @pytest.mark.skip_less_device_memory(70000) diff --git a/tests/integration/test_lists/qa/llm_function_core.txt b/tests/integration/test_lists/qa/llm_function_core.txt index f8055658d371..edcdded5cf85 100644 --- a/tests/integration/test_lists/qa/llm_function_core.txt +++ b/tests/integration/test_lists/qa/llm_function_core.txt @@ -43,6 +43,9 @@ accuracy/test_disaggregated_serving.py::TestQwen3_8B::test_chunked_prefill accuracy/test_disaggregated_serving.py::TestQwen3_8B::test_gen_first accuracy/test_disaggregated_serving.py::TestQwen3_8B::test_gen_first_kv_cache_v1 accuracy/test_disaggregated_serving.py::TestQwen3_8B::test_nixl_backend +accuracy/test_disaggregated_serving.py::TestQwen3_8_Flash_Next::test_fp8_nixl_python[no_reuse] +accuracy/test_disaggregated_serving.py::TestQwen3_8_Flash_Next::test_fp8_nixl_python[prefix_cache] +accuracy/test_disaggregated_serving.py::TestQwen3_8_Flash_Next::test_fp8_nixl_python[prefix_cache_offsets] accuracy/test_dwdp_aggregated.py::TestDwdpAggDeepSeekV3Lite::test_dwdp_agg_accuracy[mode_a_uniform] accuracy/test_dwdp_aggregated.py::TestDwdpAggDeepSeekV3Lite::test_dwdp_agg_accuracy[mode_a_uniform_contention_opt] accuracy/test_dwdp_aggregated.py::TestDwdpAggDeepSeekV3Lite::test_dwdp_agg_accuracy[mode_b_overlap] @@ -635,6 +638,7 @@ accuracy/test_llm_api_pytorch.py::TestQwen3_8_2_4T_A95B::test_nvfp4_adp16_cuteds accuracy/test_llm_api_pytorch.py::TestQwen3_8_2_4T_A95B::test_nvfp4_tp8_mtp3_trtllm accuracy/test_llm_api_pytorch.py::TestQwen3_8_Flash_Next::test_bf16_tep4_cutlass accuracy/test_llm_api_pytorch.py::TestQwen3_8_Flash_Next::test_fp8_adp4_mtp3_trtllm_ple_offload +accuracy/test_llm_api_pytorch.py::TestQwen3_8_Flash_Next::test_fp8_adp4_mtp3_prefix_cache accuracy/test_llm_api_pytorch.py::TestQwen3_8_Flash_Next::test_nvfp4_1gpu_cutedsl_ple_offload accuracy/test_llm_api_pytorch.py::TestQwen3_8_Flash_Next::test_nvfp4_adp4_mtp3_trtllm_ple_offload accuracy/test_llm_api_pytorch.py::TestSeedOss_36B::test_auto_dtype diff --git a/tests/integration/test_lists/test-db/l0_gb300.yml b/tests/integration/test_lists/test-db/l0_gb300.yml index 53e2eb5fd577..120a5e82064e 100644 --- a/tests/integration/test_lists/test-db/l0_gb300.yml +++ b/tests/integration/test_lists/test-db/l0_gb300.yml @@ -22,5 +22,6 @@ l0_gb300: - accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_fp8_block_scales[latency] # Cover nvbugs 5461712 and 5505402 - accuracy/test_llm_api_pytorch.py::TestQwen3_8_Flash_Next::test_fp8_adp4_mtp3_trtllm_ple_offload - accuracy/test_llm_api_pytorch.py::TestQwen3_8_Flash_Next::test_nvfp4_adp4_mtp3_trtllm_ple_offload + - accuracy/test_disaggregated_serving.py::TestQwen3_8_Flash_Next::test_fp8_nixl_python[prefix_cache] - unittest/_torch/thop/parallel TIMEOUT (90) - unittest/_torch/visual_gen/kernels/parallel diff --git a/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml b/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml index 0658027824da..d3142cb9b50e 100644 --- a/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml +++ b/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml @@ -91,3 +91,5 @@ l0_gb300_multi_gpus: - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=0-tp4-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-low_precision_combine=True-torch_compile=False] - accuracy/test_llm_api_pytorch.py::TestDeepSeekR1::test_nvfp4_multi_gpus[throughput_tp4] TIMEOUT (180) - accuracy/test_llm_api_pytorch.py::TestQwen3_8_Flash_Next::test_bf16_tep4_cutlass + - accuracy/test_llm_api_pytorch.py::TestQwen3_8_Flash_Next::test_fp8_adp4_mtp3_prefix_cache + - accuracy/test_disaggregated_serving.py::TestQwen3_8_Flash_Next::test_fp8_nixl_python[no_reuse] diff --git a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py index e292e1c2224d..9174796cd075 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py @@ -10,7 +10,11 @@ import torch from tensorrt_llm._torch.disaggregation.resource.kv_extractor import build_page_table_from_manager -from tensorrt_llm._torch.disaggregation.resource.page import AttentionLayerGroup, MambaLayerGroup +from tensorrt_llm._torch.disaggregation.resource.page import ( + MAMBA_CONV_ROLE, + AttentionLayerGroup, + MambaLayerGroup, +) from tensorrt_llm._torch.disaggregation.transceiver import KvCacheTransceiverV2 from tensorrt_llm._torch.modules.mamba.mamba2_metadata import Mamba2Metadata from tensorrt_llm._torch.pyexecutor._util import ( @@ -3400,7 +3404,8 @@ def test_v2_hybrid_disagg_page_table_uses_qwen3_next_conv_sections(): assert isinstance(mamba_group, MambaLayerGroup) d_conv_m1 = mgr.conv_state_shape[1] conv_elem_size = mgr.all_conv_states[0].element_size() - assert mamba_group.conv_section_bytes == [ + conv_view = next(pv for pv in mamba_group.pool_views if pv.pool_role == MAMBA_CONV_ROLE) + assert conv_view.section_bytes == [ dim * d_conv_m1 * conv_elem_size for dim in mgr.conv_section_dims ] assert mgr.conv_section_dims == [8, 8, 32] diff --git a/tests/unittest/_torch/modeling/test_qwen4_exp_support.py b/tests/unittest/_torch/modeling/test_qwen4_exp_support.py index 3196cfde67fb..5fa2f5323f8f 100644 --- a/tests/unittest/_torch/modeling/test_qwen4_exp_support.py +++ b/tests/unittest/_torch/modeling/test_qwen4_exp_support.py @@ -532,13 +532,26 @@ def test_ple_cache_layout_excludes_separate_mtp_draft() -> None: assert draft is None -def test_v2_cache_estimator_counts_ple_lifecycle_state() -> None: +@pytest.mark.parametrize( + "offsets_from_start,offsets_from_end,expected_state_slots", + [ + # Two live request slots plus one non-speculative CUDA-graph dummy slot. + ([], [], 3), + # Each snapshot rule reserves one more slot per resident sequence, and + # the PLE state is part of every one of those slots. + ([256], [0], 7), + ], + ids=["no_reuse", "snapshot_offsets"], +) +def test_v2_cache_estimator_counts_ple_lifecycle_state( + offsets_from_start, offsets_from_end, expected_state_slots +) -> None: from tensorrt_llm._torch.configs import Qwen4ExpTextConfig from tensorrt_llm._torch.pyexecutor.config_utils import extract_qwen4_exp_ple_cache_params from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import ( MambaHybridCacheManagerV2, ) - from tensorrt_llm.llmapi.llm_args import KvCacheConfig + from tensorrt_llm.llmapi.llm_args import KvCacheConfig, MambaStateConfig from tensorrt_llm.mapping import Mapping config = Qwen4ExpTextConfig.from_dict(_text_config_dict()) @@ -547,7 +560,13 @@ def test_v2_cache_estimator_counts_ple_lifecycle_state() -> None: common = { "mapping": Mapping(world_size=1, rank=0, tp_size=1, pp_size=1), "max_batch_size": 2, - "kv_cache_config": KvCacheConfig(enable_block_reuse=False), + "kv_cache_config": KvCacheConfig( + enable_block_reuse=bool(offsets_from_start or offsets_from_end), + mamba_state_config=MambaStateConfig( + additional_snapshot_offsets_from_start=offsets_from_start, + additional_snapshot_offsets_from_end=offsets_from_end, + ), + ), } with_ple = MambaHybridCacheManagerV2.get_cache_size_per_token( SimpleNamespace(pretrained_config=config, quant_config=None), **common @@ -562,8 +581,7 @@ def test_v2_cache_estimator_counts_ple_lifecycle_state() -> None: + ple.ngram_context_len * torch.int64.itemsize ) assert with_ple[0] == without_ple[0] - # Two live request slots plus one non-speculative CUDA-graph dummy slot. - assert with_ple[1] - without_ple[1] == 3 * bytes_per_slot + assert with_ple[1] - without_ple[1] == expected_state_slots * bytes_per_slot def test_ple_states_use_v2_lifecycle_buffers(monkeypatch) -> None: diff --git a/tests/unittest/disaggregated/test_bounce.py b/tests/unittest/disaggregated/test_bounce.py index 24f0122da025..9850f90c9016 100644 --- a/tests/unittest/disaggregated/test_bounce.py +++ b/tests/unittest/disaggregated/test_bounce.py @@ -735,6 +735,7 @@ def _k3_page_table() -> KVCachePageTable: pool_role=MAMBA_CONV_ROLE, mapper_kind=MapperKind.SECTIONED, bytes_per_layer=_K3_CONV_SLOT_BYTES, + section_bytes=[_K3_CONV_SLOT_BYTES // 3] * 3, ), PoolView( pool_idx=1, @@ -744,14 +745,13 @@ def _k3_page_table() -> KVCachePageTable: pool_role=MAMBA_SSM_ROLE, mapper_kind=MapperKind.INDEXED, bytes_per_layer=_K3_SSM_SLOT_BYTES, + bytes_per_head=_K3_SSM_SLOT_BYTES // 4, ), ] mamba = MambaLayerGroup( pool_group_idx=1, local_layers=local_layers, pool_views=pool_views, - conv_section_bytes=[_K3_CONV_SLOT_BYTES // 3] * 3, - ssm_bytes_per_head=_K3_SSM_SLOT_BYTES // 4, ) return KVCachePageTable( tokens_per_block=32, diff --git a/tests/unittest/disaggregated/test_extractor.py b/tests/unittest/disaggregated/test_extractor.py index 8c4948d0f523..ec0a638524cb 100644 --- a/tests/unittest/disaggregated/test_extractor.py +++ b/tests/unittest/disaggregated/test_extractor.py @@ -20,7 +20,7 @@ from tensorrt_llm._torch.disaggregation.base.region import MemRegionGroup, SpecRegion from tensorrt_llm._torch.disaggregation.resource.kv_extractor import ( KVRegionExtractorV1, - _build_v2_mamba_state_pool, + _build_mamba_pool_views, build_page_table, build_page_table_from_manager, ) @@ -35,6 +35,10 @@ get_unique_layers, ) from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import Role +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import ( + MambaHybridCacheManagerV2, + MambaRole, +) from tensorrt_llm._torch.pyexecutor.resource_manager import ( CacheTypeCpp, DataType, @@ -525,6 +529,7 @@ def _make_fake_v2_manager(attrs, role_mapper_kinds, *, num_pools=1, slot_bytes_l pp_layers=[0, 1], num_kv_heads_per_layer=[1, 1], get_disagg_role_mapper_kinds=lambda: role_mapper_kinds, + get_disagg_role_layouts=lambda: {}, ) @@ -699,20 +704,22 @@ def test_mamba_layer_group_serialization(): pool_role=MAMBA_SSM_ROLE, mapper_kind=MapperKind.INDEXED, bytes_per_layer=ssm_slot_bytes, + bytes_per_head=128, ), ] + pool_views[0].section_bytes = [512, 256, 256] mlg = MambaLayerGroup( pool_group_idx=1, local_layers=local_layers, pool_views=pool_views, - conv_section_bytes=[512, 256, 256], - ssm_bytes_per_head=128, ) d = mlg.to_dict() - assert d["conv_section_bytes"] == [512, 256, 256] + assert d["pool_views"][0]["section_bytes"] == [512, 256, 256] + assert d["pool_views"][1]["bytes_per_head"] == 128 assert d["kind"] == int(CacheKind.STATE) assert "mamba_layer_offsets" not in d + assert "conv_section_bytes" not in d restored = LayerGroup.from_dict(d) assert isinstance(restored, MambaLayerGroup) @@ -723,8 +730,10 @@ def test_mamba_layer_group_serialization(): assert restored.pool_views[1].pool_role == MAMBA_SSM_ROLE assert restored.pool_views[1].bytes_per_layer == ssm_slot_bytes assert restored.pool_views[1].mapper_kind == MapperKind.INDEXED - assert restored.conv_section_bytes == [512, 256, 256] - assert restored.ssm_bytes_per_head == 128 + assert restored.pool_views[0].section_bytes == [512, 256, 256] + assert restored.pool_views[0].bytes_per_head is None + assert restored.pool_views[1].bytes_per_head == 128 + assert restored.pool_views[1].section_bytes is None assert [(ll.local_layer_id, ll.global_layer_id) for ll in restored.local_layers] == [ (0, 10), (1, 11), @@ -736,38 +745,191 @@ def test_mamba_layer_group_serialization(): assert legacy_pool.layer_stride_bytes == legacy_pool.num_slots * legacy_pool.slot_bytes -def test_v2_mamba_state_pool_uses_affine_layer_and_slot_strides(): - num_slots = 3 - state_bytes = 64 - storage = torch.empty((num_slots, 4, state_bytes), dtype=torch.uint8) - states = [storage[:, 0, :], storage[:, 2, :]] +def _make_fake_v2_mamba_manager(*, ple_on_layer=1): + """Duck-typed MambaHybridCacheManagerV2 whose descriptors hold two life cycles. - pool = _build_v2_mamba_state_pool(states) + Life cycle 0 is recurrent (layers 0 and 1): SSM and conv state coalesce + into one 64-byte pool in ``[ssm, conv]`` order per layer, and PLE n-gram + context lives in a second, 32-byte pool on ``ple_on_layer`` only. Life + cycle 1 is attention (layer 2) with K/V in its own pool group. + """ + from types import SimpleNamespace - assert pool.base_address == states[0].data_ptr() - assert pool.slot_bytes == state_bytes - assert pool.num_slots == num_slots - assert pool.slot_stride_bytes == 4 * state_bytes - assert pool.layer_stride_bytes == 2 * state_bytes + from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import Role + def buf(layer_id, role): + return SimpleNamespace(layer_id=layer_id, role=role) -def test_v2_mamba_single_layer_pool_preserves_shared_role_footprint(): - state_bytes = 64 - storage = torch.empty((3, 2, state_bytes), dtype=torch.uint8) + recurrent_pg = SimpleNamespace( + pool_group_index=0, + num_slots=4, + pools=[ + SimpleNamespace(pool_index=0, base_address=0x10000, slot_bytes=256), + SimpleNamespace(pool_index=1, base_address=0x20000, slot_bytes=32), + ], + slot_desc=SimpleNamespace( + variants=[ + SimpleNamespace( + layer_group_id=0, + coalesced_buffers=[ + SimpleNamespace( + single_buffer_size=64, + buffer_ids=[ + buf(0, MambaRole.SSM_STATE), + buf(0, MambaRole.CONV_STATE), + buf(1, MambaRole.SSM_STATE), + buf(1, MambaRole.CONV_STATE), + ], + ), + SimpleNamespace( + single_buffer_size=32, + buffer_ids=[buf(ple_on_layer, MambaRole.PLE_NGRAM_CONTEXT)], + ), + ], + ) + ] + ), + ) + attention_pg = SimpleNamespace( + pool_group_index=1, + num_slots=8, + pools=[SimpleNamespace(pool_index=0, base_address=0x30000, slot_bytes=128)], + slot_desc=SimpleNamespace( + variants=[ + SimpleNamespace( + layer_group_id=1, + coalesced_buffers=[ + SimpleNamespace( + single_buffer_size=64, + buffer_ids=[buf(2, Role.KEY), buf(2, Role.VALUE)], + ) + ], + ) + ] + ), + ) + + manager = object.__new__(MambaHybridCacheManagerV2) + manager.impl = SimpleNamespace( + layer_grouping=((0, 1), (2,)), + init_config=SimpleNamespace( + tokens_per_block=16, + layers=[SimpleNamespace(window_size=None) for _ in range(3)], + ), + pool_group_descs=[recurrent_pg, attention_pg], + ) + manager.pp_layers = [0, 1, 2] + manager._mamba_layer_mask = [True, True, False] + manager.local_num_mamba_layers = 2 + manager.num_kv_heads_per_layer = [1, 1, 1] + # conv: [8, 4] fp16 = 64 bytes as sections [2, 2, 4] * 4 * 2 = [16, 16, 32] + manager.conv_state_shape = [8, 4] + manager.conv_state_dtype = torch.float16 + manager.conv_section_dims = [2, 2, 4] + # ssm: [2 heads, 4, 4] fp16 = 64 bytes, 32 bytes per head + manager.ssm_state_shape = [2, 4, 4] + manager.ssm_state_dtype = torch.float16 + return manager + + +def test_v2_builder_describes_recurrent_life_cycle_from_descriptors(): + """Conv, SSM and PLE views all come from pool_group_descs, geometry included.""" + from tensorrt_llm._torch.disaggregation.resource.page import ( + MAMBA_CONV_ROLE, + MAMBA_SSM_ROLE, + AttentionLayerGroup, + MambaLayerGroup, + ) + + manager = _make_fake_v2_mamba_manager() + page_table = build_page_table_from_manager(manager) + + assert len(page_table.layer_groups) == 2 + recurrent, attention = page_table.layer_groups + assert isinstance(recurrent, MambaLayerGroup) + assert isinstance(attention, AttentionLayerGroup) + # The recurrent group uses the shared descriptor pool group, not a + # dedicated tensor-derived one, and keeps V2 layer numbering. + assert recurrent.pool_group_idx == 0 + assert recurrent.slot_major_layout + assert [(ll.local_layer_id, ll.global_layer_id) for ll in recurrent.local_layers] == [ + (0, 0), + (1, 1), + ] + pools = page_table.pool_groups[recurrent.pool_group_idx].pools + assert (pools[0].base_address, pools[0].slot_bytes, pools[0].num_slots) == (0x10000, 256, 4) + + by_role = {pv.pool_role: pv for pv in recurrent.pool_views} + assert set(by_role) == { + MAMBA_SSM_ROLE, + MAMBA_CONV_ROLE, + frozenset({str(MambaRole.PLE_NGRAM_CONTEXT)}), + } - pool = _build_v2_mamba_state_pool([storage[:, 1, :]]) + ssm = by_role[MAMBA_SSM_ROLE] + assert ssm.pool_idx == 0 + assert ssm.mapper_kind == MapperKind.INDEXED + assert ssm.bytes_per_layer == 64 + assert ssm.bytes_per_head == 32 + assert ssm.section_bytes is None + assert [tuple(int(x) for x in e) for e in ssm.buffer_entries] == [(0, 0, 64), (1, 128, 64)] + + conv = by_role[MAMBA_CONV_ROLE] + assert conv.pool_idx == 0 + assert conv.mapper_kind == MapperKind.SECTIONED + assert conv.bytes_per_layer == 64 + assert conv.section_bytes == [16, 16, 32] + assert conv.bytes_per_head is None + assert [tuple(int(x) for x in e) for e in conv.buffer_entries] == [(0, 64, 64), (1, 192, 64)] + + ple = by_role[frozenset({str(MambaRole.PLE_NGRAM_CONTEXT)})] + assert ple.pool_idx == 1 + assert ple.mapper_kind == MapperKind.REPLICATED + assert ple.bytes_per_layer == 32 + assert [tuple(int(x) for x in e) for e in ple.buffer_entries] == [(1, 0, 32)] + + # extract_slot addresses layer regions inside the descriptor slot. + # (Look views up by role: PoolView.__eq__ compares numpy entries.) + extractor = KVRegionExtractorV1(page_table) + view_index = {pv.pool_role: idx for idx, pv in enumerate(recurrent.pool_views)} + conv_idx = view_index[MAMBA_CONV_ROLE] + ple_idx = view_index[frozenset({str(MambaRole.PLE_NGRAM_CONTEXT)})] + assert extractor.extract_slot(2, 0, conv_idx).memory.ptrs.tolist() == [ + 0x10000 + 2 * 256 + 64, + 0x10000 + 2 * 256 + 192, + ] + assert extractor.extract_slot(3, 0, ple_idx).memory.ptrs.tolist() == [0x20000 + 3 * 32] - assert pool.slot_bytes == state_bytes - assert pool.slot_stride_bytes == 2 * state_bytes - assert pool.layer_stride_bytes == 2 * state_bytes + # Attention life cycle is untouched by the recurrent one. + assert attention.pool_group_idx == 1 + assert [ll.global_layer_id for ll in attention.local_layers] == [2] -def test_v2_mamba_state_pool_rejects_non_affine_layer_offsets(): - storage = torch.empty((3, 6, 64), dtype=torch.uint8) - states = [storage[:, 0, :], storage[:, 2, :], storage[:, 5, :]] +def test_v2_builder_rejects_sectioned_view_in_attention_life_cycle(): + from types import SimpleNamespace - with pytest.raises(ValueError, match="uniform layer stride"): - _build_v2_mamba_state_pool(states) + from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import Role + + attrs = { + (0, Role.KEY): SimpleNamespace(pool_index=0, offset=0, size=64), + (0, Role.VALUE): SimpleNamespace(pool_index=0, offset=64, size=64), + (1, Role.KEY): SimpleNamespace(pool_index=0, offset=128, size=64), + (1, Role.VALUE): SimpleNamespace(pool_index=0, offset=192, size=64), + } + manager = _make_fake_v2_manager( + attrs, + {Role.ALL: MapperKind.INDEXED, Role.KEY: MapperKind.SECTIONED}, + slot_bytes_list=(256,), + ) + with pytest.raises(ValueError, match="no mapper for SECTIONED"): + build_page_table_from_manager(manager) + + +def test_v2_builder_rejects_section_bytes_that_do_not_cover_the_layer(): + manager = _make_fake_v2_mamba_manager() + manager.conv_section_dims = [2, 2, 2] # 48 bytes, but the conv buffer is 64 + with pytest.raises(ValueError, match="do not sum to bytes_per_layer"): + build_page_table_from_manager(manager) def test_v2_mamba_registration_uses_coalesced_physical_pool(): @@ -841,6 +1003,30 @@ def test_v2_mamba_registration_uses_coalesced_physical_pool(): ] +def test_legacy_mamba_pool_views_address_layers_past_four_gib() -> None: + """V1 pools are layer-major, so a layer offset spans the whole pool.""" + from tensorrt_llm._torch.disaggregation.resource.page import LocalLayer, PhysicalPool + + num_layers = 4 + slot_bytes = 1 << 20 + num_slots = 2048 # 2 GiB per layer: layer 2 already sits past 4 GiB + conv_pool = PhysicalPool(base_address=1000, slot_bytes=slot_bytes, num_slots=num_slots) + ssm_pool = PhysicalPool(base_address=8000, slot_bytes=slot_bytes, num_slots=num_slots) + local_layers = [ + LocalLayer(local_layer_id=lid, global_layer_id=10 + lid) for lid in range(num_layers) + ] + + conv_view, ssm_view = _build_mamba_pool_views( + conv_pool, ssm_pool, local_layers, conv_section_bytes=[1 << 18] * 4, ssm_bytes_per_head=4096 + ) + assert conv_view.section_bytes == [1 << 18] * 4 + assert ssm_view.bytes_per_head == 4096 + + assert [int(entry["offset"]) for entry in conv_view.buffer_entries] == [ + lid * num_slots * slot_bytes for lid in range(num_layers) + ] + + def test_legacy_mamba_registration_uses_layer_major_pools() -> None: import numpy as np @@ -960,6 +1146,7 @@ def test_mixed_page_table_serialization(): pool_role=MAMBA_CONV_ROLE, mapper_kind=MapperKind.SECTIONED, bytes_per_layer=1024, + section_bytes=[256, 128, 128], ), PoolView( pool_idx=1, @@ -969,14 +1156,13 @@ def test_mixed_page_table_serialization(): pool_role=MAMBA_SSM_ROLE, mapper_kind=MapperKind.INDEXED, bytes_per_layer=2048, + bytes_per_head=64, ), ] mamba_lg = MambaLayerGroup( pool_group_idx=1, local_layers=mamba_local_layers, pool_views=mamba_pool_views, - conv_section_bytes=[256, 128, 128], - ssm_bytes_per_head=64, ) page_table = KVCachePageTable( diff --git a/tests/unittest/disaggregated/test_kda_mamba_transfer.py b/tests/unittest/disaggregated/test_kda_mamba_transfer.py index 6250f9574ec0..9494849f23fd 100644 --- a/tests/unittest/disaggregated/test_kda_mamba_transfer.py +++ b/tests/unittest/disaggregated/test_kda_mamba_transfer.py @@ -221,10 +221,10 @@ def test_kda_layer_group_descriptors(enable_attention_dp): assert ssm_pv.bytes_per_layer == SSM_SLOT_BYTES # qwen3_next 3-sectioning: equal sections summing to the conv slot. - assert mlg.conv_section_bytes == [CONV_SLOT_BYTES // 3] * 3 - assert sum(mlg.conv_section_bytes) == conv_pv.bytes_per_layer - assert mlg.ssm_bytes_per_head == KDA_HEAD_DIM * KDA_HEAD_DIM * SSM_DTYPE.itemsize - assert ssm_pv.bytes_per_layer // mlg.ssm_bytes_per_head == KDA_NUM_HEADS + assert conv_pv.section_bytes == [CONV_SLOT_BYTES // 3] * 3 + assert sum(conv_pv.section_bytes) == conv_pv.bytes_per_layer + assert ssm_pv.bytes_per_head == KDA_HEAD_DIM * KDA_HEAD_DIM * SSM_DTYPE.itemsize + assert ssm_pv.bytes_per_layer // ssm_pv.bytes_per_head == KDA_NUM_HEADS # local_layers cover exactly the KDA layers (by global_layer_id). assert sorted(ll.global_layer_id for ll in mlg.local_layers) == [ @@ -356,6 +356,7 @@ def _synthetic_kda_page_table(ssm_slot_bytes: int, conv_slot_bytes: int, layer_i pool_role=MAMBA_CONV_ROLE, mapper_kind=MapperKind.SECTIONED, bytes_per_layer=conv_slot_bytes, + section_bytes=[conv_slot_bytes // 3] * 3, ), PoolView( pool_idx=1, @@ -365,14 +366,13 @@ def _synthetic_kda_page_table(ssm_slot_bytes: int, conv_slot_bytes: int, layer_i pool_role=MAMBA_SSM_ROLE, mapper_kind=MapperKind.INDEXED, bytes_per_layer=ssm_slot_bytes, + bytes_per_head=KDA_HEAD_DIM * KDA_HEAD_DIM * SSM_DTYPE.itemsize, ), ] mlg = MambaLayerGroup( pool_group_idx=0, local_layers=local_layers, pool_views=pool_views, - conv_section_bytes=[conv_slot_bytes // 3] * 3, - ssm_bytes_per_head=KDA_HEAD_DIM * KDA_HEAD_DIM * SSM_DTYPE.itemsize, ) return KVCachePageTable( tokens_per_block=8, diff --git a/tests/unittest/disaggregated/test_mamba_transfer.py b/tests/unittest/disaggregated/test_mamba_transfer.py index 44b7b8040982..f31a4ca7a17c 100644 --- a/tests/unittest/disaggregated/test_mamba_transfer.py +++ b/tests/unittest/disaggregated/test_mamba_transfer.py @@ -33,7 +33,11 @@ import tensorrt_llm.bindings import tensorrt_llm.tensorrt_llm_transfer_agent_binding # noqa: F401 from tensorrt_llm import DisaggregatedParams, Mapping, SamplingParams +from tensorrt_llm._torch.disaggregation.native.mixers.attention.peer import AttentionPolicy +from tensorrt_llm._torch.disaggregation.native.mixers.ssm import peer +from tensorrt_llm._torch.disaggregation.native.rank_info import RankInfo from tensorrt_llm._torch.disaggregation.transceiver import KvCacheTransceiverV2 +from tensorrt_llm._torch.pyexecutor.config_utils import Qwen4ExpPLECacheParams from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import ( MambaHybridCacheManagerV2, MixedMambaHybridCacheManager, @@ -208,6 +212,7 @@ def _create_managers( enable_attention_dp=False, use_v2=False, conv_state_layout="x_b_c", + with_ple=False, ): """Create Mamba hybrid cache managers for all TP ranks (PP=1). @@ -228,6 +233,17 @@ def _create_managers( if use_v2 else {} ) + if with_ple: + if not use_v2: + raise ValueError("PLE state requires the V2 cache manager") + manager_kwargs["qwen4_exp_ple_cache_params"] = Qwen4ExpPLECacheParams( + ple_layer_mask=[False, False, True, False, False], + num_ple_layers=1, + short_conv_channels=12, + short_conv_state_len=3, + ngram_context_len=2, + conv_state_dtype=torch.float32, + ) mgr = manager_cls( mamba_d_state=MAMBA_D_STATE, mamba_d_conv=MAMBA_D_CONV, @@ -276,6 +292,11 @@ def _zero_mamba_states(manager): for layer_idx in _mamba_layer_ids(manager): manager.get_conv_states(layer_idx).zero_() manager.get_ssm_states(layer_idx).zero_() + if isinstance(manager, MambaHybridCacheManagerV2): + for state in manager._ple_conv_states.values(): + state.zero_() + for state in manager._ple_ngram_contexts.values(): + state.zero_() def test_mamba_receiver_payload_bytes_matched_tp(): @@ -298,6 +319,7 @@ def test_mamba_receiver_payload_bytes_matched_tp(): conv_slot_bytes = 1024 ssm_slot_bytes = 2048 + side_slot_bytes = 128 sorted_lids = [0, 1] conv_pool = PhysicalPool(base_address=0xA000, slot_bytes=conv_slot_bytes, num_slots=8) @@ -311,6 +333,7 @@ def test_mamba_receiver_payload_bytes_matched_tp(): pool_role=MAMBA_CONV_ROLE, mapper_kind=MapperKind.SECTIONED, bytes_per_layer=conv_slot_bytes, + section_bytes=[512, 256, 256], ), PoolView( pool_idx=1, @@ -320,6 +343,14 @@ def test_mamba_receiver_payload_bytes_matched_tp(): pool_role=MAMBA_SSM_ROLE, mapper_kind=MapperKind.INDEXED, bytes_per_layer=ssm_slot_bytes, + bytes_per_head=64, + ), + PoolView( + pool_idx=2, + buffer_entries=np.array([(1, 0, side_slot_bytes)], dtype=BUFFER_ENTRY_DTYPE), + pool_role=frozenset({"ple_ngram_context"}), + mapper_kind=MapperKind.REPLICATED, + bytes_per_layer=side_slot_bytes, ), ] local_layers = [ @@ -331,21 +362,152 @@ def test_mamba_receiver_payload_bytes_matched_tp(): pool_group_idx=1, local_layers=local_layers, pool_views=pool_views, - conv_section_bytes=[512, 256, 256], - ssm_bytes_per_head=64, ) pt = KVCachePageTable( tokens_per_block=8, layer_groups=[mlg], pool_groups=[ PhysicalPoolGroup(pools=[]), # idx 0 placeholder - PhysicalPoolGroup(pools=[conv_pool, ssm_pool]), + PhysicalPoolGroup( + pools=[ + conv_pool, + ssm_pool, + PhysicalPool( + base_address=0xC000, + slot_bytes=side_slot_bytes, + num_slots=8, + ), + ] + ), ], ) got = mamba_receiver_payload_bytes(sender_page_table=pt, receiver_page_table=pt, dst_slot=3) - # Full per-rank slot bytes for 2 layers x (conv + ssm) - assert got == 2 * (conv_slot_bytes + ssm_slot_bytes) + # Standard state covers both layers; the side role exists only on layer 2. + assert got == 2 * (conv_slot_bytes + ssm_slot_bytes) + side_slot_bytes + + +def _make_rank_info(*, tp_size: int, tp_rank: int = 0) -> RankInfo: + return RankInfo( + instance_name="test", + instance_rank=tp_rank, + tp_size=tp_size, + tp_rank=tp_rank, + pp_size=1, + pp_rank=0, + layer_num_per_pp=[2], + sender_endpoints=[], + self_endpoint="", + transfer_engine_info=b"", + ) + + +def _make_side_state_page_table(*, include_side: bool = True): + from tensorrt_llm._torch.disaggregation.resource.page import ( + BUFFER_ENTRY_DTYPE, + MAMBA_CONV_ROLE, + MAMBA_SSM_ROLE, + KVCachePageTable, + LocalLayer, + MambaLayerGroup, + MapperKind, + PhysicalPool, + PhysicalPoolGroup, + PoolView, + ) + + local_layers = [ + LocalLayer(local_layer_id=0, global_layer_id=1), + LocalLayer(local_layer_id=1, global_layer_id=2), + ] + pools = [ + PhysicalPool(base_address=1000, slot_bytes=16, num_slots=4), + PhysicalPool(base_address=2000, slot_bytes=16, num_slots=4), + ] + views = [ + PoolView( + pool_idx=0, + buffer_entries=np.array([(0, 0, 16), (1, 16, 16)], dtype=BUFFER_ENTRY_DTYPE), + pool_role=MAMBA_CONV_ROLE, + mapper_kind=MapperKind.SECTIONED, + bytes_per_layer=16, + section_bytes=[4, 4, 8], + ), + PoolView( + pool_idx=1, + buffer_entries=np.array([(0, 0, 16), (1, 16, 16)], dtype=BUFFER_ENTRY_DTYPE), + pool_role=MAMBA_SSM_ROLE, + mapper_kind=MapperKind.INDEXED, + bytes_per_layer=16, + bytes_per_head=16, + ), + ] + if include_side: + pools.append(PhysicalPool(base_address=3000, slot_bytes=8, num_slots=4)) + views.append( + PoolView( + pool_idx=2, + buffer_entries=np.array([(1, 0, 8)], dtype=BUFFER_ENTRY_DTYPE), + pool_role=frozenset({"ple_ngram_context"}), + mapper_kind=MapperKind.REPLICATED, + bytes_per_layer=8, + ) + ) + return KVCachePageTable( + tokens_per_block=8, + layer_groups=[ + MambaLayerGroup( + pool_group_idx=0, + local_layers=local_layers, + pool_views=views, + ) + ], + pool_groups=[PhysicalPoolGroup(pools=pools)], + ) + + +def test_mamba_policy_rejects_incompatible_replicated_side_state(): + rank_info = _make_rank_info(tp_size=1) + pt_with_side = _make_side_state_page_table() + pt_without_side = _make_side_state_page_table(include_side=False) + + # Replicated recurrent state is validated by MambaPolicy; the attention + # policy only looks at PAGED groups and must not reject this pair. + with pytest.raises(ValueError, match="differ on overlapping layers"): + peer.MambaPolicy.validate_peer_compatible( + rank_info, rank_info, pt_with_side, pt_without_side + ) + AttentionPolicy.validate_peer_compatible(pt_with_side, pt_without_side) + + # Same roles but a different per-layer size is rejected up front. + pt_wider_side = _make_side_state_page_table() + side_view = pt_wider_side.layer_groups[0].pool_views[2] + side_view.buffer_entries = np.array([(1, 0, 16)], dtype=side_view.buffer_entries.dtype) + side_view.bytes_per_layer = 16 + with pytest.raises(ValueError, match="replicated state size differs"): + peer.MambaPolicy.validate_peer_compatible(rank_info, rank_info, pt_with_side, pt_wider_side) + + # REPLICATED state under the STATE life cycle uses the per-layer-pointer + # head-match mapper and requires equal per-layer sizes. + state_lg = pt_with_side.layer_groups[0] + common = dict( + peer_ri=rank_info, + mapper_kind=peer.MapperKind.REPLICATED, + self_layer_offsets=np.array([0], dtype=np.int64), + peer_layer_offsets=np.array([0], dtype=np.int64), + src_layer_off=0, + dst_layer_off=0, + self_lg=state_lg, + peer_lg=state_lg, + ) + with pytest.raises(ValueError, match="Replicated state size differs"): + peer.MambaPolicy(rank_info).build_mapper( + self_bytes_per_layer=8, peer_bytes_per_layer=16, **common + ) + mapper = peer.MambaPolicy(rank_info).build_mapper( + self_bytes_per_layer=8, peer_bytes_per_layer=8, **common + ) + assert isinstance(mapper, peer.MambaHeadMatchMapper) # --------------------------------------------------------------------------- @@ -472,6 +634,31 @@ def _read_actual(gen_managers, gen_request_ids) -> Dict: return actual +def _write_ple_ground_truth_to_ctx(managers, request_ids): + """Populate the replicated PLE state and return CPU reference tensors.""" + expected = {} + for req_idx, request_id in enumerate(request_ids): + for layer_idx in sorted(managers[0]._ple_conv_states): + conv_state = managers[0]._ple_conv_states[layer_idx] + ngram_context = managers[0]._ple_ngram_contexts[layer_idx] + conv = ( + torch.arange(conv_state[0].numel(), dtype=conv_state.dtype) + .reshape(conv_state[0].shape) + .add_(1000 * (req_idx + 1) + layer_idx) + ) + ngram = ( + torch.arange(ngram_context[0].numel(), dtype=ngram_context.dtype) + .reshape(ngram_context[0].shape) + .add_(100 * (req_idx + 1) + layer_idx) + ) + expected[(req_idx, layer_idx)] = (conv, ngram) + for manager in managers: + slot = _mamba_state_slot(manager, request_id) + manager._ple_conv_states[layer_idx][slot].copy_(conv) + manager._ple_ngram_contexts[layer_idx][slot].copy_(ngram) + return expected + + # --------------------------------------------------------------------------- # Main test logic # --------------------------------------------------------------------------- @@ -509,6 +696,7 @@ def run_mamba_transfer_test( gen_tp: int, use_v2: bool = False, conv_state_layout: str = "x_b_c", + with_ple: bool = False, ): """Test mamba transfer: ctx_tp -> gen_tp (PP=1, no DP).""" # -- 1. Create managers, zero mamba caches -- @@ -516,11 +704,13 @@ def run_mamba_transfer_test( ctx_tp, use_v2=use_v2, conv_state_layout=conv_state_layout, + with_ple=with_ple, ) gen_mgrs = _create_managers( gen_tp, use_v2=use_v2, conv_state_layout=conv_state_layout, + with_ple=with_ple, ) for mgr in ctx_mgrs + gen_mgrs: _zero_mamba_states(mgr) @@ -618,6 +808,7 @@ def run_mamba_transfer_test( ctx_rids, conv_state_layout, ) + expected_ple = _write_ple_ground_truth_to_ctx(ctx_mgrs, ctx_rids) if with_ple else {} # -- 6. Compute expected BEFORE transfer -- expected = _compute_expected( @@ -663,6 +854,28 @@ def run_mamba_transfer_test( ), ) + # Only the V2 manager owns PLE state. + if with_ple: + for gen_rank, manager in enumerate(gen_mgrs): + for req_idx, request_id in enumerate(gen_rids): + slot = _mamba_state_slot(manager, request_id) + for layer_idx in sorted(manager._ple_conv_states): + expected_conv, expected_ngram = expected_ple[(req_idx, layer_idx)] + torch.testing.assert_close( + manager._ple_conv_states[layer_idx][slot].cpu(), + expected_conv, + rtol=0, + atol=0, + msg=lambda msg, r=gen_rank: f"PLE conv mismatch on gen rank {r}: {msg}", + ) + torch.testing.assert_close( + manager._ple_ngram_contexts[layer_idx][slot].cpu(), + expected_ngram, + rtol=0, + atol=0, + msg=lambda msg, r=gen_rank: f"PLE n-gram mismatch on gen rank {r}: {msg}", + ) + # -- 10. Cleanup -- for mgr in ctx_mgrs + gen_mgrs: mgr.shutdown() @@ -704,3 +917,9 @@ def test_v2_mamba_transfer(ctx_tp, gen_tp, conv_state_layout): use_v2=True, conv_state_layout=conv_state_layout, ) + + +@pytest.mark.timeout(180) +def test_v2_mamba_transfer_with_replicated_ple_state(): + """Transfer PLE state once while standard Mamba state contracts from TP2 to TP1.""" + run_mamba_transfer_test(2, 1, use_v2=True, with_ple=True) diff --git a/tests/unittest/disaggregated/test_peer.py b/tests/unittest/disaggregated/test_peer.py index 4c5c2698d4d0..dd55120214c0 100644 --- a/tests/unittest/disaggregated/test_peer.py +++ b/tests/unittest/disaggregated/test_peer.py @@ -23,12 +23,18 @@ build_aux_transfer_layout, ) from tensorrt_llm._torch.disaggregation.native.mixers.attention.peer import ( + AttentionPolicy, HNDHeadMismatchMapper, IntactMapper, NHDHeadMismatchMapper, ReplicatedMapper, ) from tensorrt_llm._torch.disaggregation.native.mixers.attention.spec import AttentionInfo +from tensorrt_llm._torch.disaggregation.native.mixers.ssm.peer import ( + ConvStateMismatchMapper, + MambaHeadMatchMapper, + MambaPolicy, +) from tensorrt_llm._torch.disaggregation.native.peer import PeerOverlap, PeerRegistrar from tensorrt_llm._torch.disaggregation.native.rank_info import RankInfo from tensorrt_llm._torch.disaggregation.resource.kv_extractor import KVRegionExtractorV1 @@ -37,6 +43,7 @@ MAMBA_CONV_ROLE, MAMBA_SSM_ROLE, AttentionLayerGroup, + CacheKind, KVCachePageTable, LocalLayer, MambaLayerGroup, @@ -110,6 +117,7 @@ def make_page_table(pool_ptrs=None, block_bytes=None, global_layer_ids=None, mam pool_role=MAMBA_CONV_ROLE, mapper_kind=MapperKind.SECTIONED, bytes_per_layer=conv_slot_bytes, + section_bytes=[1024 // mamba_tp, 512 // mamba_tp, 512 // mamba_tp], ), PoolView( pool_idx=1, @@ -119,14 +127,13 @@ def make_page_table(pool_ptrs=None, block_bytes=None, global_layer_ids=None, mam pool_role=MAMBA_SSM_ROLE, mapper_kind=MapperKind.INDEXED, bytes_per_layer=ssm_slot_bytes, + bytes_per_head=64, ), ] mamba_lg = MambaLayerGroup( pool_group_idx=1, local_layers=mamba_local_layers, pool_views=mamba_pool_views, - conv_section_bytes=[1024 // mamba_tp, 512 // mamba_tp, 512 // mamba_tp], - ssm_bytes_per_head=64, ) conv_pool = PhysicalPool(base_address=0xA000, slot_bytes=conv_slot_bytes, num_slots=128) ssm_pool = PhysicalPool(base_address=0xB000, slot_bytes=ssm_slot_bytes, num_slots=128) @@ -1256,3 +1263,104 @@ def test_indexed_head_mismatch_inconsistent_slot_geometry_raises(): self_buffers_per_layer=2, peer_buffers_per_layer=2, ) + + +# --------------------------------------------------------------------------- +# Policy dispatch: the policy is chosen by CacheKind; each policy picks the +# mapper for a view from its MapperKind, including REPLICATED. +# --------------------------------------------------------------------------- + + +def _add_replicated_side_view(page_table, *, bytes_per_layer=8): + """Graft a replicated recurrent side view (e.g. PLE) onto the mamba group.""" + mamba_lg = page_table.layer_groups[1] + pool_group = page_table.pool_groups[mamba_lg.pool_group_idx] + pool_group.pools.append( + PhysicalPool(base_address=0xC000, slot_bytes=bytes_per_layer, num_slots=128) + ) + mamba_lg.pool_views.append( + PoolView( + pool_idx=len(pool_group.pools) - 1, + buffer_entries=np.array( + [(0, 0, bytes_per_layer), (1, bytes_per_layer, bytes_per_layer)], + dtype=BUFFER_ENTRY_DTYPE, + ), + pool_role=frozenset({"ple_ngram_context"}), + mapper_kind=MapperKind.REPLICATED, + bytes_per_layer=bytes_per_layer, + ) + ) + return len(mamba_lg.pool_views) - 1 + + +def test_policy_dispatch_keys_on_cache_kind_only(): + reg = _make_peer_registrar(make_rankinfo(instance_name="local")) + assert isinstance(reg._get_policy(CacheKind.PAGED), AttentionPolicy) + assert isinstance(reg._get_policy(CacheKind.STATE), MambaPolicy) + + +def test_mamba_replicated_side_view_elects_one_sender(): + """Replicated side state elects one sender per fan-in group. + + Sharded conv state on the same layer group is sent by every mamba TP rank. + """ + peer_pt = make_page_table(mamba_tp=1) + _add_replicated_side_view(peer_pt) + peer_ri = make_rankinfo(instance_name="peer", tp_size=1, page_table=peer_pt) + overlap = PeerOverlap() + + conv_ownership = [] + side_ownership = [] + for tp_rank in range(2): + self_pt = make_page_table(mamba_tp=2) + side_idx = _add_replicated_side_view(self_pt) + self_ri = make_rankinfo( + instance_name="local", tp_size=2, tp_rank=tp_rank, page_table=self_pt + ) + reg = _make_peer_registrar(self_ri) + conv_ownership.append(reg.should_send_pool(overlap, peer_ri, 1, 0)) + side_ownership.append(reg.should_send_pool(overlap, peer_ri, 1, side_idx)) + + assert conv_ownership == [True, True] + assert side_ownership == [True, False] + + +def test_mamba_replicated_side_view_copies_whole_layers_under_tp_mismatch(): + """REPLICATED state ignores mamba TP. + + The head-match mapper is used even when conv/SSM on the same group need a + resharding mapper. + """ + self_pt = make_page_table(mamba_tp=2) + self_side = _add_replicated_side_view(self_pt) + peer_pt = make_page_table(mamba_tp=1) + peer_side = _add_replicated_side_view(peer_pt) + self_ri = make_rankinfo(instance_name="local", tp_size=2, page_table=self_pt) + peer_ri = make_rankinfo(instance_name="peer", tp_size=1, page_table=peer_pt) + reg = _make_peer_registrar(self_ri) + + mapper = reg.get_kv_map(peer_ri, (1, self_side), (1, peer_side)) + assert isinstance(mapper, MambaHeadMatchMapper) + conv_mapper = reg.get_kv_map(peer_ri, (1, 0), (1, 0)) + assert isinstance(conv_mapper, ConvStateMismatchMapper) + + +def test_attention_replicated_validation_ignores_state_groups(): + """A STATE-only replicated-role mismatch is MambaPolicy's to reject.""" + self_pt = make_page_table() + _add_replicated_side_view(self_pt) + peer_pt = make_page_table() + AttentionPolicy.validate_peer_compatible(self_pt, peer_pt) + ri = make_rankinfo(instance_name="local", page_table=self_pt) + with pytest.raises(ValueError, match="MambaPolicy.*differ on overlapping layers"): + MambaPolicy.validate_peer_compatible(ri, ri, self_pt, peer_pt) + + +def test_attention_replicated_validation_rejects_missing_index_key_view(): + self_pt = make_page_table(global_layer_ids=[0]) + view = self_pt.layer_groups[0].pool_views[0] + view.pool_role = frozenset({"index_key"}) + view.mapper_kind = MapperKind.REPLICATED + peer_pt = make_page_table(global_layer_ids=[0]) + with pytest.raises(ValueError, match="AttentionPolicy.*differ on overlapping layers"): + AttentionPolicy.validate_peer_compatible(self_pt, peer_pt) From 63d217f252ea108cf4b469fa33ebfac997fccf92 Mon Sep 17 00:00:00 2001 From: TensorRT LLM AI Agent <296075020+trtllm-agent@users.noreply.github.com> Date: Tue, 15 Sep 2026 09:58:08 +0800 Subject: [PATCH 4/8] [None][test] Remove 59 closed-bug waive entries for main (#19116) Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> --- tests/integration/test_lists/waives.txt | 46 ------------------------- 1 file changed, 46 deletions(-) diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 35305a169afa..5e28af2b8dbf 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -6,9 +6,6 @@ accuracy/test_dwdp_aggregated.py::TestDwdpAggDeepSeekV3Lite::test_dwdp_agg_accur accuracy/test_kimi3.py::TestKimiK3DSpark::test_gsm8k_tep8 SKIP (https://nvbugs/6766812) accuracy/test_llm_api_pytorch.py::TestDeepSeekR1::test_nvfp4_multi_gpus[throughput_pp4_mtp] SKIP (https://nvbugs/6481323) accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_nvfp4_multi_gpus_piecewise_cuda_graph[mtp3_fp8kv_chunked] SKIP (https://nvbugs/6739553) -accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=2-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=False-enable_chunked_prefill=False-v2_kv_cache=True] SKIP (https://nvbugs/6633927) -accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=2-attention_dp=False-cuda_graph=False-overlap_scheduler=True-torch_compile=False-enable_chunked_prefill=False-v2_kv_cache=True] SKIP (https://nvbugs/6633927) -accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=2-attention_dp=False-cuda_graph=True-overlap_scheduler=False-torch_compile=False-enable_chunked_prefill=False-v2_kv_cache=False] SKIP (https://nvbugs/6633927) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[ep4-mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=True] SKIP (https://nvbugs/6402058) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[pp4-mtp_nextn=0-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=False] SKIP (https://nvbugs/6278337) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[pp4-mtp_nextn=0-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=True] SKIP (https://nvbugs/6428057) @@ -34,9 +31,6 @@ accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_4gpus[v2_kv_cache-ep4-cutl accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_4gpus[v2_kv_cache-tp4-cutlass-auto] SKIP (https://nvbugs/5596343) accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_chunked_prefill[cutlass-auto] SKIP (https://nvbugs/5596343) accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4[use_msa=True] SKIP (https://nvbugs/6601633) -accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_4gpus_hopper_w4a16 SKIP (https://nvbugs/6644472) -accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_8gpus_mtp SKIP (https://nvbugs/6581065) -accuracy/test_llm_api_pytorch.py::TestNemotronV3Ultra::test_nvfp4_4gpus_static_eplb[moe_backend=CUTLASS] SKIP (https://nvbugs/6535767) accuracy/test_llm_api_pytorch.py::TestQwen3_235B_A22B::test_fp8[latency] SKIP (https://nvbugs/6177390) accuracy/test_llm_api_pytorch.py::TestQwen3_235B_A22B::test_fp8[throughput_latency] SKIP (https://nvbugs/6177390) accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_nvfp4[tep4_latency_moe_cutlass-torch_compile=True] SKIP (https://nvbugs/6561558) @@ -48,7 +42,6 @@ accuracy/test_llm_api_pytorch.py::TestQwen3_8_Flash_Next::test_fp8_adp4_mtp3_trt accuracy/test_llm_api_pytorch_multimodal.py::TestMistralLarge3_675B::test_nvfp4_4gpus[latency_moe_trtllm] SKIP (https://nvbugs/6665906) cpp/test_multi_gpu.py::test_cache_transceiver[8proc-mooncake_kvcache-90] SKIP (https://nvbugs/5838199) cpp/test_multi_gpu.py::test_cache_transceiver[8proc-ucx_kvcache-90] SKIP (https://nvbugs/5838199) -disaggregated/test_auto_scaling.py::test_disagg_server_restart[etcd-round_robin] SKIP (https://nvbugs/6611817) disaggregated/test_disaggregated.py::test_disaggregated_cancel_large_context_requests[DeepSeek-V3-Lite-bf16] SKIP (https://nvbugs/6105768) disaggregated/test_disaggregated.py::test_disaggregated_deepseek_v3_lite_bf16_cache_aware_balance[DeepSeek-V3-Lite-bf16] SKIP (https://nvbugs/6162322) disaggregated/test_disaggregated.py::test_disaggregated_deepseek_v3_lite_bf16_conditional[DeepSeek-V3-Lite-bf16] SKIP (https://nvbugs/6162322) @@ -114,18 +107,9 @@ full:A100/llmapi/test_llm_api_pytorch_moe_lora.py::test_qwen_moe_routed_expert_m full:A100/llmapi/test_llm_api_pytorch_moe_lora.py::test_qwen_moe_routed_expert_multi_lora_varying_ranks[eager] SKIP (https://nvbugs/6758594) full:A100/llmapi/test_llm_api_qa.py::TestLlmDefaultBackend::test_llm_args_logging SKIP (https://nvbugs/6727262) full:A100/llmapi/test_llm_api_qa.py::TestLlmDefaultBackend::test_llm_args_type_default SKIP (https://nvbugs/6727262) -full:B200/accuracy/test_disaggregated_serving.py::TestNemotron3Super120B::test_ctx_dp2_gen_tp4 SKIP (https://nvbugs/6577550) full:B200/accuracy/test_epd_disagg_multimodal.py::TestVideoMMEEPD::test_disaggregated_videomme[nemotron_nano_v3_omni_fp8] SKIP (https://nvbugs/6327718) -full:B200/accuracy/test_llm_api_pytorch.py::TestDeepSeekR1::test_nvfp4_multi_gpus[latency_adp_lmtp] SKIP (https://nvbugs/6695515) -full:B200/accuracy/test_llm_api_pytorch.py::TestDeepSeekR1::test_nvfp4_multi_gpus_corner_case SKIP (https://nvbugs/6695515) -full:B200/accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_nvfp4_multi_gpus[baseline_mtp1] SKIP (https://nvbugs/6695515) -full:B200/accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_nvfp4_multi_gpus[latency] SKIP (https://nvbugs/6695515) -full:B200/accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_nvfp4_multi_gpus_chunked_prefill[latency_qsplit] SKIP (https://nvbugs/6695515) -full:B200/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=TRTLLM-mtp_nextn=2-tp4-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-low_precision_combine=False-torch_compile=False] SKIP (https://nvbugs/6695515) -full:B200/accuracy/test_llm_api_pytorch.py::TestKimiK2::test_nvfp4[4gpus] SKIP (https://nvbugs/6649733) full:B200/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_mxfp8[use_msa=False] SKIP (https://nvbugs/6424188) full:B200/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4[use_msa=False] SKIP (https://nvbugs/6424188) -full:B200/accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_4gpus_block_reuse[DEP4_MTP_ON] SKIP (https://nvbugs/6695515) full:B200/accuracy/test_llm_api_pytorch.py::TestNemotronV3Ultra::test_nvfp4_4gpus_static_eplb[moe_backend=CUTEDSL] SKIP (https://nvbugs/6566659) full:B200/accuracy/test_llm_api_pytorch.py::TestQwen3_235B_A22B::test_nvfp4[latency_moe_trtllm] SKIP (https://nvbugs/6731978) full:B200/accuracy/test_llm_api_pytorch_multimodal.py::TestExaone4_5_33B::test_auto_dtype[forced_chunked_prefill] SKIP (https://nvbugs/6597570) @@ -158,25 +142,15 @@ full:B300/llmapi/test_llm_api_pytorch_moe_lora.py::test_qwen_moe_routed_expert_m full:B300/llmapi/test_llm_api_pytorch_moe_lora.py::test_qwen_moe_routed_expert_multi_lora_varying_ranks[eager] SKIP (https://nvbugs/6758594) full:DGX_B200/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=False-enable_chunked_prefill=False-v2_kv_cache=True] SKIP (https://nvbugs/6633268) full:DGX_B200/unittest/_torch/modules/kimi_kda/test_kda_prefill_op.py::test_kda_mixer_empty_prefill SKIP (https://nvbugs/6705034) -full:DGX_B200/unittest/_torch/visual_gen/test_trtllm_serve_e2e.py::TestFlux1TextToImage SKIP (https://nvbugs/6720944) -full:DGX_B200/unittest/_torch/visual_gen/test_trtllm_serve_e2e.py::TestFlux2TextToImage SKIP (https://nvbugs/6720944) full:DGX_B200/unittest/tools/test_layer_wise_benchmarks.py::test_performance_alignment[1] SKIP (https://nvbugs/6669275) -full:DGX_H100/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=2-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=False-enable_chunked_prefill=False-v2_kv_cache=False] SKIP (https://nvbugs/6633927) -full:DGX_H100/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=2-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=True-enable_chunked_prefill=False-v2_kv_cache=True] SKIP (https://nvbugs/6633927) -full:DGX_H100/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=2-attention_dp=False-cuda_graph=True-overlap_scheduler=False-torch_compile=False-enable_chunked_prefill=False-v2_kv_cache=True] SKIP (https://nvbugs/6633927) full:DGX_H100/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_fp8_block_scales[mtp=disable-fp8kv=True-attention_dp=False-cuda_graph=True-overlap_scheduler=True-torch_compile=True] SKIP (https://nvbugs/6700265) full:DGX_H100/accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_4gpus[v2_kv_cache-dp4-triton-auto] SKIP (https://nvbugs/6766636) full:GB200/accuracy/test_dwdp_disaggregated_serving.py::TestDwdpDeepSeekV3Lite::test_dwdp_accuracy SKIP (https://nvbugs/6276923) full:GB200/accuracy/test_dwdp_disaggregated_serving.py::TestDwdpDeepSeekV3Lite::test_dwdp_accuracy_contention_opt SKIP (https://nvbugs/6276923) full:GB200/accuracy/test_dwdp_disaggregated_serving.py::TestDwdpDeepSeekV3Lite::test_dwdp_accuracy_mode_b_overlap SKIP (https://nvbugs/6276923) -full:GB200/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=0-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=False-enable_chunked_prefill=False-v2_kv_cache=False] SKIP (https://nvbugs/6525896) full:GB200/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus_online_eplb[mtp_nextn=2-moe_backend=TRTLLM] SKIP (https://nvbugs/6690084) full:GB200/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_cute_dsl_bf16_gemm[cuda_graph=True] SKIP (https://nvbugs/6525897) -full:GB200/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_mxfp8[use_msa=False] SKIP (https://nvbugs/6479471) -full:GB200/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4[use_msa=False] SKIP (https://nvbugs/6479471) -full:GB200/accuracy/test_llm_api_pytorch.py::TestNemotronV3Ultra::test_nvfp4_4gpus_static_eplb[moe_backend=TRTLLM] SKIP (https://nvbugs/6525898) full:GB200/accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_nvfp4[dep4_latency_moe_cutlass-torch_compile=True] SKIP (https://nvbugs/5929339) -full:GB200/accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_fp8_moe_dflash SKIP (https://nvbugs/6316985) full:GB200/accuracy/test_llm_api_pytorch_multimodal.py::TestExaone4_5_33B::test_auto_dtype[forced_chunked_prefill] SKIP (https://nvbugs/6597570) full:GB200/accuracy/test_llm_api_pytorch_multimodal.py::TestExaone4_5_33B::test_auto_dtype[full_budget] SKIP (https://nvbugs/6597570) full:GB200/disaggregated/test_disaggregated.py::test_disaggregated_logprobs_serving[llama-3.1-8b-instruct] SKIP (https://nvbugs/6275959) @@ -186,9 +160,7 @@ full:GB300/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_fp8_block_ full:GB300/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_fp8_block_scales_4gpus[pp4-mtp_nextn=0-fp8kv=True-attention_dp=False-cuda_graph=True-overlap_scheduler=True-torch_compile=True-sampler_async_worker=False] SKIP (https://nvbugs/6661948) full:GB300/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=0-pp4-fp8kv=False-attention_dp=False-cuda_graph=False-overlap_scheduler=False-low_precision_combine=False-torch_compile=True] SKIP (https://nvbugs/6697099) full:GB300/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_mxfp8_piecewise_cuda_graph[use_msa=False] SKIP (https://nvbugs/6714109) -full:GB300/accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_fp8_moe_dflash SKIP (https://nvbugs/6316985) full:GB300/accuracy/test_llm_api_pytorch.py::TestQwen3_5_4B::test_bf16 SKIP (https://nvbugs/6487918) -full:GB300/accuracy/test_llm_api_pytorch.py::TestQwen3_8_2_4T_A95B::test_fp8_tp16_mtp3_trtllm SKIP (https://nvbugs/6694922) full:GB300/accuracy/test_llm_api_pytorch_multimodal.py::TestExaone4_5_33B::test_auto_dtype[forced_chunked_prefill] SKIP (https://nvbugs/6597570) full:GB300/accuracy/test_llm_api_pytorch_multimodal.py::TestExaone4_5_33B::test_auto_dtype[full_budget] SKIP (https://nvbugs/6597570) full:GB300/disaggregated/test_disaggregated.py::test_disaggregated_deepseek_v3_lite_fp8_attention_dp_overlap[DeepSeek-V3-Lite-fp8] SKIP (https://nvbugs/6581064) @@ -208,17 +180,9 @@ full:H100/llmapi/test_llm_api_qa.py::TestLlmDefaultBackend::test_llm_args_loggin full:H100/llmapi/test_llm_api_qa.py::TestLlmDefaultBackend::test_llm_args_type_default SKIP (https://nvbugs/6727262) full:H100/test_e2e.py::test_openai_chat_guided_decoding[openai/gpt-oss-120b] SKIP (https://nvbugs/6768488) full:H100/test_e2e.py::test_ptp_quickstart_advanced[Qwen3-30B-A3B-Qwen3/Qwen3-30B-A3B] SKIP (https://nvbugs/6700875) -full:H100_PCIe/unittest/auto_deploy/standalone/test_standalone_package.py::TestStandalonePackage::test_run_unit_tests SKIP (https://nvbugs/6672542) full:H20/accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype[mtp_nextn=2-overlap_scheduler=False] SKIP (https://nvbugs/6345827) full:H20/accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype[mtp_nextn=2-overlap_scheduler=True] SKIP (https://nvbugs/6345827) full:H20/accuracy/test_epd_disagg_multimodal.py::TestVideoMMEEPD::test_disaggregated_videomme[nemotron_nano_v3_omni_fp8] SKIP (https://nvbugs/6327718) -full:H20/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=2-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=False-enable_chunked_prefill=False-v2_kv_cache=False] SKIP (https://nvbugs/6662724) -full:H20/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=2-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=True-enable_chunked_prefill=False-v2_kv_cache=True] SKIP (https://nvbugs/6662724) -full:H20/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=2-attention_dp=False-cuda_graph=True-overlap_scheduler=False-torch_compile=False-enable_chunked_prefill=False-v2_kv_cache=True] SKIP (https://nvbugs/6662724) -full:H20/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=False-enable_chunked_prefill=True-v2_kv_cache=False] SKIP (https://nvbugs/6662724) -full:H20/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=False-enable_chunked_prefill=True-v2_kv_cache=True] SKIP (https://nvbugs/6662724) -full:H20/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=True-enable_chunked_prefill=True-v2_kv_cache=False] SKIP (https://nvbugs/6662724) -full:H20/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=True-enable_chunked_prefill=True-v2_kv_cache=True] SKIP (https://nvbugs/6662724) full:H20/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_fp8_block_scales_4gpus[pp4-mtp_nextn=0-fp8kv=False-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=False-sampler_async_worker=False] SKIP (https://nvbugs/6692005) full:H20/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_fp8_block_scales_4gpus[tp4-mtp_nextn=0-fp8kv=False-attention_dp=True-cuda_graph=False-overlap_scheduler=False-torch_compile=False-sampler_async_worker=False] SKIP (https://nvbugs/6692005) full:H20/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_fp8_block_scales_cuda_graph_padding_4gpus[attention_dp=True-mtp_nextn=0] SKIP (https://nvbugs/6692005) @@ -228,7 +192,6 @@ full:H20/accuracy/test_llm_api_pytorch_multimodal.py::TestExaone4_5_33B::test_au full:H20/accuracy/test_llm_api_pytorch_multimodal.py::TestExaone4_5_33B::test_auto_dtype[full_budget] SKIP (https://nvbugs/6597570) full:H20/accuracy/test_llm_api_pytorch_multimodal.py::TestNanoV3Omni::test_auto_dtype[bf16] SKIP (https://nvbugs/6618649) full:H20/disaggregated/test_disaggregated.py::test_disaggregated_logprobs_serving[llama-3.1-8b-instruct] SKIP (https://nvbugs/6275959) -full:H20/disaggregated/test_disaggregated.py::test_disaggregated_overlap_gen_first[ctx_pp1-TinyLlama-1.1B-Chat-v1.0] SKIP (https://nvbugs/6692009) full:L40S/disaggregated/test_disaggregated.py::test_disaggregated_ctxpp2_genpp2[TinyLlama-1.1B-Chat-v1.0] SKIP (https://nvbugs/6728119) full:L40S/disaggregated/test_disaggregated.py::test_disaggregated_ctxpp2_gentp2[TinyLlama-1.1B-Chat-v1.0] SKIP (https://nvbugs/6728119) full:L40S/disaggregated/test_disaggregated.py::test_disaggregated_ctxtp2_genpp2[TinyLlama-1.1B-Chat-v1.0] SKIP (https://nvbugs/6728119) @@ -242,8 +205,6 @@ full:RTX_PRO_6000_Blackwell_Server_Edition/accuracy/test_disaggregated_serving.p full:RTX_PRO_6000_Blackwell_Server_Edition/accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype[mtp_nextn=2-overlap_scheduler=True] SKIP (https://nvbugs/6451323) full:RTX_PRO_6000_Blackwell_Server_Edition/accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_guided_decoding[xgrammar-mtp_nextn=0] SKIP (https://nvbugs/6616033) full:RTX_PRO_6000_Blackwell_Server_Edition/accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_guided_decoding[xgrammar-mtp_nextn=2] SKIP (https://nvbugs/6451323) -full:RTX_PRO_6000_Blackwell_Server_Edition/accuracy/test_disaggregated_serving.py::TestGPTOSS::test_auto_dtype[False] SKIP (https://nvbugs/6567731) -full:RTX_PRO_6000_Blackwell_Server_Edition/accuracy/test_disaggregated_serving.py::TestGPTOSS::test_auto_dtype[True] SKIP (https://nvbugs/6567731) full:RTX_PRO_6000_Blackwell_Server_Edition/accuracy/test_llm_api_pytorch.py::TestDeepSeekR1::test_nvfp4_multi_gpus_sm120[throughput_tp8] SKIP (https://nvbugs/6616036) full:RTX_PRO_6000_Blackwell_Server_Edition/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=0-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=False-enable_chunked_prefill=False-v2_kv_cache=True] SKIP (https://nvbugs/6313072) full:RTX_PRO_6000_Blackwell_Server_Edition/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=0-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=True-enable_chunked_prefill=False-v2_kv_cache=False] SKIP (https://nvbugs/6313072) @@ -280,18 +241,14 @@ full:RTX_PRO_6000_Blackwell_Server_Edition/accuracy/test_llm_api_pytorch.py::Tes full:RTX_PRO_6000_Blackwell_Server_Edition/accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v2_kv_cache-cutlass-one_model-overlap_scheduler] SKIP (https://nvbugs/6672360) full:RTX_PRO_6000_Blackwell_Server_Edition/accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_vswa_reuse_4gpus[v1_kv_cache-one_model] SKIP (https://nvbugs/6672360) full:RTX_PRO_6000_Blackwell_Server_Edition/accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_1gpu[v1_kv_cache-True-True-cutlass-fp8] SKIP (https://nvbugs/6626640) -full:RTX_PRO_6000_Blackwell_Server_Edition/accuracy/test_llm_api_pytorch.py::TestQwen3_235B_A22B::test_nvfp4[latency_moe_cutlass] SKIP (https://nvbugs/6470267) full:sm100/unittest/bindings SKIP (Disable for Blackwell) llmapi/test_llm_api_pytorch_bart.py::test_bart_pytorch_generate_encoder_decoder_end_to_end[bf16-kv-v1-cuda-graph-on-greedy-tp2-bart-large-cnn] SKIP (https://nvbugs/6463812) llmapi/test_llm_api_pytorch_bart.py::test_mbart_pytorch_generate_encoder_decoder_end_to_end SKIP (https://nvbugs/6758881) llmapi/test_llm_examples.py::test_llmapi_speculative_decoding_eagle3 SKIP (https://nvbugs/6075431) perf/test_perf_sanity.py::test_e2e[aggr_upload-ctx_only-gb200_deepseek-r1-fp4_8k1k_con4096_ctx1_dep4_gen1_dep16_eplb0_mtp1_ccb-NIXL] SKIP (https://nvbugs/6627795) -perf/test_perf_sanity.py::test_e2e[aggr_upload-deepseek_r1_fp4_v2_2_nodes_grace_blackwell-r1_fp4_v2_tep8_mtp3] SKIP (https://nvbugs/6668776) perf/test_perf_sanity.py::test_e2e[aggr_upload-glm5_fp4_2_nodes_grace_blackwell-glm5_fp4_tep8_mtp3_8k1k] SKIP (https://nvbugs/6746175) perf/test_perf_sanity.py::test_e2e[disagg_upload-e2e-gb300_deepseek-v4-pro-dspark_agentx_con1456_ctx3_dep8_gen1_dep16_eplb0_dspark5_ccb-NIXL] SKIP (https://nvbugs/6746175) perf/test_visual_gen_perf_sanity.py::test_visual_gen_e2e[vg_upload-flux2_blackwell-flux2_fp8_cfg1_ulysses4_teacache_on] SKIP (https://nvbugs/6759009) -test_doc.py::test_http_url_validity SKIP (https://nvbugs/6709495) -test_doc.py::test_relative_path_validity SKIP (https://nvbugs/6731971) test_e2e.py::test_openai_chat_harmony SKIP (https://nvbugs/6751484) test_e2e.py::test_ptp_quickstart_advanced[Nemotron-Nano-9B-v2-nvfp4-NVIDIA-Nemotron-Nano-9B-v2-NVFP4] SKIP (https://nvbugs/6624972) test_e2e.py::test_ptp_quickstart_bert[TRTLLM-BertForSequenceClassification-bert/bert-base-uncased-yelp-polarity] SKIP (https://nvbugs/6605819) @@ -316,9 +273,6 @@ unittest/_torch/thop/parallel/test_fp4_linear.py::test_fp4_gemm_bias_per_backend unittest/_torch/thop/parallel/test_fp4_linear.py::test_fp4_gemm_bias_per_backend[mnk2-cutlass] SKIP (https://nvbugs/6581067) unittest/_torch/thop/parallel/test_fp4_linear.py::test_fp4_gemm_bias_per_backend[mnk3-cublaslt] SKIP (https://nvbugs/6581067) unittest/_torch/thop/parallel/test_fp4_linear.py::test_fp4_gemm_bias_per_backend[mnk3-cutlass] SKIP (https://nvbugs/6581067) -unittest/_torch/visual_gen/multi_gpu/test_ring_attention.py::TestRingAttention::test_ring_invalid_mask_raises SKIP (https://nvbugs/6412105) -unittest/_torch/visual_gen/multi_gpu/test_ulysses_async.py::test_capture_smoke SKIP (https://nvbugs/6385134) -unittest/_torch/visual_gen/multi_gpu/test_ulysses_attention.py SKIP (https://nvbugs/6311866) unittest/_torch/visual_gen/test_cosmos3_t2v_offload.py::TestCosmos3Offload::test_cosmos3_offload_matches_baseline SKIP (https://nvbugs/6702264) unittest/bindings/test_transfer_agent_bindings.py::TestMooncakeFunctionalTransfer::test_mooncake_wait_in_progress_on_zero_timeout SKIP (https://nvbugs/6517836) unittest/bindings/test_transfer_agent_bindings.py::TestMooncakeFunctionalTransfer::test_mooncake_write_transfer_gpu_tensor SKIP (https://nvbugs/6517836) From e22f1e926399ae6d10b3d874b830a2e5f756eea1 Mon Sep 17 00:00:00 2001 From: Aurelien Chartier <2567591+achartier@users.noreply.github.com> Date: Mon, 14 Sep 2026 19:38:35 -0700 Subject: [PATCH 5/8] [TRTLLM-16404][fix] thread LoRA params through MoE models (#19109) Signed-off-by: Aurelien Chartier <2567591+achartier@users.noreply.github.com> --- tensorrt_llm/_torch/models/modeling_afmoe.py | 11 +- tensorrt_llm/_torch/models/modeling_glm.py | 36 +- .../_torch/models/modeling_minimaxm2.py | 7 +- .../_torch/models/modeling_minimaxm3.py | 29 +- .../_torch/models/modeling_step3p7.py | 104 +++- .../peft/lora/cuda_graph_lora_params.py | 21 +- .../llmapi/test_llm_api_pytorch_moe_lora.py | 23 +- .../_torch/modeling/test_modeling_step3p7.py | 74 +++ .../peft/test_moe_lora_cuda_graph_params.py | 11 +- .../_torch/peft/test_moe_lora_model_path.py | 455 +++++++++++++++++- 10 files changed, 726 insertions(+), 45 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_afmoe.py b/tensorrt_llm/_torch/models/modeling_afmoe.py index 09d57ee0d027..8739c4887f51 100644 --- a/tensorrt_llm/_torch/models/modeling_afmoe.py +++ b/tensorrt_llm/_torch/models/modeling_afmoe.py @@ -198,6 +198,7 @@ def __init__( overridden_tp_size=1 if self.enable_attention_dp else None, reduce_output=False, layer_idx=layer_idx, + is_shared_expert=True, ) else: self.shared_experts = None @@ -215,6 +216,7 @@ def forward( self, hidden_states: torch.Tensor, attn_metadata: AttentionMetadata, + lora_params: Optional[dict] = None, ) -> torch.Tensor: all_rank_num_tokens = attn_metadata.all_rank_num_tokens router_logits = self.gate(hidden_states) @@ -224,10 +226,11 @@ def forward( router_logits, all_rank_num_tokens=all_rank_num_tokens, use_dp_padding=False, + lora_params=lora_params, ) if self.shared_experts is not None: - shared_output = self.shared_experts(hidden_states) + shared_output = self.shared_experts(hidden_states, lora_params=lora_params) final_output = shared_output.add_(routed_output) else: final_output = routed_output @@ -356,6 +359,7 @@ def forward( hidden_states: torch.Tensor, attn_metadata: AttentionMetadata, residual: Optional[torch.Tensor], + lora_params: Optional[dict] = None, **kwargs, ) -> torch.Tensor: if residual is None: @@ -368,6 +372,7 @@ def forward( position_ids=position_ids, hidden_states=hidden_states, attn_metadata=attn_metadata, + lora_params=lora_params, **kwargs, ) hidden_states = self.post_attention_layernorm(hidden_states) @@ -375,9 +380,9 @@ def forward( hidden_states, residual = self.pre_mlp_layernorm(hidden_states, residual) if self.moe_enabled: - hidden_states = self.mlp(hidden_states, attn_metadata) + hidden_states = self.mlp(hidden_states, attn_metadata, lora_params=lora_params) else: - hidden_states = self.mlp(hidden_states) + hidden_states = self.mlp(hidden_states, lora_params=lora_params) hidden_states = self.post_mlp_layernorm(hidden_states) diff --git a/tensorrt_llm/_torch/models/modeling_glm.py b/tensorrt_llm/_torch/models/modeling_glm.py index 5d282145060b..38669e007ed9 100644 --- a/tensorrt_llm/_torch/models/modeling_glm.py +++ b/tensorrt_llm/_torch/models/modeling_glm.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + import inspect import math import os @@ -316,6 +319,8 @@ def __init__( config=model_config, overridden_tp_size=shared_tp_size, reduce_output=False, + layer_idx=layer_idx, + is_shared_expert=True, ) self.allreduce = AllReduce( @@ -375,7 +380,12 @@ def _get_experts_quant_config(model_config, layer_idx: int) -> QuantConfig: ) def compute_routed_output( - self, hidden_states, hidden_states_fp4, all_rank_num_tokens, do_finalize + self, + hidden_states, + hidden_states_fp4, + all_rank_num_tokens, + do_finalize, + lora_params, ): # max-throughput use_dp_padding = False @@ -396,6 +406,7 @@ def compute_routed_output( output_dtype=hidden_states.dtype, all_rank_num_tokens=all_rank_num_tokens, use_dp_padding=use_dp_padding, + lora_params=lora_params, ) return routed_output @@ -407,13 +418,15 @@ def forward( all_rank_num_tokens: Optional[list[int]] = None, final_all_reduce_params: Optional[AllReduceParams] = None, do_finalize: Optional[bool] = True, + lora_params: Optional[dict] = None, ) -> torch.Tensor: if not do_finalize: assert not self.use_dp def _compute_shared_output(): shared_output = self.shared_experts( - hidden_states_fp4 if hidden_states_fp4 is not None else hidden_states + hidden_states_fp4 if hidden_states_fp4 is not None else hidden_states, + lora_params=lora_params, ) if self.shared_output_scale is not None: shared_output *= self.shared_output_scale @@ -421,7 +434,11 @@ def _compute_shared_output(): def _compute_routed_output(): routed_output = self.compute_routed_output( - hidden_states, hidden_states_fp4, all_rank_num_tokens, do_finalize + hidden_states, + hidden_states_fp4, + all_rank_num_tokens, + do_finalize, + lora_params, ) return routed_output @@ -565,6 +582,7 @@ def __init__( config=model_config, overridden_tp_size=self.mlp_tp_size, reduce_output=True, + layer_idx=layer_idx, ) self.input_layernorm = RMSNorm( @@ -645,6 +663,7 @@ def forward( attn_metadata: AttentionMetadata, residual: torch.Tensor, spec_metadata: Optional[SpecMetadata] = None, + lora_params: Optional[dict] = None, **kwargs, ) -> Tuple[torch.Tensor, torch.Tensor]: if residual is None: @@ -656,6 +675,7 @@ def forward( hidden_states=hidden_states, attn_metadata=attn_metadata, all_reduce_params=AllReduceParams(enable_allreduce=not (self.disable_attn_allreduce)), + lora_params=lora_params, **kwargs, ) if isinstance(self.mlp, Glm4MoE): @@ -666,6 +686,7 @@ def forward( attn_metadata=attn_metadata, residual=residual, spec_metadata=spec_metadata, + lora_params=lora_params, ) else: if spec_metadata is not None and spec_metadata.is_layer_capture(self.layer_idx): @@ -675,6 +696,7 @@ def forward( hidden_states=hidden_states, residual=residual, spec_metadata=spec_metadata, + lora_params=lora_params, ) def forward_MoE( @@ -683,6 +705,7 @@ def forward_MoE( attn_metadata: AttentionMetadata, residual: torch.Tensor, spec_metadata: Optional[SpecMetadata] = None, + lora_params: Optional[dict] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: def _run_MoE(hidden_states, hidden_states_fp4, do_finalize): return self.mlp( @@ -695,6 +718,7 @@ def _run_MoE(hidden_states, hidden_states_fp4, do_finalize): ) ), do_finalize=do_finalize, + lora_params=lora_params, ) if self.fusion_config.PRE_MOE_FUSION: @@ -772,6 +796,7 @@ def forward_mlp( hidden_states: torch.Tensor, residual: torch.Tensor, spec_metadata: Optional[SpecMetadata] = None, + lora_params: Optional[dict] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: if self.fusion_config.PRE_MLP_FUSION: act_fp4, act_sf, residual = self.allreduce( @@ -795,6 +820,7 @@ def forward_mlp( final_all_reduce_params=AllReduceParams( enable_allreduce=not (self.fusion_config.POST_MLP_FUSION or self.mlp_tp_size == 1) ), + lora_params=lora_params, ) if self.fusion_config.POST_MLP_FUSION: @@ -872,6 +898,7 @@ def forward( embed_tokens: Embedding, attn_metadata: AttentionMetadata, all_rank_num_tokens: Optional[List[int]] = None, + lora_params: Optional[dict] = None, **kwargs, ) -> torch.Tensor: def norm_embeds(): @@ -907,6 +934,7 @@ def norm_hidden(): hidden_states=hidden_states, attn_metadata=attn_metadata, all_reduce_params=AllReduceParams(enable_allreduce=not (self.disable_attn_allreduce)), + lora_params=lora_params, **kwargs, ) @@ -933,6 +961,7 @@ def norm_hidden(): self.fusion_config.POST_MOE_FUSION or self.mapping.tp_size == 1 ) ), + lora_params=lora_params, ) if self.fusion_config.POST_MOE_FUSION: @@ -1009,6 +1038,7 @@ def forward( attn_metadata=attn_metadata, residual=residual, spec_metadata=spec_metadata, + **kwargs, ) return hidden_states diff --git a/tensorrt_llm/_torch/models/modeling_minimaxm2.py b/tensorrt_llm/_torch/models/modeling_minimaxm2.py index 8563dbeee1ae..2b359512f325 100644 --- a/tensorrt_llm/_torch/models/modeling_minimaxm2.py +++ b/tensorrt_llm/_torch/models/modeling_minimaxm2.py @@ -104,6 +104,7 @@ def forward( self, hidden_states: torch.Tensor, attn_metadata: AttentionMetadata, + lora_params: Optional[dict] = None, ) -> torch.Tensor: all_rank_num_tokens = attn_metadata.all_rank_num_tokens hidden_states_f32 = hidden_states.to(torch.float32) @@ -113,6 +114,7 @@ def forward( router_logits, all_rank_num_tokens=all_rank_num_tokens, use_dp_padding=False, + lora_params=lora_params, ) return final_hidden_states @@ -313,6 +315,7 @@ def forward( hidden_states: torch.Tensor, attn_metadata: AttentionMetadata, residual: Optional[torch.Tensor], + lora_params: Optional[dict] = None, **kwargs, ) -> torch.Tensor: if residual is None: @@ -326,12 +329,13 @@ def forward( position_ids=position_ids, hidden_states=hidden_states, attn_metadata=attn_metadata, + lora_params=lora_params, **kwargs, ) # Fully Connected hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) - hidden_states = self.block_sparse_moe(hidden_states, attn_metadata) + hidden_states = self.block_sparse_moe(hidden_states, attn_metadata, lora_params=lora_params) return hidden_states, residual @@ -391,6 +395,7 @@ def forward( hidden_states=hidden_states, attn_metadata=attn_metadata, residual=residual, + **kwargs, ) hidden_states, _ = self.norm(hidden_states, residual) diff --git a/tensorrt_llm/_torch/models/modeling_minimaxm3.py b/tensorrt_llm/_torch/models/modeling_minimaxm3.py index 2495e8dd446b..2a79aa5868f5 100644 --- a/tensorrt_llm/_torch/models/modeling_minimaxm3.py +++ b/tensorrt_llm/_torch/models/modeling_minimaxm3.py @@ -241,6 +241,7 @@ def _minimax_m3_swiglu_oai(gate_up: torch.Tensor, *, alpha: float, limit: float) def _build_swiglu_oai_dense_mlp( model_config: "ModelConfig[PretrainedConfig]", intermediate_size: int, + layer_idx: Optional[int] = None, *, is_shared_expert: bool = False, ) -> GatedMLP: @@ -290,6 +291,7 @@ def _build_swiglu_oai_dense_mlp( overridden_tp_size=1 if enable_adp else None, reduce_output=reduce_output, is_shared_expert=is_shared_expert, + layer_idx=layer_idx, ) @@ -473,6 +475,7 @@ def __init__( self.shared_experts = _build_swiglu_oai_dense_mlp( model_config=model_config, intermediate_size=shared_intermediate, + layer_idx=layer_idx, is_shared_expert=True, ) else: @@ -501,6 +504,7 @@ def forward( hidden_states: torch.Tensor, attn_metadata: AttentionMetadata, final_all_reduce_params: Optional[AllReduceParams] = None, + lora_params: Optional[dict] = None, ) -> torch.Tensor: all_rank_num_tokens = attn_metadata.all_rank_num_tokens @@ -511,10 +515,11 @@ def _compute_routed_output(): router_logits, all_rank_num_tokens=all_rank_num_tokens, use_dp_padding=False, + lora_params=lora_params, ) def _compute_shared_output(): - return self.shared_experts(hidden_states) + return self.shared_experts(hidden_states, lora_params=lora_params) if self.shared_experts is None: result = _compute_routed_output() @@ -1687,7 +1692,9 @@ def __init__( getattr(config, "dense_intermediate_size", config.intermediate_size) ) self.mlp = _build_swiglu_oai_dense_mlp( - model_config=model_config, intermediate_size=dense_intermediate + model_config=model_config, + intermediate_size=dense_intermediate, + layer_idx=layer_idx, ) self.block_sparse_moe = None @@ -1747,6 +1754,7 @@ def forward( attn_metadata: AttentionMetadata, residual: Optional[torch.Tensor], spec_metadata: Optional[SpecMetadata] = None, + lora_params: Optional[dict] = None, **kwargs, ) -> torch.Tensor: # Layer-0 prologue only. For every subsequent layer the input_layernorm @@ -1769,6 +1777,7 @@ def forward( hidden_states=hidden_states, attn_metadata=attn_metadata, all_reduce_params=attn_all_reduce_params, + lora_params=lora_params, **kwargs, ) @@ -1779,11 +1788,18 @@ def forward( self.post_feed_forward_fusion = False if self.block_sparse_moe is not None: hidden_states, residual = self.forward_MoE( - hidden_states, attn_metadata, residual, spec_metadata=spec_metadata + hidden_states, + attn_metadata, + residual, + spec_metadata=spec_metadata, + lora_params=lora_params, ) else: hidden_states, residual = self.forward_mlp( - hidden_states, residual, spec_metadata=spec_metadata + hidden_states, + residual, + spec_metadata=spec_metadata, + lora_params=lora_params, ) return hidden_states, residual @@ -1859,6 +1875,7 @@ def forward_MoE( attn_metadata: AttentionMetadata, residual: torch.Tensor, spec_metadata: Optional[SpecMetadata] = None, + lora_params: Optional[dict] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: hidden_states, residual = self._apply_pre_feed_forward_norm(hidden_states, residual) @@ -1866,6 +1883,7 @@ def forward_MoE( hidden_states, attn_metadata, final_all_reduce_params=self._feed_forward_all_reduce_params(), + lora_params=lora_params, ) if spec_metadata is not None and spec_metadata.is_layer_capture(self.layer_idx): @@ -1878,12 +1896,14 @@ def forward_mlp( hidden_states: torch.Tensor, residual: torch.Tensor, spec_metadata: Optional[SpecMetadata] = None, + lora_params: Optional[dict] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: hidden_states, residual = self._apply_pre_feed_forward_norm(hidden_states, residual) hidden_states = self.mlp( hidden_states, final_all_reduce_params=self._feed_forward_all_reduce_params(), + lora_params=lora_params, ) if spec_metadata is not None and spec_metadata.is_layer_capture(self.layer_idx): @@ -1969,6 +1989,7 @@ def forward( attn_metadata=attn_metadata, residual=residual, spec_metadata=spec_metadata, + **kwargs, ) # When setup_aliases has chained the final norm into the last decoder diff --git a/tensorrt_llm/_torch/models/modeling_step3p7.py b/tensorrt_llm/_torch/models/modeling_step3p7.py index 191e6a658e5e..12cde6e4078e 100644 --- a/tensorrt_llm/_torch/models/modeling_step3p7.py +++ b/tensorrt_llm/_torch/models/modeling_step3p7.py @@ -58,9 +58,11 @@ from ..modules.gated_mlp import GatedMLP from ..modules.linear import Linear, TensorParallelMode from ..modules.rms_norm import RMSNorm -from ..moe.fused_moe import create_moe +from ..moe.fused_moe import SwigluActivation, create_moe from ..moe.fused_moe.interface import MoEWeightLoadingMode from ..moe.fused_moe.routing import MiniMaxM2MoeRoutingMethod +from ..peft.lora.layer import LoraLayer +from ..peft.lora.validation import has_moe_lora_targets from ..speculative import SpecMetadata from ..utils import AuxStreamType, create_lm_head_tp_mapping from .modeling_speculative import SpecDecOneEngineForCausalLM, _slice_spec_position_ids @@ -577,13 +579,14 @@ def forward( attn_metadata: AttentionMetadata, attention_mask=PredefinedAttentionMask.CAUSAL, attention_window_size: Optional[int] = None, + lora_params: Optional[dict] = None, **kwargs, ) -> torch.Tensor: """Step3p7 attention with module-side QK-norm + RoPE + head gate. Bypasses the base ``Attention.forward`` to apply the per-head output gate (``sigmoid(g_proj(hidden))``) between the attention backend and - ``o_proj``. Helix CP, LoRA, and attention sinks are not plumbed here. + ``o_proj``. Helix CP and attention sinks are not plumbed here. """ effective_window = ( attention_window_size if attention_window_size is not None else self.sliding_window @@ -605,7 +608,16 @@ def forward( if position_ids is not None and self.layer_idx >= num_text_layers: position_ids = position_ids.clamp_min(0) - qkv = self.qkv_proj(hidden_states) + if bool(lora_params): + qkv = LoraLayer.forward_with_base( + lambda: self.qkv_proj(hidden_states), + (self.splitted_qkv_lora, self.fused_qkv_lora), + hidden_states, + lora_params, + self.layer_idx, + ) + else: + qkv = self.qkv_proj(hidden_states) q, k, v = self.apply_rope(qkv, None, None, position_ids) q, k, v = self.convert_qkv(q, k, v) @@ -619,7 +631,7 @@ def forward( kwargs.get("attention_mask_data"), mrope_config=kwargs.get("mrope_config"), attention_sinks=None, - has_lora=False, + has_lora=bool(lora_params), ) if self.use_head_wise_gate: @@ -630,7 +642,11 @@ def forward( attn_output = attn_output * gate.unsqueeze(-1).sigmoid() attn_output = attn_output.view(*orig_shape) - return self.o_proj(attn_output) + return self.o_proj( + attn_output, + lora_params=lora_params, + layer_idx=self.layer_idx, + ) # --------------------------------------------------------------------------- @@ -669,13 +685,47 @@ def __init__( ) self.swiglu_limit = swiglu_limit - def forward(self, hidden_states: torch.Tensor, **kwargs) -> torch.Tensor: + def forward( + self, + hidden_states: torch.Tensor, + all_rank_num_tokens=None, + final_all_reduce_params: Optional[AllReduceParams] = None, + lora_params: Optional[dict] = None, + **kwargs, + ) -> torch.Tensor: if self.swiglu_limit is None: - return super().forward(hidden_states, **kwargs) - gate, up = self.gate_up_proj(hidden_states).chunk(2, dim=-1) + return super().forward( + hidden_states, + all_rank_num_tokens=all_rank_num_tokens, + final_all_reduce_params=final_all_reduce_params, + lora_params=lora_params, + **kwargs, + ) + if bool(lora_params): + assert self.layer_idx is not None, "layer_idx is required for lora" + if self._uneven_tp_blocks_lora: + raise NotImplementedError( + "LoRA is not supported with uneven TP for ClampedGatedMLP " + "(intermediate_size not divisible by tp_size)." + ) + gate_up = LoraLayer.forward_with_base( + lambda: self.gate_up_proj(hidden_states), + (self.splitted_gate_up_lora, self.fused_gate_up_lora), + hidden_states, + lora_params, + self.layer_idx, + ) + else: + gate_up = self.gate_up_proj(hidden_states) + gate, up = gate_up.chunk(2, dim=-1) gate = torch.nn.functional.silu(gate).clamp(max=self.swiglu_limit) up = up.clamp(min=-self.swiglu_limit, max=self.swiglu_limit) - return self.down_proj(gate * up) + return self.down_proj( + gate * up, + all_reduce_params=final_all_reduce_params, + lora_params=lora_params, + layer_idx=self.layer_idx, + ) class Step3p7MoE(nn.Module): @@ -715,11 +765,12 @@ def __init__( ) self.router_bias = Step3p7RouterBiasHolder(self.num_experts) + # Keep routing weights normalized but unscaled. The common output scale + # in forward applies the checkpoint factor exactly once for every backend. routing_method = Step3p7MoeRoutingMethod( top_k=self.top_k, num_experts=self.num_experts, callable_router_bias=lambda: self.router_bias.router_bias, - routed_scaling_factor=self.routed_scaling_factor, ) ( @@ -727,6 +778,7 @@ def __init__( self._use_python_clamp, self._python_path_reason, ) = _select_python_expert_path(model_config, text_config, layer_idx) + self._moe_lora_enabled = has_moe_lora_targets(model_config.lora_config) self.experts = create_moe( num_experts=self.num_experts, @@ -739,6 +791,7 @@ def __init__( model_config=model_config, layer_idx=layer_idx, weight_loading_mode=MoEWeightLoadingMode.VANILLA, + activation=SwigluActivation(clamp=self._routed_swiglu_limit), ) if self._use_python_clamp: @@ -884,6 +937,7 @@ def forward( hidden_states: torch.Tensor, attn_metadata: AttentionMetadata, all_reduce_params: Optional[AllReduceParams] = None, + lora_params: Optional[dict] = None, **kwargs, ) -> torch.Tensor: assert hidden_states.shape[-1] == self.hidden_size @@ -893,16 +947,18 @@ def forward( gate_input = h.to(torch.float32) if self.need_fp32_gate else h router_logits = self.gate(gate_input) - if self._use_python_clamp and self._clamp_weights_loaded: - # Python path's routing.apply already scales topk_weights. - return self._python_clamped_moe_forward(h, router_logits).view(orig_shape) - - routed = self.experts( - h, - router_logits, - all_rank_num_tokens=attn_metadata.all_rank_num_tokens, - use_dp_padding=False, - ) + if self._use_python_clamp and self._clamp_weights_loaded and not self._moe_lora_enabled: + # The Python loop consumes the same unscaled routing weights as the + # separated CUTLASS path, then joins the common output scaling below. + routed = self._python_clamped_moe_forward(h, router_logits) + else: + routed = self.experts( + h, + router_logits, + all_rank_num_tokens=attn_metadata.all_rank_num_tokens, + use_dp_padding=False, + lora_params=lora_params, + ) # Step3p7 uses the generic MiniMax2 metadata with routeScale=1.0, so apply # ``routed_scaling_factor`` to the MoE output here (mathematically # equivalent to scaling each topk weight). @@ -1010,6 +1066,7 @@ def forward( hidden_states: torch.Tensor, attn_metadata: AttentionMetadata, residual: Optional[torch.Tensor], + lora_params: Optional[dict] = None, **kwargs, ) -> Tuple[torch.Tensor, torch.Tensor]: if residual is None: @@ -1021,17 +1078,18 @@ def forward( position_ids=position_ids, hidden_states=hidden_states, attn_metadata=attn_metadata, + lora_params=lora_params, **kwargs, ) hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) if self.moe is not None: - routed = self.moe(hidden_states, attn_metadata) - shared = self.share_expert(hidden_states) + routed = self.moe(hidden_states, attn_metadata, lora_params=lora_params) + shared = self.share_expert(hidden_states, lora_params=lora_params) hidden_states = routed + shared if self.allreduce is not None: hidden_states = self.allreduce(hidden_states) else: - hidden_states = self.mlp(hidden_states) + hidden_states = self.mlp(hidden_states, lora_params=lora_params) return hidden_states, residual diff --git a/tensorrt_llm/_torch/peft/lora/cuda_graph_lora_params.py b/tensorrt_llm/_torch/peft/lora/cuda_graph_lora_params.py index df9c681c80fd..5cc3a476c4c2 100644 --- a/tensorrt_llm/_torch/peft/lora/cuda_graph_lora_params.py +++ b/tensorrt_llm/_torch/peft/lora/cuda_graph_lora_params.py @@ -132,8 +132,8 @@ def __init__( # layer/module actually has pointers for: an adapter need not cover every # MoE layer, and the grouped-GEMM problem builder reads an unmasked # rank > 0 on a null-pointer slot as an active rank-sized GEMM. - self._moe_slot_cache: Dict[ - Tuple[int, int], Tuple[torch.Tensor, torch.Tensor, torch.Tensor] + self._moe_slot_cache: dict[ + tuple[int, int], tuple[torch.Tensor, torch.Tensor, torch.Tensor] ] = {} # The pointer table is packed as PTR_DTYPE and relies on Tensor.copy_'s @@ -342,7 +342,7 @@ def zero_out_weight_pointers(slot_id: int): def _moe_slot_entry( self, layer_idx: int, module_id: int - ) -> Optional[Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None: """Return this (layer, module)'s cached MoE slot buffers, allocating them on first use. Never refreshes their contents. @@ -353,7 +353,13 @@ def _moe_slot_entry( would hide a missing refresh from any test that reads back through `get_moe_slot_inputs`. - Returns None if (layer_idx, module_id) carries no LoRA modules. + Args: + layer_idx: Model layer whose cached slot buffers are requested. + module_id: Routed-expert LoRA module within the layer. + + Returns: + The cached rank, pointer, and scratch-mask tensors, or None if the + layer and module carry no LoRA weights. """ key = self.layer_module2key.get((layer_idx, module_id)) layer_param = self.layer_params.get(key) if key is not None else None @@ -400,6 +406,13 @@ def _refresh_moe_slot_cache(self, layer_idx: int, module_id: int) -> None: The cached pinned buffers are updated in place to keep their addresses stable (the captured H2D copy reads them by address at replay). + + Args: + layer_idx: Model layer whose cached slot buffers are refreshed. + module_id: Routed-expert LoRA module within the layer. + + Returns: + None. """ key = self.layer_module2key[(layer_idx, module_id)] layer_param = self.layer_params[key] diff --git a/tests/integration/defs/llmapi/test_llm_api_pytorch_moe_lora.py b/tests/integration/defs/llmapi/test_llm_api_pytorch_moe_lora.py index 12cf2dcbb70c..789d83b3fe37 100644 --- a/tests/integration/defs/llmapi/test_llm_api_pytorch_moe_lora.py +++ b/tests/integration/defs/llmapi/test_llm_api_pytorch_moe_lora.py @@ -108,14 +108,14 @@ def randn(rows, cols, std=0.02): def _run_routed_expert_multi_lora( model_dir: str, - lora_paths: list, + lora_paths: list[str], *, max_rank: int, - target_modules: list, - trtllm_modules_to_hf_modules: dict, - cuda_graph_config, + target_modules: list[str], + trtllm_modules_to_hf_modules: dict[str, str], + cuda_graph_config: CudaGraphConfig | None, preallocate_all_adapters: bool = True, - peft_cache_config=None, + peft_cache_config: PeftCacheConfig | None = None, ) -> None: """Serve a MoE checkpoint with routed-expert LoRA and assert it applies. @@ -128,6 +128,19 @@ def _run_routed_expert_multi_lora( calibration. With a CUDA graph the decode takes the slot-indexed input schema; without one it takes the per-request schema. Both feed the same grouped-GEMM LoRA core. + + Args: + model_dir: Path to the base model checkpoint. + lora_paths: Paths to routed-expert LoRA adapters. + max_rank: Maximum adapter rank accepted by the cache. + target_modules: TensorRT-LLM LoRA module names to enable. + trtllm_modules_to_hf_modules: TensorRT-LLM to Hugging Face module mapping. + cuda_graph_config: CUDA graph configuration, or None for eager execution. + preallocate_all_adapters: Whether to reserve every adapter slot up front. + peft_cache_config: Optional explicit PEFT cache configuration. + + Returns: + None. """ cache_config = {} if preallocate_all_adapters: diff --git a/tests/unittest/_torch/modeling/test_modeling_step3p7.py b/tests/unittest/_torch/modeling/test_modeling_step3p7.py index 913091f20cc0..7a564b9a208f 100644 --- a/tests/unittest/_torch/modeling/test_modeling_step3p7.py +++ b/tests/unittest/_torch/modeling/test_modeling_step3p7.py @@ -39,6 +39,7 @@ import os import types import unittest +from unittest.mock import MagicMock, patch import pytest import torch @@ -357,6 +358,79 @@ def test_nvfp4_dequant_batched_round_trips_constant_values(self): self.assertTrue(torch.all(out3[0] == 1.0)) # 2.0 * 1.0 * 0.5 = 1.0 self.assertTrue(torch.all(out3[1] == 0.5)) # 2.0 * 1.0 * 0.25 = 0.5 + def test_moe_scaling_and_lora_clamp_path_selection(self): + """Python and LoRA-capable expert paths apply routed scaling once.""" + import tensorrt_llm._torch.models.modeling_step3p7 as step3p7_module + from tensorrt_llm._torch.models.modeling_step3p7 import Step3p7MoE + + scaling_factor = 3.0 + text_config = types.SimpleNamespace( + hidden_size=4, + moe_num_experts=2, + moe_top_k=2, + moe_intermediate_size=8, + moe_router_scaling_factor=scaling_factor, + need_fp32_gate=False, + torch_dtype=torch.float32, + ) + model_config = types.SimpleNamespace( + pretrained_config=text_config, + mapping=types.SimpleNamespace(enable_attention_dp=True, tp_size=1), + lora_config=object(), + ) + + for moe_lora_enabled in (False, True): + with self.subTest(moe_lora_enabled=moe_lora_enabled): + experts = MagicMock() + with ( + patch.object(step3p7_module, "Linear", return_value=MagicMock()), + patch.object( + step3p7_module, + "_select_python_expert_path", + return_value=(1.0, True, "clamp"), + ), + patch.object( + step3p7_module, + "has_moe_lora_targets", + return_value=moe_lora_enabled, + ), + patch.object(step3p7_module, "create_moe", return_value=experts) as create_moe, + patch.object(Step3p7MoE, "_allocate_clamp_buffers"), + ): + moe = Step3p7MoE(model_config, layer_idx=0, aux_stream_dict={}) + + routing_method = create_moe.call_args.kwargs["routing_method"] + moe.router_bias.router_bias.data.zero_() + moe.gate.return_value = torch.zeros(1, text_config.moe_num_experts) + moe._clamp_weights_loaded = True + python_output = MagicMock(return_value=torch.ones(1, text_config.hidden_size)) + moe._python_clamped_moe_forward = python_output + + def separated_experts(hidden_states, logits, **kwargs): + del kwargs + _, routing_weights = routing_method.apply(logits) + return routing_weights.sum(dim=-1, keepdim=True).expand_as(hidden_states) + + experts.side_effect = separated_experts + hidden_states = torch.ones(1, text_config.hidden_size) + lora_params = {"active": True} if moe_lora_enabled else None + output = moe( + hidden_states, + types.SimpleNamespace(all_rank_num_tokens=[1]), + lora_params=lora_params, + ) + + torch.testing.assert_close( + output, + torch.full_like(hidden_states, scaling_factor), + ) + if moe_lora_enabled: + python_output.assert_not_called() + experts.assert_called_once() + else: + python_output.assert_called_once() + experts.assert_not_called() + def test_mtp_head_normalizes_before_output_projection(self): """Step3p7 MTP applies shared-head norm only when producing draft logits.""" from tensorrt_llm._torch.model_config import ModelConfig diff --git a/tests/unittest/_torch/peft/test_moe_lora_cuda_graph_params.py b/tests/unittest/_torch/peft/test_moe_lora_cuda_graph_params.py index 6e9bac4c0414..8d935124036e 100644 --- a/tests/unittest/_torch/peft/test_moe_lora_cuda_graph_params.py +++ b/tests/unittest/_torch/peft/test_moe_lora_cuda_graph_params.py @@ -102,7 +102,16 @@ def __init__(self, layer_idx): ) -def _make_params(max_lora_size=2, max_rank=8, max_batch_size=2, layer_idxs=(0,)): +def _make_params( + max_lora_size: int = 2, + max_rank: int = 8, + max_batch_size: int = 2, + layer_idxs: tuple[int, ...] = (0,), +) -> tuple[ + CudaGraphLoraParams, + CudaGraphLoraParams.LoraLayerKey, + tuple[int, int, int], +]: """Build a CudaGraphLoraParams carrying one MoE layer per entry in `layer_idxs`, each with all three routed-expert modules (fc1=moe_h_to_4h, gated=moe_gate, fc2=moe_4h_to_h).""" diff --git a/tests/unittest/_torch/peft/test_moe_lora_model_path.py b/tests/unittest/_torch/peft/test_moe_lora_model_path.py index bed5faff8c44..a45a11d86061 100644 --- a/tests/unittest/_torch/peft/test_moe_lora_model_path.py +++ b/tests/unittest/_torch/peft/test_moe_lora_model_path.py @@ -17,12 +17,25 @@ """ from types import SimpleNamespace -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import pytest import torch +import tensorrt_llm._torch.models.modeling_glm as glm_module +import tensorrt_llm._torch.models.modeling_minimaxm3 as minimaxm3_module +import tensorrt_llm._torch.models.modeling_step3p7 as step3p7_module +from tensorrt_llm._torch.models.modeling_afmoe import AfmoeDecoderLayer, AfmoeMoE +from tensorrt_llm._torch.models.modeling_glm import Glm4DecoderLayer, Glm4MoE +from tensorrt_llm._torch.models.modeling_minimaxm2 import MiniMaxM2DecoderLayer, MiniMaxM2MoE +from tensorrt_llm._torch.models.modeling_minimaxm3 import MiniMaxM3DecoderLayer, MiniMaxM3MoE from tensorrt_llm._torch.models.modeling_mixtral import MixtralMoE +from tensorrt_llm._torch.models.modeling_step3p7 import ( + ClampedGatedMLP, + Step3p7Attention, + Step3p7DecoderLayer, + Step3p7MoE, +) from tensorrt_llm._torch.moe.fused_moe.configurable_moe import ConfigurableMoE from tensorrt_llm._torch.moe.fused_moe.fused_moe_cutlass import CutlassFusedMoE from tensorrt_llm._torch.moe.fused_moe.fused_moe_deepgemm import DeepGemmFusedMoE @@ -70,6 +83,446 @@ def test_mixtral_moe_forward_passes_lora_params_to_routed_experts(): ) +@pytest.mark.parametrize( + "moe_cls", + [AfmoeMoE, MiniMaxM2MoE, MiniMaxM3MoE], +) +def test_model_moe_forward_passes_lora_params_to_routed_experts( + moe_cls: type[torch.nn.Module], +) -> None: + """MoE model wrappers must forward lora_params to routed experts.""" + num_tokens, hidden_dim = 4, 8 + hidden_states = torch.randn(num_tokens, hidden_dim) + router_logits = torch.randn(num_tokens, 2) + experts = MagicMock(return_value=torch.randn_like(hidden_states)) + fake_self = SimpleNamespace( + gate=MagicMock(return_value=router_logits), + experts=experts, + shared_experts=None, + allreduce=None, + ) + if moe_cls is MiniMaxM2MoE: + fake_self.hidden_dim = hidden_dim + + moe_cls.forward( + fake_self, + hidden_states, + SimpleNamespace(all_rank_num_tokens=[num_tokens]), + lora_params=_LORA_PARAMS_SENTINEL, + ) + + experts.assert_called_once() + assert experts.call_args.kwargs.get("lora_params") is _LORA_PARAMS_SENTINEL + + +def test_glm_moe_forward_passes_lora_params_to_routed_experts() -> None: + """Glm4MoE.compute_routed_output must forward lora_params to experts.""" + num_tokens, hidden_dim = 4, 8 + hidden_states = torch.randn(num_tokens, hidden_dim) + experts = MagicMock(return_value=torch.randn_like(hidden_states)) + fake_self = SimpleNamespace( + use_dp=False, + gate=MagicMock(return_value=torch.randn(num_tokens, 2)), + experts=experts, + ) + + Glm4MoE.compute_routed_output( + fake_self, + hidden_states, + hidden_states_fp4=None, + all_rank_num_tokens=[num_tokens], + do_finalize=True, + lora_params=_LORA_PARAMS_SENTINEL, + ) + + experts.assert_called_once() + assert experts.call_args.kwargs.get("lora_params") is _LORA_PARAMS_SENTINEL + + +def test_step3p7_moe_forward_passes_lora_params_to_routed_experts() -> None: + """Step3p7MoE.forward must forward lora_params to routed experts.""" + num_tokens, hidden_dim = 4, 8 + hidden_states = torch.randn(num_tokens, hidden_dim) + experts = MagicMock(return_value=torch.randn_like(hidden_states)) + fake_self = SimpleNamespace( + hidden_size=hidden_dim, + need_fp32_gate=False, + gate=MagicMock(return_value=torch.randn(num_tokens, 2)), + experts=experts, + _use_python_clamp=False, + _moe_lora_enabled=True, + routed_scaling_factor=1.0, + ) + + Step3p7MoE.forward( + fake_self, + hidden_states, + SimpleNamespace(all_rank_num_tokens=[num_tokens]), + lora_params=_LORA_PARAMS_SENTINEL, + ) + + experts.assert_called_once() + assert experts.call_args.kwargs.get("lora_params") is _LORA_PARAMS_SENTINEL + + +def _make_decoder_layer(decoder_cls: type[torch.nn.Module]) -> tuple[torch.nn.Module, MagicMock]: + """Build a lightweight decoder whose real forward reaches a mocked MoE.""" + decoder = decoder_cls.__new__(decoder_cls) + torch.nn.Module.__init__(decoder) + + hidden_states = torch.randn(4, 8) + residual = torch.randn_like(hidden_states) + decoder.input_layernorm = MagicMock(return_value=(hidden_states, residual)) + decoder.self_attn = MagicMock(return_value=hidden_states) + decoder.post_attention_layernorm = MagicMock(return_value=(hidden_states, residual)) + + moe = MagicMock(return_value=hidden_states) + if decoder_cls is AfmoeDecoderLayer: + decoder.pre_mlp_layernorm = MagicMock(return_value=(hidden_states, residual)) + decoder.post_mlp_layernorm = MagicMock(return_value=hidden_states) + decoder.moe_enabled = True + decoder.mlp = moe + elif decoder_cls is MiniMaxM2DecoderLayer: + decoder.block_sparse_moe = moe + elif decoder_cls is MiniMaxM3DecoderLayer: + decoder.pre_feed_forward_fusion = False + decoder.post_feed_forward_fusion = False + decoder.next_layer_layernorm = None + decoder.block_sparse_moe = moe + elif decoder_cls is Glm4DecoderLayer: + glm_moe = Glm4MoE.__new__(Glm4MoE) + torch.nn.Module.__init__(glm_moe) + glm_moe.forward = moe + decoder.mlp = glm_moe + decoder.disable_attn_allreduce = False + decoder.layer_idx = 0 + decoder.fusion_config = SimpleNamespace(PRE_MOE_FUSION=False, POST_MOE_FUSION=False) + decoder.mapping = SimpleNamespace(tp_size=1, is_multi_node=lambda: True) + decoder.next_layer_layernorm = None + elif decoder_cls is Step3p7DecoderLayer: + decoder.moe = moe + decoder.share_expert = MagicMock(return_value=torch.zeros_like(hidden_states)) + decoder.allreduce = None + else: + raise AssertionError(f"Unsupported decoder class: {decoder_cls.__name__}") + + return decoder, moe + + +@pytest.mark.parametrize( + "decoder_cls", + [ + AfmoeDecoderLayer, + MiniMaxM2DecoderLayer, + MiniMaxM3DecoderLayer, + Glm4DecoderLayer, + Step3p7DecoderLayer, + ], +) +def test_decoder_layer_passes_lora_params_to_moe( + decoder_cls: type[torch.nn.Module], +) -> None: + """Decoder layers must preserve request LoRA state at the MoE boundary.""" + decoder, moe = _make_decoder_layer(decoder_cls) + hidden_states = torch.randn(4, 8) + + decoder.forward( + position_ids=torch.arange(4), + hidden_states=hidden_states, + attn_metadata=SimpleNamespace(all_rank_num_tokens=[4]), + residual=torch.randn_like(hidden_states), + lora_params=_LORA_PARAMS_SENTINEL, + ) + + moe.assert_called_once() + assert moe.call_args.kwargs.get("lora_params") is _LORA_PARAMS_SENTINEL + if decoder_cls is Step3p7DecoderLayer: + assert decoder.share_expert.call_args.kwargs.get("lora_params") is _LORA_PARAMS_SENTINEL + + +def test_afmoe_dense_mlp_receives_lora_params() -> None: + """AFMoE dense layers must forward request LoRA state to GatedMLP.""" + hidden_states = torch.randn(2, 4) + residual = torch.randn_like(hidden_states) + mlp = MagicMock(return_value=hidden_states) + fake_self = SimpleNamespace( + input_layernorm=MagicMock(return_value=(hidden_states, residual)), + self_attn=MagicMock(return_value=hidden_states), + post_attention_layernorm=MagicMock(return_value=hidden_states), + pre_mlp_layernorm=MagicMock(return_value=(hidden_states, residual)), + moe_enabled=False, + mlp=mlp, + post_mlp_layernorm=MagicMock(return_value=hidden_states), + ) + + AfmoeDecoderLayer.forward( + fake_self, + position_ids=torch.arange(2), + hidden_states=hidden_states, + attn_metadata=SimpleNamespace(), + residual=residual, + lora_params=_LORA_PARAMS_SENTINEL, + ) + + assert mlp.call_args.kwargs["lora_params"] is _LORA_PARAMS_SENTINEL + + +def test_glm_dense_mlp_receives_lora_params() -> None: + """GLM dense layers must forward request LoRA state to GatedMLP.""" + hidden_states = torch.randn(2, 4) + residual = torch.randn_like(hidden_states) + mlp = MagicMock(return_value=hidden_states) + fake_self = SimpleNamespace( + fusion_config=SimpleNamespace(PRE_MLP_FUSION=False, POST_MLP_FUSION=False), + post_attention_layernorm=MagicMock(return_value=(hidden_states, residual)), + mlp=mlp, + mlp_tp_size=1, + next_layer_layernorm=None, + ) + + Glm4DecoderLayer.forward_mlp( + fake_self, + hidden_states, + residual, + lora_params=_LORA_PARAMS_SENTINEL, + ) + + assert mlp.call_args.kwargs["lora_params"] is _LORA_PARAMS_SENTINEL + + +def test_minimax_m3_dense_mlp_receives_lora_params() -> None: + """MiniMax-M3 dense layers must forward request LoRA state to GatedMLP.""" + hidden_states = torch.randn(2, 4) + residual = torch.randn_like(hidden_states) + mlp = MagicMock(return_value=hidden_states) + fake_self = SimpleNamespace( + _apply_pre_feed_forward_norm=MagicMock(return_value=(hidden_states, residual)), + mlp=mlp, + _feed_forward_all_reduce_params=MagicMock(return_value=None), + _apply_next_layer_layernorm=MagicMock(return_value=(hidden_states, residual)), + ) + + MiniMaxM3DecoderLayer.forward_mlp( + fake_self, + hidden_states, + residual, + lora_params=_LORA_PARAMS_SENTINEL, + ) + + assert mlp.call_args.kwargs["lora_params"] is _LORA_PARAMS_SENTINEL + + +def test_step3p7_dense_mlp_receives_lora_params() -> None: + """Step3p7 dense layers must forward request LoRA state to GatedMLP.""" + hidden_states = torch.randn(2, 4) + residual = torch.randn_like(hidden_states) + mlp = MagicMock(return_value=hidden_states) + fake_self = SimpleNamespace( + input_layernorm=MagicMock(return_value=(hidden_states, residual)), + self_attn=MagicMock(return_value=hidden_states), + post_attention_layernorm=MagicMock(return_value=(hidden_states, residual)), + moe=None, + mlp=mlp, + ) + + Step3p7DecoderLayer.forward( + fake_self, + position_ids=torch.arange(2), + hidden_states=hidden_states, + attn_metadata=SimpleNamespace(), + residual=residual, + lora_params=_LORA_PARAMS_SENTINEL, + ) + + assert mlp.call_args.kwargs["lora_params"] is _LORA_PARAMS_SENTINEL + + +def test_feed_forward_moe_wrappers_combine_routed_and_shared_lora() -> None: + """AFMoE, GLM, and MiniMax-M3 must retain both adapted MoE branches.""" + hidden_states = torch.zeros(2, 4) + routed = torch.ones_like(hidden_states) + shared = torch.full_like(hidden_states, 2.0) + + af_shared = MagicMock(return_value=shared.clone()) + af_self = SimpleNamespace( + gate=MagicMock(return_value=torch.zeros(2, 2)), + experts=MagicMock(return_value=routed.clone()), + shared_experts=af_shared, + allreduce=None, + ) + af_output = AfmoeMoE.forward( + af_self, + hidden_states, + SimpleNamespace(all_rank_num_tokens=[2]), + lora_params=_LORA_PARAMS_SENTINEL, + ) + + def run_both(routed_fn, shared_fn, *_args, **_kwargs): + return routed_fn(), shared_fn() + + glm_shared = MagicMock(return_value=shared.clone()) + glm_self = SimpleNamespace( + use_dp=True, + shared_experts=glm_shared, + shared_output_scale=None, + compute_routed_output=MagicMock(return_value=routed.clone()), + event_dict=MagicMock(), + aux_stream=object(), + mapping=SimpleNamespace(tp_size=1), + top_k=1, + ) + with patch.object(glm_module, "maybe_execute_in_parallel", side_effect=run_both): + glm_output = Glm4MoE.forward( + glm_self, + hidden_states, + all_rank_num_tokens=[2], + lora_params=_LORA_PARAMS_SENTINEL, + ) + + m3_shared = MagicMock(return_value=shared.clone()) + m3_self = SimpleNamespace( + gate=MagicMock(return_value=torch.zeros(2, 2)), + experts=MagicMock(return_value=routed.clone()), + shared_experts=m3_shared, + event_dict=MagicMock(), + aux_stream=object(), + allreduce=None, + ) + with patch.object(minimaxm3_module, "maybe_execute_in_parallel", side_effect=run_both): + m3_output = MiniMaxM3MoE.forward( + m3_self, + hidden_states, + SimpleNamespace(all_rank_num_tokens=[2]), + lora_params=_LORA_PARAMS_SENTINEL, + ) + + for output in (af_output, glm_output, m3_output): + torch.testing.assert_close(output, torch.full_like(hidden_states, 3.0)) + for shared_expert in (af_shared, glm_shared, m3_shared): + assert shared_expert.call_args.kwargs["lora_params"] is _LORA_PARAMS_SENTINEL + + +@pytest.mark.parametrize("gate_value", [10.0, -10.0]) +def test_step3p7_clamped_mlp_applies_gate_up_and_down_lora(gate_value: float) -> None: + """The clamped path must preserve both projection adapter contributions.""" + hidden_states = torch.zeros(1, 2) + gate_up_lora = torch.tensor([[gate_value, gate_value, 2.0, 2.0]]) + down_lora = torch.full_like(hidden_states, 3.0) + + def forward_with_base(base_forward, lora_layers, x, params, layer_idx): + assert params is _LORA_PARAMS_SENTINEL + assert layer_idx == 4 + assert lora_layers == ("split_gate_up", "fused_gate_up") + return base_forward() + gate_up_lora + + def down_projection(x, *, all_reduce_params, lora_params, layer_idx): + assert all_reduce_params == "reduce" + assert lora_params is _LORA_PARAMS_SENTINEL + assert layer_idx == 4 + return x + down_lora + + fake_self = SimpleNamespace( + swiglu_limit=5.0, + layer_idx=4, + _uneven_tp_blocks_lora=False, + gate_up_proj=MagicMock(return_value=torch.zeros_like(gate_up_lora)), + splitted_gate_up_lora="split_gate_up", + fused_gate_up_lora="fused_gate_up", + down_proj=MagicMock(side_effect=down_projection), + ) + + with patch.object(step3p7_module.LoraLayer, "forward_with_base", side_effect=forward_with_base): + output = ClampedGatedMLP.forward( + fake_self, + hidden_states, + final_all_reduce_params="reduce", + lora_params=_LORA_PARAMS_SENTINEL, + ) + + clamped_gate = torch.nn.functional.silu(torch.tensor(gate_value)).clamp(max=5.0) + expected = torch.full_like( + hidden_states, + clamped_gate * 2 + 3, + ) + torch.testing.assert_close(output, expected) + + +def test_step3p7_moe_lora_uses_clamp_capable_expert_path() -> None: + """Step3p7 must not bypass routed-expert LoRA through its Python path.""" + hidden_states = torch.randn(4, 8) + experts = MagicMock(return_value=torch.randn_like(hidden_states)) + python_clamped_forward = MagicMock(return_value=torch.randn_like(hidden_states)) + fake_self = SimpleNamespace( + hidden_size=hidden_states.shape[-1], + need_fp32_gate=False, + gate=MagicMock(return_value=torch.randn(4, 2)), + experts=experts, + _use_python_clamp=True, + _clamp_weights_loaded=True, + _moe_lora_enabled=True, + _python_clamped_moe_forward=python_clamped_forward, + routed_scaling_factor=1.0, + ) + + Step3p7MoE.forward( + fake_self, + hidden_states, + SimpleNamespace(all_rank_num_tokens=[4]), + lora_params=_LORA_PARAMS_SENTINEL, + ) + + python_clamped_forward.assert_not_called() + assert experts.call_args.kwargs.get("lora_params") is _LORA_PARAMS_SENTINEL + + +def test_step3p7_attention_applies_lora_to_qkv_and_output() -> None: + """Step3p7 attention must retain non-zero QKV and output LoRA contributions.""" + hidden_states = torch.zeros(1, 4) + lora_params = {"enabled_modules": {"attn_qkv", "attn_dense"}} + qkv_lora_contribution = torch.ones_like(hidden_states) + output_lora_contribution = torch.full_like(hidden_states, 2.0) + + def forward_with_base(base_forward, lora_layers, x, params, layer_idx): + assert params is lora_params + assert layer_idx == 3 + assert lora_layers == ("split_qkv_lora", "fused_qkv_lora") + return base_forward() + qkv_lora_contribution + + def output_projection(x, *, lora_params, layer_idx): + assert lora_params is not None + assert lora_params["enabled_modules"] == {"attn_qkv", "attn_dense"} + assert layer_idx == 3 + return x + output_lora_contribution + + fake_self = SimpleNamespace( + rope_fusion=True, + text_config=SimpleNamespace(num_hidden_layers=8), + layer_idx=3, + sliding_window=None, + qkv_proj=MagicMock(return_value=torch.zeros_like(hidden_states)), + splitted_qkv_lora="split_qkv_lora", + fused_qkv_lora="fused_qkv_lora", + apply_rope=MagicMock(side_effect=lambda qkv, _k, _v, _positions: (qkv, qkv, qkv)), + convert_qkv=MagicMock(side_effect=lambda q, k, v: (q, k, v)), + forward_impl=MagicMock(side_effect=lambda q, _k, _v, *_args, **_kwargs: q), + use_head_wise_gate=False, + o_proj=MagicMock(side_effect=output_projection), + ) + + with patch.object(step3p7_module.LoraLayer, "forward_with_base", side_effect=forward_with_base): + output = Step3p7Attention.forward( + fake_self, + position_ids=torch.arange(1), + hidden_states=hidden_states, + attn_metadata=SimpleNamespace(), + lora_params=lora_params, + ) + + torch.testing.assert_close(output, torch.full_like(hidden_states, 3.0)) + assert fake_self.forward_impl.call_args.kwargs["has_lora"] is True + assert fake_self.o_proj.call_args.kwargs["lora_params"] is lora_params + + def test_configurable_moe_forward_impl_forwards_lora_params_to_scheduler(): """ConfigurableMoE.forward_impl must forward lora_params to the scheduler so routed-expert MoE LoRA is not dropped.""" From 5cb1c9500b77354a340bb10c0f55afca889df2dd Mon Sep 17 00:00:00 2001 From: Mgluhovskoi Date: Mon, 14 Sep 2026 19:39:01 -0700 Subject: [PATCH 6/8] [None][fix] compose telemetry capture policies (#18978) Signed-off-by: Maxim Gluhovskoi --- docs/source/_ext/llmapi_config_telemetry.py | 34 +- docs/source/developer-guide/telemetry.md | 609 ++++----- tensorrt_llm/llmapi/llm_args.py | 24 +- tensorrt_llm/usage/config.py | 26 +- .../usage/llm_args_golden_manifest.json | 1169 +++++------------ tensorrt_llm/usage/llmapi_config.py | 655 ++++----- tensorrt_llm/usage/schemas/README.md | 40 +- tests/unittest/usage/test_config.py | 46 +- .../usage/test_llmapi_config_capture.py | 606 +++++---- .../test_llmapi_config_telemetry_docs.py | 146 +- 10 files changed, 1431 insertions(+), 1924 deletions(-) diff --git a/docs/source/_ext/llmapi_config_telemetry.py b/docs/source/_ext/llmapi_config_telemetry.py index 0b0bf9e63427..69b85c3d9858 100644 --- a/docs/source/_ext/llmapi_config_telemetry.py +++ b/docs/source/_ext/llmapi_config_telemetry.py @@ -34,12 +34,16 @@ for the wire schema. **No PII or free-form fields are captured.** LLM API configuration capture is -*type-driven*: fields whose type is categorical (`Literal`/`Enum`/`bool`) or -numeric (`int`/`float`), plus safe collections of those, are captured -automatically. Free-form `str`/`Any`/`Path`/`dict`/`Callable` are never captured -unless a field carries an explicit allowlist (`TelemetryField.categorical(...)`), -and any field may opt out with `telemetry=False`. Every captured field is listed -below; the runtime can capture nothing absent from this list. +automatic for `bool`, `int`, finite `float`, `Literal`, `Enum`, supported unions, +and homogeneous sequences. Unsafe scalar `str`, `Any`, and `object` branches +require `TelemetryField.categorical(...)`; paths, mappings, callables, and +unsupported structures always fail closed. Use `telemetry=False` to exclude a +field. The runtime can capture nothing absent from the list below. + +`capture_policy` branches separated by `|` are tried independently; `enum[X]` +requires the exact enum type `X`. The categorical domain lists tokens from +`Literal`/`Enum` annotations or explicit `allowed_values`; it does not restrict +`bool`, `int`, or `float` branches. If the manifest check fails, run `python3 scripts/generate_llm_args_golden_manifest.py`, then commit `tensorrt_llm/usage/llm_args_golden_manifest.json`; new fields require telemetry/privacy CODEOWNER approval. @@ -55,20 +59,24 @@ def _escape(text: str) -> str: return text.replace("|", "\\|").replace("\n", " ") -def _format_values(values: list[str]) -> str: - return ", ".join(f"`{_escape(v)}`" for v in values) if values else "" +def _format_values(values: list[object]) -> str: + def format_value(value: object) -> str: + text = value if isinstance(value, str) else json.dumps(value) + return f"`{_escape(text)}`" + + return ", ".join(format_value(value) for value in values) def _table(rows: list[dict]) -> str: lines = [ - "| Captured key | Annotation | Kind | Converter | Allowed values |", - "|--------------|------------|------|-----------|----------------|", + "| Captured key | Capture policy | Kind | Categorical domain |", + "|--------------|----------------|------|--------------------|", ] for row in rows: lines.append( - f"| `{_escape(row['path'])}` | `{_escape(row['annotation'])}` | " - f"`{_escape(row['kind'])}` | {_escape(row['converter'])} | " - f"{_format_values(row['allowed_values'])} |" + f"| `{_escape(row['path'])}` | `{_escape(row['capture_policy'])}` | " + f"`{_escape(row['kind'])}` | " + f"{_format_values(row.get('allowed_values', []))} |" ) return "\n".join(lines) diff --git a/docs/source/developer-guide/telemetry.md b/docs/source/developer-guide/telemetry.md index 9a67bfa2c5dd..c512b450d3a9 100644 --- a/docs/source/developer-guide/telemetry.md +++ b/docs/source/developer-guide/telemetry.md @@ -11,12 +11,16 @@ for the user-facing collection and opt-out overview, and the for the wire schema. **No PII or free-form fields are captured.** LLM API configuration capture is -*type-driven*: fields whose type is categorical (`Literal`/`Enum`/`bool`) or -numeric (`int`/`float`), plus safe collections of those, are captured -automatically. Free-form `str`/`Any`/`Path`/`dict`/`Callable` are never captured -unless a field carries an explicit allowlist (`TelemetryField.categorical(...)`), -and any field may opt out with `telemetry=False`. Every captured field is listed -below; the runtime can capture nothing absent from this list. +automatic for `bool`, `int`, finite `float`, `Literal`, `Enum`, supported unions, +and homogeneous sequences. Unsafe scalar `str`, `Any`, and `object` branches +require `TelemetryField.categorical(...)`; paths, mappings, callables, and +unsupported structures always fail closed. Use `telemetry=False` to exclude a +field. The runtime can capture nothing absent from the list below. + +`capture_policy` branches separated by `|` are tried independently; `enum[X]` +requires the exact enum type `X`. The categorical domain lists tokens from +`Literal`/`Enum` annotations or explicit `allowed_values`; it does not restrict +`bool`, `int`, or `float` branches. If the manifest check fails, run `python3 scripts/generate_llm_args_golden_manifest.py`, then commit `tensorrt_llm/usage/llm_args_golden_manifest.json`; new fields require telemetry/privacy CODEOWNER approval. @@ -28,292 +32,309 @@ unset or when the safety sanitizer rejects the runtime value. ### `TorchLlmArgs` -294 captured fields. +302 captured fields. -| Captured key | Annotation | Kind | Converter | Allowed values | -|--------------|------------|------|-----------|----------------| -| `allreduce_strategy` | `Optional[Literal['AUTO', 'NCCL', 'UB', 'MINLATENCY', 'ONESHOT', 'TWOSHOT', 'LOWPRECISION', 'MNNVL', 'NCCL_SYMMETRIC']]` | `categorical` | | `AUTO`, `NCCL`, `UB`, `MINLATENCY`, `ONESHOT`, `TWOSHOT`, `LOWPRECISION`, `MNNVL`, `NCCL_SYMMETRIC` | -| `attention_dp_config.batching_wait_iters` | `` | `value` | | | -| `attention_dp_config.enable_balance` | `` | `value` | | | -| `attention_dp_config.enable_kv_cache_aware_routing` | `` | `value` | | | -| `attention_dp_config.kv_cache_routing_account_for_in_transfer` | `` | `value` | | | -| `attention_dp_config.kv_cache_routing_cold_start_warmup` | `` | `value` | | | -| `attention_dp_config.kv_cache_routing_conversation_affinity` | `` | `value` | | | -| `attention_dp_config.kv_cache_routing_fair_share_multiplier` | `` | `value` | | | -| `attention_dp_config.kv_cache_routing_load_balance_weight` | `` | `value` | | | -| `attention_dp_config.kv_cache_routing_match_rate_threshold` | `` | `value` | | | -| `attention_dp_config.kv_cache_routing_max_sessions` | `` | `value` | | | -| `attention_dp_config.kv_cache_routing_new_conv_placement` | `Literal['round_robin', 'least_queued']` | `categorical` | | `round_robin`, `least_queued` | -| `attention_dp_config.timeout_iters` | `` | `value` | | | -| `attn_backend` | `` | `categorical` | allowlist | `VANILLA`, `TRTLLM`, `FLASHINFER` | -| `backend` | `Literal['pytorch']` | `categorical` | | `pytorch` | -| `batch_wait_max_tokens_ratio` | `` | `value` | | | -| `batch_wait_timeout_iters` | `` | `value` | | | -| `batch_wait_timeout_ms` | `` | `value` | | | -| `cache_transceiver_config.backend` | `Optional[Literal['DEFAULT', 'UCX', 'NIXL', 'MOONCAKE', 'MPI']]` | `categorical` | | `DEFAULT`, `UCX`, `NIXL`, `MOONCAKE`, `MPI` | -| `cache_transceiver_config.kv_cache_bounce_size_mb` | `` | `value` | | | -| `cache_transceiver_config.kv_transfer_poll_interval_ms` | `Optional[Annotated[int, Gt(gt=0)]]` | `value` | | | -| `cache_transceiver_config.kv_transfer_sender_future_timeout_ms` | `Optional[Annotated[int, Gt(gt=0)]]` | `value` | | | -| `cache_transceiver_config.kv_transfer_timeout_ms` | `Optional[Annotated[int, Gt(gt=0)]]` | `value` | | | -| `cache_transceiver_config.max_tokens_in_buffer` | `Optional[int]` | `value` | | | -| `cache_transceiver_config.transceiver_runtime` | `Optional[Literal['CPP', 'PYTHON', 'auto']]` | `categorical` | | `CPP`, `PYTHON`, `auto` | -| `context_parallel_size` | `` | `value` | | | -| `cp_config.cp_type` | `` | `categorical` | | `ULYSSES`, `RING`, `HELIX` | -| `cp_config.fifo_version` | `Optional[int]` | `value` | | | -| `cp_config.tokens_per_block` | `Optional[int]` | `value` | | | -| `cp_config.use_nccl_for_alltoall` | `Optional[bool]` | `value` | | | -| `cuda_graph_config.batch_sizes` | `Optional[List[int]]` | `value` | | | -| `cuda_graph_config.enable_padding` | `` | `value` | | | -| `cuda_graph_config.max_batch_size` | `` | `value` | | | -| `cuda_graph_config.max_num_token` | `` | `value` | | | -| `cuda_graph_config.max_seq_len` | `` | `value` | | | -| `cuda_graph_config.mode` | `Literal['decode']` | `categorical` | | `decode`, `encode` | -| `cuda_graph_config.num_tokens` | `Optional[List[Annotated[int, Gt(gt=0)]]]` | `value` | | | -| `cuda_graph_config.seq_lens` | `Optional[List[Annotated[int, Gt(gt=0)]]]` | `value` | | | -| `disable_mm_encoder` | `` | `value` | | | -| `disable_overlap_scheduler` | `` | `value` | | | -| `dtype` | `` | `categorical` | allowlist | `auto`, `float16`, `bfloat16`, `float32` | -| `dwdp_config.contention_opt` | `` | `value` | | | -| `dwdp_config.dwdp_size` | `` | `value` | | | -| `dwdp_config.num_experts_per_worker` | `` | `value` | | | -| `dwdp_config.num_groups` | `` | `value` | | | -| `dwdp_config.num_prefetch_experts` | `` | `value` | | | -| `enable_attention_dp` | `` | `value` | | | -| `enable_autotuner` | `` | `value` | | | -| `enable_chunked_prefill` | `` | `value` | | | -| `enable_early_first_token_response` | `` | `value` | | | -| `enable_encoder_decoder_mixed_cuda_graph` | `` | `value` | | | -| `enable_energy_metrics` | `` | `value` | | | -| `enable_iter_perf_stats` | `` | `value` | | | -| `enable_iter_req_stats` | `` | `value` | | | -| `enable_layerwise_nvtx_marker` | `` | `value` | | | -| `enable_lm_head_tp_in_adp` | `` | `value` | | | -| `enable_lora` | `` | `value` | | | -| `enable_low_latency_host_dispatch` | `` | `value` | | | -| `enable_min_latency` | `` | `value` | | | -| `enable_resource_governor` | `` | `value` | | | -| `enable_speculative_beam_history_d2h` | `` | `value` | | | -| `encode_only` | `` | `value` | | | -| `encoder_cuda_graph_config.batch_sizes` | `Optional[List[int]]` | `value` | | | -| `encoder_cuda_graph_config.enable_padding` | `` | `value` | | | -| `encoder_cuda_graph_config.max_batch_size` | `` | `value` | | | -| `encoder_cuda_graph_config.max_num_token` | `` | `value` | | | -| `encoder_cuda_graph_config.max_seq_len` | `` | `value` | | | -| `encoder_cuda_graph_config.mode` | `Literal['encode']` | `categorical` | | `encode` | -| `encoder_cuda_graph_config.num_tokens` | `Optional[List[Annotated[int, Gt(gt=0)]]]` | `value` | | | -| `encoder_cuda_graph_config.seq_lens` | `Optional[List[Annotated[int, Gt(gt=0)]]]` | `value` | | | -| `encoder_max_batch_size` | `Optional[int]` | `value` | | | -| `encoder_max_num_tokens` | `Optional[int]` | `value` | | | -| `force_dynamic_quantization` | `` | `value` | | | -| `garbage_collection_gen0_threshold` | `` | `value` | | | -| `gather_generation_logits` | `` | `value` | | | -| `generation_config` | `Literal['auto', 'trtllm']` | `categorical` | | `auto`, `trtllm` | -| `gms_config.mode` | `Literal['auto', 'rw', 'ro']` | `categorical` | | `auto`, `rw`, `ro` | -| `gpus_per_node` | `Optional[int]` | `value` | | | -| `guided_decoding_backend` | `Optional[Literal['xgrammar', 'llguidance']]` | `categorical` | | `xgrammar`, `llguidance` | -| `iter_stats_max_iterations` | `Optional[int]` | `value` | | | -| `kv_cache_compression_config.algorithm` | `Literal['triattention']` | `categorical` | | `triattention` | -| `kv_cache_compression_config.beta` | `` | `value` | | | -| `kv_cache_compression_config.budget` | `` | `value` | | | -| `kv_cache_compression_config.eviction_mode` | `Literal['union', 'per_head', 'per_layer_perhead']` | `categorical` | | `union`, `per_head`, `per_layer_perhead` | -| `kv_cache_compression_config.normalize_scores` | `` | `value` | | | -| `kv_cache_config.attention_dp_events_gather_period_ms` | `` | `value` | | | -| `kv_cache_config.avg_seq_len` | `Optional[Annotated[int, Gt(gt=0)]]` | `value` | | | -| `kv_cache_config.block_reuse_config.max_num_turns` | `` | `value` | | | -| `kv_cache_config.block_reuse_config.policy` | `Literal['all_reusable', 'per_request', 'per_conversation']` | `categorical` | | `all_reusable`, `per_request`, `per_conversation` | -| `kv_cache_config.copy_on_partial_reuse` | `` | `value` | | | -| `kv_cache_config.cross_kv_cache_fraction` | `Optional[float]` | `value` | | | -| `kv_cache_config.disk_cache_size` | `Optional[Annotated[int, Ge(ge=0)]]` | `value` | | | -| `kv_cache_config.disk_prefetch_num_reqs` | `` | `value` | | | -| `kv_cache_config.dtype` | `` | `categorical` | allowlist | `auto`, `float16`, `bfloat16`, `float32`, `fp8`, `nvfp4` | -| `kv_cache_config.enable_block_reuse` | `` | `value` | | | -| `kv_cache_config.enable_kv_pool_rebalance` | `` | `value` | | | -| `kv_cache_config.enable_partial_reuse` | `` | `value` | | | -| `kv_cache_config.enable_swa_scratch_reuse` | `` | `value` | | | -| `kv_cache_config.event_buffer_max_size` | `` | `value` | | | -| `kv_cache_config.fp8_context_mla_kv_len_cap` | `Optional[int]` | `value` | | | -| `kv_cache_config.free_gpu_memory_fraction` | `Optional[float]` | `value` | | | -| `kv_cache_config.host_cache_size` | `Optional[int]` | `value` | | | -| `kv_cache_config.iteration_stats_interval` | `` | `value` | | | -| `kv_cache_config.kv_cache_event_hash_algo` | `Literal['auto', 'v1_block_key', 'v2_sha256', 'v2_sha256_64']` | `categorical` | | `auto`, `v1_block_key`, `v2_sha256`, `v2_sha256_64` | -| `kv_cache_config.mamba_ssm_cache_dtype` | `Literal['auto', 'float16', 'bfloat16', 'float32']` | `categorical` | | `auto`, `float16`, `bfloat16`, `float32` | -| `kv_cache_config.mamba_ssm_philox_rounds` | `` | `value` | | | -| `kv_cache_config.mamba_ssm_stochastic_rounding` | `` | `value` | | | -| `kv_cache_config.mamba_state_config.periodic_snapshot_interval` | `` | `value` | | | -| `kv_cache_config.max_attention_window` | `Optional[List[int]]` | `value` | | | -| `kv_cache_config.max_gpu_total_bytes` | `` | `value` | | | -| `kv_cache_config.max_tokens` | `Optional[int]` | `value` | | | -| `kv_cache_config.max_util_for_resume` | `` | `value` | | | -| `kv_cache_config.pool_ratio` | `Optional[List[float]]` | `value` | | | -| `kv_cache_config.secondary_offload_min_priority` | `Optional[int]` | `value` | | | -| `kv_cache_config.sink_token_length` | `Optional[int]` | `value` | | | -| `kv_cache_config.tokens_per_block` | `` | `value` | | | -| `kv_cache_config.use_kv_cache_manager_v2` | `Union[bool, Literal['auto']]` | `value` | | `auto` | -| `kv_cache_config.use_uvm` | `` | `value` | | | -| `kv_connector_config.connector` | `Optional[str]` | `categorical` | allowlist | `lmcache`, `lmcache-mp`, `kvbm` | -| `layer_wise_benchmarks_config.calibration_layer_indices` | `Optional[List[int]]` | `value` | | | -| `layer_wise_benchmarks_config.calibration_mode` | `Literal['NONE', 'MARK', 'COLLECT']` | `categorical` | | `NONE`, `MARK`, `COLLECT` | -| `load_format` | `Union[str, tensorrt_llm.llmapi.llm_args.LoadFormat]` | `categorical` | allowlist | `auto`, `dummy`, `vision_only`, `gms` | -| `lora_config.lora_ckpt_source` | `Literal['hf', 'nemo']` | `categorical` | | `hf`, `nemo` | -| `lora_config.max_cpu_loras` | `Optional[int]` | `value` | | | -| `lora_config.max_lora_rank` | `` | `value` | | | -| `lora_config.max_loras` | `Optional[int]` | `value` | | | -| `lora_config.swap_gate_up_proj_lora_b_weight` | `` | `value` | | | -| `max_batch_size` | `Optional[int]` | `value` | | | -| `max_beam_width` | `Optional[int]` | `value` | | | -| `max_input_len` | `Optional[int]` | `value` | | | -| `max_num_tokens` | `Optional[int]` | `value` | | | -| `max_seq_len` | `Optional[int]` | `value` | | | -| `max_stats_len` | `` | `value` | | | -| `mm_encoder_only` | `` | `value` | | | -| `moe_cluster_parallel_size` | `Optional[int]` | `value` | | | -| `moe_config.backend` | `Literal['AUTO', 'CUTLASS', 'CUTEDSL', 'WIDEEP', 'TRTLLM', 'DEEPGEMM', 'DENSEGEMM', 'VANILLA', 'TRITON', 'MARLIN', 'MEGAMOE_DEEPGEMM', 'MEGAMOE_CUTEDSL']` | `categorical` | | `AUTO`, `CUTLASS`, `CUTEDSL`, `WIDEEP`, `TRTLLM`, `DEEPGEMM`, `DENSEGEMM`, `VANILLA`, `TRITON`, `MARLIN`, `MEGAMOE_DEEPGEMM`, `MEGAMOE_CUTEDSL` | -| `moe_config.disable_finalize_fusion` | `` | `value` | | | -| `moe_config.max_num_tokens` | `Optional[int]` | `value` | | | -| `moe_config.use_low_precision_moe_combine` | `` | `value` | | | -| `moe_expert_parallel_size` | `Optional[int]` | `value` | | | -| `moe_tensor_parallel_size` | `Optional[int]` | `value` | | | -| `multimodal_config.encoder_cache_max_bytes` | `` | `value` | | | -| `multimodal_config.encoder_scheduling_policy` | `` | `categorical` | | `DISABLED`, `DEFAULT`, `EAGER` | -| `multimodal_config.encoder_side_stream_max_ahead` | `` | `value` | | | -| `multimodal_config.video_pruning_rate` | `Optional[float]` | `value` | | | -| `mx_config.preshard_strategy` | `` | `categorical` | allowlist | `per_module` | -| `mx_config.server_query_timeout_s` | `Optional[Annotated[int, Ge(ge=0)]]` | `value` | | | -| `num_postprocess_workers` | `` | `value` | | | -| `num_serve_frontends` | `` | `value` | | | -| `nvfp4_gemm_config.allowed_backends` | `List[Literal['cutlass', 'cublaslt', 'cutedsl', 'cuda_core', 'marlin']]` | `value` | | `cutlass`, `cublaslt`, `cutedsl`, `cuda_core`, `marlin` | -| `orchestrator_type` | `Optional[Literal['rpc', 'ray']]` | `categorical` | | `rpc`, `ray` | -| `peft_cache_config.device_cache_percent` | `` | `value` | | | -| `peft_cache_config.host_cache_size` | `` | `value` | | | -| `peft_cache_config.max_adapter_size` | `` | `value` | | | -| `peft_cache_config.max_pages_per_block_device` | `` | `value` | | | -| `peft_cache_config.max_pages_per_block_host` | `` | `value` | | | -| `peft_cache_config.num_copy_streams` | `` | `value` | | | -| `peft_cache_config.num_device_module_layer` | `` | `value` | | | -| `peft_cache_config.num_ensure_workers` | `` | `value` | | | -| `peft_cache_config.num_host_module_layer` | `` | `value` | | | -| `peft_cache_config.num_put_workers` | `` | `value` | | | -| `peft_cache_config.optimal_adapter_size` | `` | `value` | | | -| `perf_metrics_max_requests` | `` | `value` | | | -| `pipeline_parallel_size` | `` | `value` | | | -| `pp_partition` | `Optional[List[int]]` | `value` | | | -| `prefill_capture_num_tokens` | `Optional[List[int]]` | `value` | | | -| `prefill_cuda_graph_backend` | `` | `categorical` | allowlist | `disabled`, `piecewise`, `breakable` | -| `print_iter_log` | `` | `value` | | | -| `prometheus_metrics_config.e2e_request_latency_buckets` | `Optional[List[float]]` | `value` | | | -| `prometheus_metrics_config.request_decode_time_buckets` | `Optional[List[float]]` | `value` | | | -| `prometheus_metrics_config.request_inference_time_buckets` | `Optional[List[float]]` | `value` | | | -| `prometheus_metrics_config.request_prefill_time_buckets` | `Optional[List[float]]` | `value` | | | -| `prometheus_metrics_config.request_queue_time_buckets` | `Optional[List[float]]` | `value` | | | -| `prometheus_metrics_config.time_per_output_token_buckets` | `Optional[List[float]]` | `value` | | | -| `prometheus_metrics_config.time_to_first_token_buckets` | `Optional[List[float]]` | `value` | | | -| `ray_placement_config.defer_workers_init` | `` | `value` | | | -| `ray_placement_config.per_worker_gpu_share` | `Optional[float]` | `value` | | | -| `ray_placement_config.placement_bundle_indices` | `Optional[List[List[int]]]` | `value` | | | -| `reasoning_parser` | `Optional[str]` | `categorical` | allowlist | `auto`, `deepseek-r1`, `poolside_v1`, `laguna`, `qwen3`, `qwen3_5`, `minimax_m2`, `minimax_m2_append_think`, `nano-v3`, `gemma4`, `kimi_k2`, `kimi_k25` | -| `reorder_policy_config.policy_args.agent_inflight_seq_num` | `` | `value` | | | -| `reorder_policy_config.policy_args.agent_percentage` | `` | `value` | | | -| `reorder_policy_config.policy_name` | `Optional[Literal['AgentTree']]` | `categorical` | | `AgentTree` | -| `request_stats_max_iterations` | `Optional[int]` | `value` | | | -| `return_perf_metrics` | `` | `value` | | | -| `sampler_force_async_worker` | `` | `value` | | | -| `scheduler_config.capacity_scheduler_policy` | `` | `categorical` | | `MAX_UTILIZATION`, `GUARANTEED_NO_EVICT`, `STATIC_BATCH` | -| `scheduler_config.context_chunking_policy` | `Optional[tensorrt_llm.llmapi.llm_args.ContextChunkingPolicy]` | `categorical` | | `FIRST_COME_FIRST_SERVED`, `EQUAL_PROGRESS`, `FORCE_CHUNK` | -| `scheduler_config.dynamic_batch_config.dynamic_batch_moving_average_window` | `` | `value` | | | -| `scheduler_config.dynamic_batch_config.enable_batch_size_tuning` | `` | `value` | | | -| `scheduler_config.dynamic_batch_config.enable_max_num_tokens_tuning` | `` | `value` | | | -| `scheduler_config.enable_prefix_aware_scheduling` | `` | `value` | | | -| `scheduler_config.use_python_scheduler` | `` | `value` | | | -| `scheduler_config.waiting_queue_policy` | `` | `categorical` | | `fcfs`, `priority` | -| `skip_tokenizer_init` | `` | `value` | | | -| `sparse_attention_config.algorithm` | `Literal['dsa']` | `categorical` | | `dsa`, `deepseek_v4`, `minimax_m3`, `rocket`, `skip_softmax` | -| `sparse_attention_config.compress_ratios` | `List[int]` | `value` | | | -| `sparse_attention_config.enable_heuristic_topk` | `` | `value` | | | -| `sparse_attention_config.implementation` | `Literal['triton', 'msa']` | `categorical` | | `triton`, `msa` | -| `sparse_attention_config.index_head_dim` | `Optional[int]` | `value` | | | -| `sparse_attention_config.index_n_heads` | `Optional[int]` | `value` | | | -| `sparse_attention_config.index_share_for_mtp_iteration` | `Optional[bool]` | `value` | | | -| `sparse_attention_config.index_topk` | `Optional[int]` | `value` | | | -| `sparse_attention_config.indexer_k_dtype` | `Literal['fp8', 'fp4']` | `categorical` | | `fp8`, `fp4` | -| `sparse_attention_config.indexer_kv_dtype` | `Literal['bf16', 'fp8']` | `categorical` | | `bf16`, `fp8` | -| `sparse_attention_config.indexer_max_chunk_size` | `Optional[int]` | `value` | | | -| `sparse_attention_config.indexer_rope_interleave` | `` | `value` | | | -| `sparse_attention_config.kernel_size` | `Optional[int]` | `value` | | | -| `sparse_attention_config.kt_cache_dtype` | `Optional[str]` | `categorical` | allowlist | `bfloat16`, `float8_e5m2` | -| `sparse_attention_config.num_attention_heads` | `Optional[int]` | `value` | | | -| `sparse_attention_config.num_key_value_heads` | `Optional[int]` | `value` | | | -| `sparse_attention_config.page_size` | `Optional[int]` | `value` | | | -| `sparse_attention_config.prompt_budget` | `Optional[int]` | `value` | | | -| `sparse_attention_config.q_split_threshold` | `` | `value` | | | -| `sparse_attention_config.seq_len_threshold` | `Optional[int]` | `value` | | | -| `sparse_attention_config.skip_indexer_for_short_seqs` | `` | `value` | | | -| `sparse_attention_config.sparse_block_size` | `` | `value` | | | -| `sparse_attention_config.sparse_disable_index_value` | `` | `value` | | | -| `sparse_attention_config.sparse_index_dim` | `` | `value` | | | -| `sparse_attention_config.sparse_init_blocks` | `` | `value` | | | -| `sparse_attention_config.sparse_local_blocks` | `` | `value` | | | -| `sparse_attention_config.sparse_num_index_heads` | `` | `value` | | | -| `sparse_attention_config.sparse_score_type` | `Literal['max']` | `categorical` | | `max` | -| `sparse_attention_config.sparse_topk_blocks` | `` | `value` | | | -| `sparse_attention_config.topk` | `Optional[int]` | `value` | | | -| `sparse_attention_config.topr` | `Union[int, float, NoneType]` | `value` | | | -| `sparse_attention_config.use_cute_dsl_paged_mqa_logits` | `` | `value` | | | -| `sparse_attention_config.use_cute_dsl_topk` | `` | `value` | | | -| `sparse_attention_config.window_size` | `` | `value` | | | -| `speculative_config.acceptance_rate_threshold` | `Optional[float]` | `value` | | | -| `speculative_config.acceptance_rate_window_size` | `Optional[Annotated[int, Ge(ge=0)]]` | `value` | | | -| `speculative_config.advanced_sampling_mode` | `` | `categorical` | | `full`, `no_topk`, `no_topp`, `no_topk_no_topp` | -| `speculative_config.allow_advanced_sampling` | `` | `value` | | | -| `speculative_config.attention_backend` | `Literal['VANILLA', 'TRTLLM']` | `categorical` | | `VANILLA`, `TRTLLM` | -| `speculative_config.begin_thinking_phase_token` | `` | `value` | | | -| `speculative_config.block_size` | `Optional[Annotated[int, Gt(gt=0)]]` | `value` | | | -| `speculative_config.decoding_type` | `Literal['AUTO']` | `categorical` | | `AUTO`, `DFlash`, `DSpark`, `Draft_Target`, `Eagle3`, `Eagle`, `MTP`, `NGram`, `PARD`, `SA`, `SaveState`, `User_Provided` | -| `speculative_config.dynamic_tree_max_topK` | `Optional[int]` | `value` | | | -| `speculative_config.eagle3_layers_to_capture` | `Optional[Set[int]]` | `value` | | | -| `speculative_config.eagle3_model_arch` | `Literal['llama3', 'mistral_large3']` | `categorical` | | `llama3`, `mistral_large3` | -| `speculative_config.eagle_choices` | `Optional[List[List[int]]]` | `value` | | | -| `speculative_config.enable_global_pool` | `` | `value` | | | -| `speculative_config.end_thinking_phase_token` | `` | `value` | | | -| `speculative_config.global_pool_size` | `Optional[Annotated[int, Gt(gt=0)]]` | `value` | | | -| `speculative_config.is_keep_all` | `` | `value` | | | -| `speculative_config.is_public_pool` | `` | `value` | | | -| `speculative_config.is_use_oldest` | `` | `value` | | | -| `speculative_config.markov_head_type` | `Optional[Literal['vanilla', 'gated', 'rnn']]` | `categorical` | | `vanilla`, `gated`, `rnn` | -| `speculative_config.markov_rank` | `Optional[int]` | `value` | | | -| `speculative_config.mask_token_id` | `Optional[int]` | `value` | | | -| `speculative_config.max_concurrency` | `Optional[Annotated[int, Gt(gt=0)]]` | `value` | | | -| `speculative_config.max_draft_len` | `Optional[Annotated[int, Ge(ge=0)]]` | `value` | | | -| `speculative_config.max_matching_ngram_size` | `` | `value` | | | -| `speculative_config.max_ngram_size` | `` | `value` | | | -| `speculative_config.max_non_leaves_per_layer` | `Optional[int]` | `value` | | | -| `speculative_config.max_total_draft_tokens` | `Optional[int]` | `value` | | | -| `speculative_config.max_verification_set_size` | `` | `value` | | | -| `speculative_config.max_window_size` | `` | `value` | | | -| `speculative_config.num_eagle_layers` | `Optional[int]` | `value` | | | -| `speculative_config.num_nextn_predict_layers` | `Optional[int]` | `value` | | | -| `speculative_config.relaxed_delta` | `` | `value` | | | -| `speculative_config.relaxed_topk` | `` | `value` | | | -| `speculative_config.sa_config.enable_global_pool` | `` | `value` | | | -| `speculative_config.sa_config.threshold` | `` | `value` | | | -| `speculative_config.target_layer_ids` | `Optional[List[int]]` | `value` | | | -| `speculative_config.use_dynamic_tree` | `Optional[bool]` | `value` | | | -| `speculative_config.use_mtp_vanilla` | `` | `value` | | | -| `speculative_config.use_rejection_sampling` | `` | `value` | | | -| `speculative_config.use_relaxed_acceptance_for_thinking` | `` | `value` | | | -| `speculative_config.write_interval` | `` | `value` | | | -| `stream_interval` | `` | `value` | | | -| `telemetry_config.disabled` | `` | `value` | | | -| `telemetry_config.usage_context` | `` | `categorical` | | `unknown`, `llm_class`, `cli_serve`, `cli_bench`, `cli_eval`, `disaggregated` | -| `tensor_parallel_size` | `` | `value` | | | -| `tokenizer_mode` | `Literal['auto', 'slow']` | `categorical` | | `auto`, `slow` | -| `torch_compile_config.capture_num_tokens` | `Optional[List[Annotated[int, Gt(gt=0)]]]` | `value` | | | -| `torch_compile_config.enable_fullgraph` | `` | `value` | | | -| `torch_compile_config.enable_inductor` | `` | `value` | | | -| `torch_compile_config.enable_piecewise_cuda_graph` | `` | `value` | | | -| `torch_compile_config.enable_userbuffers` | `` | `value` | | | -| `torch_compile_config.max_num_streams` | `` | `value` | | | -| `trust_remote_code` | `` | `value` | | | -| `use_cute_dsl_bf16_bmm` | `` | `value` | | | -| `use_cute_dsl_bf16_gemm` | `` | `value` | | | -| `use_cute_dsl_blockscaling_bmm` | `` | `value` | | | -| `use_cute_dsl_blockscaling_mm` | `` | `value` | | | +| Captured key | Capture policy | Kind | Categorical domain | +|--------------|----------------|------|--------------------| +| `allreduce_strategy` | `literal\|none` | `categorical` | `AUTO`, `NCCL`, `UB`, `MINLATENCY`, `ONESHOT`, `TWOSHOT`, `LOWPRECISION`, `MNNVL`, `NCCL_SYMMETRIC` | +| `attention_dp_config.batching_wait_iters` | `int` | `value` | | +| `attention_dp_config.enable_balance` | `bool` | `value` | | +| `attention_dp_config.enable_kv_cache_aware_routing` | `bool` | `value` | | +| `attention_dp_config.kv_cache_routing_account_for_in_transfer` | `bool` | `value` | | +| `attention_dp_config.kv_cache_routing_cold_start_warmup` | `bool` | `value` | | +| `attention_dp_config.kv_cache_routing_conversation_affinity` | `bool` | `value` | | +| `attention_dp_config.kv_cache_routing_fair_share_multiplier` | `float` | `value` | | +| `attention_dp_config.kv_cache_routing_load_balance_weight` | `float` | `value` | | +| `attention_dp_config.kv_cache_routing_match_rate_threshold` | `float` | `value` | | +| `attention_dp_config.kv_cache_routing_max_sessions` | `int` | `value` | | +| `attention_dp_config.kv_cache_routing_new_conv_placement` | `literal` | `categorical` | `round_robin`, `least_queued` | +| `attention_dp_config.timeout_iters` | `int` | `value` | | +| `attn_backend` | `allowlist` | `categorical` | `VANILLA`, `TRTLLM`, `FLASHINFER` | +| `backend` | `literal` | `categorical` | `pytorch` | +| `batch_wait_max_tokens_ratio` | `float` | `value` | | +| `batch_wait_timeout_iters` | `int` | `value` | | +| `batch_wait_timeout_ms` | `float` | `value` | | +| `cache_transceiver_config.backend` | `literal\|none` | `categorical` | `DEFAULT`, `UCX`, `NIXL`, `MOONCAKE`, `MPI` | +| `cache_transceiver_config.enable_pipelined_transfer` | `bool` | `value` | | +| `cache_transceiver_config.kv_cache_bounce_size_mb` | `int` | `value` | | +| `cache_transceiver_config.kv_transfer_poll_interval_ms` | `int\|none` | `value` | | +| `cache_transceiver_config.kv_transfer_sender_future_timeout_ms` | `int\|none` | `value` | | +| `cache_transceiver_config.kv_transfer_timeout_ms` | `int\|none` | `value` | | +| `cache_transceiver_config.max_tokens_in_buffer` | `int\|none` | `value` | | +| `cache_transceiver_config.transceiver_runtime` | `literal\|none` | `categorical` | `CPP`, `PYTHON`, `auto` | +| `checkpoint_io_policy` | `literal` | `categorical` | `auto`, `native`, `rank_striped_read_ahead` | +| `context_parallel_size` | `int` | `value` | | +| `cp_config.cp_type` | `enum[CpType]` | `categorical` | `ULYSSES`, `RING`, `HELIX` | +| `cp_config.fifo_version` | `int\|none` | `value` | | +| `cp_config.tokens_per_block` | `int\|none` | `value` | | +| `cp_config.use_nccl_for_alltoall` | `bool\|none` | `value` | | +| `cuda_graph_config.batch_sizes` | `list[int]\|none` | `value` | | +| `cuda_graph_config.enable_padding` | `bool` | `value` | | +| `cuda_graph_config.max_batch_size` | `int` | `value` | | +| `cuda_graph_config.max_num_token` | `int` | `value` | | +| `cuda_graph_config.max_seq_len` | `int` | `value` | | +| `cuda_graph_config.mode` | `literal` | `categorical` | `decode`, `encode` | +| `cuda_graph_config.num_tokens` | `list[int]\|none` | `value` | | +| `cuda_graph_config.seq_lens` | `list[int]\|none` | `value` | | +| `disable_mm_encoder` | `bool` | `value` | | +| `disable_overlap_scheduler` | `bool` | `value` | | +| `dtype` | `allowlist` | `categorical` | `auto`, `float16`, `bfloat16`, `float32` | +| `dwdp_config.contention_opt` | `bool` | `value` | | +| `dwdp_config.dwdp_size` | `int` | `value` | | +| `dwdp_config.num_experts_per_worker` | `int` | `value` | | +| `dwdp_config.num_groups` | `int` | `value` | | +| `dwdp_config.num_prefetch_experts` | `int` | `value` | | +| `enable_attention_dp` | `bool` | `value` | | +| `enable_autotuner` | `bool` | `value` | | +| `enable_chunked_prefill` | `bool` | `value` | | +| `enable_early_first_token_response` | `bool` | `value` | | +| `enable_encoder_decoder_mixed_cuda_graph` | `bool` | `value` | | +| `enable_energy_metrics` | `bool` | `value` | | +| `enable_in_graph_sampling` | `bool` | `value` | | +| `enable_iter_perf_stats` | `bool` | `value` | | +| `enable_iter_req_stats` | `bool` | `value` | | +| `enable_layerwise_nvtx_marker` | `bool` | `value` | | +| `enable_lm_head_tp_in_adp` | `bool` | `value` | | +| `enable_lora` | `bool` | `value` | | +| `enable_low_latency_host_dispatch` | `bool` | `value` | | +| `enable_min_latency` | `bool` | `value` | | +| `enable_mla_skip_correction` | `bool` | `value` | | +| `enable_resource_governor` | `bool` | `value` | | +| `enable_speculative_beam_history_d2h` | `bool` | `value` | | +| `encode_only` | `bool` | `value` | | +| `encoder_cuda_graph_config.batch_sizes` | `list[int]\|none` | `value` | | +| `encoder_cuda_graph_config.enable_padding` | `bool` | `value` | | +| `encoder_cuda_graph_config.max_batch_size` | `int` | `value` | | +| `encoder_cuda_graph_config.max_num_token` | `int` | `value` | | +| `encoder_cuda_graph_config.max_seq_len` | `int` | `value` | | +| `encoder_cuda_graph_config.mode` | `literal` | `categorical` | `encode` | +| `encoder_cuda_graph_config.num_tokens` | `list[int]\|none` | `value` | | +| `encoder_cuda_graph_config.seq_lens` | `list[int]\|none` | `value` | | +| `encoder_max_batch_size` | `int\|none` | `value` | | +| `encoder_max_num_tokens` | `int\|none` | `value` | | +| `force_dynamic_quantization` | `bool` | `value` | | +| `garbage_collection_gen0_threshold` | `int` | `value` | | +| `gather_generation_logits` | `bool` | `value` | | +| `generation_config` | `literal` | `categorical` | `auto`, `trtllm` | +| `gms_config.mode` | `literal` | `categorical` | `auto`, `rw`, `ro` | +| `gpus_per_node` | `int\|none` | `value` | | +| `guided_decoding_backend` | `literal\|none` | `categorical` | `xgrammar`, `llguidance` | +| `iter_stats_max_iterations` | `int\|none` | `value` | | +| `kv_cache_compression_config.algorithm` | `literal` | `categorical` | `quantization_for_cold_page`, `triattention` | +| `kv_cache_compression_config.beta` | `int` | `value` | | +| `kv_cache_compression_config.budget` | `int` | `value` | | +| `kv_cache_compression_config.eviction_mode` | `literal` | `categorical` | `union`, `per_head`, `per_layer_perhead` | +| `kv_cache_compression_config.normalize_scores` | `bool` | `value` | | +| `kv_cache_compression_config.quant` | `literal` | `categorical` | `nvfp4` | +| `kv_cache_config.attention_dp_events_gather_period_ms` | `int` | `value` | | +| `kv_cache_config.avg_seq_len` | `int\|none` | `value` | | +| `kv_cache_config.block_reuse_config.max_num_turns` | `int` | `value` | | +| `kv_cache_config.block_reuse_config.policy` | `literal` | `categorical` | `all_reusable`, `per_request`, `per_conversation` | +| `kv_cache_config.copy_on_partial_reuse` | `bool` | `value` | | +| `kv_cache_config.cross_kv_cache_fraction` | `float\|none` | `value` | | +| `kv_cache_config.disk_cache_size` | `int\|none` | `value` | | +| `kv_cache_config.disk_prefetch_num_reqs` | `int` | `value` | | +| `kv_cache_config.dtype` | `allowlist` | `categorical` | `auto`, `float16`, `bfloat16`, `float32`, `fp8`, `fp8_ds_mla`, `nvfp4` | +| `kv_cache_config.enable_block_reuse` | `bool` | `value` | | +| `kv_cache_config.enable_kv_pool_rebalance` | `bool` | `value` | | +| `kv_cache_config.enable_partial_reuse` | `bool` | `value` | | +| `kv_cache_config.enable_swa_scratch_reuse` | `bool` | `value` | | +| `kv_cache_config.event_buffer_max_size` | `int` | `value` | | +| `kv_cache_config.fp8_context_mla_kv_len_cap` | `int\|none` | `value` | | +| `kv_cache_config.free_gpu_memory_fraction` | `float\|none` | `value` | | +| `kv_cache_config.host_cache_size` | `int\|none` | `value` | | +| `kv_cache_config.iteration_stats_interval` | `int` | `value` | | +| `kv_cache_config.kv_cache_event_hash_algo` | `literal` | `categorical` | `auto`, `v1_block_key`, `v2_sha256`, `v2_sha256_64` | +| `kv_cache_config.kv_events_config.buffer_steps` | `int` | `value` | | +| `kv_cache_config.kv_events_config.enable_kv_cache_events` | `bool` | `value` | | +| `kv_cache_config.kv_events_config.hwm` | `int` | `value` | | +| `kv_cache_config.kv_events_config.max_queue_size` | `int` | `value` | | +| `kv_cache_config.kv_events_config.publisher` | `literal\|none` | `categorical` | `null`, `zmq` | +| `kv_cache_config.mamba_ssm_cache_dtype` | `literal` | `categorical` | `auto`, `float16`, `bfloat16`, `float32` | +| `kv_cache_config.mamba_ssm_philox_rounds` | `int` | `value` | | +| `kv_cache_config.mamba_ssm_stochastic_rounding` | `bool` | `value` | | +| `kv_cache_config.mamba_state_config.enable_branch_snapshot` | `bool` | `value` | | +| `kv_cache_config.mamba_state_config.periodic_snapshot_interval` | `int` | `value` | | +| `kv_cache_config.max_attention_window` | `list[int]\|none` | `value` | | +| `kv_cache_config.max_gpu_total_bytes` | `int` | `value` | | +| `kv_cache_config.max_tokens` | `int\|none` | `value` | | +| `kv_cache_config.max_util_for_resume` | `float` | `value` | | +| `kv_cache_config.pool_ratio` | `list[float]\|none` | `value` | | +| `kv_cache_config.secondary_offload_min_priority` | `int\|none` | `value` | | +| `kv_cache_config.sink_token_length` | `int\|none` | `value` | | +| `kv_cache_config.tokens_per_block` | `int` | `value` | | +| `kv_cache_config.use_kv_cache_manager_v2` | `bool\|literal` | `categorical` | `auto` | +| `kv_cache_config.use_uvm` | `bool` | `value` | | +| `kv_connector_config.connector` | `allowlist\|none` | `categorical` | `lmcache`, `lmcache-mp`, `kvbm` | +| `layer_wise_benchmarks_config.calibration_layer_indices` | `list[int]\|none` | `value` | | +| `layer_wise_benchmarks_config.calibration_mode` | `literal` | `categorical` | `NONE`, `MARK`, `COLLECT` | +| `load_format` | `allowlist\|enum[LoadFormat]` | `categorical` | `auto`, `dummy`, `vision_only`, `gms`, `AUTO`, `DUMMY`, `VISION_ONLY`, `GMS` | +| `lora_config.cuda_graph_specialize_lora` | `bool` | `value` | | +| `lora_config.lora_ckpt_source` | `literal` | `categorical` | `hf`, `nemo` | +| `lora_config.max_cpu_loras` | `int\|none` | `value` | | +| `lora_config.max_lora_rank` | `int` | `value` | | +| `lora_config.max_loras` | `int\|none` | `value` | | +| `lora_config.overlap_lora_and_base` | `bool` | `value` | | +| `lora_config.swap_gate_up_proj_lora_b_weight` | `bool` | `value` | | +| `max_batch_size` | `int\|none` | `value` | | +| `max_beam_width` | `int\|none` | `value` | | +| `max_input_len` | `int\|none` | `value` | | +| `max_num_tokens` | `int\|none` | `value` | | +| `max_seq_len` | `int\|none` | `value` | | +| `max_stats_len` | `int` | `value` | | +| `mla_skip_correction_threshold` | `float` | `value` | | +| `mm_encoder_only` | `bool` | `value` | | +| `moe_cluster_parallel_size` | `int\|none` | `value` | | +| `moe_config.backend` | `literal` | `categorical` | `AUTO`, `CUTLASS`, `CUTEDSL`, `TRTLLM`, `DEEPGEMM`, `DENSEGEMM`, `VANILLA`, `TRITON`, `MARLIN`, `MEGAMOE_DEEPGEMM`, `MEGAMOE_CUTEDSL` | +| `moe_config.disable_finalize_fusion` | `bool` | `value` | | +| `moe_config.max_num_tokens` | `int\|none` | `value` | | +| `moe_config.use_low_precision_moe_combine` | `bool` | `value` | | +| `moe_expert_parallel_size` | `int\|none` | `value` | | +| `moe_tensor_parallel_size` | `int\|none` | `value` | | +| `multimodal_config.encoder_cache_max_bytes` | `int` | `value` | | +| `multimodal_config.encoder_scheduling_policy` | `enum[MultimodalEncoderSchedulingPolicy]` | `categorical` | `DISABLED`, `DEFAULT`, `EAGER` | +| `multimodal_config.encoder_side_stream_max_ahead` | `int` | `value` | | +| `multimodal_config.video_pruning_rate` | `float\|none` | `value` | | +| `mx_config.preshard_strategy` | `allowlist` | `categorical` | `per_module` | +| `mx_config.server_query_timeout_s` | `int\|none` | `value` | | +| `num_postprocess_workers` | `int` | `value` | | +| `num_serve_frontends` | `int` | `value` | | +| `nvfp4_gemm_config.allowed_backends` | `list[literal]` | `categorical` | `cutlass`, `cublaslt`, `cutedsl`, `cuda_core`, `marlin` | +| `orchestrator_type` | `literal\|none` | `categorical` | `rpc`, `ray` | +| `peft_cache_config.device_cache_percent` | `float` | `value` | | +| `peft_cache_config.host_cache_size` | `int` | `value` | | +| `peft_cache_config.max_adapter_size` | `int` | `value` | | +| `peft_cache_config.max_pages_per_block_device` | `int` | `value` | | +| `peft_cache_config.max_pages_per_block_host` | `int` | `value` | | +| `peft_cache_config.num_copy_streams` | `int` | `value` | | +| `peft_cache_config.num_device_module_layer` | `int` | `value` | | +| `peft_cache_config.num_ensure_workers` | `int` | `value` | | +| `peft_cache_config.num_host_module_layer` | `int` | `value` | | +| `peft_cache_config.num_put_workers` | `int` | `value` | | +| `peft_cache_config.optimal_adapter_size` | `int` | `value` | | +| `perf_metrics_max_requests` | `int` | `value` | | +| `pipeline_parallel_size` | `int` | `value` | | +| `pp_partition` | `list[int]\|none` | `value` | | +| `prefill_capture_num_tokens` | `list[int]\|none` | `value` | | +| `prefill_cuda_graph_backend` | `enum[PrefillCudaGraphBackend]` | `categorical` | `disabled`, `piecewise`, `breakable` | +| `print_iter_log` | `bool` | `value` | | +| `prometheus_metrics_config.e2e_request_latency_buckets` | `list[float]\|none` | `value` | | +| `prometheus_metrics_config.request_decode_time_buckets` | `list[float]\|none` | `value` | | +| `prometheus_metrics_config.request_inference_time_buckets` | `list[float]\|none` | `value` | | +| `prometheus_metrics_config.request_prefill_time_buckets` | `list[float]\|none` | `value` | | +| `prometheus_metrics_config.request_queue_time_buckets` | `list[float]\|none` | `value` | | +| `prometheus_metrics_config.time_per_output_token_buckets` | `list[float]\|none` | `value` | | +| `prometheus_metrics_config.time_to_first_token_buckets` | `list[float]\|none` | `value` | | +| `ray_placement_config.defer_workers_init` | `bool` | `value` | | +| `ray_placement_config.per_worker_gpu_share` | `float\|none` | `value` | | +| `ray_placement_config.placement_bundle_indices` | `list[list[int]]\|none` | `value` | | +| `reasoning_parser` | `allowlist\|none` | `categorical` | `auto`, `deepseek-r1`, `poolside_v1`, `laguna`, `qwen3`, `qwen3_5`, `minimax_m2`, `minimax_m2_append_think`, `nano-v3`, `gemma4`, `kimi_k2`, `kimi_k25` | +| `reorder_policy_config.policy_args.agent_inflight_seq_num` | `int` | `value` | | +| `reorder_policy_config.policy_args.agent_percentage` | `float` | `value` | | +| `reorder_policy_config.policy_name` | `literal\|none` | `categorical` | `AgentTree` | +| `request_stats_max_iterations` | `int\|none` | `value` | | +| `return_perf_metrics` | `bool` | `value` | | +| `sampler_force_async_worker` | `bool` | `value` | | +| `scheduler_config.capacity_scheduler_policy` | `enum[CapacitySchedulerPolicy]` | `categorical` | `MAX_UTILIZATION`, `GUARANTEED_NO_EVICT`, `STATIC_BATCH` | +| `scheduler_config.context_chunking_policy` | `enum[ContextChunkingPolicy]\|none` | `categorical` | `FIRST_COME_FIRST_SERVED`, `EQUAL_PROGRESS`, `FORCE_CHUNK` | +| `scheduler_config.dynamic_batch_config.dynamic_batch_moving_average_window` | `int` | `value` | | +| `scheduler_config.dynamic_batch_config.enable_batch_size_tuning` | `bool` | `value` | | +| `scheduler_config.dynamic_batch_config.enable_max_num_tokens_tuning` | `bool` | `value` | | +| `scheduler_config.enable_prefix_aware_scheduling` | `bool` | `value` | | +| `scheduler_config.use_python_scheduler` | `bool` | `value` | | +| `scheduler_config.waiting_queue_policy` | `enum[WaitingQueuePolicy]` | `categorical` | `fcfs`, `priority` | +| `skip_tokenizer_init` | `bool` | `value` | | +| `sparse_attention_config.algorithm` | `literal` | `categorical` | `dsa`, `deepseek_v4`, `minimax_m3`, `qsa`, `rocket`, `skip_softmax` | +| `sparse_attention_config.compress_ratios` | `list[int]` | `value` | | +| `sparse_attention_config.enable_heuristic_topk` | `bool` | `value` | | +| `sparse_attention_config.implementation` | `literal` | `categorical` | `triton`, `msa` | +| `sparse_attention_config.index_head_dim` | `int\|none` | `value` | | +| `sparse_attention_config.index_n_heads` | `int\|none` | `value` | | +| `sparse_attention_config.index_share_for_mtp_iteration` | `bool\|none` | `value` | | +| `sparse_attention_config.index_topk` | `int\|none` | `value` | | +| `sparse_attention_config.indexer_k_dtype` | `literal` | `categorical` | `fp8`, `fp4` | +| `sparse_attention_config.indexer_kv_dtype` | `literal` | `categorical` | `bf16`, `fp8` | +| `sparse_attention_config.indexer_max_chunk_size` | `int\|none` | `value` | | +| `sparse_attention_config.indexer_rope_interleave` | `bool` | `value` | | +| `sparse_attention_config.kernel_size` | `int\|none` | `value` | | +| `sparse_attention_config.kt_cache_dtype` | `allowlist\|none` | `categorical` | `bfloat16`, `float8_e5m2` | +| `sparse_attention_config.num_attention_heads` | `int\|none` | `value` | | +| `sparse_attention_config.num_key_value_heads` | `int\|none` | `value` | | +| `sparse_attention_config.page_size` | `int\|none` | `value` | | +| `sparse_attention_config.prompt_budget` | `int\|none` | `value` | | +| `sparse_attention_config.q_split_threshold` | `int` | `value` | | +| `sparse_attention_config.seq_len_threshold` | `int\|none` | `value` | | +| `sparse_attention_config.skip_indexer_for_short_seqs` | `bool` | `value` | | +| `sparse_attention_config.sparse_block_size` | `int` | `value` | | +| `sparse_attention_config.sparse_disable_index_value` | `bool` | `value` | | +| `sparse_attention_config.sparse_index_dim` | `int` | `value` | | +| `sparse_attention_config.sparse_init_blocks` | `int` | `value` | | +| `sparse_attention_config.sparse_local_blocks` | `int` | `value` | | +| `sparse_attention_config.sparse_num_index_heads` | `int` | `value` | | +| `sparse_attention_config.sparse_score_type` | `literal` | `categorical` | `max` | +| `sparse_attention_config.sparse_topk_blocks` | `int` | `value` | | +| `sparse_attention_config.target_sparsity` | `float\|none` | `value` | | +| `sparse_attention_config.threshold_scale_factor` | `float\|none` | `value` | | +| `sparse_attention_config.topk` | `int\|none` | `value` | | +| `sparse_attention_config.topr` | `float\|int\|none` | `value` | | +| `sparse_attention_config.use_cute_dsl_paged_mqa_logits` | `bool` | `value` | | +| `sparse_attention_config.use_cute_dsl_topk` | `bool` | `value` | | +| `sparse_attention_config.use_gvr_emission` | `bool` | `value` | | +| `sparse_attention_config.use_self_sampling_topk` | `bool` | `value` | | +| `sparse_attention_config.window_size` | `int\|none` | `value` | | +| `speculative_config.acceptance_rate_threshold` | `float\|none` | `value` | | +| `speculative_config.acceptance_rate_window_size` | `int\|none` | `value` | | +| `speculative_config.advanced_sampling_mode` | `enum[AdvancedSamplingMode]` | `categorical` | `full`, `no_topk`, `no_topp`, `no_topk_no_topp` | +| `speculative_config.allow_advanced_sampling` | `bool` | `value` | | +| `speculative_config.attention_backend` | `literal` | `categorical` | `VANILLA`, `TRTLLM`, `FA4` | +| `speculative_config.begin_thinking_phase_token` | `int` | `value` | | +| `speculative_config.block_size` | `int\|none` | `value` | | +| `speculative_config.decoding_type` | `literal` | `categorical` | `AUTO`, `DFlash`, `DSpark`, `Draft_Target`, `Eagle3`, `Eagle`, `MTP`, `NGram`, `PARD`, `SA`, `SaveState`, `User_Provided` | +| `speculative_config.dynamic_tree_max_topK` | `int\|none` | `value` | | +| `speculative_config.eagle3_layers_to_capture` | `none\|set[int]` | `value` | | +| `speculative_config.eagle3_model_arch` | `literal` | `categorical` | `llama3`, `mistral_large3` | +| `speculative_config.eagle_choices` | `list[list[int]]\|none` | `value` | | +| `speculative_config.enable_global_pool` | `bool` | `value` | | +| `speculative_config.enable_penalty` | `bool` | `value` | | +| `speculative_config.end_thinking_phase_token` | `int` | `value` | | +| `speculative_config.global_pool_size` | `int\|none` | `value` | | +| `speculative_config.is_keep_all` | `bool` | `value` | | +| `speculative_config.is_public_pool` | `bool` | `value` | | +| `speculative_config.is_use_oldest` | `bool` | `value` | | +| `speculative_config.markov_head_type` | `literal\|none` | `categorical` | `vanilla`, `gated`, `rnn` | +| `speculative_config.markov_rank` | `int\|none` | `value` | | +| `speculative_config.mask_token_id` | `int\|none` | `value` | | +| `speculative_config.max_concurrency` | `int\|none` | `value` | | +| `speculative_config.max_draft_len` | `int\|none` | `value` | | +| `speculative_config.max_matching_ngram_size` | `int` | `value` | | +| `speculative_config.max_non_leaves_per_layer` | `int\|none` | `value` | | +| `speculative_config.max_total_draft_tokens` | `int\|none` | `value` | | +| `speculative_config.num_eagle_layers` | `int\|none` | `value` | | +| `speculative_config.num_nextn_predict_layers` | `int\|none` | `value` | | +| `speculative_config.relaxed_delta` | `float` | `value` | | +| `speculative_config.relaxed_topk` | `int` | `value` | | +| `speculative_config.sa_config.enable_global_pool` | `bool` | `value` | | +| `speculative_config.sa_config.threshold` | `int` | `value` | | +| `speculative_config.target_layer_ids` | `list[int]\|none` | `value` | | +| `speculative_config.use_dynamic_tree` | `bool\|none` | `value` | | +| `speculative_config.use_mtp_vanilla` | `bool` | `value` | | +| `speculative_config.use_rejection_sampling` | `bool` | `value` | | +| `speculative_config.use_relaxed_acceptance_for_thinking` | `bool` | `value` | | +| `speculative_config.write_interval` | `int` | `value` | | +| `stream_interval` | `int` | `value` | | +| `telemetry_config.disabled` | `bool` | `value` | | +| `telemetry_config.usage_context` | `enum[UsageContext]` | `categorical` | `unknown`, `llm_class`, `cli_serve`, `cli_bench`, `cli_eval`, `disaggregated` | +| `tensor_parallel_size` | `int` | `value` | | +| `tokenizer_mode` | `literal` | `categorical` | `auto`, `slow` | +| `torch_compile_config.capture_num_tokens` | `list[int]\|none` | `value` | | +| `torch_compile_config.enable_fullgraph` | `bool` | `value` | | +| `torch_compile_config.enable_inductor` | `bool` | `value` | | +| `torch_compile_config.enable_piecewise_cuda_graph` | `bool` | `value` | | +| `torch_compile_config.enable_userbuffers` | `bool` | `value` | | +| `torch_compile_config.max_num_streams` | `int` | `value` | | +| `trust_remote_code` | `bool` | `value` | | +| `use_cute_dsl_bf16_bmm` | `bool` | `value` | | +| `use_cute_dsl_bf16_gemm` | `bool` | `value` | | +| `use_cute_dsl_blockscaling_bmm` | `bool` | `value` | | +| `use_cute_dsl_blockscaling_mm` | `bool` | `value` | | +| `use_fine_grained_sync` | `bool` | `value` | | diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 2d6f08e0e812..59c5030de54c 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -110,8 +110,9 @@ def Field(default: Any = ..., - "prototype": Not yet stable and subject to breaking changes; intended for experimentation only. telemetry: Optional field-local telemetry override for LLM API config capture. Type-safe fields (categorical/numeric) auto-enroll; pass - telemetry=TelemetryField.categorical(...) to opt a free-form str/Any field - in via an allowlist, or telemetry=False to opt a type-safe field out. + telemetry=TelemetryField.categorical(...) to opt an otherwise unsafe + categorical branch in with exact allowed values, or telemetry=False + to opt a type-safe field out. **kwargs: All other arguments passed to the original Pydantic Field Returns: @@ -134,7 +135,7 @@ def Field(default: Any = ..., if isinstance(telemetry, TelemetryField): telemetry_metadata = telemetry.as_json_schema_extra() elif telemetry is True: - telemetry_metadata = {"kind": "value"} + telemetry_metadata = {} elif isinstance(telemetry, dict): telemetry_metadata = dict(telemetry) else: @@ -3868,17 +3869,11 @@ def supports_speculative_decoding(self) -> bool: return False -_KV_CACHE_COMPRESSION_ALGORITHM_TELEMETRY = TelemetryField.categorical( - "quantization_for_cold_page", "triattention") - - class ColdPageQuantizationCompressionConfig(KvCacheCompressionConfig): """Quantize Host and Disk KV pages without changing the active GPU cache.""" algorithm: Literal["quantization_for_cold_page"] = Field( - default="quantization_for_cold_page", - telemetry=False, - ) + default="quantization_for_cold_page") quant: Literal["nvfp4"] = Field( default="nvfp4", description="Quantization format stored in the compressed cache tier.") @@ -3909,10 +3904,7 @@ class TriAttentionKvCacheCompressionConfig(KvCacheCompressionConfig): changes_physical_kv_length: ClassVar[bool] = True - algorithm: Literal["triattention"] = Field( - default="triattention", - telemetry=_KV_CACHE_COMPRESSION_ALGORITHM_TELEMETRY, - ) + algorithm: Literal["triattention"] = Field(default="triattention") eviction_mode: Literal["union", "per_head", "per_layer_perhead"] = Field( default="union", description= @@ -5804,9 +5796,7 @@ def validate_mla_skip_correction_config(self) -> 'TorchLlmArgs': default=PrefillCudaGraphBackend.DISABLED, description="CUDA graph implementation used for prefill requests. " "Defaults to disabled.", - status="prototype", - telemetry=TelemetryField.categorical("disabled", "piecewise", - "breakable")) + status="prototype") prefill_capture_num_tokens: Optional[List[int]] = Field( default=None, diff --git a/tensorrt_llm/usage/config.py b/tensorrt_llm/usage/config.py index c3b38b777160..639e6edbcd2c 100644 --- a/tensorrt_llm/usage/config.py +++ b/tensorrt_llm/usage/config.py @@ -24,7 +24,7 @@ from dataclasses import dataclass from enum import Enum -from typing import Any, Literal, Optional +from typing import Any from pydantic import BaseModel, ConfigDict, Field @@ -52,31 +52,15 @@ class UsageContext(str, Enum): class TelemetryField: """Field-local opt-in metadata for LLM API config telemetry capture.""" - kind: Literal["value", "categorical"] = "value" - converter: Optional[Literal["allowlist"]] = None - allowed_values: Optional[tuple[Any, ...]] = None + allowed_values: tuple[Any, ...] @classmethod def categorical(cls, *allowed_values: Any) -> "TelemetryField": - """Build a categorical allowlist field from the recognized values. - - Shorthand for the common bare-string allowlist case: marks the field - categorical and pins capture to the explicit allowed values via the - allowlist converter. - """ - return cls( - kind="categorical", - converter="allowlist", - allowed_values=tuple(allowed_values), - ) + """Allow capture only for the specified values and their exact types.""" + return cls(allowed_values=allowed_values) def as_json_schema_extra(self) -> dict[str, Any]: - data: dict[str, Any] = {"kind": self.kind} - if self.converter is not None: - data["converter"] = self.converter - if self.allowed_values is not None: - data["allowed_values"] = list(self.allowed_values) - return data + return {"allowed_values": list(self.allowed_values)} class TelemetryConfig(_StrictUsageBaseModel): diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index aef78b3a38ce..e4f0bcc10f7e 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -12,78 +12,57 @@ "MNNVL", "NCCL_SYMMETRIC" ], - "annotation": "Optional[Literal['AUTO', 'NCCL', 'UB', 'MINLATENCY', 'ONESHOT', 'TWOSHOT', 'LOWPRECISION', 'MNNVL', 'NCCL_SYMMETRIC']]", - "converter": "", + "capture_policy": "literal|none", "kind": "categorical", "path": "allreduce_strategy" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "int", "kind": "value", "path": "attention_dp_config.batching_wait_iters" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "attention_dp_config.enable_balance" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "attention_dp_config.enable_kv_cache_aware_routing" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "attention_dp_config.kv_cache_routing_account_for_in_transfer" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "attention_dp_config.kv_cache_routing_cold_start_warmup" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "attention_dp_config.kv_cache_routing_conversation_affinity" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "float", "kind": "value", "path": "attention_dp_config.kv_cache_routing_fair_share_multiplier" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "float", "kind": "value", "path": "attention_dp_config.kv_cache_routing_load_balance_weight" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "float", "kind": "value", "path": "attention_dp_config.kv_cache_routing_match_rate_threshold" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "int", "kind": "value", "path": "attention_dp_config.kv_cache_routing_max_sessions" }, @@ -92,15 +71,12 @@ "round_robin", "least_queued" ], - "annotation": "Literal['round_robin', 'least_queued']", - "converter": "", + "capture_policy": "literal", "kind": "categorical", "path": "attention_dp_config.kv_cache_routing_new_conv_placement" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "int", "kind": "value", "path": "attention_dp_config.timeout_iters" }, @@ -110,8 +86,7 @@ "TRTLLM", "FLASHINFER" ], - "annotation": "", - "converter": "allowlist", + "capture_policy": "allowlist", "kind": "categorical", "path": "attn_backend" }, @@ -119,29 +94,22 @@ "allowed_values": [ "pytorch" ], - "annotation": "Literal['pytorch']", - "converter": "", + "capture_policy": "literal", "kind": "categorical", "path": "backend" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "float", "kind": "value", "path": "batch_wait_max_tokens_ratio" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "int", "kind": "value", "path": "batch_wait_timeout_iters" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "float", "kind": "value", "path": "batch_wait_timeout_ms" }, @@ -153,50 +121,37 @@ "MOONCAKE", "MPI" ], - "annotation": "Optional[Literal['DEFAULT', 'UCX', 'NIXL', 'MOONCAKE', 'MPI']]", - "converter": "", + "capture_policy": "literal|none", "kind": "categorical", "path": "cache_transceiver_config.backend" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "cache_transceiver_config.enable_pipelined_transfer" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "int", "kind": "value", "path": "cache_transceiver_config.kv_cache_bounce_size_mb" }, { - "allowed_values": [], - "annotation": "Optional[Annotated[int, Gt(gt=0)]]", - "converter": "", + "capture_policy": "int|none", "kind": "value", "path": "cache_transceiver_config.kv_transfer_poll_interval_ms" }, { - "allowed_values": [], - "annotation": "Optional[Annotated[int, Gt(gt=0)]]", - "converter": "", + "capture_policy": "int|none", "kind": "value", "path": "cache_transceiver_config.kv_transfer_sender_future_timeout_ms" }, { - "allowed_values": [], - "annotation": "Optional[Annotated[int, Gt(gt=0)]]", - "converter": "", + "capture_policy": "int|none", "kind": "value", "path": "cache_transceiver_config.kv_transfer_timeout_ms" }, { - "allowed_values": [], - "annotation": "Optional[int]", - "converter": "", + "capture_policy": "int|none", "kind": "value", "path": "cache_transceiver_config.max_tokens_in_buffer" }, @@ -206,8 +161,7 @@ "PYTHON", "auto" ], - "annotation": "Optional[Literal['CPP', 'PYTHON', 'auto']]", - "converter": "", + "capture_policy": "literal|none", "kind": "categorical", "path": "cache_transceiver_config.transceiver_runtime" }, @@ -217,15 +171,12 @@ "native", "rank_striped_read_ahead" ], - "annotation": "Literal['auto', 'native', 'rank_striped_read_ahead']", - "converter": "", + "capture_policy": "literal", "kind": "categorical", "path": "checkpoint_io_policy" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "int", "kind": "value", "path": "context_parallel_size" }, @@ -235,64 +186,47 @@ "RING", "HELIX" ], - "annotation": "", - "converter": "", + "capture_policy": "enum[CpType]", "kind": "categorical", "path": "cp_config.cp_type" }, { - "allowed_values": [], - "annotation": "Optional[int]", - "converter": "", + "capture_policy": "int|none", "kind": "value", "path": "cp_config.fifo_version" }, { - "allowed_values": [], - "annotation": "Optional[int]", - "converter": "", + "capture_policy": "int|none", "kind": "value", "path": "cp_config.tokens_per_block" }, { - "allowed_values": [], - "annotation": "Optional[bool]", - "converter": "", + "capture_policy": "bool|none", "kind": "value", "path": "cp_config.use_nccl_for_alltoall" }, { - "allowed_values": [], - "annotation": "Optional[List[int]]", - "converter": "", + "capture_policy": "list[int]|none", "kind": "value", "path": "cuda_graph_config.batch_sizes" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "cuda_graph_config.enable_padding" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "int", "kind": "value", "path": "cuda_graph_config.max_batch_size" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "int", "kind": "value", "path": "cuda_graph_config.max_num_token" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "int", "kind": "value", "path": "cuda_graph_config.max_seq_len" }, @@ -301,36 +235,27 @@ "decode", "encode" ], - "annotation": "Literal['decode']", - "converter": "", + "capture_policy": "literal", "kind": "categorical", "path": "cuda_graph_config.mode" }, { - "allowed_values": [], - "annotation": "Optional[List[Annotated[int, Gt(gt=0)]]]", - "converter": "", + "capture_policy": "list[int]|none", "kind": "value", "path": "cuda_graph_config.num_tokens" }, { - "allowed_values": [], - "annotation": "Optional[List[Annotated[int, Gt(gt=0)]]]", - "converter": "", + "capture_policy": "list[int]|none", "kind": "value", "path": "cuda_graph_config.seq_lens" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "disable_mm_encoder" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "disable_overlap_scheduler" }, @@ -341,204 +266,147 @@ "bfloat16", "float32" ], - "annotation": "", - "converter": "allowlist", + "capture_policy": "allowlist", "kind": "categorical", "path": "dtype" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "dwdp_config.contention_opt" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "int", "kind": "value", "path": "dwdp_config.dwdp_size" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "int", "kind": "value", "path": "dwdp_config.num_experts_per_worker" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "int", "kind": "value", "path": "dwdp_config.num_groups" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "int", "kind": "value", "path": "dwdp_config.num_prefetch_experts" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "enable_attention_dp" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "enable_autotuner" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "enable_chunked_prefill" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "enable_early_first_token_response" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "enable_encoder_decoder_mixed_cuda_graph" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "enable_energy_metrics" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "enable_in_graph_sampling" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "enable_iter_perf_stats" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "enable_iter_req_stats" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "enable_layerwise_nvtx_marker" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "enable_lm_head_tp_in_adp" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "enable_lora" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "enable_low_latency_host_dispatch" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "enable_min_latency" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "enable_mla_skip_correction" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "enable_resource_governor" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "enable_speculative_beam_history_d2h" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "encode_only" }, { - "allowed_values": [], - "annotation": "Optional[List[int]]", - "converter": "", + "capture_policy": "list[int]|none", "kind": "value", "path": "encoder_cuda_graph_config.batch_sizes" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "encoder_cuda_graph_config.enable_padding" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "int", "kind": "value", "path": "encoder_cuda_graph_config.max_batch_size" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "int", "kind": "value", "path": "encoder_cuda_graph_config.max_num_token" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "int", "kind": "value", "path": "encoder_cuda_graph_config.max_seq_len" }, @@ -546,57 +414,42 @@ "allowed_values": [ "encode" ], - "annotation": "Literal['encode']", - "converter": "", + "capture_policy": "literal", "kind": "categorical", "path": "encoder_cuda_graph_config.mode" }, { - "allowed_values": [], - "annotation": "Optional[List[Annotated[int, Gt(gt=0)]]]", - "converter": "", + "capture_policy": "list[int]|none", "kind": "value", "path": "encoder_cuda_graph_config.num_tokens" }, { - "allowed_values": [], - "annotation": "Optional[List[Annotated[int, Gt(gt=0)]]]", - "converter": "", + "capture_policy": "list[int]|none", "kind": "value", "path": "encoder_cuda_graph_config.seq_lens" }, { - "allowed_values": [], - "annotation": "Optional[int]", - "converter": "", + "capture_policy": "int|none", "kind": "value", "path": "encoder_max_batch_size" }, { - "allowed_values": [], - "annotation": "Optional[int]", - "converter": "", + "capture_policy": "int|none", "kind": "value", "path": "encoder_max_num_tokens" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "force_dynamic_quantization" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "int", "kind": "value", "path": "garbage_collection_gen0_threshold" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "gather_generation_logits" }, @@ -605,8 +458,7 @@ "auto", "trtllm" ], - "annotation": "Literal['auto', 'trtllm']", - "converter": "", + "capture_policy": "literal", "kind": "categorical", "path": "generation_config" }, @@ -616,15 +468,12 @@ "rw", "ro" ], - "annotation": "Literal['auto', 'rw', 'ro']", - "converter": "", + "capture_policy": "literal", "kind": "categorical", "path": "gms_config.mode" }, { - "allowed_values": [], - "annotation": "Optional[int]", - "converter": "", + "capture_policy": "int|none", "kind": "value", "path": "gpus_per_node" }, @@ -633,15 +482,12 @@ "xgrammar", "llguidance" ], - "annotation": "Optional[Literal['xgrammar', 'llguidance']]", - "converter": "", + "capture_policy": "literal|none", "kind": "categorical", "path": "guided_decoding_backend" }, { - "allowed_values": [], - "annotation": "Optional[int]", - "converter": "", + "capture_policy": "int|none", "kind": "value", "path": "iter_stats_max_iterations" }, @@ -650,22 +496,17 @@ "quantization_for_cold_page", "triattention" ], - "annotation": "Literal['triattention']", - "converter": "allowlist", + "capture_policy": "literal", "kind": "categorical", "path": "kv_cache_compression_config.algorithm" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "int", "kind": "value", "path": "kv_cache_compression_config.beta" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "int", "kind": "value", "path": "kv_cache_compression_config.budget" }, @@ -675,15 +516,12 @@ "per_head", "per_layer_perhead" ], - "annotation": "Literal['union', 'per_head', 'per_layer_perhead']", - "converter": "", + "capture_policy": "literal", "kind": "categorical", "path": "kv_cache_compression_config.eviction_mode" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "kv_cache_compression_config.normalize_scores" }, @@ -691,29 +529,22 @@ "allowed_values": [ "nvfp4" ], - "annotation": "Literal['nvfp4']", - "converter": "", + "capture_policy": "literal", "kind": "categorical", "path": "kv_cache_compression_config.quant" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "int", "kind": "value", "path": "kv_cache_config.attention_dp_events_gather_period_ms" }, { - "allowed_values": [], - "annotation": "Optional[Annotated[int, Gt(gt=0)]]", - "converter": "", + "capture_policy": "int|none", "kind": "value", "path": "kv_cache_config.avg_seq_len" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "int", "kind": "value", "path": "kv_cache_config.block_reuse_config.max_num_turns" }, @@ -723,36 +554,27 @@ "per_request", "per_conversation" ], - "annotation": "Literal['all_reusable', 'per_request', 'per_conversation']", - "converter": "", + "capture_policy": "literal", "kind": "categorical", "path": "kv_cache_config.block_reuse_config.policy" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "kv_cache_config.copy_on_partial_reuse" }, { - "allowed_values": [], - "annotation": "Optional[float]", - "converter": "", + "capture_policy": "float|none", "kind": "value", "path": "kv_cache_config.cross_kv_cache_fraction" }, { - "allowed_values": [], - "annotation": "Optional[Annotated[int, Ge(ge=0)]]", - "converter": "", + "capture_policy": "int|none", "kind": "value", "path": "kv_cache_config.disk_cache_size" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "int", "kind": "value", "path": "kv_cache_config.disk_prefetch_num_reqs" }, @@ -766,71 +588,52 @@ "fp8_ds_mla", "nvfp4" ], - "annotation": "", - "converter": "allowlist", + "capture_policy": "allowlist", "kind": "categorical", "path": "kv_cache_config.dtype" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "kv_cache_config.enable_block_reuse" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "kv_cache_config.enable_kv_pool_rebalance" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "kv_cache_config.enable_partial_reuse" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "kv_cache_config.enable_swa_scratch_reuse" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "int", "kind": "value", "path": "kv_cache_config.event_buffer_max_size" }, { - "allowed_values": [], - "annotation": "Optional[int]", - "converter": "", + "capture_policy": "int|none", "kind": "value", "path": "kv_cache_config.fp8_context_mla_kv_len_cap" }, { - "allowed_values": [], - "annotation": "Optional[float]", - "converter": "", + "capture_policy": "float|none", "kind": "value", "path": "kv_cache_config.free_gpu_memory_fraction" }, { - "allowed_values": [], - "annotation": "Optional[int]", - "converter": "", + "capture_policy": "int|none", "kind": "value", "path": "kv_cache_config.host_cache_size" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "int", "kind": "value", "path": "kv_cache_config.iteration_stats_interval" }, @@ -841,36 +644,27 @@ "v2_sha256", "v2_sha256_64" ], - "annotation": "Literal['auto', 'v1_block_key', 'v2_sha256', 'v2_sha256_64']", - "converter": "", + "capture_policy": "literal", "kind": "categorical", "path": "kv_cache_config.kv_cache_event_hash_algo" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "int", "kind": "value", "path": "kv_cache_config.kv_events_config.buffer_steps" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "kv_cache_config.kv_events_config.enable_kv_cache_events" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "int", "kind": "value", "path": "kv_cache_config.kv_events_config.hwm" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "int", "kind": "value", "path": "kv_cache_config.kv_events_config.max_queue_size" }, @@ -879,8 +673,7 @@ "null", "zmq" ], - "annotation": "Optional[Literal['null', 'zmq']]", - "converter": "", + "capture_policy": "literal|none", "kind": "categorical", "path": "kv_cache_config.kv_events_config.publisher" }, @@ -891,92 +684,67 @@ "bfloat16", "float32" ], - "annotation": "Literal['auto', 'float16', 'bfloat16', 'float32']", - "converter": "", + "capture_policy": "literal", "kind": "categorical", "path": "kv_cache_config.mamba_ssm_cache_dtype" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "int", "kind": "value", "path": "kv_cache_config.mamba_ssm_philox_rounds" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "kv_cache_config.mamba_ssm_stochastic_rounding" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "kv_cache_config.mamba_state_config.enable_branch_snapshot" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "int", "kind": "value", "path": "kv_cache_config.mamba_state_config.periodic_snapshot_interval" }, { - "allowed_values": [], - "annotation": "Optional[List[int]]", - "converter": "", + "capture_policy": "list[int]|none", "kind": "value", "path": "kv_cache_config.max_attention_window" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "int", "kind": "value", "path": "kv_cache_config.max_gpu_total_bytes" }, { - "allowed_values": [], - "annotation": "Optional[int]", - "converter": "", + "capture_policy": "int|none", "kind": "value", "path": "kv_cache_config.max_tokens" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "float", "kind": "value", "path": "kv_cache_config.max_util_for_resume" }, { - "allowed_values": [], - "annotation": "Optional[List[float]]", - "converter": "", + "capture_policy": "list[float]|none", "kind": "value", "path": "kv_cache_config.pool_ratio" }, { - "allowed_values": [], - "annotation": "Optional[int]", - "converter": "", + "capture_policy": "int|none", "kind": "value", "path": "kv_cache_config.secondary_offload_min_priority" }, { - "allowed_values": [], - "annotation": "Optional[int]", - "converter": "", + "capture_policy": "int|none", "kind": "value", "path": "kv_cache_config.sink_token_length" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "int", "kind": "value", "path": "kv_cache_config.tokens_per_block" }, @@ -984,15 +752,12 @@ "allowed_values": [ "auto" ], - "annotation": "Union[bool, Literal['auto']]", - "converter": "", - "kind": "value", + "capture_policy": "bool|literal", + "kind": "categorical", "path": "kv_cache_config.use_kv_cache_manager_v2" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "kv_cache_config.use_uvm" }, @@ -1002,15 +767,12 @@ "lmcache-mp", "kvbm" ], - "annotation": "Optional[str]", - "converter": "allowlist", + "capture_policy": "allowlist|none", "kind": "categorical", "path": "kv_connector_config.connector" }, { - "allowed_values": [], - "annotation": "Optional[List[int]]", - "converter": "", + "capture_policy": "list[int]|none", "kind": "value", "path": "layer_wise_benchmarks_config.calibration_layer_indices" }, @@ -1020,8 +782,7 @@ "MARK", "COLLECT" ], - "annotation": "Literal['NONE', 'MARK', 'COLLECT']", - "converter": "", + "capture_policy": "literal", "kind": "categorical", "path": "layer_wise_benchmarks_config.calibration_mode" }, @@ -1030,17 +791,18 @@ "auto", "dummy", "vision_only", - "gms" + "gms", + "AUTO", + "DUMMY", + "VISION_ONLY", + "GMS" ], - "annotation": "Union[str, tensorrt_llm.llmapi.llm_args.LoadFormat]", - "converter": "allowlist", + "capture_policy": "allowlist|enum[LoadFormat]", "kind": "categorical", "path": "load_format" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "lora_config.cuda_graph_specialize_lora" }, @@ -1049,106 +811,77 @@ "hf", "nemo" ], - "annotation": "Literal['hf', 'nemo']", - "converter": "", + "capture_policy": "literal", "kind": "categorical", "path": "lora_config.lora_ckpt_source" }, { - "allowed_values": [], - "annotation": "Optional[int]", - "converter": "", + "capture_policy": "int|none", "kind": "value", "path": "lora_config.max_cpu_loras" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "int", "kind": "value", "path": "lora_config.max_lora_rank" }, { - "allowed_values": [], - "annotation": "Optional[int]", - "converter": "", + "capture_policy": "int|none", "kind": "value", "path": "lora_config.max_loras" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "lora_config.overlap_lora_and_base" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "lora_config.swap_gate_up_proj_lora_b_weight" }, { - "allowed_values": [], - "annotation": "Optional[int]", - "converter": "", + "capture_policy": "int|none", "kind": "value", "path": "max_batch_size" }, { - "allowed_values": [], - "annotation": "Optional[int]", - "converter": "", + "capture_policy": "int|none", "kind": "value", "path": "max_beam_width" }, { - "allowed_values": [], - "annotation": "Optional[int]", - "converter": "", + "capture_policy": "int|none", "kind": "value", "path": "max_input_len" }, { - "allowed_values": [], - "annotation": "Optional[int]", - "converter": "", + "capture_policy": "int|none", "kind": "value", "path": "max_num_tokens" }, { - "allowed_values": [], - "annotation": "Optional[int]", - "converter": "", + "capture_policy": "int|none", "kind": "value", "path": "max_seq_len" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "int", "kind": "value", "path": "max_stats_len" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "float", "kind": "value", "path": "mla_skip_correction_threshold" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "mm_encoder_only" }, { - "allowed_values": [], - "annotation": "Optional[int]", - "converter": "", + "capture_policy": "int|none", "kind": "value", "path": "moe_cluster_parallel_size" }, @@ -1166,50 +899,37 @@ "MEGAMOE_DEEPGEMM", "MEGAMOE_CUTEDSL" ], - "annotation": "Literal['AUTO', 'CUTLASS', 'CUTEDSL', 'TRTLLM', 'DEEPGEMM', 'DENSEGEMM', 'VANILLA', 'TRITON', 'MARLIN', 'MEGAMOE_DEEPGEMM', 'MEGAMOE_CUTEDSL']", - "converter": "", + "capture_policy": "literal", "kind": "categorical", "path": "moe_config.backend" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "moe_config.disable_finalize_fusion" }, { - "allowed_values": [], - "annotation": "Optional[int]", - "converter": "", + "capture_policy": "int|none", "kind": "value", "path": "moe_config.max_num_tokens" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "moe_config.use_low_precision_moe_combine" }, { - "allowed_values": [], - "annotation": "Optional[int]", - "converter": "", + "capture_policy": "int|none", "kind": "value", "path": "moe_expert_parallel_size" }, { - "allowed_values": [], - "annotation": "Optional[int]", - "converter": "", + "capture_policy": "int|none", "kind": "value", "path": "moe_tensor_parallel_size" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "int", "kind": "value", "path": "multimodal_config.encoder_cache_max_bytes" }, @@ -1219,22 +939,17 @@ "DEFAULT", "EAGER" ], - "annotation": "", - "converter": "", + "capture_policy": "enum[MultimodalEncoderSchedulingPolicy]", "kind": "categorical", "path": "multimodal_config.encoder_scheduling_policy" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "int", "kind": "value", "path": "multimodal_config.encoder_side_stream_max_ahead" }, { - "allowed_values": [], - "annotation": "Optional[float]", - "converter": "", + "capture_policy": "float|none", "kind": "value", "path": "multimodal_config.video_pruning_rate" }, @@ -1242,29 +957,22 @@ "allowed_values": [ "per_module" ], - "annotation": "", - "converter": "allowlist", + "capture_policy": "allowlist", "kind": "categorical", "path": "mx_config.preshard_strategy" }, { - "allowed_values": [], - "annotation": "Optional[Annotated[int, Ge(ge=0)]]", - "converter": "", + "capture_policy": "int|none", "kind": "value", "path": "mx_config.server_query_timeout_s" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "int", "kind": "value", "path": "num_postprocess_workers" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "int", "kind": "value", "path": "num_serve_frontends" }, @@ -1276,9 +984,8 @@ "cuda_core", "marlin" ], - "annotation": "List[Literal['cutlass', 'cublaslt', 'cutedsl', 'cuda_core', 'marlin']]", - "converter": "", - "kind": "value", + "capture_policy": "list[literal]", + "kind": "categorical", "path": "nvfp4_gemm_config.allowed_backends" }, { @@ -1286,113 +993,82 @@ "rpc", "ray" ], - "annotation": "Optional[Literal['rpc', 'ray']]", - "converter": "", + "capture_policy": "literal|none", "kind": "categorical", "path": "orchestrator_type" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "float", "kind": "value", "path": "peft_cache_config.device_cache_percent" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "int", "kind": "value", "path": "peft_cache_config.host_cache_size" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "int", "kind": "value", "path": "peft_cache_config.max_adapter_size" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "int", "kind": "value", "path": "peft_cache_config.max_pages_per_block_device" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "int", "kind": "value", "path": "peft_cache_config.max_pages_per_block_host" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "int", "kind": "value", "path": "peft_cache_config.num_copy_streams" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "int", "kind": "value", "path": "peft_cache_config.num_device_module_layer" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "int", "kind": "value", "path": "peft_cache_config.num_ensure_workers" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "int", "kind": "value", "path": "peft_cache_config.num_host_module_layer" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "int", "kind": "value", "path": "peft_cache_config.num_put_workers" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "int", "kind": "value", "path": "peft_cache_config.optimal_adapter_size" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "int", "kind": "value", "path": "perf_metrics_max_requests" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "int", "kind": "value", "path": "pipeline_parallel_size" }, { - "allowed_values": [], - "annotation": "Optional[List[int]]", - "converter": "", + "capture_policy": "list[int]|none", "kind": "value", "path": "pp_partition" }, { - "allowed_values": [], - "annotation": "Optional[List[int]]", - "converter": "", + "capture_policy": "list[int]|none", "kind": "value", "path": "prefill_capture_num_tokens" }, @@ -1402,85 +1078,62 @@ "piecewise", "breakable" ], - "annotation": "", - "converter": "allowlist", + "capture_policy": "enum[PrefillCudaGraphBackend]", "kind": "categorical", "path": "prefill_cuda_graph_backend" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "print_iter_log" }, { - "allowed_values": [], - "annotation": "Optional[List[float]]", - "converter": "", + "capture_policy": "list[float]|none", "kind": "value", "path": "prometheus_metrics_config.e2e_request_latency_buckets" }, { - "allowed_values": [], - "annotation": "Optional[List[float]]", - "converter": "", + "capture_policy": "list[float]|none", "kind": "value", "path": "prometheus_metrics_config.request_decode_time_buckets" }, { - "allowed_values": [], - "annotation": "Optional[List[float]]", - "converter": "", + "capture_policy": "list[float]|none", "kind": "value", "path": "prometheus_metrics_config.request_inference_time_buckets" }, { - "allowed_values": [], - "annotation": "Optional[List[float]]", - "converter": "", + "capture_policy": "list[float]|none", "kind": "value", "path": "prometheus_metrics_config.request_prefill_time_buckets" }, { - "allowed_values": [], - "annotation": "Optional[List[float]]", - "converter": "", + "capture_policy": "list[float]|none", "kind": "value", "path": "prometheus_metrics_config.request_queue_time_buckets" }, { - "allowed_values": [], - "annotation": "Optional[List[float]]", - "converter": "", + "capture_policy": "list[float]|none", "kind": "value", "path": "prometheus_metrics_config.time_per_output_token_buckets" }, { - "allowed_values": [], - "annotation": "Optional[List[float]]", - "converter": "", + "capture_policy": "list[float]|none", "kind": "value", "path": "prometheus_metrics_config.time_to_first_token_buckets" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "ray_placement_config.defer_workers_init" }, { - "allowed_values": [], - "annotation": "Optional[float]", - "converter": "", + "capture_policy": "float|none", "kind": "value", "path": "ray_placement_config.per_worker_gpu_share" }, { - "allowed_values": [], - "annotation": "Optional[List[List[int]]]", - "converter": "", + "capture_policy": "list[list[int]]|none", "kind": "value", "path": "ray_placement_config.placement_bundle_indices" }, @@ -1499,22 +1152,17 @@ "kimi_k2", "kimi_k25" ], - "annotation": "Optional[str]", - "converter": "allowlist", + "capture_policy": "allowlist|none", "kind": "categorical", "path": "reasoning_parser" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "int", "kind": "value", "path": "reorder_policy_config.policy_args.agent_inflight_seq_num" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "float", "kind": "value", "path": "reorder_policy_config.policy_args.agent_percentage" }, @@ -1522,29 +1170,22 @@ "allowed_values": [ "AgentTree" ], - "annotation": "Optional[Literal['AgentTree']]", - "converter": "", + "capture_policy": "literal|none", "kind": "categorical", "path": "reorder_policy_config.policy_name" }, { - "allowed_values": [], - "annotation": "Optional[int]", - "converter": "", + "capture_policy": "int|none", "kind": "value", "path": "request_stats_max_iterations" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "return_perf_metrics" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "sampler_force_async_worker" }, @@ -1554,8 +1195,7 @@ "GUARANTEED_NO_EVICT", "STATIC_BATCH" ], - "annotation": "", - "converter": "", + "capture_policy": "enum[CapacitySchedulerPolicy]", "kind": "categorical", "path": "scheduler_config.capacity_scheduler_policy" }, @@ -1565,43 +1205,32 @@ "EQUAL_PROGRESS", "FORCE_CHUNK" ], - "annotation": "Optional[tensorrt_llm.llmapi.llm_args.ContextChunkingPolicy]", - "converter": "", + "capture_policy": "enum[ContextChunkingPolicy]|none", "kind": "categorical", "path": "scheduler_config.context_chunking_policy" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "int", "kind": "value", "path": "scheduler_config.dynamic_batch_config.dynamic_batch_moving_average_window" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "scheduler_config.dynamic_batch_config.enable_batch_size_tuning" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "scheduler_config.dynamic_batch_config.enable_max_num_tokens_tuning" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "scheduler_config.enable_prefix_aware_scheduling" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "scheduler_config.use_python_scheduler" }, @@ -1610,15 +1239,12 @@ "fcfs", "priority" ], - "annotation": "", - "converter": "", + "capture_policy": "enum[WaitingQueuePolicy]", "kind": "categorical", "path": "scheduler_config.waiting_queue_policy" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "skip_tokenizer_init" }, @@ -1631,22 +1257,17 @@ "rocket", "skip_softmax" ], - "annotation": "Literal['dsa']", - "converter": "", + "capture_policy": "literal", "kind": "categorical", "path": "sparse_attention_config.algorithm" }, { - "allowed_values": [], - "annotation": "List[int]", - "converter": "", + "capture_policy": "list[int]", "kind": "value", "path": "sparse_attention_config.compress_ratios" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "sparse_attention_config.enable_heuristic_topk" }, @@ -1655,36 +1276,27 @@ "triton", "msa" ], - "annotation": "Literal['triton', 'msa']", - "converter": "", + "capture_policy": "literal", "kind": "categorical", "path": "sparse_attention_config.implementation" }, { - "allowed_values": [], - "annotation": "Optional[int]", - "converter": "", + "capture_policy": "int|none", "kind": "value", "path": "sparse_attention_config.index_head_dim" }, { - "allowed_values": [], - "annotation": "Optional[int]", - "converter": "", + "capture_policy": "int|none", "kind": "value", "path": "sparse_attention_config.index_n_heads" }, { - "allowed_values": [], - "annotation": "Optional[bool]", - "converter": "", + "capture_policy": "bool|none", "kind": "value", "path": "sparse_attention_config.index_share_for_mtp_iteration" }, { - "allowed_values": [], - "annotation": "Optional[int]", - "converter": "", + "capture_policy": "int|none", "kind": "value", "path": "sparse_attention_config.index_topk" }, @@ -1693,8 +1305,7 @@ "fp8", "fp4" ], - "annotation": "Literal['fp8', 'fp4']", - "converter": "", + "capture_policy": "literal", "kind": "categorical", "path": "sparse_attention_config.indexer_k_dtype" }, @@ -1703,29 +1314,22 @@ "bf16", "fp8" ], - "annotation": "Literal['bf16', 'fp8']", - "converter": "", + "capture_policy": "literal", "kind": "categorical", "path": "sparse_attention_config.indexer_kv_dtype" }, { - "allowed_values": [], - "annotation": "Optional[int]", - "converter": "", + "capture_policy": "int|none", "kind": "value", "path": "sparse_attention_config.indexer_max_chunk_size" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "sparse_attention_config.indexer_rope_interleave" }, { - "allowed_values": [], - "annotation": "Optional[int]", - "converter": "", + "capture_policy": "int|none", "kind": "value", "path": "sparse_attention_config.kernel_size" }, @@ -1734,99 +1338,72 @@ "bfloat16", "float8_e5m2" ], - "annotation": "Optional[str]", - "converter": "allowlist", + "capture_policy": "allowlist|none", "kind": "categorical", "path": "sparse_attention_config.kt_cache_dtype" }, { - "allowed_values": [], - "annotation": "Optional[int]", - "converter": "", + "capture_policy": "int|none", "kind": "value", "path": "sparse_attention_config.num_attention_heads" }, { - "allowed_values": [], - "annotation": "Optional[int]", - "converter": "", + "capture_policy": "int|none", "kind": "value", "path": "sparse_attention_config.num_key_value_heads" }, { - "allowed_values": [], - "annotation": "Optional[int]", - "converter": "", + "capture_policy": "int|none", "kind": "value", "path": "sparse_attention_config.page_size" }, { - "allowed_values": [], - "annotation": "Optional[int]", - "converter": "", + "capture_policy": "int|none", "kind": "value", "path": "sparse_attention_config.prompt_budget" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "int", "kind": "value", "path": "sparse_attention_config.q_split_threshold" }, { - "allowed_values": [], - "annotation": "Optional[int]", - "converter": "", + "capture_policy": "int|none", "kind": "value", "path": "sparse_attention_config.seq_len_threshold" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "sparse_attention_config.skip_indexer_for_short_seqs" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "int", "kind": "value", "path": "sparse_attention_config.sparse_block_size" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "sparse_attention_config.sparse_disable_index_value" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "int", "kind": "value", "path": "sparse_attention_config.sparse_index_dim" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "int", "kind": "value", "path": "sparse_attention_config.sparse_init_blocks" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "int", "kind": "value", "path": "sparse_attention_config.sparse_local_blocks" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "int", "kind": "value", "path": "sparse_attention_config.sparse_num_index_heads" }, @@ -1834,78 +1411,67 @@ "allowed_values": [ "max" ], - "annotation": "Literal['max']", - "converter": "", + "capture_policy": "literal", "kind": "categorical", "path": "sparse_attention_config.sparse_score_type" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "int", "kind": "value", "path": "sparse_attention_config.sparse_topk_blocks" }, { - "allowed_values": [], - "annotation": "Optional[int]", - "converter": "", + "capture_policy": "float|none", + "kind": "value", + "path": "sparse_attention_config.target_sparsity" + }, + { + "capture_policy": "float|none", + "kind": "value", + "path": "sparse_attention_config.threshold_scale_factor" + }, + { + "capture_policy": "int|none", "kind": "value", "path": "sparse_attention_config.topk" }, { - "allowed_values": [], - "annotation": "Union[int, float, NoneType]", - "converter": "", + "capture_policy": "float|int|none", "kind": "value", "path": "sparse_attention_config.topr" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "sparse_attention_config.use_cute_dsl_paged_mqa_logits" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "sparse_attention_config.use_cute_dsl_topk" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "sparse_attention_config.use_gvr_emission" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "sparse_attention_config.use_self_sampling_topk" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "int|none", "kind": "value", "path": "sparse_attention_config.window_size" }, { - "allowed_values": [], - "annotation": "Optional[float]", - "converter": "", + "capture_policy": "float|none", "kind": "value", "path": "speculative_config.acceptance_rate_threshold" }, { - "allowed_values": [], - "annotation": "Optional[Annotated[int, Ge(ge=0)]]", - "converter": "", + "capture_policy": "int|none", "kind": "value", "path": "speculative_config.acceptance_rate_window_size" }, @@ -1916,15 +1482,12 @@ "no_topp", "no_topk_no_topp" ], - "annotation": "", - "converter": "", + "capture_policy": "enum[AdvancedSamplingMode]", "kind": "categorical", "path": "speculative_config.advanced_sampling_mode" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "speculative_config.allow_advanced_sampling" }, @@ -1934,22 +1497,17 @@ "TRTLLM", "FA4" ], - "annotation": "Literal['VANILLA', 'TRTLLM', 'FA4']", - "converter": "", + "capture_policy": "literal", "kind": "categorical", "path": "speculative_config.attention_backend" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "int", "kind": "value", "path": "speculative_config.begin_thinking_phase_token" }, { - "allowed_values": [], - "annotation": "Optional[Annotated[int, Gt(gt=0)]]", - "converter": "", + "capture_policy": "int|none", "kind": "value", "path": "speculative_config.block_size" }, @@ -1968,22 +1526,17 @@ "SaveState", "User_Provided" ], - "annotation": "Literal['AUTO']", - "converter": "", + "capture_policy": "literal", "kind": "categorical", "path": "speculative_config.decoding_type" }, { - "allowed_values": [], - "annotation": "Optional[int]", - "converter": "", + "capture_policy": "int|none", "kind": "value", "path": "speculative_config.dynamic_tree_max_topK" }, { - "allowed_values": [], - "annotation": "Optional[Set[int]]", - "converter": "", + "capture_policy": "none|set[int]", "kind": "value", "path": "speculative_config.eagle3_layers_to_capture" }, @@ -1992,64 +1545,47 @@ "llama3", "mistral_large3" ], - "annotation": "Literal['llama3', 'mistral_large3']", - "converter": "", + "capture_policy": "literal", "kind": "categorical", "path": "speculative_config.eagle3_model_arch" }, { - "allowed_values": [], - "annotation": "Optional[List[List[int]]]", - "converter": "", + "capture_policy": "list[list[int]]|none", "kind": "value", "path": "speculative_config.eagle_choices" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "speculative_config.enable_global_pool" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "speculative_config.enable_penalty" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "int", "kind": "value", "path": "speculative_config.end_thinking_phase_token" }, { - "allowed_values": [], - "annotation": "Optional[Annotated[int, Gt(gt=0)]]", - "converter": "", + "capture_policy": "int|none", "kind": "value", "path": "speculative_config.global_pool_size" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "speculative_config.is_keep_all" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "speculative_config.is_public_pool" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "speculative_config.is_use_oldest" }, @@ -2059,155 +1595,112 @@ "gated", "rnn" ], - "annotation": "Optional[Literal['vanilla', 'gated', 'rnn']]", - "converter": "", + "capture_policy": "literal|none", "kind": "categorical", "path": "speculative_config.markov_head_type" }, { - "allowed_values": [], - "annotation": "Optional[int]", - "converter": "", + "capture_policy": "int|none", "kind": "value", "path": "speculative_config.markov_rank" }, { - "allowed_values": [], - "annotation": "Optional[int]", - "converter": "", + "capture_policy": "int|none", "kind": "value", "path": "speculative_config.mask_token_id" }, { - "allowed_values": [], - "annotation": "Optional[Annotated[int, Gt(gt=0)]]", - "converter": "", + "capture_policy": "int|none", "kind": "value", "path": "speculative_config.max_concurrency" }, { - "allowed_values": [], - "annotation": "Optional[Annotated[int, Ge(ge=0)]]", - "converter": "", + "capture_policy": "int|none", "kind": "value", "path": "speculative_config.max_draft_len" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "int", "kind": "value", "path": "speculative_config.max_matching_ngram_size" }, { - "allowed_values": [], - "annotation": "Optional[int]", - "converter": "", + "capture_policy": "int|none", "kind": "value", "path": "speculative_config.max_non_leaves_per_layer" }, { - "allowed_values": [], - "annotation": "Optional[int]", - "converter": "", + "capture_policy": "int|none", "kind": "value", "path": "speculative_config.max_total_draft_tokens" }, { - "allowed_values": [], - "annotation": "Optional[int]", - "converter": "", + "capture_policy": "int|none", "kind": "value", "path": "speculative_config.num_eagle_layers" }, { - "allowed_values": [], - "annotation": "Optional[int]", - "converter": "", + "capture_policy": "int|none", "kind": "value", "path": "speculative_config.num_nextn_predict_layers" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "float", "kind": "value", "path": "speculative_config.relaxed_delta" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "int", "kind": "value", "path": "speculative_config.relaxed_topk" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "speculative_config.sa_config.enable_global_pool" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "int", "kind": "value", "path": "speculative_config.sa_config.threshold" }, { - "allowed_values": [], - "annotation": "Optional[List[int]]", - "converter": "", + "capture_policy": "list[int]|none", "kind": "value", "path": "speculative_config.target_layer_ids" }, { - "allowed_values": [], - "annotation": "Optional[bool]", - "converter": "", + "capture_policy": "bool|none", "kind": "value", "path": "speculative_config.use_dynamic_tree" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "speculative_config.use_mtp_vanilla" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "speculative_config.use_rejection_sampling" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "speculative_config.use_relaxed_acceptance_for_thinking" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "int", "kind": "value", "path": "speculative_config.write_interval" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "int", "kind": "value", "path": "stream_interval" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "telemetry_config.disabled" }, @@ -2220,15 +1713,12 @@ "cli_eval", "disaggregated" ], - "annotation": "", - "converter": "", + "capture_policy": "enum[UsageContext]", "kind": "categorical", "path": "telemetry_config.usage_context" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "int", "kind": "value", "path": "tensor_parallel_size" }, @@ -2237,92 +1727,67 @@ "auto", "slow" ], - "annotation": "Literal['auto', 'slow']", - "converter": "", + "capture_policy": "literal", "kind": "categorical", "path": "tokenizer_mode" }, { - "allowed_values": [], - "annotation": "Optional[List[Annotated[int, Gt(gt=0)]]]", - "converter": "", + "capture_policy": "list[int]|none", "kind": "value", "path": "torch_compile_config.capture_num_tokens" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "torch_compile_config.enable_fullgraph" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "torch_compile_config.enable_inductor" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "torch_compile_config.enable_piecewise_cuda_graph" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "torch_compile_config.enable_userbuffers" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "int", "kind": "value", "path": "torch_compile_config.max_num_streams" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "trust_remote_code" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "use_cute_dsl_bf16_bmm" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "use_cute_dsl_bf16_gemm" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "use_cute_dsl_blockscaling_bmm" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "use_cute_dsl_blockscaling_mm" }, { - "allowed_values": [], - "annotation": "", - "converter": "", + "capture_policy": "bool", "kind": "value", "path": "use_fine_grained_sync" } diff --git a/tensorrt_llm/usage/llmapi_config.py b/tensorrt_llm/usage/llmapi_config.py index 25c661af8d61..a8124b5bc8d5 100644 --- a/tensorrt_llm/usage/llmapi_config.py +++ b/tensorrt_llm/usage/llmapi_config.py @@ -31,7 +31,7 @@ from tensorrt_llm.usage.config import TelemetryField CAPTURE_VERSION = "2" -FIELD_POLICY_VERSION = "2" +FIELD_POLICY_VERSION = "3" API_CONTRACT_VERSION = "0.2.0" CAPTURE_SOURCE = "effective_validated_llm_args" @@ -43,7 +43,6 @@ _TELEMETRY_EXTRA_KEY = "telemetry" _TRTLLM_JSON_SCHEMA_EXTRA_ATTR = "_trtllm_json_schema_extra" -_APPROVED_CONVERTERS = {"allowlist"} # Per-sequence cap, applied recursively so each inner list of a nested # List[List[int]] is bounded independently. 256 sits above the longest realistic @@ -72,10 +71,6 @@ def _is_pydantic_model(value: Any) -> bool: return isinstance(value, BaseModel) -def _none_type() -> type[None]: - return type(None) - - def _unwrap_annotated(annotation: Any) -> Any: while get_origin(annotation) is Annotated: annotation = get_args(annotation)[0] @@ -86,32 +81,16 @@ def _is_union(annotation: Any) -> bool: return get_origin(annotation) in {Union, types.UnionType} -def _unwrap_optional(annotation: Any) -> Any: - annotation = _unwrap_annotated(annotation) - if not _is_union(annotation): - return annotation - branches = [arg for arg in get_args(annotation) if arg is not _none_type()] - if len(branches) == 1: - return branches[0] - return annotation - - def _is_literal(annotation: Any) -> bool: return get_origin(annotation) is Literal def _is_enum_annotation(annotation: Any) -> bool: - try: - return isinstance(annotation, type) and issubclass(annotation, Enum) - except TypeError: - return False + return isinstance(annotation, type) and issubclass(annotation, Enum) def _is_path_annotation(annotation: Any) -> bool: - try: - return isinstance(annotation, type) and issubclass(annotation, Path) - except TypeError: - return False + return isinstance(annotation, type) and issubclass(annotation, Path) def _is_callable_annotation(annotation: Any) -> bool: @@ -119,51 +98,13 @@ def _is_callable_annotation(annotation: Any) -> bool: return annotation is Callable or origin is Callable -def _is_safe_annotation_branch(annotation: Any) -> bool: - annotation = _unwrap_annotated(annotation) - origin = get_origin(annotation) - - if annotation is Any: - return False - if annotation is _none_type(): - return True - if _is_literal(annotation): - return True - if _is_enum_annotation(annotation): - return True - if annotation in {bool, int, float}: - return True - if annotation is str or annotation is object: - return False - if _is_path_annotation(annotation) or _is_callable_annotation(annotation): - return False - if _is_union(annotation): - return all(_is_safe_annotation_branch(arg) for arg in get_args(annotation)) - if origin in {list, tuple, set}: - args = get_args(annotation) - return bool(args) and all(_is_safe_annotation_branch(arg) for arg in args) - if origin is dict: - return False - return False - - -def _union_needs_converter(annotation: Any) -> bool: - annotation = _unwrap_annotated(annotation) - if not _is_union(annotation): - return False - branches = [arg for arg in get_args(annotation) if arg is not _none_type()] - if len(branches) <= 1: - return False - return not all(_is_safe_annotation_branch(arg) for arg in branches) - - def _normalize_metadata(metadata: Any) -> dict[str, Any] | None: if metadata is None: return None if metadata is False: return {"exclude": True} if metadata is True: - return {"kind": "value"} + return {} if isinstance(metadata, TelemetryField): return metadata.as_json_schema_extra() if isinstance(metadata, dict): @@ -180,81 +121,230 @@ def _get_telemetry_metadata(field_info: Any) -> dict[str, Any] | None: return _normalize_metadata(json_schema_extra.get(_TELEMETRY_EXTRA_KEY)) -def _converter_is_approved(metadata: dict[str, Any]) -> bool: - return metadata.get("converter") in _APPROVED_CONVERTERS - - def _is_explicit_exclude(metadata: dict[str, Any] | None) -> bool: return bool(metadata) and metadata.get("exclude") is True -def _metadata_has_allowlist(metadata: dict[str, Any]) -> bool: - return _converter_is_approved(metadata) or metadata.get("allowed_values") is not None +@dataclass(frozen=True) +class _CapturePolicy: + """Compiled runtime sanitizer policy for one annotated field.""" + policy_type: str + runtime_type: Any = None + allowed_values: tuple[Any, ...] = () + branches: tuple["_CapturePolicy", ...] = () -def derive_kind(annotation: Any, metadata: dict[str, Any]) -> str: - """Derive a telemetry field's kind from its annotation and metadata. - - Categorical iff the Optional-unwrapped annotation is a Literal or Enum, or it - carries an allowlist; otherwise a plain value. Any registered kind is ignored - so kind stays annotation-driven. - """ - if _metadata_has_allowlist(metadata): - return "categorical" - unwrapped = _unwrap_optional(annotation) - if _is_literal(unwrapped) or _is_enum_annotation(unwrapped): - return "categorical" - return "value" +@dataclass(frozen=True) +class _PolicyVariant: + """Policy for a field as declared by one reachable model arm.""" -def _annotation_is_capture_safe(annotation: Any) -> bool: - return _is_safe_annotation_branch(annotation) + owner: type[BaseModel] + policy: _CapturePolicy @dataclass(frozen=True) class _ManifestEntry: path: str - annotation: Any # real type object — needed by the sanitizer kind: str - converter: str - allowed_values: tuple[str, ...] - metadata: dict[str, Any] # normalized metadata (allowlist) for the sanitizer + capture_types: tuple[str, ...] + allowed_values: tuple[Any, ...] + variants: tuple[_PolicyVariant, ...] -def _domain_values(annotation: Any, metadata: dict[str, Any]) -> list[str]: - """Compute a field's allowed-value domain. +def _is_safe_scalar(value: Any) -> bool: + if value is None: + return True + if type(value) not in {bool, int, float, str}: + return False + return not isinstance(value, float) or math.isfinite(value) - Explicit allowlist wins; else Literal args AND Enum members found anywhere in - the annotation tree, order-preserving + deduped. Mirrors what _sanitize_value - would emit for an Enum (str value, else name). - """ - allowlist = metadata.get("allowed_values") if metadata else None - if isinstance(allowlist, (list, tuple, set)): - return [str(v) for v in allowlist] - values: list[str] = [] +def _typed_equal(left: Any, right: Any) -> bool: + return type(left) is type(right) and left == right - def rec(ann: Any) -> None: - ann = _unwrap_annotated(ann) - if _is_literal(ann): - for v in get_args(ann): - if isinstance(v, (bool, int, float, str)) or v is None: - values.append(str(v)) - return - if _is_enum_annotation(ann): - for member in ann: - values.append(member.value if isinstance(member.value, str) else member.name) - return - if _is_union(ann) or get_origin(ann) in {list, tuple, set}: - for arg in get_args(ann): - rec(arg) - rec(annotation) - seen: list[str] = [] - for v in values: - if v not in seen: - seen.append(v) - return seen +def _dedupe_typed(values: list[Any]) -> tuple[Any, ...]: + deduped: list[Any] = [] + for value in values: + if not any(_typed_equal(value, seen) for seen in deduped): + deduped.append(value) + return tuple(deduped) + + +def _policy_key(policy: _CapturePolicy) -> tuple[Any, ...]: + """Typed structural identity for a compiled capture policy.""" + return ( + policy.policy_type, + policy.runtime_type, + tuple((type(value), value) for value in policy.allowed_values), + tuple(_policy_key(branch) for branch in policy.branches), + ) + + +def _combine_policies(policies: list[_CapturePolicy]) -> _CapturePolicy | None: + branches: list[_CapturePolicy] = [] + for policy in policies: + candidates = policy.branches if policy.policy_type == "union" else (policy,) + for candidate in candidates: + if not any(_policy_key(candidate) == _policy_key(branch) for branch in branches): + branches.append(candidate) + if not branches: + return None + if len(branches) == 1: + return branches[0] + return _CapturePolicy("union", branches=tuple(branches)) + + +def _compile_sequence_policy(annotation: Any) -> _CapturePolicy | None: + origin = get_origin(annotation) + args = get_args(annotation) + if origin in {list, set} and len(args) == 1: + element_annotations = args + elif origin is tuple and args: + element_annotations = args[:1] if len(args) == 2 and args[1] is Ellipsis else args + else: + return None + + element_policies: list[_CapturePolicy] = [] + for element_annotation in element_annotations: + policy = _compile_annotation_policy(element_annotation) + if policy is None: + return None + element_policies.append(policy) + if any( + _policy_key(policy) != _policy_key(element_policies[0]) for policy in element_policies[1:] + ): + return None + return _CapturePolicy("sequence", runtime_type=origin, branches=(element_policies[0],)) + + +def _compile_annotation_policy(annotation: Any) -> _CapturePolicy | None: + """Compile only capture-safe annotation branches into a sanitizer policy.""" + annotation = _unwrap_annotated(annotation) + + if annotation is Any or annotation in {str, object}: + return None + if annotation is types.NoneType: + return _CapturePolicy("none") + if _is_path_annotation(annotation) or _is_callable_annotation(annotation): + return None + if _is_literal(annotation): + values = list(get_args(annotation)) + if not values or not all(_is_safe_scalar(value) for value in values): + return None + return _CapturePolicy("literal", allowed_values=_dedupe_typed(values)) + if _is_enum_annotation(annotation): + return _CapturePolicy("enum", runtime_type=annotation) + if annotation in {bool, int, float}: + return _CapturePolicy(annotation.__name__) + if _is_union(annotation): + policies: list[_CapturePolicy] = [] + for branch in get_args(annotation): + policy = _compile_annotation_policy(branch) + if policy is not None: + policies.append(policy) + # Optional does not make an otherwise unsafe field capturable by itself. + if not any(policy.policy_type != "none" for policy in policies): + return None + return _combine_policies(policies) + + return _compile_sequence_policy(annotation) + + +def _annotation_allows_none(annotation: Any) -> bool: + annotation = _unwrap_annotated(annotation) + return annotation is types.NoneType or ( + _is_union(annotation) + and any(_unwrap_annotated(branch) is types.NoneType for branch in get_args(annotation)) + ) + + +def _annotation_accepts_allowed_value(annotation: Any, value: Any) -> bool: + """Whether an unsafe annotation branch can own an explicit scalar token.""" + annotation = _unwrap_annotated(annotation) + if annotation is Any or annotation is object: + return _is_safe_scalar(value) + if annotation is str: + return type(value) is str + if _is_path_annotation(annotation) or _is_callable_annotation(annotation): + return False + if _is_union(annotation): + return any( + _annotation_accepts_allowed_value(branch, value) for branch in get_args(annotation) + ) + return False + + +def _explicit_allowlist_policy(annotation: Any, metadata: dict[str, Any]) -> _CapturePolicy | None: + if "allowed_values" not in metadata: + return None + allowed_values = metadata["allowed_values"] + if not isinstance(allowed_values, (list, tuple, set)) or not allowed_values: + raise ValueError("telemetry allowed_values must be a non-empty sequence") + values = list(allowed_values) + if isinstance(allowed_values, set): + values.sort(key=lambda value: f"{type(value).__name__}:{_canonical_json(value)}") + if not all(_is_safe_scalar(value) for value in values): + raise ValueError("telemetry allowed_values must contain only finite JSON scalars") + if not all(_annotation_accepts_allowed_value(annotation, value) for value in values): + raise ValueError( + "telemetry allowed_values must belong to an unsafe scalar annotation branch" + ) + return _CapturePolicy("allowlist", allowed_values=_dedupe_typed(values)) + + +def _compile_field_policy(annotation: Any, metadata: dict[str, Any]) -> _CapturePolicy | None: + annotation_policy = _compile_annotation_policy(annotation) + policies: list[_CapturePolicy] = [] + allowlist_policy = _explicit_allowlist_policy(annotation, metadata) + if allowlist_policy is not None: + policies.append(allowlist_policy) + if annotation_policy is None and _annotation_allows_none(annotation): + policies.append(_CapturePolicy("none")) + if annotation_policy is not None: + policies.append(annotation_policy) + return _combine_policies(policies) + + +def _enum_output(member: Enum) -> str: + return member.value if isinstance(member.value, str) else member.name + + +def _policy_allowed_values(policy: _CapturePolicy) -> tuple[Any, ...]: + values: list[Any] = [] + if policy.policy_type in {"literal", "allowlist"}: + values.extend(policy.allowed_values) + elif policy.policy_type == "enum": + values.extend(_enum_output(member) for member in policy.runtime_type) + elif policy.policy_type in {"union", "sequence"}: + for branch in policy.branches: + values.extend(_policy_allowed_values(branch)) + return _dedupe_typed(values) + + +def _policy_signature(policy: _CapturePolicy) -> str: + if policy.policy_type == "sequence": + return f"{policy.runtime_type.__name__}[{_policy_signature(policy.branches[0])}]" + if policy.policy_type == "union": + signatures = sorted({_policy_signature(branch) for branch in policy.branches}) + return "|".join(signatures) + if policy.policy_type == "enum": + return f"enum[{policy.runtime_type.__name__}]" + return policy.policy_type + + +def _policy_capture_types(policy: _CapturePolicy) -> tuple[str, ...]: + branches = policy.branches if policy.policy_type == "union" else (policy,) + return tuple(sorted({_policy_signature(branch) for branch in branches})) + + +def _policy_kind(policy: _CapturePolicy) -> str: + if policy.policy_type in {"literal", "enum", "allowlist"}: + return "categorical" + if any(_policy_kind(branch) == "categorical" for branch in policy.branches): + return "categorical" + return "value" def _nested_models(annotation: Any) -> list[type]: @@ -282,32 +372,16 @@ def rec(ann: Any) -> None: return deduped -def _defining_class(cls: type, field_name: str) -> str: - for klass in cls.__mro__: - if field_name in getattr(klass, "__annotations__", {}): - return f"{klass.__name__}.{field_name}" - return f"{cls.__name__}.{field_name}" - - -def _field_is_selected(annotation: Any, metadata: dict[str, Any] | None) -> bool: - if _is_explicit_exclude(metadata): - return False - if _annotation_is_capture_safe(annotation): - return True - if metadata is not None and _converter_is_approved(metadata): - return True - return False - - def build_capture_manifest(model_cls: type[BaseModel]) -> list[_ManifestEntry]: """Walk real type objects and emit the complete capturable manifest. The single source of truth. Type-safe annotations auto-enroll; str/Any allowlist escape hatches opt in; telemetry=False opts out. Recurses into statically reachable nested BaseModels with a cycle guard. Collapses - duplicate keys (shared union-arm base fields): keeps the first by - (key, defining_class), unions allowed_values across arms, and FAILS if two - arms give a key a different kind. + duplicate keys shared by nested model union arms while retaining an + owner-specific compiled policy for every arm. Display domains and capture + types are merged, but runtime sanitization selects only the active arm's + policy. Conflicting kinds fail manifest construction. """ rows: list[dict[str, Any]] = [] @@ -318,72 +392,83 @@ def walk(cls: type, prefix: str, stack: tuple) -> None: key = f"{prefix}.{fname}" if prefix else fname ann = finfo.annotation meta = _get_telemetry_metadata(finfo) - if _field_is_selected(ann, meta): - normalized = meta if (meta and not _is_explicit_exclude(meta)) else {} + if _is_explicit_exclude(meta): + continue + policy = _compile_field_policy(ann, meta or {}) + if policy is not None: rows.append( { "key": key, - "defining": _defining_class(cls, fname), - "annotation": ann, - "kind": derive_kind(ann, normalized), - "converter": str(normalized.get("converter", "")), - "allowed": _domain_values(ann, normalized), - "metadata": normalized, + "owner": cls, + "policy": policy, + "kind": _policy_kind(policy), } ) - if not _is_explicit_exclude(meta): - for sub in _nested_models(ann): - walk(sub, key, (*stack, cls)) + for sub in _nested_models(ann): + walk(sub, key, (*stack, cls)) walk(model_cls, "", ()) - rows.sort(key=lambda r: (r["key"], r["defining"])) - first: dict[str, dict] = {} - union_allowed: dict[str, list[str]] = {} + rows.sort(key=lambda r: (r["key"], r["owner"].__module__, r["owner"].__qualname__)) + grouped: dict[str, dict[str, Any]] = {} for r in rows: - if r["key"] not in first: - first[r["key"]] = r - elif first[r["key"]]["kind"] != r["kind"]: + key = r["key"] + if key not in grouped: + grouped[key] = {"kind": r["kind"], "variants": []} + elif grouped[key]["kind"] != r["kind"]: raise ValueError( - f"telemetry manifest: key '{r['key']}' has conflicting kinds " - f"across union arms: {first[r['key']]['kind']} vs {r['kind']}" + f"telemetry manifest: key '{key}' has conflicting kinds " + f"across model arms: {grouped[key]['kind']} vs {r['kind']}" + ) + variants: list[_PolicyVariant] = grouped[key]["variants"] + matching = [variant for variant in variants if variant.owner is r["owner"]] + if matching: + if any(_policy_key(variant.policy) != _policy_key(r["policy"]) for variant in matching): + raise ValueError( + f"telemetry manifest: key '{key}' has conflicting policies " + f"for model arm {r['owner'].__qualname__}" + ) + continue + variants.append(_PolicyVariant(owner=r["owner"], policy=r["policy"])) + + entries = [] + for key, group in grouped.items(): + variants = tuple(group["variants"]) + allowed_values: list[Any] = [] + capture_types: set[str] = set() + for variant in variants: + allowed_values.extend(_policy_allowed_values(variant.policy)) + capture_types.update(_policy_capture_types(variant.policy)) + entries.append( + _ManifestEntry( + path=key, + kind=group["kind"], + capture_types=tuple(sorted(capture_types)), + allowed_values=_dedupe_typed(allowed_values), + variants=variants, ) - seen = union_allowed.setdefault(r["key"], []) - for v in r["allowed"]: - if v not in seen: - seen.append(v) - - entries = [ - _ManifestEntry( - path=key, - annotation=r["annotation"], - kind=r["kind"], - converter=r["converter"], - allowed_values=tuple(union_allowed[key]), - metadata=r["metadata"], ) - for key, r in first.items() - ] entries.sort(key=lambda e: e.path) return entries -def manifest_rows(model_cls: type[BaseModel]) -> list[dict[str, Any]]: - """Serializable, human-legible projection of build_capture_manifest. - - Used by the committed golden, the docs renderer, and the - capture_manifest_digest. - """ - return [ - { - "path": e.path, - "annotation": _annotation_repr(e.annotation), - "kind": e.kind, - "converter": e.converter, - "allowed_values": list(e.allowed_values), +def _manifest_rows(entries: list[_ManifestEntry]) -> list[dict[str, Any]]: + rows = [] + for entry in entries: + row = { + "path": entry.path, + "kind": entry.kind, + "capture_policy": "|".join(entry.capture_types), } - for e in build_capture_manifest(model_cls) - ] + if entry.allowed_values: + row["allowed_values"] = list(entry.allowed_values) + rows.append(row) + return rows + + +def manifest_rows(model_cls: type[BaseModel]) -> list[dict[str, Any]]: + """Serializable projection used by the golden, docs, and manifest digest.""" + return _manifest_rows(build_capture_manifest(model_cls)) def golden_manifest() -> dict[str, list[dict[str, Any]]]: @@ -394,124 +479,74 @@ def golden_manifest() -> dict[str, list[dict[str, Any]]]: } -def _sanitize_allowlist(value: Any, metadata: dict[str, Any]) -> tuple[bool, Any]: - """Capture value only when it matches field-owned allowed values.""" - allowed_values = metadata.get("allowed_values") - if not isinstance(allowed_values, (list, tuple, set)): - return False, None +def _sanitize_allowed_value(value: Any, allowed_values: tuple[Any, ...]) -> tuple[bool, Any]: + """Capture only a finite scalar matching an exact typed allowed value.""" candidates = [value] if isinstance(value, Enum): - candidates.append(value.value) - candidates.append(value.name) - candidates.append(value.name.lower()) - + enum_output = _enum_output(value) + candidates.extend((value.value, enum_output)) + if isinstance(enum_output, str): + candidates.append(enum_output.lower()) for candidate in candidates: - if candidate in allowed_values and ( - isinstance(candidate, (bool, int, float, str)) or candidate is None + if _is_safe_scalar(candidate) and any( + _typed_equal(candidate, allowed) for allowed in allowed_values ): return True, candidate return False, None -def _sanitize_literal(value: Any, annotation: Any) -> tuple[bool, Any]: - """Capture only values declared by Literal annotation.""" - allowed = get_args(annotation) - if value in allowed: - return True, value - return False, None - - -def _sanitize_sequence( - value: Any, - annotation: Any, - metadata: dict[str, Any], - state: _CaptureState | None = None, -) -> tuple[bool, Any]: - """Sanitize homogeneous sequence values. Reject whole sequence on one bad item. - - Captured items are capped at MAX_SEQ_ITEMS. The cap rides the recursive - _sanitize_value call, so each inner list of a nested sequence is bounded - independently. When any sequence is clipped, state.sequence_truncated is - set so the metadata reports the truncation honestly. - """ - annotation = _unwrap_optional(annotation) - origin = get_origin(annotation) - if origin not in {list, tuple, set}: - element_annotation = Any - else: - args = get_args(annotation) - # Homogeneous only. tuple[int, str] uses first annotation and rejects on - # mismatch. Fail closed if future telemetry field adds heterogeneous tuple. - element_annotation = args[0] if args else Any - - sanitized = [] - for item in value: - item_safe, item_value = _sanitize_value(item, element_annotation, metadata, state) - if not item_safe: - return False, None - sanitized.append(item_value) - if origin is set: - sanitized = sorted(sanitized, key=_canonical_json) - if len(sanitized) > MAX_SEQ_ITEMS: - sanitized = sanitized[:MAX_SEQ_ITEMS] - if state is not None: - state.sequence_truncated = True - return True, sanitized - - -def _sanitize_value( +def _sanitize_policy( value: Any, - annotation: Any, - metadata: dict[str, Any], + policy: _CapturePolicy, state: _CaptureState | None = None, ) -> tuple[bool, Any]: - """Return telemetry-safe primitive value, else exclude it. - - Bare strings are unsafe unless Literal or allowlist-converted. Exclusion is - visible through unsafe_excluded metadata. - """ - has_converter = _converter_is_approved(metadata) - if _union_needs_converter(annotation) and not has_converter: - return False, None - if not has_converter and not _annotation_is_capture_safe(annotation): + """Try the compiled policy branches and return one telemetry-safe value.""" + if policy.policy_type == "union": + for branch in policy.branches: + is_safe, sanitized = _sanitize_policy(value, branch, state) + if is_safe: + return True, sanitized return False, None - annotation = _unwrap_optional(annotation) - - # None: Optional field unset -> capture as null, regardless of converter. - # Must precede the allowlist branch, else None on an Optional allowlist field - # fails the allowlist and falsely flips unsafe_excluded. - if value is None: - return True, None - - if has_converter: - return _sanitize_allowlist(value, metadata) - if isinstance(value, Enum): - enum_value = value.value - if isinstance(enum_value, str): - return True, enum_value - # Enum.name always str. Prefer stable names for int-valued enums. - return True, value.name - # bool before int. Python bool is int subclass; keep True/False not 1/0. - if isinstance(value, bool): - return True, value - if isinstance(value, int) and not isinstance(value, bool): - return True, value - if isinstance(value, float): - # Reject nan/inf. json.dumps emits the bare NaN/Infinity tokens for - # non-finite floats, which are invalid JSON and break downstream - # parsing and digest stability. - if not math.isfinite(value): + if policy.policy_type == "none": + return (True, None) if value is None else (False, None) + if policy.policy_type == "bool": + return (True, value) if type(value) is bool else (False, None) + if policy.policy_type == "int": + return (True, value) if type(value) is int else (False, None) + if policy.policy_type == "float": + # The Python numeric tower permits an int where float is annotated, + # and Pydantic does not validate every default. Normalize such values + # to the annotation's JSON number shape, while still excluding bool. + if type(value) not in {int, float}: return False, None - return True, value - if isinstance(value, str): - if _is_literal(annotation): - return _sanitize_literal(value, annotation) - # Bare str can be path/secret/user text. Require Literal or allowlist. - return False, None - if isinstance(value, Path): - return False, None - if isinstance(value, (list, tuple, set)): - return _sanitize_sequence(value, annotation, metadata, state) + try: + normalized = float(value) + except OverflowError: + return False, None + return (True, normalized) if math.isfinite(normalized) else (False, None) + if policy.policy_type in {"literal", "allowlist"}: + return _sanitize_allowed_value(value, policy.allowed_values) + if policy.policy_type == "enum": + if type(value) is not policy.runtime_type: + return False, None + return True, _enum_output(value) + if policy.policy_type == "sequence": + if type(value) is not policy.runtime_type: + return False, None + element_policy = policy.branches[0] + sanitized = [] + for item in value: + item_safe, item_value = _sanitize_policy(item, element_policy, state) + if not item_safe: + return False, None + sanitized.append(item_value) + if policy.runtime_type is set: + sanitized.sort(key=_canonical_json) + if len(sanitized) > MAX_SEQ_ITEMS: + sanitized = sanitized[:MAX_SEQ_ITEMS] + if state is not None: + state.sequence_truncated = True + return True, sanitized return False, None @@ -533,27 +568,34 @@ def _schema_digest(model_cls: type[BaseModel]) -> str: return _digest({"class": model_cls.__name__, "fields": schema_fields}) -def _resolve_path(instance: BaseModel, path: str) -> tuple[bool, Any]: +def _resolve_path(instance: BaseModel, path: str) -> tuple[bool, Any, BaseModel | None]: """Resolve a dotted manifest path against a live instance. - Returns (present, value). Skips when a parent segment is missing/None or is - not a pydantic model (unset config, or a discriminated-union arm that isn't - the active one). A present leaf whose value is None resolves as (True, None). + Returns (present, value, leaf owner). Skips when a parent segment is + missing/None or is not a pydantic model (unset config, or a union arm that + is not active). The owner selects that arm's compiled capture policy. """ segments = path.split(".") obj: Any = instance for seg in segments[:-1]: if not _is_pydantic_model(obj): - return False, None + return False, None, None if seg not in obj.__class__.model_fields: - return False, None + return False, None, None obj = getattr(obj, seg, None) if obj is None: - return False, None + return False, None, None leaf = segments[-1] if not _is_pydantic_model(obj) or leaf not in obj.__class__.model_fields: - return False, None - return True, getattr(obj, leaf, None) + return False, None, None + return True, getattr(obj, leaf, None), obj + + +def _policy_for_owner(entry: _ManifestEntry, owner: BaseModel) -> _CapturePolicy | None: + return next( + (variant.policy for variant in entry.variants if type(owner) is variant.owner), + None, + ) def _truncate_to_budget(values: dict[str, Any]) -> tuple[dict[str, Any], str]: @@ -608,10 +650,13 @@ def collect_llm_api_config_payloads(llm_args: Any) -> tuple[str, str]: entries = build_capture_manifest(cls) state = _CaptureState() for entry in entries: - present, value = _resolve_path(llm_args, entry.path) - if not present: + present, value, owner = _resolve_path(llm_args, entry.path) + if not present or owner is None: + continue + policy = _policy_for_owner(entry, owner) + if policy is None: continue - is_safe, sanitized = _sanitize_value(value, entry.annotation, entry.metadata, state) + is_safe, sanitized = _sanitize_policy(value, policy, state) if is_safe: state.values[entry.path] = sanitized else: @@ -623,7 +668,7 @@ def collect_llm_api_config_payloads(llm_args: Any) -> tuple[str, str]: state.values, config_json = _truncate_to_budget(state.values) state.payload_truncated = True - rows = manifest_rows(cls) + rows = _manifest_rows(entries) metadata = { "api_contract_version": API_CONTRACT_VERSION, "args_class": cls.__name__, diff --git a/tensorrt_llm/usage/schemas/README.md b/tensorrt_llm/usage/schemas/README.md index 2bb4b4dc2010..5b8e0565c94d 100644 --- a/tensorrt_llm/usage/schemas/README.md +++ b/tensorrt_llm/usage/schemas/README.md @@ -225,28 +225,28 @@ flags such as LoRA/speculative decoding have explicit safe config fields. The `llmApiConfigJson` field is a JSON-serialized dict containing a type-driven subset of the validated, effective LLM API configuration. Capture is -**type-driven**: a field is captured automatically when its type is categorical -(`Literal`/`Enum`/`bool`) or numeric (`int`/`float`), or a safe collection of -those. Free-form `str`/`Any`/`Path`/`dict`/`Callable` are not captured unless the -field carries an explicit allowlist (`TelemetryField.categorical(...)`). Any field -can opt out with `telemetry=False`. +automatic for `bool`, `int`, finite `float`, `Literal`, `Enum`, supported unions, +and homogeneous sequences. Unsafe scalar `str`, `Any`, and `object` branches +require `TelemetryField.categorical(...)`; paths, mappings, callables, and +unsupported structures always fail closed. Use `telemetry=False` to exclude a +field. Captured values must be safe primitives. Raw strings are excluded unless the -field is a `Literal[...]` or uses an explicit `allowlist` converter. Paths, +field is a `Literal[...]`/`Enum` or matches explicit `allowed_values`. Paths, tokenizer locations, dicts, objects, callables, raw `Any` values, non-finite floats (`nan`/`inf`), and unsafe or heterogeneous sequences are excluded. +Union branches are compiled and sanitized independently: explicit +`allowed_values` opt in only otherwise unsafe scalar branches and do not filter +safe numeric, boolean, `Literal`, or `Enum` branches in the same union. Captured sequences are capped at a fixed length and any clipping is reported in `llmApiConfigMetaJson`. Exclusion is fail-closed: the value is omitted instead of being serialized, and `llmApiConfigMetaJson` reports whether any resolved field was excluded as unsafe. -The table below is a non-exhaustive set of examples for readers building -dashboards. The exhaustive source of truth is -`tensorrt_llm/usage/llm_args_golden_manifest.json` (regenerated from -`build_capture_manifest`), after the safety sanitizer has excluded unsafe values. -Use `llmApiConfigMetaJson` digests and field counts to track the exact capture -manifest for a given release. The rendered documentation generates the -exhaustive field table at docs build time under **Developer Guide > Telemetry**. +The table below gives common dashboard examples. The committed manifest is the +canonical list of `TorchLlmArgs` capturable paths, merged policies, and +categorical domains. `llmApiConfigMetaJson` identifies that manifest by digest. +The docs build renders the full table under **Developer Guide > Telemetry**. | Key | Description | |-----|-------------| @@ -256,7 +256,7 @@ exhaustive field table at docs build time under **Developer Guide > Telemetry**. | `moe_expert_parallel_size` | MoE expert parallelism degree (None/unset when runtime decides). | | `moe_tensor_parallel_size` | MoE tensor parallelism degree (None/unset when runtime decides). | | `moe_cluster_parallel_size` | MoE cluster parallelism degree (None/unset when runtime decides). | -| `backend` | Execution backend. Captured as the `Literal["pytorch"]` value on the PyTorch args, and through an explicit allowlist (`pytorch`, `tensorrt`, `_autodeploy`) on the base/TRT args. | +| `backend` | Execution backend. Captured as the `Literal["pytorch"]` value on the PyTorch args, and through explicit allowed values (`pytorch`, `tensorrt`, `_autodeploy`) on the base/TRT args. | | `dtype` | Model dtype, captured through an explicit allowlist. | | `load_format` | Weight load format, captured as a low-cardinality enum/string value. | | `quant_config.quant_algo` | Quantization algorithm, captured as a closed `QuantAlgo` enum value (TRT args only). Empty/absent when unquantized. | @@ -321,9 +321,10 @@ Checklist for adding an LLM API config capture field inside `llmApiConfigJson`: exclusion sentinel keeps a categorical/numeric field out of capture. 4. **Do not capture unsafe data.** No model/tokenizer/file paths, prompts, outputs, secrets/tokens/URLs/hostnames, free-form user strings, raw - dict/object payloads, or callables. The sanitizer fails closed regardless: - bare `str`, `Any`, `object`, `Path`, `dict`, callables, permissive unions, and - non-finite floats are dropped unless an approved `allowlist` converter applies. + dict/object payloads, or callables. Paths, mappings, callables, arbitrary + non-scalars, heterogeneous sequences, and non-finite floats always fail + closed. A `str`/`Any`/`object` scalar branch is emitted only when it exactly + matches the field's finite `allowed_values`. 5. **`tests/unittest/usage/test_llmapi_config_capture.py`** — Add behavior coverage: assert the value is captured, and for a categorical bare-string field assert that an out-of-allowlist value is redacted (dropped) while an @@ -338,8 +339,9 @@ Checklist for adding an LLM API config capture field inside `llmApiConfigJson`: important enough for dashboard users to know by name. Dashboard note: payloads carry `capture_version` and `field_policy_version` in -`llmApiConfigMetaJson`. During release adoption, v1 (opt-in) and v2 (type-driven) -payloads coexist in the same index — **bucket by these before aggregating** +`llmApiConfigMetaJson`. During release adoption, v1 (opt-in), v2 (initial +type-driven), and v3 (composed branch-policy) payloads coexist in the same index +— **bucket by these before aggregating** `captured_field_count` or any `llmApiConfigJson.`. ### Conventions diff --git a/tests/unittest/usage/test_config.py b/tests/unittest/usage/test_config.py index 6680f8ca6a03..115e8ec57e2c 100644 --- a/tests/unittest/usage/test_config.py +++ b/tests/unittest/usage/test_config.py @@ -22,18 +22,6 @@ class TestTelemetryConfigLocation: """Verify TelemetryConfig and UsageContext live in tensorrt_llm.usage.config.""" - def test_import_telemetry_config_from_usage_config(self): - """TelemetryConfig must be importable from tensorrt_llm.usage.config.""" - from tensorrt_llm.usage import config - - assert hasattr(config, "TelemetryConfig") - - def test_import_usage_context_from_usage_config(self): - """UsageContext must be importable from tensorrt_llm.usage.config.""" - from tensorrt_llm.usage import config - - assert hasattr(config, "UsageContext") - def test_telemetry_config_defaults(self): """TelemetryConfig defaults: disabled=False, usage_context=UNKNOWN.""" from tensorrt_llm.usage import config @@ -95,38 +83,20 @@ def test_telemetry_config_rejects_extra_fields(self): class TestBackwardCompatibility: """Verify types are still importable from llm_args for backward compat.""" - def test_telemetry_config_importable_from_llm_args(self): - """TelemetryConfig must still be importable from llm_args.""" - from tensorrt_llm.llmapi import llm_args + def test_telemetry_types_preserve_legacy_imports(self) -> None: + from tensorrt_llm.llmapi.llm_args import TelemetryConfig as LegacyTelemetryConfig + from tensorrt_llm.llmapi.llm_args import UsageContext as LegacyUsageContext + from tensorrt_llm.usage.config import TelemetryConfig, UsageContext - assert hasattr(llm_args, "TelemetryConfig") - - def test_usage_context_importable_from_llm_args(self): - """UsageContext must still be importable from llm_args.""" - from tensorrt_llm.llmapi import llm_args - - assert hasattr(llm_args, "UsageContext") - - def test_same_types_both_locations(self): - """Types from both locations must be the same class.""" - from tensorrt_llm.llmapi import llm_args - from tensorrt_llm.usage import config - - assert config.TelemetryConfig is llm_args.TelemetryConfig - assert config.UsageContext is llm_args.UsageContext + assert LegacyTelemetryConfig is TelemetryConfig + assert LegacyUsageContext is UsageContext class TestFieldTelemetryMetadata: """Verify llm_args.Field telemetry metadata handling.""" def test_telemetry_false_records_exclude_marker(self): - """Field(telemetry=False) records an honored exclude sentinel. - - Under type-driven auto-enroll, telemetry=False is no longer a no-op: it - is the explicit opt-out for a type-safe-but-sensitive field, recorded as - json_schema_extra['telemetry'] = {"exclude": True} and honored by - build_capture_manifest's selection rule. - """ + """Field(telemetry=False) records the explicit exclude sentinel.""" from tensorrt_llm.llmapi import llm_args field = llm_args.Field(default=0, telemetry=False) @@ -154,7 +124,5 @@ def test_categorical_builds_allowlist_metadata(self): field = TelemetryField.categorical("a", "b") assert field.as_json_schema_extra() == { - "kind": "categorical", - "converter": "allowlist", "allowed_values": ["a", "b"], } diff --git a/tests/unittest/usage/test_llmapi_config_capture.py b/tests/unittest/usage/test_llmapi_config_capture.py index ff558552ebbf..9798459343dd 100644 --- a/tests/unittest/usage/test_llmapi_config_capture.py +++ b/tests/unittest/usage/test_llmapi_config_capture.py @@ -29,25 +29,24 @@ ) from tensorrt_llm.llmapi.utils import StrictBaseModel from tensorrt_llm.usage import usage_lib +from tensorrt_llm.usage.config import TelemetryField from tensorrt_llm.usage.llmapi_config import collect_llm_api_config_payloads pytestmark = pytest.mark.cpu_only class _NestedConfig(StrictBaseModel): - marked: int = Field(default=7, telemetry={"kind": "value"}) - unmarked: int = Field(default=11) + safe_plain_int: int = 7 + safe_field_int: int = Field(default=11) class _ExampleConfig(StrictBaseModel): - safe_marked: int = Field(default=3, telemetry={"kind": "value"}) - safe_unmarked: int = Field(default=5) - private_path: str = Field(default="/customer/private/model", telemetry={"kind": "value"}) - mode: Literal["auto", "slow"] = Field(default="auto", telemetry={"kind": "categorical"}) + safe_plain_int: int = 3 + safe_field_int: int = Field(default=5) + private_path: str = "/customer/private/model" + mode: Literal["auto", "slow"] = "auto" nested: _NestedConfig = Field(default_factory=_NestedConfig) - unsafe_union: Optional[Union[str, Path]] = Field( - default="/customer/tokenizer", telemetry={"kind": "categorical"} - ) + unsafe_union: Optional[Union[str, Path]] = "/customer/tokenizer" def _loads_payloads(args) -> tuple[dict, dict]: @@ -56,17 +55,15 @@ def _loads_payloads(args) -> tuple[dict, dict]: def test_collect_llm_api_config_uses_type_driven_autoenroll_and_safety_vetoes(): - # Renamed from ..._uses_strict_opt_in_...: under auto-enroll, unmarked - # type-safe ints (safe_unmarked, nested.unmarked) are now captured; bare - # str / Union[str,Path] without an approved allowlist remain uncapturable. + """Safe annotations auto-enroll while bare strings and paths stay excluded.""" config, meta = _loads_payloads(_ExampleConfig()) assert config == { "mode": "auto", - "nested.marked": 7, - "nested.unmarked": 11, - "safe_marked": 3, - "safe_unmarked": 5, + "nested.safe_field_int": 11, + "nested.safe_plain_int": 7, + "safe_field_int": 5, + "safe_plain_int": 3, } assert "private_path" not in config # bare str, no allowlist -> not capturable assert "unsafe_union" not in config # Union[str,Path], no allowlist -> not capturable @@ -79,87 +76,57 @@ def test_collect_llm_api_config_uses_type_driven_autoenroll_and_safety_vetoes(): assert meta["payload_truncated"] is False -def test_collect_llm_api_config_allows_approved_string_converters_only(): - # The union_backend / union_path fixtures below are Union[str, Path] solely - # to exercise the value-fail-closed allowlist seam: union_path defaults to a - # Path, which is dropped because it is not an allowlisted scalar, while - # union_backend's allowlisted str is captured. No production telemetry field - # is Union[str, Path]; the only real Union allowlist fields are - # Union[str, Enum] (load_format). See CR-E (declined). +def test_collect_llm_api_config_allows_only_explicit_categorical_values(): + """Explicit strings capture while arbitrary strings and Path values fail closed.""" + class _StringConfig(StrictBaseModel): backend: Optional[str] = Field( default="pytorch", - telemetry={ - "kind": "categorical", - "converter": "allowlist", - "allowed_values": ["pytorch", "tensorrt"], - }, + telemetry=TelemetryField.categorical("pytorch", "tensorrt"), ) unsafe_backend: Optional[str] = Field( default="file:///customer/private", - telemetry={ - "kind": "categorical", - "converter": "allowlist", - "allowed_values": ["pytorch", "tensorrt"], - }, + telemetry=TelemetryField.categorical("pytorch", "tensorrt"), ) unconverted: Optional[str] = Field( default="arbitrary-user-string", telemetry={"kind": "categorical"} ) union_backend: Union[str, Path] = Field( default="tensorrt", - telemetry={ - "kind": "categorical", - "converter": "allowlist", - "allowed_values": ["pytorch", "tensorrt"], - }, + telemetry=TelemetryField.categorical("pytorch", "tensorrt"), ) union_path: Union[str, Path] = Field( default=Path("/customer/private"), - telemetry={ - "kind": "categorical", - "converter": "allowlist", - "allowed_values": ["pytorch", "tensorrt"], - }, + telemetry=TelemetryField.categorical("pytorch", "tensorrt"), ) config, meta = _loads_payloads(_StringConfig()) assert config == {"backend": "pytorch", "union_backend": "tensorrt"} assert meta["captured_field_count"] == 2 - # unconverted (kind=categorical, no converter) is not capturable -> not in - # manifest; unsafe_backend + union_path resolve but fail the allowlist sanitizer. + # unconverted has no allowed values and is not capturable; unsafe_backend and + # union_path resolve but fail the explicit-value policy. assert meta["capturable_field_count"] == 4 assert meta["excluded_field_count"] == 2 assert meta["unsafe_excluded"] is True -def test_sanitize_allowlist_is_value_fail_closed_for_non_scalars(): - """The allowlist only emits scalars, even if a non-scalar is allowlisted. - - Documents the CR-E decision (declined): the sanitizer is value-fail-closed, - so a Path or arbitrary object cannot leak through the allowlist converter - even if it were placed in allowed_values. _sanitize_allowlist returns a - candidate only when it is BOTH in allowed_values AND a scalar - (bool/int/float/str) or None. Excluding Union-with-Any/Path from allowlist - eligibility at the type level is therefore unnecessary for safety, and a - coarse rule would also break legitimate Union[str, Enum] allowlist fields - such as load_format (verified captured elsewhere). - """ - from tensorrt_llm.usage import llmapi_config +def test_non_scalar_allowed_value_fails_manifest_build_loudly(): + """Explicit domains accept only finite JSON scalars.""" + from tensorrt_llm.usage.llmapi_config import build_capture_manifest secret = Path("/customer/secret") - metadata = {"converter": "allowlist", "allowed_values": [secret]} - # A Path placed in the allowlist is still rejected: not a scalar. - assert llmapi_config._sanitize_allowlist(secret, metadata) == (False, None) - # An allowlisted scalar is captured. - str_metadata = {"converter": "allowlist", "allowed_values": ["pytorch"]} - assert llmapi_config._sanitize_allowlist("pytorch", str_metadata) == (True, "pytorch") + + class _InvalidDomainConfig(StrictBaseModel): + value: Any = Field(default=secret, telemetry=TelemetryField.categorical(secret)) + + with pytest.raises(ValueError, match="finite JSON scalars"): + build_capture_manifest(_InvalidDomainConfig) def test_collect_llm_api_config_walks_only_declared_pydantic_fields(): class _DeclaredFieldsOnlyConfig(StrictBaseModel): - safe_value: int = Field(default=3, telemetry={"kind": "value"}) + safe_value: int = 3 @property def leaked_value(self): @@ -174,17 +141,13 @@ def leaked_value(self): def test_collect_llm_api_config_rejects_unsafe_annotations_even_for_safe_values(): class _UnsafeAnnotationConfig(StrictBaseModel): - safe_value: int = Field(default=3, telemetry={"kind": "value"}) - raw_any: Any = Field(default=11, telemetry={"kind": "value"}) - object_like: object = Field(default=True, telemetry={"kind": "value"}) - raw_dict: dict[str, Any] = Field(default_factory=dict, telemetry={"kind": "value"}) + safe_value: int = 3 + raw_any: Any = 11 + object_like: object = True + raw_dict: dict[str, Any] = Field(default_factory=dict) converted_any: Any = Field( default="known", - telemetry={ - "kind": "categorical", - "converter": "allowlist", - "allowed_values": ["known"], - }, + telemetry=TelemetryField.categorical("known"), ) config, meta = _loads_payloads(_UnsafeAnnotationConfig()) @@ -276,10 +239,8 @@ class _PrecisionMode(Enum): FP8 = "fp8" class _EnumConfig(StrictBaseModel): - mode: _LoadMode = Field(default=_LoadMode.AUTO, telemetry={"kind": "categorical"}) - precision: _PrecisionMode = Field( - default=_PrecisionMode.FP8, telemetry={"kind": "categorical"} - ) + mode: _LoadMode = _LoadMode.AUTO + precision: _PrecisionMode = _PrecisionMode.FP8 config, meta = _loads_payloads(_EnumConfig()) @@ -289,7 +250,7 @@ class _EnumConfig(StrictBaseModel): def test_collect_llm_api_config_keeps_bool_values_boolean(): class _BoolConfig(StrictBaseModel): - enabled: bool = Field(default=True, telemetry={"kind": "value"}) + enabled: bool = True config, _ = _loads_payloads(_BoolConfig()) @@ -297,130 +258,239 @@ class _BoolConfig(StrictBaseModel): assert type(config["enabled"]) is bool -def test_collect_llm_api_config_rejects_non_finite_floats(): - """Non-finite floats (nan/inf) are excluded; finite floats are captured. - - json.dumps emits the bare NaN/Infinity tokens for non-finite floats, which - are invalid JSON and break downstream parsing and digest stability. Guard - the float branch with math.isfinite so a marked float set to inf/nan is - dropped (unsafe_excluded) while a finite float is still captured. - """ - +def test_float_policy_normalizes_integer_defaults_but_rejects_bool(): class _FloatConfig(StrictBaseModel): - finite: float = Field(default=0.5, telemetry={"kind": "value"}) - infinite: float = Field(default=float("inf"), telemetry={"kind": "value"}) - not_a_number: float = Field(default=float("nan"), telemetry={"kind": "value"}) + value: float = 0 config, meta = _loads_payloads(_FloatConfig()) + assert config == {"value": 0.0} + assert type(config["value"]) is float + assert meta["unsafe_excluded"] is False - assert config == {"finite": 0.5} - assert "infinite" not in config - assert "not_a_number" not in config - assert meta["excluded_field_count"] == 2 + invalid = _FloatConfig.model_construct(value=True) + config, meta = _loads_payloads(invalid) + assert config == {} assert meta["unsafe_excluded"] is True -def test_collect_llm_api_config_rejects_non_finite_floats_in_sequence(): - """A single non-finite float poisons the whole marked sequence. +def test_bool_literal_union_composes_branches_without_filtering_bool(): + from tensorrt_llm.usage.llmapi_config import build_capture_manifest - The sequence sanitizer fails closed on one bad item, so a list containing - inf/nan is dropped entirely rather than emitting invalid JSON tokens. - """ + for value in (True, False, "auto"): + config, meta = _loads_payloads(KvCacheConfig(use_kv_cache_manager_v2=value)) + assert config["use_kv_cache_manager_v2"] == value + assert meta["unsafe_excluded"] is False - class _FloatSeqConfig(StrictBaseModel): - finite_buckets: list[float] = Field( - default_factory=lambda: [0.1, 0.5, 1.0], telemetry={"kind": "value"} - ) - poisoned_buckets: list[float] = Field( - default_factory=lambda: [0.1, float("inf"), 1.0], telemetry={"kind": "value"} - ) + invalid = KvCacheConfig() + invalid.use_kv_cache_manager_v2 = "private-value" + config, meta = _loads_payloads(invalid) + assert "use_kv_cache_manager_v2" not in config + assert meta["unsafe_excluded"] is True - config, meta = _loads_payloads(_FloatSeqConfig()) + entry = next( + item + for item in build_capture_manifest(KvCacheConfig) + if item.path == "use_kv_cache_manager_v2" + ) + assert entry.kind == "categorical" + assert entry.capture_types == ("bool", "literal") + assert entry.allowed_values == ("auto",) - assert config == {"finite_buckets": [0.1, 0.5, 1.0]} - assert "poisoned_buckets" not in config - assert meta["unsafe_excluded"] is True +def test_int_literal_union_preserves_typed_branch_distinctions(): + class _IntOrAutoConfig(StrictBaseModel): + value: int | Literal["auto"] = "auto" -def test_collect_llm_api_config_caps_long_sequences_and_flags_truncation(): - """A marked sequence longer than MAX_SEQ_ITEMS is clipped and flagged. + for value in (128, "auto"): + config, meta = _loads_payloads(_IntOrAutoConfig(value=value)) + assert config["value"] == value + assert meta["unsafe_excluded"] is False - llmApiConfigJson is unbounded on the wire and the reporter is fail-silent, - so a pathological user-sized list could silently drop the whole payload. - Cap captured sequences to MAX_SEQ_ITEMS and record a single honest - sequence_truncated boolean in the metadata. - """ - from tensorrt_llm.usage import llmapi_config + for value in ("private-value", True): + invalid = _IntOrAutoConfig.model_construct(value=value) + config, meta = _loads_payloads(invalid) + assert "value" not in config + assert meta["unsafe_excluded"] is True - cap = llmapi_config.MAX_SEQ_ITEMS - class _LongSeqConfig(StrictBaseModel): - values: list[int] = Field( - default_factory=lambda: list(range(cap + 50)), telemetry={"kind": "value"} - ) +def test_equal_valued_literal_union_branches_preserve_python_types(): + from tensorrt_llm.usage.llmapi_config import build_capture_manifest - config, meta = _loads_payloads(_LongSeqConfig()) + class _TypedLiteralConfig(StrictBaseModel): + value: Literal[1] | Literal[True] = 1 + values: list[Literal[0] | Literal[False]] = Field(default_factory=lambda: [0, False]) - assert len(config["values"]) == cap - assert config["values"] == list(range(cap)) - assert meta["sequence_truncated"] is True + entry = next( + item for item in build_capture_manifest(_TypedLiteralConfig) if item.path == "value" + ) + assert [type(value) for value in entry.allowed_values] == [int, bool] + for value in (1, True): + config, meta = _loads_payloads(_TypedLiteralConfig.model_construct(value=value)) + assert type(config["value"]) is type(value) + assert meta["unsafe_excluded"] is False -def test_collect_llm_api_config_caps_nested_inner_sequences(): - """Each inner list of a nested List[List[int]] is capped independently.""" - from tensorrt_llm.usage import llmapi_config + config, meta = _loads_payloads(_TypedLiteralConfig()) + assert [type(value) for value in config["values"]] == [int, bool] + assert meta["unsafe_excluded"] is False - cap = llmapi_config.MAX_SEQ_ITEMS - class _NestedSeqConfig(StrictBaseModel): - rows: list[list[int]] = Field( - default_factory=lambda: [list(range(cap + 10)), list(range(cap + 20))], - telemetry={"kind": "value"}, +def test_literal_mamba_cache_dtype_needs_no_explicit_allowlist(): + allowed = ("auto", "float16", "bfloat16", "float32") + for value in allowed: + config, meta = _loads_payloads(KvCacheConfig(mamba_ssm_cache_dtype=value)) + assert config["mamba_ssm_cache_dtype"] == value + assert meta["unsafe_excluded"] is False + + invalid = KvCacheConfig() + invalid.mamba_ssm_cache_dtype = "private-value" + config, meta = _loads_payloads(invalid) + assert "mamba_ssm_cache_dtype" not in config + assert meta["unsafe_excluded"] is True + + +def test_explicit_allowlist_applies_only_to_unsafe_union_branch(): + class _MixedConfig(StrictBaseModel): + value: int | str = Field( + default="auto", + telemetry=TelemetryField.categorical("auto"), ) - config, meta = _loads_payloads(_NestedSeqConfig()) + for value in (128, "auto"): + config, meta = _loads_payloads(_MixedConfig(value=value)) + assert config["value"] == value + assert meta["unsafe_excluded"] is False - assert len(config["rows"]) == 2 - assert all(len(inner) == cap for inner in config["rows"]) - assert config["rows"][0] == list(range(cap)) - assert meta["sequence_truncated"] is True + for value in ("private-value", True): + invalid = _MixedConfig.model_construct(value=value) + config, meta = _loads_payloads(invalid) + assert "value" not in config + assert meta["unsafe_excluded"] is True -def test_collect_llm_api_config_caps_outer_nested_sequence(): - """The outer list of a nested List[List[int]] is also capped.""" - from tensorrt_llm.usage import llmapi_config +def test_explicit_allowlist_membership_is_type_exact(): + class _NumericDomainConfig(StrictBaseModel): + value: Any = Field(default=1, telemetry=TelemetryField.categorical(1)) + + config, _ = _loads_payloads(_NumericDomainConfig()) + assert config == {"value": 1} + + invalid = _NumericDomainConfig(value=True) + config, meta = _loads_payloads(invalid) + assert config == {} + assert meta["unsafe_excluded"] is True - cap = llmapi_config.MAX_SEQ_ITEMS - class _WideNestedSeqConfig(StrictBaseModel): - rows: list[list[int]] = Field( - default_factory=lambda: [[0, 1] for _ in range(cap + 30)], - telemetry={"kind": "value"}, +def test_explicit_allowlist_cannot_broaden_a_safe_annotation(): + from tensorrt_llm.usage.llmapi_config import build_capture_manifest + + class _LiteralConfig(StrictBaseModel): + value: Literal["auto"] = Field( + default="auto", + telemetry=TelemetryField.categorical("private-value"), ) - config, meta = _loads_payloads(_WideNestedSeqConfig()) + with pytest.raises(ValueError, match="unsafe scalar annotation branch"): + build_capture_manifest(_LiteralConfig) - assert len(config["rows"]) == cap - assert meta["sequence_truncated"] is True +def test_nested_union_sequence_is_sanitized_branch_by_branch(): + class _NestedUnionConfig(StrictBaseModel): + values: Optional[list[int | Literal["auto"]]] = None + + config, meta = _loads_payloads(_NestedUnionConfig(values=[128, "auto"])) + assert config == {"values": [128, "auto"]} + assert meta["unsafe_excluded"] is False + + invalid = _NestedUnionConfig.model_construct(values=[True]) + config, meta = _loads_payloads(invalid) + assert config == {} + assert meta["unsafe_excluded"] is True -def test_collect_llm_api_config_small_sequence_not_truncated(): - """A sequence within the cap is captured whole and the flag stays false.""" - class _SmallSeqConfig(StrictBaseModel): - values: list[int] = Field(default_factory=lambda: [1, 2, 3], telemetry={"kind": "value"}) +def test_homogeneous_tuple_and_set_sequences_remain_capturable(): + class _SequenceConfig(StrictBaseModel): + sizes: tuple[int, ...] = (3, 1) + fixed_sizes: tuple[int, int] = (5, 2) + fixed_modes: tuple[Literal["auto"], Literal["auto"]] = ("auto", "auto") + modes: set[Literal["auto", "manual"]] = {"manual", "auto"} - config, meta = _loads_payloads(_SmallSeqConfig()) + config, meta = _loads_payloads(_SequenceConfig()) + assert config == { + "fixed_modes": ["auto", "auto"], + "fixed_sizes": [5, 2], + "modes": ["auto", "manual"], + "sizes": [3, 1], + } + assert meta["unsafe_excluded"] is False - assert config["values"] == [1, 2, 3] - assert meta["sequence_truncated"] is False +def test_collect_llm_api_config_rejects_non_finite_floats(): + """Non-finite floats are excluded while finite floats are captured.""" + + class _FloatConfig(StrictBaseModel): + finite: float = 0.5 + infinite: float = float("inf") + not_a_number: float = float("nan") -def test_collect_llm_api_config_failure_meta_has_truncation_key(): - """Failure metadata carries sequence_truncated for shape parity.""" + config, meta = _loads_payloads(_FloatConfig()) + + assert config == {"finite": 0.5} + assert "infinite" not in config + assert "not_a_number" not in config + assert meta["excluded_field_count"] == 2 + assert meta["unsafe_excluded"] is True + + +def test_collect_llm_api_config_rejects_non_finite_floats_in_sequence(): + """One non-finite item excludes the entire sequence.""" + + class _FloatSeqConfig(StrictBaseModel): + finite_buckets: list[float] = Field(default_factory=lambda: [0.1, 0.5, 1.0]) + poisoned_buckets: list[float] = Field(default_factory=lambda: [0.1, float("inf"), 1.0]) + + config, meta = _loads_payloads(_FloatSeqConfig()) + + assert config == {"finite_buckets": [0.1, 0.5, 1.0]} + assert "poisoned_buckets" not in config + assert meta["unsafe_excluded"] is True + + +def test_collect_llm_api_config_caps_sequences_recursively_and_flags_truncation(): from tensorrt_llm.usage import llmapi_config - meta = llmapi_config._failure_meta(args_class="Foo") + cap = llmapi_config.MAX_SEQ_ITEMS + + class _SequenceConfig(StrictBaseModel): + flat: list[int] + inner: list[list[int]] + outer: list[list[int]] + + config, meta = _loads_payloads( + _SequenceConfig( + flat=list(range(cap + 50)), + inner=[list(range(cap + 10)), list(range(cap + 20))], + outer=[[0, 1] for _ in range(cap + 30)], + ) + ) + assert config["flat"] == list(range(cap)) + assert len(config["inner"]) == 2 + assert all(row == list(range(cap)) for row in config["inner"]) + assert len(config["outer"]) == cap + assert meta["sequence_truncated"] is True + + exact = list(range(cap)) + exact_outer = [[0, 1] for _ in range(cap)] + config, meta = _loads_payloads(_SequenceConfig(flat=exact, inner=[exact], outer=exact_outer)) + assert config == {"flat": exact, "inner": [exact], "outer": exact_outer} + assert meta["sequence_truncated"] is False + + config, meta = _loads_payloads( + _SequenceConfig(flat=[1, 2, 3], inner=[[1], [2]], outer=[[0, 1]]) + ) + assert config["flat"] == [1, 2, 3] + assert config["inner"] == [[1], [2]] + assert config["outer"] == [[0, 1]] assert meta["sequence_truncated"] is False @@ -430,9 +500,10 @@ def test_failure_meta_uses_new_contract_keys_and_versions(): meta = rc._failure_meta(args_class="X") assert meta["capture_version"] == "2" assert meta["api_contract_version"] == "0.2.0" - assert meta["field_policy_version"] == "2" + assert meta["field_policy_version"] == "3" assert meta["excluded_field_count"] == 0 # renamed from the old marked-count key assert meta["payload_truncated"] is False + assert meta["sequence_truncated"] is False # The pre-migration keys must be gone from the new contract; assert by literal # so a regression that reintroduces them fails loudly. assert "excluded_marked_field_count" not in meta @@ -441,21 +512,19 @@ def test_failure_meta_uses_new_contract_keys_and_versions(): def test_collect_llm_api_config_rejects_heterogeneous_tuples(): class _TupleConfig(StrictBaseModel): - pair: tuple[int, Literal["safe"]] = Field(default=(1, "safe"), telemetry={"kind": "value"}) + pair: tuple[int, Literal["safe"]] = (1, "safe") + typed_pair: tuple[Literal[1], Literal[True]] = (1, True) config, meta = _loads_payloads(_TupleConfig()) assert config == {} - assert meta["excluded_field_count"] == 1 - assert meta["unsafe_excluded"] is True + assert meta["capturable_field_count"] == 0 + assert meta["excluded_field_count"] == 0 + assert meta["unsafe_excluded"] is False def test_collect_llm_api_config_derives_manifest_kind_from_annotation(): - """Manifest 'kind' is derived per D1, not taken from the registered value. - - Categorical iff (Optional-unwrapped) annotation is Literal/Enum OR an - allowlist is present; otherwise 'value'. The registered kind is ignored. - """ + """Stored kind metadata cannot override the kind derived from compiled policies.""" class _Mode(Enum): AUTO = "auto" @@ -470,7 +539,6 @@ class _KindConfig(StrictBaseModel): default="x", telemetry={ "kind": "value", - "converter": "allowlist", "allowed_values": ["x", "y"], }, ) @@ -517,50 +585,6 @@ def raise_runtime_error(*_args, **_kwargs): collect_llm_api_config_payloads(_ExampleConfig()) -def test_collect_llm_api_config_captures_expanded_value_fields(): - """Representative newly-marked value fields are captured on TorchLlmArgs.""" - args = TorchLlmArgs( - model="/customer/private/Llama", - skip_tokenizer_init=True, - moe_expert_parallel_size=2, - moe_tensor_parallel_size=1, - moe_cluster_parallel_size=1, - num_postprocess_workers=3, - stream_interval=4, - trust_remote_code=True, - ) - - config, meta = _loads_payloads(args) - - assert config["moe_expert_parallel_size"] == 2 - assert config["moe_tensor_parallel_size"] == 1 - assert config["moe_cluster_parallel_size"] == 1 - assert config["num_postprocess_workers"] == 3 - assert config["stream_interval"] == 4 - assert config["trust_remote_code"] is True - # backend on TorchLlmArgs is the Literal["pytorch"] override -> value capture. - assert config["backend"] == "pytorch" - assert meta["capture_succeeded"] is True - - -def test_collect_llm_api_config_captures_nested_config_value_fields(): - """Newly-marked nested-config value fields are captured via recursion.""" - from tensorrt_llm.llmapi.llm_args import MoeConfig - - args = TorchLlmArgs( - model="/customer/private/Llama", - skip_tokenizer_init=True, - moe_config=MoeConfig(max_num_tokens=8192, disable_finalize_fusion=True), - ) - - config, _ = _loads_payloads(args) - - assert config["moe_config.max_num_tokens"] == 8192 - assert config["moe_config.disable_finalize_fusion"] is True - # MoeConfig.backend is a Literal -> derived categorical, value captured. - assert config["moe_config.backend"] == "AUTO" - - def test_field_wrapper_preserves_callable_json_schema_extra_with_metadata(): """Preserve callable json_schema_extra when adding status/telemetry metadata. @@ -586,7 +610,7 @@ def mark_schema(schema: dict[str, object]) -> dict[str, object]: "original": True, "returned": True, "status": "beta", - "telemetry": {"kind": "value"}, + "telemetry": {}, } class _CallableExtraConfig(StrictBaseModel): @@ -602,21 +626,12 @@ class _CallableExtraConfig(StrictBaseModel): def test_collect_llm_api_config_captures_none_on_optional_allowlist_field(): - """None on an Optional allowlist field is captured as null, not excluded. - - Regression: the None check must precede the allowlist branch, else a None - default (e.g. reasoning_parser, TrtLlmArgs.backend) fails the allowlist and - permanently flips unsafe_excluded on default configs. - """ + """None is captured through the independent policy on Optional fields.""" class _C(StrictBaseModel): backend: Optional[str] = Field( default=None, - telemetry={ - "kind": "categorical", - "converter": "allowlist", - "allowed_values": ["pytorch", "tensorrt"], - }, + telemetry=TelemetryField.categorical("pytorch", "tensorrt"), ) config, meta = _loads_payloads(_C()) @@ -627,13 +642,7 @@ class _C(StrictBaseModel): def test_collect_llm_api_config_captures_transceiver_runtime_categorical(): - """transceiver_runtime is a single Optional[Literal] categorical. - - Regression: a two-branch Literal union (Optional[Literal['CPP','PYTHON']] - | Literal['auto']) would not unwrap to a top-level Literal, degrading the - manifest kind to 'value' and making the sanitizer reject all three - strings. Every allowed value plus None must be captured, not excluded. - """ + """Every Optional transceiver-runtime Literal value is captured.""" from tensorrt_llm.llmapi.llm_args import CacheTransceiverConfig for runtime in ("CPP", "PYTHON", "auto", None): @@ -652,12 +661,7 @@ def test_collect_llm_api_config_captures_transceiver_runtime_categorical(): def test_collect_llm_api_config_redacts_out_of_allowlist_categorical_str(): - """An out-of-allowlist value on a categorical bare-str field is dropped. - - reasoning_parser is captured via TelemetryField.categorical mirroring the - ReasoningParserFactory registry; any value outside that recognized domain - (e.g. injected free-form text) must be excluded, not captured. - """ + """The reasoning-parser allowlist rejects arbitrary strings.""" args = TorchLlmArgs( model="/customer/private/Llama", skip_tokenizer_init=True, @@ -679,12 +683,7 @@ def test_collect_llm_api_config_redacts_out_of_allowlist_categorical_str(): def test_collect_llm_api_config_captures_gms_load_format(): - """load_format=GMS is captured as 'gms' (was dropped before the allowlist fix). - - LoadFormat.GMS is a real, accepted value (convert_load_format maps the - string 'gms' to the enum), but it was missing from the load_format telemetry - allowlist, so GMS deployments were silently excluded from llmApiConfigJson. - """ + """The explicit load-format policy captures GMS as ``gms``.""" args = TorchLlmArgs( model="/customer/private/Llama", skip_tokenizer_init=True, @@ -718,13 +717,7 @@ def _walk_captured_keys(model) -> set[str]: def test_collect_llm_api_config_captures_decoding_type_for_every_arm(): - """The speculative discriminator decoding_type is captured for every arm. - - decoding_type is the single most valuable categorical (it identifies which - speculative mode is active). The runtime collector walks the concrete active - arm's model_fields, so marking it on only one arm drops it for the others. - Assert representative non-UserProvided arms capture it. - """ + """Every reachable decoding_type Literal is captured from its active model arm.""" from tensorrt_llm.llmapi.llm_args import ( AutoDecodingConfig, MTPDecodingConfig, @@ -742,12 +735,7 @@ def test_collect_llm_api_config_captures_decoding_type_for_every_arm(): def test_collect_llm_api_config_captures_max_total_draft_tokens_for_every_arm(): - """max_total_draft_tokens is a safe value field on the shared base. - - Marking it only on a single override (SaveHiddenStates) drops it for every - other arm at runtime even though the doc-gen union-collapse advertises it. - Mark it on DecodingBaseConfig so all arms capture it. - """ + """Inherited safe fields remain capturable on concrete model arms.""" from tensorrt_llm.llmapi.llm_args import MTPDecodingConfig mtp = MTPDecodingConfig(num_nextn_predict_layers=1) @@ -755,21 +743,67 @@ def test_collect_llm_api_config_captures_max_total_draft_tokens_for_every_arm(): def test_collect_llm_api_config_captures_sparse_algorithm_for_every_arm(): - """The sparse-attention discriminator algorithm is captured for every arm. - - algorithm identifies the active sparse algorithm (rocket, dsa, skip_softmax). - Marking it on only one arm drops it for the others at runtime even though the - doc-gen union-collapse advertises sparse_attention_config.algorithm. - """ + """Every reachable sparse algorithm Literal is captured from its active model arm.""" from tensorrt_llm.llmapi.llm_args import ( DeepSeekSparseAttentionConfig, + DeepSeekV4SparseAttentionConfig, + MiniMaxM3SparseAttentionConfig, + QSASparseAttentionConfig, RocketSparseAttentionConfig, SkipSoftmaxAttentionConfig, ) + from tensorrt_llm.usage.llmapi_config import build_capture_manifest - assert _walk_captured_keys(RocketSparseAttentionConfig()) >= {"algorithm"} - assert _walk_captured_keys(DeepSeekSparseAttentionConfig()) >= {"algorithm"} - assert _walk_captured_keys(SkipSoftmaxAttentionConfig()) >= {"algorithm"} + configs = ( + QSASparseAttentionConfig(), + DeepSeekSparseAttentionConfig(), + DeepSeekV4SparseAttentionConfig(), + RocketSparseAttentionConfig(), + SkipSoftmaxAttentionConfig(), + MiniMaxM3SparseAttentionConfig(), + ) + expected = {config.algorithm for config in configs} + entry = next( + item + for item in build_capture_manifest(TorchLlmArgs) + if item.path == "sparse_attention_config.algorithm" + ) + assert set(entry.allowed_values) == expected + assert entry.capture_types == ("literal",) + + for sparse_config in configs: + args = TorchLlmArgs( + model="/customer/private/Llama", + skip_tokenizer_init=True, + sparse_attention_config=sparse_config, + ) + captured, meta = _loads_payloads(args) + assert captured["sparse_attention_config.algorithm"] == sparse_config.algorithm + assert meta["unsafe_excluded"] is False + + +@pytest.mark.parametrize("field_name", ["target_sparsity", "threshold_scale_factor"]) +def test_sparse_scalar_union_branches_capture_float_but_not_mapping(field_name): + from tensorrt_llm.llmapi.llm_args import SkipSoftmaxAttentionConfig + + def capture(value): + sparse_config = SkipSoftmaxAttentionConfig(**{field_name: value}) + args = TorchLlmArgs( + model="/customer/private/Llama", + skip_tokenizer_init=True, + sparse_attention_config=sparse_config, + ) + return _loads_payloads(args) + + path = f"sparse_attention_config.{field_name}" + config, meta = capture(0.5) + assert config[path] == 0.5 + assert type(config[path]) is float + assert meta["unsafe_excluded"] is False + + config, meta = capture({"decode": 0.5}) + assert path not in config + assert meta["unsafe_excluded"] is True def test_background_reporter_keeps_initial_report_when_config_capture_fails( @@ -815,16 +849,6 @@ def raise_capture_error(_): assert params["featuresJson"] -def test_field_wrapper_records_explicit_exclude_marker(): - from tensorrt_llm.usage.llmapi_config import _get_telemetry_metadata - - class _C(StrictBaseModel): - a: int = Field(default=1, telemetry=False) - - meta = _get_telemetry_metadata(_C.model_fields["a"]) - assert meta == {"exclude": True} - - def test_collect_llm_api_config_honors_explicit_exclude_sentinel(): class _ExcludeConfig(StrictBaseModel): kept: int = Field(default=1) @@ -836,6 +860,21 @@ class _ExcludeConfig(StrictBaseModel): assert meta["capturable_field_count"] == 1 +def test_unknown_subclass_cannot_inherit_a_base_capture_policy(): + class _BaseArm(StrictBaseModel): + sensitive: int = 1 + + class _ChildArm(_BaseArm): + sensitive: int = Field(default=2, telemetry=False) + + class _Root(StrictBaseModel): + arm: _BaseArm + + config, meta = _loads_payloads(_Root(arm=_ChildArm())) + assert "arm.sensitive" not in config + assert meta["capturable_field_count"] == 1 + + def test_collect_llm_api_config_honors_raw_json_schema_extra_exclude(): # Cross-module models use bare pydantic Field with json_schema_extra={"telemetry": ...}. # A raw {"telemetry": False} must be honored as an exclude, like the wrapper telemetry=False. @@ -850,15 +889,6 @@ class _RawExcludeConfig(StrictBaseModel): assert config == {"kept": 1} -def test_runtime_keys_are_subset_of_manifest_for_fixture(): - from tensorrt_llm.usage.llmapi_config import build_capture_manifest - - inst = _ExampleConfig() - manifest_paths = {e.path for e in build_capture_manifest(_ExampleConfig)} - config, _ = _loads_payloads(inst) - assert set(config) <= manifest_paths - - def test_manifest_excludes_loosely_typed_model_children(): # B-1 regression: moe_config.load_balancer is Optional[Union[object, str]]; # a validator coerces it into a MoeLoadBalancerConfig at runtime, but the diff --git a/tests/unittest/usage/test_llmapi_config_telemetry_docs.py b/tests/unittest/usage/test_llmapi_config_telemetry_docs.py index 025f362ddf54..f21a058f6d11 100644 --- a/tests/unittest/usage/test_llmapi_config_telemetry_docs.py +++ b/tests/unittest/usage/test_llmapi_config_telemetry_docs.py @@ -68,9 +68,7 @@ def _sample_manifest() -> dict[str, list[dict[str, object]]]: { "path": "flag", "kind": "value", - "converter": "", - "annotation": "", - "allowed_values": [], + "capture_policy": "bool", } ], } @@ -288,8 +286,34 @@ def test_build_capture_manifest_matches_committed_golden(): _assert_committed_manifest_current(golden_manifest()) +def test_golden_manifest_uses_compact_semantic_policies(): + from tensorrt_llm.usage.llmapi_config import golden_manifest + + by_path = {row["path"]: row for row in golden_manifest()["TorchLlmArgs"]} + bool_row = by_path["enable_chunked_prefill"] + enum_row = by_path["prefill_cuda_graph_backend"] + literal_row = by_path["kv_cache_config.mamba_ssm_cache_dtype"] + + assert bool_row == { + "path": "enable_chunked_prefill", + "kind": "value", + "capture_policy": "bool", + } + assert literal_row["capture_policy"] == "literal" + assert literal_row["allowed_values"] == [ + "auto", + "float16", + "bfloat16", + "float32", + ] + assert enum_row["capture_policy"] == "enum[PrefillCudaGraphBackend]" + assert enum_row["allowed_values"] == ["disabled", "piecewise", "breakable"] + assert "annotation" not in literal_row + assert "converter" not in literal_row + + def test_kv_cache_compression_discriminator_captures_both_algorithms() -> None: - """The shared allowlist captures either compression discriminator.""" + """Each arm's Literal composes without duplicated telemetry metadata.""" from tensorrt_llm.llmapi.llm_args import ( ColdPageQuantizationCompressionConfig, TorchLlmArgs, @@ -305,20 +329,15 @@ def test_kv_cache_compression_discriminator_captures_both_algorithms() -> None: for item in build_capture_manifest(TorchLlmArgs) if item.path == "kv_cache_compression_config.algorithm" ) - assert repr(entry.annotation) == "typing.Literal['triattention']" - assert entry.converter == "allowlist" + assert entry.capture_types == ("literal",) assert set(entry.allowed_values) == { "quantization_for_cold_page", "triattention", } cold_field = ColdPageQuantizationCompressionConfig.model_fields["algorithm"] - assert cold_field.json_schema_extra["telemetry"] == {"exclude": True} + assert not cold_field.json_schema_extra tri_field = TriAttentionKvCacheCompressionConfig.model_fields["algorithm"] - assert tri_field.json_schema_extra["telemetry"] == { - "kind": "categorical", - "converter": "allowlist", - "allowed_values": ["quantization_for_cold_page", "triattention"], - } + assert not tri_field.json_schema_extra private_paths = ( "/private/modelopt-scales", @@ -349,67 +368,6 @@ def test_kv_cache_compression_discriminator_captures_both_algorithms() -> None: assert metadata["capture_succeeded"] is True -def test_load_generator_does_not_leak_sys_modules(): - """_load_generator must not leak its temporary module into sys.modules. - - The loader needs the module registered while exec_module runs (frozen - dataclasses resolve their module via sys.modules), but it must restore the - prior state on both success and failure. - """ - name = "llmapi_config_telemetry" - sys.modules.pop(name, None) - - _load_generator() - assert name not in sys.modules - - import importlib.util as _util - - real_spec_from_file_location = _util.spec_from_file_location - - def _boom(*args, **kwargs): - spec = real_spec_from_file_location(*args, **kwargs) - - class _BoomLoader: - name = spec.name - - def create_module(self, spec): - return None - - def exec_module(self, module): - raise RuntimeError("synthetic load failure") - - spec.loader = _BoomLoader() - return spec - - _util.spec_from_file_location = _boom - try: - try: - _load_generator() - except RuntimeError: - pass - assert name not in sys.modules - finally: - _util.spec_from_file_location = real_spec_from_file_location - - -def test_domain_values_cover_literal_and_enum(): - from enum import Enum - from typing import Literal, Optional - - from tensorrt_llm.usage import llmapi_config as rc - - class _Color(str, Enum): - RED = "red" - BLUE = "blue" - - assert rc._domain_values(Optional[Literal["a", "b"]], {}) == ["a", "b"] - assert rc._domain_values(Optional[_Color], {}) == ["red", "blue"] - assert rc._domain_values(int, {"converter": "allowlist", "allowed_values": ["x", "y"]}) == [ - "x", - "y", - ] - - def _small_models(): from enum import Enum @@ -434,8 +392,6 @@ class Root(BaseModel): allow: str = Field( default="a", telemetry={ - "kind": "categorical", - "converter": "allowlist", "allowed_values": ["a", "b"], }, ) @@ -466,7 +422,7 @@ def test_build_capture_manifest_kinds_and_domains(): assert by_path["mode"].kind == "categorical" assert list(by_path["mode"].allowed_values) == ["a", "b"] # Enum domain assert by_path["allow"].kind == "categorical" - assert by_path["allow"].converter == "allowlist" + assert by_path["allow"].capture_types == ("allowlist",) assert list(by_path["allow"].allowed_values) == ["a", "b"] @@ -478,6 +434,9 @@ def test_renderer_emits_table_from_committed_golden(tmp_path): assert "## LLM API Configuration Fields" in text assert "python3 scripts/generate_llm_args_golden_manifest.py" in text assert "explicitly marked" not in text # opt-in prose must be gone + assert "Capture policy" in text + assert "Categorical domain" in text + assert "Converter" not in text assert "`backend`" in text # a known captured key renders @@ -504,6 +463,41 @@ class Root(BaseModel): build_capture_manifest(Root) +def test_shared_path_policies_are_scoped_to_the_active_union_arm(): + from typing import Literal, Union + + from pydantic import BaseModel + + from tensorrt_llm.usage.llmapi_config import ( + build_capture_manifest, + collect_llm_api_config_payloads, + ) + + class ArmA(BaseModel): + tag: Literal["a"] = "a" + shared: Literal["a-only"] = "a-only" + + class ArmB(BaseModel): + tag: Literal["b"] = "b" + shared: Literal["b-only"] = "b-only" + + class Root(BaseModel): + arm: Union[ArmA, ArmB] = ArmA() + + entry = next(item for item in build_capture_manifest(Root) if item.path == "arm.shared") + assert entry.allowed_values == ("a-only", "b-only") + + arm_b = Root(arm=ArmB()) + config_json, _ = collect_llm_api_config_payloads(arm_b) + assert json.loads(config_json)["arm.shared"] == "b-only" + + invalid_arm_a = Root() + invalid_arm_a.arm.shared = "b-only" + config_json, metadata_json = collect_llm_api_config_payloads(invalid_arm_a) + assert "arm.shared" not in json.loads(config_json) + assert json.loads(metadata_json)["unsafe_excluded"] is True + + def test_build_capture_manifest_cycle_guard_terminates_on_self_reference(): from typing import Optional From c10d6efb7078d57404458928552ad03bf6a585b8 Mon Sep 17 00:00:00 2001 From: TensorRT LLM <90828364+tensorrt-cicd@users.noreply.github.com> Date: Tue, 15 Sep 2026 03:19:26 +0000 Subject: [PATCH 7/8] [None][infra] Check in most recent lock file from nightly pipeline Signed-off-by: TensorRT LLM <90828364+tensorrt-cicd@users.noreply.github.com> --- security_scanning/examples/apps/poetry.lock | 24 +- .../examples/ray_orchestrator/poetry.lock | 122 +++--- security_scanning/metadata.json | 4 +- security_scanning/poetry.lock | 402 ++++++++++-------- security_scanning/pyproject.toml | 8 +- 5 files changed, 313 insertions(+), 247 deletions(-) diff --git a/security_scanning/examples/apps/poetry.lock b/security_scanning/examples/apps/poetry.lock index 0ed92427d406..ce44505ef153 100644 --- a/security_scanning/examples/apps/poetry.lock +++ b/security_scanning/examples/apps/poetry.lock @@ -78,15 +78,15 @@ files = [ [[package]] name = "httpcore2" -version = "2.12.0" +version = "2.13.0" description = "A minimal low-level HTTP client." optional = false python-versions = ">=3.10" groups = ["main"] markers = "sys_platform != \"emscripten\"" files = [ - {file = "httpcore2-2.12.0-py3-none-any.whl", hash = "sha256:7e04258ce01013d7d615e5b910a3b27fac937d7a95038227e79652b4ba3b4ceb"}, - {file = "httpcore2-2.12.0.tar.gz", hash = "sha256:9293522bba0aa7c4c8e9e3f040c16575bd8868e155a77fa30c7a9085a5eae648"}, + {file = "httpcore2-2.13.0-py3-none-any.whl", hash = "sha256:35ae5be347aa40467b4a5dc032ac67ebb6d27189fc97e8cebcf99616f6a1bb9e"}, + {file = "httpcore2-2.13.0.tar.gz", hash = "sha256:2adc8be4fb285fbcd6d894298db3b52c177e74b6674eda3a76bd36be3292a3db"}, ] [package.dependencies] @@ -97,30 +97,30 @@ truststore = ">=0.10" asyncio = ["anyio (>=4.5.0,<5.0)"] http2 = ["h2 (>=3,<5)"] socks = ["socksio (==1.*)"] -trio = ["trio (>=0.33.0,<1.0)"] +trio = ["trio (>=0.34.0,<1.0)"] [[package]] name = "httpx2" -version = "2.12.0" +version = "2.13.0" description = "The next generation HTTP client." optional = false python-versions = ">=3.10" groups = ["main"] files = [ - {file = "httpx2-2.12.0-py3-none-any.whl", hash = "sha256:cc8b6eecb8661c146b8f89a60e97456ee086e91a784ed31ac450c3a9e613dd36"}, - {file = "httpx2-2.12.0.tar.gz", hash = "sha256:7631fe9887a8a2275f4a2540e053aa670fcc50742864a9ae7c66e609fdcf12cf"}, + {file = "httpx2-2.13.0-py3-none-any.whl", hash = "sha256:fc12720cedf72faa26cca6b4ca394e05c894e7d7933fc45cafe767960804e49a"}, + {file = "httpx2-2.13.0.tar.gz", hash = "sha256:81bd07dc67a3701729ef1f777a3c00c915d4539604fdb5afd327f8682f6b7b44"}, ] [package.dependencies] anyio = {version = ">=4.10", markers = "sys_platform != \"emscripten\""} -httpcore2 = {version = "2.12.0", markers = "sys_platform != \"emscripten\""} +httpcore2 = {version = "2.13.0", markers = "sys_platform != \"emscripten\""} httpx2-jsfetch = {version = "*", markers = "sys_platform == \"emscripten\" and python_version >= \"3.12\""} idna = ">=3.18" truststore = {version = ">=0.10", markers = "sys_platform != \"emscripten\""} typing-extensions = {version = ">=4.5.0", markers = "python_version < \"3.13\""} [package.extras] -brotli = ["brotli (>=1.2.0) ; platform_python_implementation == \"CPython\"", "brotlicffi (>=1.2.0.0) ; platform_python_implementation != \"CPython\""] +brotli = ["brotli (>=1.2.0) ; platform_python_implementation == \"CPython\"", "brotlicffi (>=1.2.0.2) ; platform_python_implementation != \"CPython\""] cli = ["click (>=8.4.2)", "pygments (==2.*)", "rich (>=10,<16)"] http2 = ["h2 (>=3,<5)"] socks = ["socksio (==1.*)"] @@ -286,14 +286,14 @@ files = [ [[package]] name = "openai" -version = "3.13.0" +version = "3.14.0" description = "The official Python library for the openai API" optional = false python-versions = ">=3.10" groups = ["main"] files = [ - {file = "openai-3.13.0-py3-none-any.whl", hash = "sha256:e35b1f6fe99245e86e37504d9fad1ad2a363807307c424232c1d849bd0666c8e"}, - {file = "openai-3.13.0.tar.gz", hash = "sha256:a8f87a9b3b9c08eb446d68bd0a80e8ec907c4c35fdea63f4265c7b34b2de3a60"}, + {file = "openai-3.14.0-py3-none-any.whl", hash = "sha256:232a85a1c0ff6820534630fbfdb99313990ed80e66f39892a3abcff4130e2b16"}, + {file = "openai-3.14.0.tar.gz", hash = "sha256:714ba70b91ee1e78e75263694bebfeaca004432032685333a46d0e0206114bf9"}, ] [package.dependencies] diff --git a/security_scanning/examples/ray_orchestrator/poetry.lock b/security_scanning/examples/ray_orchestrator/poetry.lock index 9de23a50a448..f02d31cbc8db 100644 --- a/security_scanning/examples/ray_orchestrator/poetry.lock +++ b/security_scanning/examples/ray_orchestrator/poetry.lock @@ -800,14 +800,14 @@ files = [ [[package]] name = "google-api-core" -version = "2.36.0" +version = "2.37.0" description = "Google API client core library" optional = false python-versions = ">=3.10" groups = ["main"] files = [ - {file = "google_api_core-2.36.0-py3-none-any.whl", hash = "sha256:e4d0b179260727ea5c42222426d9199285214dbef7bf48f8b16600c7f9a78944"}, - {file = "google_api_core-2.36.0.tar.gz", hash = "sha256:32779307b52e64c9a9592a3621de6281676ecaeea299fe8524e4637ab7ac2531"}, + {file = "google_api_core-2.37.0-py3-none-any.whl", hash = "sha256:d84042a0034cce9c4304e17d2e5e9966c2fcb1a851b73799291f687a67470c8b"}, + {file = "google_api_core-2.37.0.tar.gz", hash = "sha256:cf58f220aa797f1ffdda52194c4c7d72efeced09297c2976529f8135f6d85b9a"}, ] [package.dependencies] @@ -873,70 +873,80 @@ grpc = ["grpcio (>=1.59.0,<2.0.0)"] [[package]] name = "grpcio" -version = "1.83.1" +version = "1.84.0" description = "HTTP/2-based RPC framework" optional = false python-versions = ">=3.10" groups = ["main"] files = [ - {file = "grpcio-1.83.1-cp310-cp310-linux_armv7l.whl", hash = "sha256:1fea1ae4795d4790579995a4dd5e20e7494d358e29a340e8368dab9723264328"}, - {file = "grpcio-1.83.1-cp310-cp310-macosx_11_0_universal2.whl", hash = "sha256:b7ace1f740b36fcd451a1bb96f71ee7650e60b308822baeb66a023965bc27f4b"}, - {file = "grpcio-1.83.1-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:2a141f7bfc1601a0942405a8af6334ab21ba1dd0fa49b8427686df7beebd374d"}, - {file = "grpcio-1.83.1-cp310-cp310-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:c7e9e19413d43077d5a5c77b02ff82610209088e8f98da929347bc03d4c848d1"}, - {file = "grpcio-1.83.1-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:b59eaaeeb03dde0a2708095fb50f1afa94f11dc1b459bb7790b53bfb8cf95153"}, - {file = "grpcio-1.83.1-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:4e7c1468cf37cca17ab18bc8072901eed8daeb81685589ccd07988e5a750ee67"}, - {file = "grpcio-1.83.1-cp310-cp310-musllinux_1_2_i686.whl", hash = "sha256:a4a87dc86b0393257a11eb11e911c4c3456cbacd1c1ab9e9441060d9a3ad126b"}, - {file = "grpcio-1.83.1-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:d0dda8af248f6971555e1d4425f64864ce4e7369c5f8ef57c3e82a9bef77e22f"}, - {file = "grpcio-1.83.1-cp310-cp310-win32.whl", hash = "sha256:0f736f8359cf7cb8d0914a290999765a4342b0c35f01adc6e3ba24598f9d62b7"}, - {file = "grpcio-1.83.1-cp310-cp310-win_amd64.whl", hash = "sha256:7d43e3bd2b7d749c2dbd41c2cc83d550c3343d299a19acbbba9e37ad8c11fa8e"}, - {file = "grpcio-1.83.1-cp311-cp311-linux_armv7l.whl", hash = "sha256:907a5e5afb31f7a46376afc1a1edddd7afa00a74bbbc5b78979bbc34479581f6"}, - {file = "grpcio-1.83.1-cp311-cp311-macosx_11_0_universal2.whl", hash = "sha256:547645f02499c972f3edec9be4db9997f1d03df307c1c199772342ed6d8b3c6d"}, - {file = "grpcio-1.83.1-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:34f1841fc6d1d76f8a2d74177eafa2d1ec7d7e039633488c9fcc1b375a1fc165"}, - {file = "grpcio-1.83.1-cp311-cp311-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:05ba265193fbd9f63355311ec7567bba32a72aeb8e9fd7b3443e4fcad87b0750"}, - {file = "grpcio-1.83.1-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:5cce1d9fe2887239f054dc9c314597e04f33d2e6bd3150a91c4946d7e5be5d98"}, - {file = "grpcio-1.83.1-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:f732feb060ef57c1a040c24cee072ba9fab99bd0a7d2c916ef3f1c4d84b98974"}, - {file = "grpcio-1.83.1-cp311-cp311-musllinux_1_2_i686.whl", hash = "sha256:145b0050d24eb38accd9dc7ae09a3c09b8e7330159f3cfb46b1dba8711d50c42"}, - {file = "grpcio-1.83.1-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:e844cdb25c3c93c7572e0a37137c12305efea493be4eb65801b3ee93f180c186"}, - {file = "grpcio-1.83.1-cp311-cp311-win32.whl", hash = "sha256:0d07661944477517b12a239e18720c8d9038f80a62f2c56260fae80327f43d2a"}, - {file = "grpcio-1.83.1-cp311-cp311-win_amd64.whl", hash = "sha256:e572da3e247b28a98f46636d33c756e81ffb0f5def96c231ba45332333060595"}, - {file = "grpcio-1.83.1-cp312-cp312-linux_armv7l.whl", hash = "sha256:5acd14c6ddf047de62cbf8745b11103ea91abbf57d1b8edd5395ccd9fcd13abb"}, - {file = "grpcio-1.83.1-cp312-cp312-macosx_11_0_universal2.whl", hash = "sha256:16138031a47b771860a16a975b53087f4fd5bbdbb2c03a188c5d90ad65d2bdae"}, - {file = "grpcio-1.83.1-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:5ccc26715fd4defca5e129e280dd883b1737b65045ec50ffe22ce42104089519"}, - {file = "grpcio-1.83.1-cp312-cp312-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:b74f2a1d9ab1dfa3e263ef33d581613679b78d0884babf11671af26e45570ead"}, - {file = "grpcio-1.83.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:72578aa07a4008f17521ef52debcc3acfd1e2c5426243bc3ffb56a38bfe610b7"}, - {file = "grpcio-1.83.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:c12e1fc59c6dc26d10d9144453ddc6cbfe4cd4c31e874ed2d0132f88e685eb8b"}, - {file = "grpcio-1.83.1-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:4910b62f7d12197160bfb7de06d876d64dd12d43483e8292f98f49ca09b628d9"}, - {file = "grpcio-1.83.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:9e703effe3ae779925c82ac24fdb82cf4105e1096810151ed9501c5f34546b9c"}, - {file = "grpcio-1.83.1-cp312-cp312-win32.whl", hash = "sha256:a2aea8bd6e0a34f12cbaddb7bb70bec836818789fa5c7ab7572c6b745396a2d4"}, - {file = "grpcio-1.83.1-cp312-cp312-win_amd64.whl", hash = "sha256:583bf2e8255040a4a312f9572dfe62a05271437b149550e1a536d5c47d2d1e8a"}, - {file = "grpcio-1.83.1-cp313-cp313-linux_armv7l.whl", hash = "sha256:8d228e253b77865efcbdd7b5894ca882c9e0ea98c02b7d20582e61ded8dfd4b5"}, - {file = "grpcio-1.83.1-cp313-cp313-macosx_11_0_universal2.whl", hash = "sha256:0468b627f2987c9a77f7580030207cbd85457ffe52998beff4f0b5c38c58a72c"}, - {file = "grpcio-1.83.1-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:a6a282e81530cead60bbd752cc04950a57f224379e9821495d6a35bd5ce9b1f4"}, - {file = "grpcio-1.83.1-cp313-cp313-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:947d945f52e8ecf3cafd2bb7113502a16ccfda3e12c854443094de32d83ad432"}, - {file = "grpcio-1.83.1-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:55656318d5dd387077396dffb929171ca3966e24bfead9a6c5dba9f889062cb4"}, - {file = "grpcio-1.83.1-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:9daf5acf4fc9d5f5627229969c2580a91e511779d76e4ccdeb9f4770f05d8bc2"}, - {file = "grpcio-1.83.1-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:7b94174cbca93316888f805efbeb08f1c020f7b7493d2d50cc4f6b64ebb7e8bd"}, - {file = "grpcio-1.83.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:65c5a7210911ffe0f67b1cdc5308f9854b6d1f1b345e3e49ab7cac1ba50fa346"}, - {file = "grpcio-1.83.1-cp313-cp313-win32.whl", hash = "sha256:179368d9361854616ce6f397d4716e07480129652752fcbcfc5a7260455ad6f2"}, - {file = "grpcio-1.83.1-cp313-cp313-win_amd64.whl", hash = "sha256:2e57af456385491a76e13c4aada8c8f43a8e47051e06ea97a9dbe2a49654e6db"}, - {file = "grpcio-1.83.1-cp314-cp314-linux_armv7l.whl", hash = "sha256:8b3c87ca908296bf125f841d3e1a2225a2b39aaa8ed7a57e7ccde465ee519bab"}, - {file = "grpcio-1.83.1-cp314-cp314-macosx_11_0_universal2.whl", hash = "sha256:c0f3f20c90e72a171917ae65706500b096a1c3eb5f162c3ce702a2e25635f132"}, - {file = "grpcio-1.83.1-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:81bbf35a46bf8cad2dfbb2eccc19c711befb58b288acb534bbcd0d74283202a6"}, - {file = "grpcio-1.83.1-cp314-cp314-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:215cec07d11176507387bda4bf2751816e880f9bff8dc1ca524bfbb8ed8f2fad"}, - {file = "grpcio-1.83.1-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:abce7d43ec29cd39230fa8339de1a07643b55adc412a454850fbd875349950ff"}, - {file = "grpcio-1.83.1-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:e256f95a40e3b0183a98556fb7164d24b97eeb353123ccabfcba94712b35ee2a"}, - {file = "grpcio-1.83.1-cp314-cp314-musllinux_1_2_i686.whl", hash = "sha256:2110059146fb0ea216e1ffddb29377b5cc2fd412a5b0a92e102616bd5edf18c2"}, - {file = "grpcio-1.83.1-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:20d944d967843f8183f9f23d5916388362e5f8eeeae855bbe4354d906dc9f31b"}, - {file = "grpcio-1.83.1-cp314-cp314-win32.whl", hash = "sha256:623c87c6d4a1cb30d82c4e896f95477050f2e01b4a1f8cf91ff2b1abdf89c457"}, - {file = "grpcio-1.83.1-cp314-cp314-win_amd64.whl", hash = "sha256:47e6934ad38779271e2e7cc5f78a63a407cf3d98114c65c1fdbcd3f5a716f29b"}, - {file = "grpcio-1.83.1.tar.gz", hash = "sha256:9cee6fcbf2eb57c4b49451787bfa87be8efc1ca02a0b327dd4b54d44502e362b"}, + {file = "grpcio-1.84.0-cp310-cp310-linux_armv7l.whl", hash = "sha256:71fd60e6e426d293d0a2f685115ad0a0845117602cf13605a4be7524fb5f7bba"}, + {file = "grpcio-1.84.0-cp310-cp310-macosx_11_0_universal2.whl", hash = "sha256:8e1a45d174b6b8589f51dce1cea804aa6c1f72c9c80cba91ae2caabeb6d90540"}, + {file = "grpcio-1.84.0-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:efb29f8633bf6630dc89de4fe0353ac3d7e4b70ef7b6e29fb40f00e68c127fa5"}, + {file = "grpcio-1.84.0-cp310-cp310-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:d0fdd25faece8a1f95e8a3a8006e29701b5cf8dadb4a8132e68f3134637004a5"}, + {file = "grpcio-1.84.0-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:393d8a78bff6731ecc5ad2151a821f8fbc1709b137ebb9c25a4ef399fbdcc914"}, + {file = "grpcio-1.84.0-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:fc66cb50c93554b86db0b6625ab5c6e9051dbf8847c08d93c84918e02e413fb7"}, + {file = "grpcio-1.84.0-cp310-cp310-musllinux_1_2_i686.whl", hash = "sha256:455ed6083353b8e938f1d58c765eab2fbb165731e5b507be30fee344915a2a11"}, + {file = "grpcio-1.84.0-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:3d6a82c4fc6c85f2fb7572c86bdb86f84c97b6580e5f6599f711800bac48a5d8"}, + {file = "grpcio-1.84.0-cp310-cp310-win32.whl", hash = "sha256:8e3f508d0e9e6236ba2f08d56e33355e434e785e813149a1b8477d3edf69779d"}, + {file = "grpcio-1.84.0-cp310-cp310-win_amd64.whl", hash = "sha256:ed2c1493c44d0932f1e55fdb5d1ead658c68288ec5d51b8c4928422d98633ef9"}, + {file = "grpcio-1.84.0-cp311-cp311-linux_armv7l.whl", hash = "sha256:4aaeceeb7fa7d824c322d1ec3208c8495c88478a927295553235435fc49043ad"}, + {file = "grpcio-1.84.0-cp311-cp311-macosx_11_0_universal2.whl", hash = "sha256:06619ba1515e5ee69fb2a514e95dd8be05ce74cb3928d5b34f87f87c86fe3c27"}, + {file = "grpcio-1.84.0-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:158c1c11cfb61b4849c3caf4d52de6f5ecd376e14446feb4a90dc95a90d616f5"}, + {file = "grpcio-1.84.0-cp311-cp311-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:a9383401d9f116f98cacd4eba6c505a6edb80ba65badfc8e8ed8ae64983bcc44"}, + {file = "grpcio-1.84.0-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:bd8ea8eb3817b226057cc1c0e7ec4b378dcda52043b972b6ff12b1152178967d"}, + {file = "grpcio-1.84.0-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:756ea5c2da00fa65c930284892d2a9706828704ca3ba40b4c51c4834eb39fcfd"}, + {file = "grpcio-1.84.0-cp311-cp311-musllinux_1_2_i686.whl", hash = "sha256:28d2609691da93051e998495108bbddd2a9f7a561253bae94828d81290f30c15"}, + {file = "grpcio-1.84.0-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:27b8b36200a9fbee6e120246f4a8a41657549107ef19fb2c819c4b2fd524f39a"}, + {file = "grpcio-1.84.0-cp311-cp311-win32.whl", hash = "sha256:465eef3d17e59ad22a556fc0138f7c7c799df426734344daec42c797d49fda99"}, + {file = "grpcio-1.84.0-cp311-cp311-win_amd64.whl", hash = "sha256:f9a456bdbed52a01c9ab8423bdebab04a5363c78676edc55ab9b58bd13bdf9e1"}, + {file = "grpcio-1.84.0-cp312-cp312-linux_armv7l.whl", hash = "sha256:b5c6f20d657ae09ae4e30d9d3a21edd13f1219d58cc6f999b9d1bb63be9c1baa"}, + {file = "grpcio-1.84.0-cp312-cp312-macosx_11_0_universal2.whl", hash = "sha256:406583b4e8fb2282ebd392e12b963e601c1f82e07125a8c2cb5b144e7e024796"}, + {file = "grpcio-1.84.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:fbdbcd06986ede3ce584083b1dc2afe6808e8943e5cf50ad11183c03aceda25a"}, + {file = "grpcio-1.84.0-cp312-cp312-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:23e6e8e8a75cff88e0a793bfd3becea03a13e2763ae90c1ff573bc19ca5b429a"}, + {file = "grpcio-1.84.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:b44f0a0fc7bc6677d38cc80bca1a32814ce6c8f200fb8b3c1a61c9d77eaefbf3"}, + {file = "grpcio-1.84.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:210e4c32f907045eb8158273e60c6ab69a3947697df6245dbda381f26c59485b"}, + {file = "grpcio-1.84.0-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:a71d24f40b0cc6798feaa978c7411dc1135b7018e9fc0442db611c139bf58344"}, + {file = "grpcio-1.84.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:f6c972474ce691aca74e58d17625450cef153dc4760364cadeb167983ea6d589"}, + {file = "grpcio-1.84.0-cp312-cp312-win32.whl", hash = "sha256:0d532ade4486dad9b302ffa4d4683d67561051c26d17c4023322845e9fa10140"}, + {file = "grpcio-1.84.0-cp312-cp312-win_amd64.whl", hash = "sha256:49717e857899f4136d7657bf5aded61ac479110a075438290923a4d86af7cd02"}, + {file = "grpcio-1.84.0-cp313-cp313-linux_armv7l.whl", hash = "sha256:209414080da8c20af94df1395b635da52dd57b5edc9e917e1deca0dc1c4bb55e"}, + {file = "grpcio-1.84.0-cp313-cp313-macosx_11_0_universal2.whl", hash = "sha256:e41c3993eee896c617dbd8a505085d28b6e84a0445ed9a1f40f95808473cf678"}, + {file = "grpcio-1.84.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:fff5ef3fe1bba7d6147e5f19e01e5e122ac2c076486887ddcb8d42e663400fbe"}, + {file = "grpcio-1.84.0-cp313-cp313-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:b8c62888c3e49debf37ad9773e3c02f77b0c1e811f8fb0962f2b6c3bbab5b97a"}, + {file = "grpcio-1.84.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:986e9751d416d7a6eaa2fecdac38da63153d63a4b340ba7d624889c490451500"}, + {file = "grpcio-1.84.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:5933a052946873d01a42119a05420d669bdca436aeba2d1851988ccb12b421c0"}, + {file = "grpcio-1.84.0-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:e094dd21f077af8194923fc263cad872eaa1802bb0156fd7e5ae18e99cd86715"}, + {file = "grpcio-1.84.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:08735e3d08d24ab3132cf87e2e5dea8746cabcc7d676c2b0b7362f195feef9d9"}, + {file = "grpcio-1.84.0-cp313-cp313-win32.whl", hash = "sha256:70bb4ce8be0c5606bec259cbd7152374470396413b7863a658a08c849e6b29ff"}, + {file = "grpcio-1.84.0-cp313-cp313-win_amd64.whl", hash = "sha256:b61692f0069b3eee2fc8a3a1b7f6c044df9e03fede6ce69b3ca832e1c39f26c5"}, + {file = "grpcio-1.84.0-cp314-cp314-linux_armv7l.whl", hash = "sha256:026d757df86c5b7a41de8200b9a2cda454aaa5004cb0c7e3374c66eb82f61499"}, + {file = "grpcio-1.84.0-cp314-cp314-macosx_11_0_universal2.whl", hash = "sha256:3de427b05f244ba2c2a9bdc67e7a6731c8340811524ecc4435466549f8af1d17"}, + {file = "grpcio-1.84.0-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:e90e3bdf7b5eac005fef631adae9cafde16f922def207b80a7c46b253c18ad20"}, + {file = "grpcio-1.84.0-cp314-cp314-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:e88d304f094f4937bc27ec6a435e218a084168f11ec630c8d5d39b431d08d81d"}, + {file = "grpcio-1.84.0-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:57dc36a5ab0e676f5f6e171de2917fd0aef73f32a9aaf23956bfe19997a30bd1"}, + {file = "grpcio-1.84.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:5deda5b4bf62769eb98c119cca43d40e1231e34846b19db5cdea821d446a2253"}, + {file = "grpcio-1.84.0-cp314-cp314-musllinux_1_2_i686.whl", hash = "sha256:9bab4cf571653a8afffb83ce21aa27b51dfe629b526b7b6adec35491fe1fc2ea"}, + {file = "grpcio-1.84.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:c5559b492007dc09b4de9b95dab05f0b5e53547aad230cf07e46c7dd017a3be5"}, + {file = "grpcio-1.84.0-cp314-cp314-win32.whl", hash = "sha256:2c024da73b296f040b8360e60bd73a659b230093684a438da0e1260f34cc724e"}, + {file = "grpcio-1.84.0-cp314-cp314-win_amd64.whl", hash = "sha256:800b7e00d92553313c0463c200087930aa78678ec1d528193aeb50906f55989b"}, + {file = "grpcio-1.84.0-cp315-cp315-linux_armv7l.whl", hash = "sha256:47ecf0d9b81d981f07b61bd89eced9d2582f5eaacc3aaa36ad27f81aef70a27f"}, + {file = "grpcio-1.84.0-cp315-cp315-macosx_10_15_universal2.whl", hash = "sha256:61386101ecaa096b694d0dd278caf99a56aeec78440cc17e918eef0b50f2d567"}, + {file = "grpcio-1.84.0-cp315-cp315-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:f6d178ba6dc8e82976c184b65fddde172d054c17237993a3e083efe4f134d55b"}, + {file = "grpcio-1.84.0-cp315-cp315-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:15bb76489e337fc492685c9758e2fd4d4ab516b901ad830dc5a91987decf00be"}, + {file = "grpcio-1.84.0-cp315-cp315-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:82da34ae4f639c73ac46e521e00c0a49bf86f717b9fb1f405f133e98731e38dc"}, + {file = "grpcio-1.84.0-cp315-cp315-musllinux_1_2_aarch64.whl", hash = "sha256:9b73836ba0e16fcbb57c31cf6cbc2907c8d8c790b83679df454b74bd15e0be04"}, + {file = "grpcio-1.84.0-cp315-cp315-musllinux_1_2_i686.whl", hash = "sha256:42959bd50dd660ffc3f2a9bec15a6da4f9aaa0dda555d59ff2d2e80b908456a8"}, + {file = "grpcio-1.84.0-cp315-cp315-musllinux_1_2_x86_64.whl", hash = "sha256:659728f20fc7a0933ed7b1945435e31014b97ab8a5a7edcbaa70da4794aeb191"}, + {file = "grpcio-1.84.0-cp315-cp315-win32.whl", hash = "sha256:edb6f87fc60ff438557291501b3e16c7a77c3b01a52d782cf276dccc7c5dd89c"}, + {file = "grpcio-1.84.0-cp315-cp315-win_amd64.whl", hash = "sha256:4119efa6519871719ad81f33bc95ab87857dcb1c5801f30a6e592f2c41164169"}, + {file = "grpcio-1.84.0.tar.gz", hash = "sha256:19aaf172fc2edbefccce3f6e92c5150975dbe56c45744e9e87cf72ebdf85bfbe"}, ] [package.dependencies] typing-extensions = ">=4.12,<5.0" [package.extras] -protobuf = ["grpcio-tools (>=1.83.1)"] +protobuf = ["grpcio-tools (>=1.84.0)"] [[package]] name = "idna" diff --git a/security_scanning/metadata.json b/security_scanning/metadata.json index 77c01a9c132d..d48da909fd4d 100644 --- a/security_scanning/metadata.json +++ b/security_scanning/metadata.json @@ -1,4 +1,4 @@ { - "commit_hash": "eefca665cabbbbf86e8d1e22bd42036efca6bcca", - "timestamp": "2026-09-14T02:49:19Z" + "commit_hash": "5cb1c9500b77354a340bb10c0f55afca889df2dd", + "timestamp": "2026-09-15T02:49:24Z" } diff --git a/security_scanning/poetry.lock b/security_scanning/poetry.lock index 8ab419a7289b..9f4507294e93 100644 --- a/security_scanning/poetry.lock +++ b/security_scanning/poetry.lock @@ -1124,55 +1124,56 @@ tileiras = ["cuda-toolkit[nvcc,nvvm,tileiras] (>=13.2,<13.5)"] [[package]] name = "cuda-toolkit" -version = "13.0.2" +version = "13.0.3" description = "CUDA Toolkit meta-package" optional = false python-versions = "*" groups = ["main"] markers = "platform_system == \"Linux\"" files = [ - {file = "cuda_toolkit-13.0.2-py2.py3-none-any.whl", hash = "sha256:b198824cf2f54003f50d64ada3a0f184b42ca0846c1c94192fa269ecd97a66eb"}, + {file = "cuda_toolkit-13.0.3.0-py2.py3-none-any.whl", hash = "sha256:d693caaa261214ddd7dbb60d68e71cbed884e68c2be7509778f3051da0b91c3f"}, ] [package.dependencies] -nvidia-cuda-cupti = {version = "==13.0.85.*", optional = true, markers = "(sys_platform == \"linux\" or sys_platform == \"win32\") and extra == \"cupti\""} -nvidia-cuda-nvrtc = {version = "==13.0.88.*", optional = true, markers = "(sys_platform == \"linux\" or sys_platform == \"win32\") and extra == \"nvrtc\""} -nvidia-cuda-runtime = {version = "==13.0.96.*", optional = true, markers = "(sys_platform == \"linux\" or sys_platform == \"win32\") and extra == \"cudart\""} -nvidia-cufft = {version = "==12.0.0.61.*", optional = true, markers = "(sys_platform == \"linux\" or sys_platform == \"win32\") and extra == \"cufft\""} -nvidia-cufile = {version = "==1.15.1.6.*", optional = true, markers = "sys_platform == \"linux\" and extra == \"cufile\""} -nvidia-curand = {version = "==10.4.0.35.*", optional = true, markers = "(sys_platform == \"linux\" or sys_platform == \"win32\") and extra == \"curand\""} -nvidia-cusolver = {version = "==12.0.4.66.*", optional = true, markers = "(sys_platform == \"linux\" or sys_platform == \"win32\") and extra == \"cusolver\""} -nvidia-cusparse = {version = "==12.6.3.3.*", optional = true, markers = "(sys_platform == \"linux\" or sys_platform == \"win32\") and extra == \"cusparse\""} -nvidia-nvjitlink = {version = "==13.0.88.*", optional = true, markers = "(sys_platform == \"linux\" or sys_platform == \"win32\") and extra == \"nvjitlink\""} -nvidia-nvtx = {version = "==13.0.85.*", optional = true, markers = "(sys_platform == \"linux\" or sys_platform == \"win32\") and extra == \"nvtx\""} +nvidia-cublas = {version = "==13.1.1.3.*", optional = true, markers = "(sys_platform == \"linux\" or sys_platform == \"win32\") and (sys_platform == \"linux\" or platform_machine == \"x86_64\") and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") and (extra == \"cublas\" or extra == \"cusolver\")"} +nvidia-cuda-cupti = {version = "==13.0.85.*", optional = true, markers = "(sys_platform == \"linux\" or sys_platform == \"win32\") and (sys_platform == \"linux\" or platform_machine == \"x86_64\") and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") and extra == \"cupti\""} +nvidia-cuda-nvrtc = {version = "==13.0.88.*", optional = true, markers = "(sys_platform == \"linux\" or sys_platform == \"win32\") and (sys_platform == \"linux\" or platform_machine == \"x86_64\") and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") and (extra == \"cublas\" or extra == \"nvrtc\")"} +nvidia-cuda-runtime = {version = "==13.0.96.*", optional = true, markers = "(sys_platform == \"linux\" or sys_platform == \"win32\") and (sys_platform == \"linux\" or platform_machine == \"x86_64\") and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") and extra == \"cudart\""} +nvidia-cufft = {version = "==12.0.0.61.*", optional = true, markers = "(sys_platform == \"linux\" or sys_platform == \"win32\") and (sys_platform == \"linux\" or platform_machine == \"x86_64\") and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") and extra == \"cufft\""} +nvidia-cufile = {version = "==1.15.1.6.*", optional = true, markers = "sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") and extra == \"cufile\""} +nvidia-curand = {version = "==10.4.0.35.*", optional = true, markers = "(sys_platform == \"linux\" or sys_platform == \"win32\") and (sys_platform == \"linux\" or platform_machine == \"x86_64\") and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") and extra == \"curand\""} +nvidia-cusolver = {version = "==12.0.4.66.*", optional = true, markers = "(sys_platform == \"linux\" or sys_platform == \"win32\") and (sys_platform == \"linux\" or platform_machine == \"x86_64\") and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") and extra == \"cusolver\""} +nvidia-cusparse = {version = "==12.6.3.3.*", optional = true, markers = "(sys_platform == \"linux\" or sys_platform == \"win32\") and (sys_platform == \"linux\" or platform_machine == \"x86_64\") and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") and (extra == \"cusolver\" or extra == \"cusparse\")"} +nvidia-nvjitlink = {version = ">=13.0.88,<14", optional = true, markers = "(sys_platform == \"linux\" or sys_platform == \"win32\") and (sys_platform == \"linux\" or platform_machine == \"x86_64\") and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") and (extra == \"cufft\" or extra == \"cusolver\" or extra == \"cusparse\" or extra == \"nvjitlink\")"} +nvidia-nvtx = {version = "==13.0.85.*", optional = true, markers = "(sys_platform == \"linux\" or sys_platform == \"win32\") and (sys_platform == \"linux\" or platform_machine == \"x86_64\") and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") and extra == \"nvtx\""} [package.extras] -all = ["nvidia-cublas (==13.1.0.3.*) ; sys_platform == \"linux\" or sys_platform == \"win32\"", "nvidia-cuda-cccl (==13.0.85.*) ; sys_platform == \"linux\" or sys_platform == \"win32\"", "nvidia-cuda-crt (==13.0.88.*) ; sys_platform == \"linux\" or sys_platform == \"win32\"", "nvidia-cuda-culibos (==13.0.85.*) ; sys_platform == \"linux\"", "nvidia-cuda-cupti (==13.0.85.*) ; sys_platform == \"linux\" or sys_platform == \"win32\"", "nvidia-cuda-cuxxfilt (==13.0.85.*) ; sys_platform == \"linux\" or sys_platform == \"win32\"", "nvidia-cuda-nvcc (==13.0.88.*) ; sys_platform == \"linux\" or sys_platform == \"win32\"", "nvidia-cuda-nvrtc (==13.0.88.*) ; sys_platform == \"linux\" or sys_platform == \"win32\"", "nvidia-cuda-opencl (==13.0.85.*) ; sys_platform == \"linux\" or sys_platform == \"win32\"", "nvidia-cuda-profiler-api (==13.0.85.*) ; sys_platform == \"linux\" or sys_platform == \"win32\"", "nvidia-cuda-runtime (==13.0.96.*) ; sys_platform == \"linux\" or sys_platform == \"win32\"", "nvidia-cuda-sanitizer-api (==13.0.85.*) ; sys_platform == \"linux\" or sys_platform == \"win32\"", "nvidia-cufft (==12.0.0.61.*) ; sys_platform == \"linux\" or sys_platform == \"win32\"", "nvidia-cufile (==1.15.1.6.*) ; sys_platform == \"linux\"", "nvidia-curand (==10.4.0.35.*) ; sys_platform == \"linux\" or sys_platform == \"win32\"", "nvidia-cusolver (==12.0.4.66.*) ; sys_platform == \"linux\" or sys_platform == \"win32\"", "nvidia-cusparse (==12.6.3.3.*) ; sys_platform == \"linux\" or sys_platform == \"win32\"", "nvidia-npp (==13.0.1.2.*) ; sys_platform == \"linux\" or sys_platform == \"win32\"", "nvidia-nvfatbin (==13.0.85.*) ; sys_platform == \"linux\" or sys_platform == \"win32\"", "nvidia-nvjitlink (==13.0.88.*) ; sys_platform == \"linux\" or sys_platform == \"win32\"", "nvidia-nvjpeg (==13.0.1.86.*) ; sys_platform == \"linux\" or sys_platform == \"win32\"", "nvidia-nvml-dev (==13.0.87.*) ; sys_platform == \"linux\" or sys_platform == \"win32\"", "nvidia-nvptxcompiler (==13.0.88.*) ; sys_platform == \"linux\" or sys_platform == \"win32\"", "nvidia-nvtx (==13.0.85.*) ; sys_platform == \"linux\" or sys_platform == \"win32\"", "nvidia-nvvm (==13.0.88.*) ; sys_platform == \"linux\" or sys_platform == \"win32\""] -cccl = ["nvidia-cuda-cccl (==13.0.85.*) ; sys_platform == \"linux\" or sys_platform == \"win32\""] -crt = ["nvidia-cuda-crt (==13.0.88.*) ; sys_platform == \"linux\" or sys_platform == \"win32\""] -cublas = ["nvidia-cublas (==13.1.0.3.*) ; sys_platform == \"linux\" or sys_platform == \"win32\""] -cudart = ["nvidia-cuda-runtime (==13.0.96.*) ; sys_platform == \"linux\" or sys_platform == \"win32\""] -cufft = ["nvidia-cufft (==12.0.0.61.*) ; sys_platform == \"linux\" or sys_platform == \"win32\""] -cufile = ["nvidia-cufile (==1.15.1.6.*) ; sys_platform == \"linux\""] -culibos = ["nvidia-cuda-culibos (==13.0.85.*) ; sys_platform == \"linux\""] -cupti = ["nvidia-cuda-cupti (==13.0.85.*) ; sys_platform == \"linux\" or sys_platform == \"win32\""] -curand = ["nvidia-curand (==10.4.0.35.*) ; sys_platform == \"linux\" or sys_platform == \"win32\""] -cusolver = ["nvidia-cusolver (==12.0.4.66.*) ; sys_platform == \"linux\" or sys_platform == \"win32\""] -cusparse = ["nvidia-cusparse (==12.6.3.3.*) ; sys_platform == \"linux\" or sys_platform == \"win32\""] -cuxxfilt = ["nvidia-cuda-cuxxfilt (==13.0.85.*) ; sys_platform == \"linux\" or sys_platform == \"win32\""] -npp = ["nvidia-npp (==13.0.1.2.*) ; sys_platform == \"linux\" or sys_platform == \"win32\""] -nvcc = ["nvidia-cuda-nvcc (==13.0.88.*) ; sys_platform == \"linux\" or sys_platform == \"win32\""] -nvfatbin = ["nvidia-nvfatbin (==13.0.85.*) ; sys_platform == \"linux\" or sys_platform == \"win32\""] -nvjitlink = ["nvidia-nvjitlink (==13.0.88.*) ; sys_platform == \"linux\" or sys_platform == \"win32\""] -nvjpeg = ["nvidia-nvjpeg (==13.0.1.86.*) ; sys_platform == \"linux\" or sys_platform == \"win32\""] -nvml = ["nvidia-nvml-dev (==13.0.87.*) ; sys_platform == \"linux\" or sys_platform == \"win32\""] -nvptxcompiler = ["nvidia-nvptxcompiler (==13.0.88.*) ; sys_platform == \"linux\" or sys_platform == \"win32\""] -nvrtc = ["nvidia-cuda-nvrtc (==13.0.88.*) ; sys_platform == \"linux\" or sys_platform == \"win32\""] -nvtx = ["nvidia-nvtx (==13.0.85.*) ; sys_platform == \"linux\" or sys_platform == \"win32\""] -nvvm = ["nvidia-nvvm (==13.0.88.*) ; sys_platform == \"linux\" or sys_platform == \"win32\""] -opencl = ["nvidia-cuda-opencl (==13.0.85.*) ; sys_platform == \"linux\" or sys_platform == \"win32\""] -profiler = ["nvidia-cuda-profiler-api (==13.0.85.*) ; sys_platform == \"linux\" or sys_platform == \"win32\""] -sanitizer = ["nvidia-cuda-sanitizer-api (==13.0.85.*) ; sys_platform == \"linux\" or sys_platform == \"win32\""] +all = ["nvidia-cublas (==13.1.1.3.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\"", "nvidia-cuda-cccl (==13.0.85.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\"", "nvidia-cuda-crt (==13.0.88.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\"", "nvidia-cuda-culibos (==13.0.85.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\")", "nvidia-cuda-cupti (==13.0.85.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\"", "nvidia-cuda-cuxxfilt (==13.0.85.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\"", "nvidia-cuda-nvcc (==13.0.88.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\"", "nvidia-cuda-nvrtc (==13.0.88.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\"", "nvidia-cuda-opencl (==13.0.85.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\"", "nvidia-cuda-profiler-api (==13.0.85.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\"", "nvidia-cuda-runtime (==13.0.96.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\"", "nvidia-cuda-sanitizer-api (==13.0.85.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\"", "nvidia-cufft (==12.0.0.61.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\"", "nvidia-cufile (==1.15.1.6.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\")", "nvidia-curand (==10.4.0.35.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\"", "nvidia-cusolver (==12.0.4.66.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\"", "nvidia-cusparse (==12.6.3.3.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\"", "nvidia-npp (==13.0.1.2.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\"", "nvidia-nvfatbin (==13.0.85.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\"", "nvidia-nvjitlink (>=13.0.88,<14) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\"", "nvidia-nvjpeg (==13.0.1.86.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\"", "nvidia-nvml-dev (==13.0.87.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\"", "nvidia-nvptxcompiler (==13.0.88.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\"", "nvidia-nvtx (==13.0.85.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\"", "nvidia-nvvm (==13.0.88.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\""] +cccl = ["nvidia-cuda-cccl (==13.0.85.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\""] +crt = ["nvidia-cuda-crt (==13.0.88.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\""] +cublas = ["nvidia-cublas (==13.1.1.3.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\"", "nvidia-cuda-nvrtc (==13.0.88.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\""] +cudart = ["nvidia-cuda-runtime (==13.0.96.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\""] +cufft = ["nvidia-cufft (==12.0.0.61.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\"", "nvidia-nvjitlink (>=13.0.88,<14) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\""] +cufile = ["nvidia-cufile (==1.15.1.6.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\")"] +culibos = ["nvidia-cuda-culibos (==13.0.85.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\")"] +cupti = ["nvidia-cuda-cupti (==13.0.85.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\""] +curand = ["nvidia-curand (==10.4.0.35.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\""] +cusolver = ["nvidia-cublas (==13.1.1.3.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\"", "nvidia-cusolver (==12.0.4.66.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\"", "nvidia-cusparse (==12.6.3.3.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\"", "nvidia-nvjitlink (>=13.0.88,<14) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\""] +cusparse = ["nvidia-cusparse (==12.6.3.3.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\"", "nvidia-nvjitlink (>=13.0.88,<14) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\""] +cuxxfilt = ["nvidia-cuda-cuxxfilt (==13.0.85.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\""] +npp = ["nvidia-npp (==13.0.1.2.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\""] +nvcc = ["nvidia-cuda-crt (==13.0.88.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\"", "nvidia-cuda-nvcc (==13.0.88.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\"", "nvidia-cuda-runtime (==13.0.96.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\"", "nvidia-nvvm (==13.0.88.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\""] +nvfatbin = ["nvidia-nvfatbin (==13.0.85.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\""] +nvjitlink = ["nvidia-nvjitlink (>=13.0.88,<14) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\""] +nvjpeg = ["nvidia-nvjpeg (==13.0.1.86.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\""] +nvml = ["nvidia-nvml-dev (==13.0.87.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\""] +nvptxcompiler = ["nvidia-nvptxcompiler (==13.0.88.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\""] +nvrtc = ["nvidia-cuda-nvrtc (==13.0.88.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\""] +nvtx = ["nvidia-nvtx (==13.0.85.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\""] +nvvm = ["nvidia-nvvm (==13.0.88.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\""] +opencl = ["nvidia-cuda-opencl (==13.0.85.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\""] +profiler = ["nvidia-cuda-profiler-api (==13.0.85.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\""] +sanitizer = ["nvidia-cuda-sanitizer-api (==13.0.85.*) ; sys_platform == \"linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\") or (sys_platform == \"linux\" or sys_platform == \"win32\") and platform_machine == \"x86_64\""] [[package]] name = "datasets" @@ -1669,70 +1670,80 @@ test = ["coverage", "pytest (>=7,<8.1)", "pytest-cov", "pytest-mock (>=3)"] [[package]] name = "grpcio" -version = "1.83.1" +version = "1.84.0" description = "HTTP/2-based RPC framework" optional = false python-versions = ">=3.10" groups = ["main"] files = [ - {file = "grpcio-1.83.1-cp310-cp310-linux_armv7l.whl", hash = "sha256:1fea1ae4795d4790579995a4dd5e20e7494d358e29a340e8368dab9723264328"}, - {file = "grpcio-1.83.1-cp310-cp310-macosx_11_0_universal2.whl", hash = "sha256:b7ace1f740b36fcd451a1bb96f71ee7650e60b308822baeb66a023965bc27f4b"}, - {file = "grpcio-1.83.1-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:2a141f7bfc1601a0942405a8af6334ab21ba1dd0fa49b8427686df7beebd374d"}, - {file = "grpcio-1.83.1-cp310-cp310-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:c7e9e19413d43077d5a5c77b02ff82610209088e8f98da929347bc03d4c848d1"}, - {file = "grpcio-1.83.1-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:b59eaaeeb03dde0a2708095fb50f1afa94f11dc1b459bb7790b53bfb8cf95153"}, - {file = "grpcio-1.83.1-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:4e7c1468cf37cca17ab18bc8072901eed8daeb81685589ccd07988e5a750ee67"}, - {file = "grpcio-1.83.1-cp310-cp310-musllinux_1_2_i686.whl", hash = "sha256:a4a87dc86b0393257a11eb11e911c4c3456cbacd1c1ab9e9441060d9a3ad126b"}, - {file = "grpcio-1.83.1-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:d0dda8af248f6971555e1d4425f64864ce4e7369c5f8ef57c3e82a9bef77e22f"}, - {file = "grpcio-1.83.1-cp310-cp310-win32.whl", hash = "sha256:0f736f8359cf7cb8d0914a290999765a4342b0c35f01adc6e3ba24598f9d62b7"}, - {file = "grpcio-1.83.1-cp310-cp310-win_amd64.whl", hash = "sha256:7d43e3bd2b7d749c2dbd41c2cc83d550c3343d299a19acbbba9e37ad8c11fa8e"}, - {file = "grpcio-1.83.1-cp311-cp311-linux_armv7l.whl", hash = "sha256:907a5e5afb31f7a46376afc1a1edddd7afa00a74bbbc5b78979bbc34479581f6"}, - {file = "grpcio-1.83.1-cp311-cp311-macosx_11_0_universal2.whl", hash = "sha256:547645f02499c972f3edec9be4db9997f1d03df307c1c199772342ed6d8b3c6d"}, - {file = "grpcio-1.83.1-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:34f1841fc6d1d76f8a2d74177eafa2d1ec7d7e039633488c9fcc1b375a1fc165"}, - {file = "grpcio-1.83.1-cp311-cp311-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:05ba265193fbd9f63355311ec7567bba32a72aeb8e9fd7b3443e4fcad87b0750"}, - {file = "grpcio-1.83.1-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:5cce1d9fe2887239f054dc9c314597e04f33d2e6bd3150a91c4946d7e5be5d98"}, - {file = "grpcio-1.83.1-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:f732feb060ef57c1a040c24cee072ba9fab99bd0a7d2c916ef3f1c4d84b98974"}, - {file = "grpcio-1.83.1-cp311-cp311-musllinux_1_2_i686.whl", hash = "sha256:145b0050d24eb38accd9dc7ae09a3c09b8e7330159f3cfb46b1dba8711d50c42"}, - {file = "grpcio-1.83.1-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:e844cdb25c3c93c7572e0a37137c12305efea493be4eb65801b3ee93f180c186"}, - {file = "grpcio-1.83.1-cp311-cp311-win32.whl", hash = "sha256:0d07661944477517b12a239e18720c8d9038f80a62f2c56260fae80327f43d2a"}, - {file = "grpcio-1.83.1-cp311-cp311-win_amd64.whl", hash = "sha256:e572da3e247b28a98f46636d33c756e81ffb0f5def96c231ba45332333060595"}, - {file = "grpcio-1.83.1-cp312-cp312-linux_armv7l.whl", hash = "sha256:5acd14c6ddf047de62cbf8745b11103ea91abbf57d1b8edd5395ccd9fcd13abb"}, - {file = "grpcio-1.83.1-cp312-cp312-macosx_11_0_universal2.whl", hash = "sha256:16138031a47b771860a16a975b53087f4fd5bbdbb2c03a188c5d90ad65d2bdae"}, - {file = "grpcio-1.83.1-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:5ccc26715fd4defca5e129e280dd883b1737b65045ec50ffe22ce42104089519"}, - {file = "grpcio-1.83.1-cp312-cp312-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:b74f2a1d9ab1dfa3e263ef33d581613679b78d0884babf11671af26e45570ead"}, - {file = "grpcio-1.83.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:72578aa07a4008f17521ef52debcc3acfd1e2c5426243bc3ffb56a38bfe610b7"}, - {file = "grpcio-1.83.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:c12e1fc59c6dc26d10d9144453ddc6cbfe4cd4c31e874ed2d0132f88e685eb8b"}, - {file = "grpcio-1.83.1-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:4910b62f7d12197160bfb7de06d876d64dd12d43483e8292f98f49ca09b628d9"}, - {file = "grpcio-1.83.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:9e703effe3ae779925c82ac24fdb82cf4105e1096810151ed9501c5f34546b9c"}, - {file = "grpcio-1.83.1-cp312-cp312-win32.whl", hash = "sha256:a2aea8bd6e0a34f12cbaddb7bb70bec836818789fa5c7ab7572c6b745396a2d4"}, - {file = "grpcio-1.83.1-cp312-cp312-win_amd64.whl", hash = "sha256:583bf2e8255040a4a312f9572dfe62a05271437b149550e1a536d5c47d2d1e8a"}, - {file = "grpcio-1.83.1-cp313-cp313-linux_armv7l.whl", hash = "sha256:8d228e253b77865efcbdd7b5894ca882c9e0ea98c02b7d20582e61ded8dfd4b5"}, - {file = "grpcio-1.83.1-cp313-cp313-macosx_11_0_universal2.whl", hash = "sha256:0468b627f2987c9a77f7580030207cbd85457ffe52998beff4f0b5c38c58a72c"}, - {file = "grpcio-1.83.1-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:a6a282e81530cead60bbd752cc04950a57f224379e9821495d6a35bd5ce9b1f4"}, - {file = "grpcio-1.83.1-cp313-cp313-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:947d945f52e8ecf3cafd2bb7113502a16ccfda3e12c854443094de32d83ad432"}, - {file = "grpcio-1.83.1-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:55656318d5dd387077396dffb929171ca3966e24bfead9a6c5dba9f889062cb4"}, - {file = "grpcio-1.83.1-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:9daf5acf4fc9d5f5627229969c2580a91e511779d76e4ccdeb9f4770f05d8bc2"}, - {file = "grpcio-1.83.1-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:7b94174cbca93316888f805efbeb08f1c020f7b7493d2d50cc4f6b64ebb7e8bd"}, - {file = "grpcio-1.83.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:65c5a7210911ffe0f67b1cdc5308f9854b6d1f1b345e3e49ab7cac1ba50fa346"}, - {file = "grpcio-1.83.1-cp313-cp313-win32.whl", hash = "sha256:179368d9361854616ce6f397d4716e07480129652752fcbcfc5a7260455ad6f2"}, - {file = "grpcio-1.83.1-cp313-cp313-win_amd64.whl", hash = "sha256:2e57af456385491a76e13c4aada8c8f43a8e47051e06ea97a9dbe2a49654e6db"}, - {file = "grpcio-1.83.1-cp314-cp314-linux_armv7l.whl", hash = "sha256:8b3c87ca908296bf125f841d3e1a2225a2b39aaa8ed7a57e7ccde465ee519bab"}, - {file = "grpcio-1.83.1-cp314-cp314-macosx_11_0_universal2.whl", hash = "sha256:c0f3f20c90e72a171917ae65706500b096a1c3eb5f162c3ce702a2e25635f132"}, - {file = "grpcio-1.83.1-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:81bbf35a46bf8cad2dfbb2eccc19c711befb58b288acb534bbcd0d74283202a6"}, - {file = "grpcio-1.83.1-cp314-cp314-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:215cec07d11176507387bda4bf2751816e880f9bff8dc1ca524bfbb8ed8f2fad"}, - {file = "grpcio-1.83.1-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:abce7d43ec29cd39230fa8339de1a07643b55adc412a454850fbd875349950ff"}, - {file = "grpcio-1.83.1-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:e256f95a40e3b0183a98556fb7164d24b97eeb353123ccabfcba94712b35ee2a"}, - {file = "grpcio-1.83.1-cp314-cp314-musllinux_1_2_i686.whl", hash = "sha256:2110059146fb0ea216e1ffddb29377b5cc2fd412a5b0a92e102616bd5edf18c2"}, - {file = "grpcio-1.83.1-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:20d944d967843f8183f9f23d5916388362e5f8eeeae855bbe4354d906dc9f31b"}, - {file = "grpcio-1.83.1-cp314-cp314-win32.whl", hash = "sha256:623c87c6d4a1cb30d82c4e896f95477050f2e01b4a1f8cf91ff2b1abdf89c457"}, - {file = "grpcio-1.83.1-cp314-cp314-win_amd64.whl", hash = "sha256:47e6934ad38779271e2e7cc5f78a63a407cf3d98114c65c1fdbcd3f5a716f29b"}, - {file = "grpcio-1.83.1.tar.gz", hash = "sha256:9cee6fcbf2eb57c4b49451787bfa87be8efc1ca02a0b327dd4b54d44502e362b"}, + {file = "grpcio-1.84.0-cp310-cp310-linux_armv7l.whl", hash = "sha256:71fd60e6e426d293d0a2f685115ad0a0845117602cf13605a4be7524fb5f7bba"}, + {file = "grpcio-1.84.0-cp310-cp310-macosx_11_0_universal2.whl", hash = "sha256:8e1a45d174b6b8589f51dce1cea804aa6c1f72c9c80cba91ae2caabeb6d90540"}, + {file = "grpcio-1.84.0-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:efb29f8633bf6630dc89de4fe0353ac3d7e4b70ef7b6e29fb40f00e68c127fa5"}, + {file = "grpcio-1.84.0-cp310-cp310-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:d0fdd25faece8a1f95e8a3a8006e29701b5cf8dadb4a8132e68f3134637004a5"}, + {file = "grpcio-1.84.0-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:393d8a78bff6731ecc5ad2151a821f8fbc1709b137ebb9c25a4ef399fbdcc914"}, + {file = "grpcio-1.84.0-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:fc66cb50c93554b86db0b6625ab5c6e9051dbf8847c08d93c84918e02e413fb7"}, + {file = "grpcio-1.84.0-cp310-cp310-musllinux_1_2_i686.whl", hash = "sha256:455ed6083353b8e938f1d58c765eab2fbb165731e5b507be30fee344915a2a11"}, + {file = "grpcio-1.84.0-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:3d6a82c4fc6c85f2fb7572c86bdb86f84c97b6580e5f6599f711800bac48a5d8"}, + {file = "grpcio-1.84.0-cp310-cp310-win32.whl", hash = "sha256:8e3f508d0e9e6236ba2f08d56e33355e434e785e813149a1b8477d3edf69779d"}, + {file = "grpcio-1.84.0-cp310-cp310-win_amd64.whl", hash = "sha256:ed2c1493c44d0932f1e55fdb5d1ead658c68288ec5d51b8c4928422d98633ef9"}, + {file = "grpcio-1.84.0-cp311-cp311-linux_armv7l.whl", hash = "sha256:4aaeceeb7fa7d824c322d1ec3208c8495c88478a927295553235435fc49043ad"}, + {file = "grpcio-1.84.0-cp311-cp311-macosx_11_0_universal2.whl", hash = "sha256:06619ba1515e5ee69fb2a514e95dd8be05ce74cb3928d5b34f87f87c86fe3c27"}, + {file = "grpcio-1.84.0-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:158c1c11cfb61b4849c3caf4d52de6f5ecd376e14446feb4a90dc95a90d616f5"}, + {file = "grpcio-1.84.0-cp311-cp311-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:a9383401d9f116f98cacd4eba6c505a6edb80ba65badfc8e8ed8ae64983bcc44"}, + {file = "grpcio-1.84.0-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:bd8ea8eb3817b226057cc1c0e7ec4b378dcda52043b972b6ff12b1152178967d"}, + {file = "grpcio-1.84.0-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:756ea5c2da00fa65c930284892d2a9706828704ca3ba40b4c51c4834eb39fcfd"}, + {file = "grpcio-1.84.0-cp311-cp311-musllinux_1_2_i686.whl", hash = "sha256:28d2609691da93051e998495108bbddd2a9f7a561253bae94828d81290f30c15"}, + {file = "grpcio-1.84.0-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:27b8b36200a9fbee6e120246f4a8a41657549107ef19fb2c819c4b2fd524f39a"}, + {file = "grpcio-1.84.0-cp311-cp311-win32.whl", hash = "sha256:465eef3d17e59ad22a556fc0138f7c7c799df426734344daec42c797d49fda99"}, + {file = "grpcio-1.84.0-cp311-cp311-win_amd64.whl", hash = "sha256:f9a456bdbed52a01c9ab8423bdebab04a5363c78676edc55ab9b58bd13bdf9e1"}, + {file = "grpcio-1.84.0-cp312-cp312-linux_armv7l.whl", hash = "sha256:b5c6f20d657ae09ae4e30d9d3a21edd13f1219d58cc6f999b9d1bb63be9c1baa"}, + {file = "grpcio-1.84.0-cp312-cp312-macosx_11_0_universal2.whl", hash = "sha256:406583b4e8fb2282ebd392e12b963e601c1f82e07125a8c2cb5b144e7e024796"}, + {file = "grpcio-1.84.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:fbdbcd06986ede3ce584083b1dc2afe6808e8943e5cf50ad11183c03aceda25a"}, + {file = "grpcio-1.84.0-cp312-cp312-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:23e6e8e8a75cff88e0a793bfd3becea03a13e2763ae90c1ff573bc19ca5b429a"}, + {file = "grpcio-1.84.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:b44f0a0fc7bc6677d38cc80bca1a32814ce6c8f200fb8b3c1a61c9d77eaefbf3"}, + {file = "grpcio-1.84.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:210e4c32f907045eb8158273e60c6ab69a3947697df6245dbda381f26c59485b"}, + {file = "grpcio-1.84.0-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:a71d24f40b0cc6798feaa978c7411dc1135b7018e9fc0442db611c139bf58344"}, + {file = "grpcio-1.84.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:f6c972474ce691aca74e58d17625450cef153dc4760364cadeb167983ea6d589"}, + {file = "grpcio-1.84.0-cp312-cp312-win32.whl", hash = "sha256:0d532ade4486dad9b302ffa4d4683d67561051c26d17c4023322845e9fa10140"}, + {file = "grpcio-1.84.0-cp312-cp312-win_amd64.whl", hash = "sha256:49717e857899f4136d7657bf5aded61ac479110a075438290923a4d86af7cd02"}, + {file = "grpcio-1.84.0-cp313-cp313-linux_armv7l.whl", hash = "sha256:209414080da8c20af94df1395b635da52dd57b5edc9e917e1deca0dc1c4bb55e"}, + {file = "grpcio-1.84.0-cp313-cp313-macosx_11_0_universal2.whl", hash = "sha256:e41c3993eee896c617dbd8a505085d28b6e84a0445ed9a1f40f95808473cf678"}, + {file = "grpcio-1.84.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:fff5ef3fe1bba7d6147e5f19e01e5e122ac2c076486887ddcb8d42e663400fbe"}, + {file = "grpcio-1.84.0-cp313-cp313-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:b8c62888c3e49debf37ad9773e3c02f77b0c1e811f8fb0962f2b6c3bbab5b97a"}, + {file = "grpcio-1.84.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:986e9751d416d7a6eaa2fecdac38da63153d63a4b340ba7d624889c490451500"}, + {file = "grpcio-1.84.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:5933a052946873d01a42119a05420d669bdca436aeba2d1851988ccb12b421c0"}, + {file = "grpcio-1.84.0-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:e094dd21f077af8194923fc263cad872eaa1802bb0156fd7e5ae18e99cd86715"}, + {file = "grpcio-1.84.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:08735e3d08d24ab3132cf87e2e5dea8746cabcc7d676c2b0b7362f195feef9d9"}, + {file = "grpcio-1.84.0-cp313-cp313-win32.whl", hash = "sha256:70bb4ce8be0c5606bec259cbd7152374470396413b7863a658a08c849e6b29ff"}, + {file = "grpcio-1.84.0-cp313-cp313-win_amd64.whl", hash = "sha256:b61692f0069b3eee2fc8a3a1b7f6c044df9e03fede6ce69b3ca832e1c39f26c5"}, + {file = "grpcio-1.84.0-cp314-cp314-linux_armv7l.whl", hash = "sha256:026d757df86c5b7a41de8200b9a2cda454aaa5004cb0c7e3374c66eb82f61499"}, + {file = "grpcio-1.84.0-cp314-cp314-macosx_11_0_universal2.whl", hash = "sha256:3de427b05f244ba2c2a9bdc67e7a6731c8340811524ecc4435466549f8af1d17"}, + {file = "grpcio-1.84.0-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:e90e3bdf7b5eac005fef631adae9cafde16f922def207b80a7c46b253c18ad20"}, + {file = "grpcio-1.84.0-cp314-cp314-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:e88d304f094f4937bc27ec6a435e218a084168f11ec630c8d5d39b431d08d81d"}, + {file = "grpcio-1.84.0-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:57dc36a5ab0e676f5f6e171de2917fd0aef73f32a9aaf23956bfe19997a30bd1"}, + {file = "grpcio-1.84.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:5deda5b4bf62769eb98c119cca43d40e1231e34846b19db5cdea821d446a2253"}, + {file = "grpcio-1.84.0-cp314-cp314-musllinux_1_2_i686.whl", hash = "sha256:9bab4cf571653a8afffb83ce21aa27b51dfe629b526b7b6adec35491fe1fc2ea"}, + {file = "grpcio-1.84.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:c5559b492007dc09b4de9b95dab05f0b5e53547aad230cf07e46c7dd017a3be5"}, + {file = "grpcio-1.84.0-cp314-cp314-win32.whl", hash = "sha256:2c024da73b296f040b8360e60bd73a659b230093684a438da0e1260f34cc724e"}, + {file = "grpcio-1.84.0-cp314-cp314-win_amd64.whl", hash = "sha256:800b7e00d92553313c0463c200087930aa78678ec1d528193aeb50906f55989b"}, + {file = "grpcio-1.84.0-cp315-cp315-linux_armv7l.whl", hash = "sha256:47ecf0d9b81d981f07b61bd89eced9d2582f5eaacc3aaa36ad27f81aef70a27f"}, + {file = "grpcio-1.84.0-cp315-cp315-macosx_10_15_universal2.whl", hash = "sha256:61386101ecaa096b694d0dd278caf99a56aeec78440cc17e918eef0b50f2d567"}, + {file = "grpcio-1.84.0-cp315-cp315-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:f6d178ba6dc8e82976c184b65fddde172d054c17237993a3e083efe4f134d55b"}, + {file = "grpcio-1.84.0-cp315-cp315-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:15bb76489e337fc492685c9758e2fd4d4ab516b901ad830dc5a91987decf00be"}, + {file = "grpcio-1.84.0-cp315-cp315-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:82da34ae4f639c73ac46e521e00c0a49bf86f717b9fb1f405f133e98731e38dc"}, + {file = "grpcio-1.84.0-cp315-cp315-musllinux_1_2_aarch64.whl", hash = "sha256:9b73836ba0e16fcbb57c31cf6cbc2907c8d8c790b83679df454b74bd15e0be04"}, + {file = "grpcio-1.84.0-cp315-cp315-musllinux_1_2_i686.whl", hash = "sha256:42959bd50dd660ffc3f2a9bec15a6da4f9aaa0dda555d59ff2d2e80b908456a8"}, + {file = "grpcio-1.84.0-cp315-cp315-musllinux_1_2_x86_64.whl", hash = "sha256:659728f20fc7a0933ed7b1945435e31014b97ab8a5a7edcbaa70da4794aeb191"}, + {file = "grpcio-1.84.0-cp315-cp315-win32.whl", hash = "sha256:edb6f87fc60ff438557291501b3e16c7a77c3b01a52d782cf276dccc7c5dd89c"}, + {file = "grpcio-1.84.0-cp315-cp315-win_amd64.whl", hash = "sha256:4119efa6519871719ad81f33bc95ab87857dcb1c5801f30a6e592f2c41164169"}, + {file = "grpcio-1.84.0.tar.gz", hash = "sha256:19aaf172fc2edbefccce3f6e92c5150975dbe56c45744e9e87cf72ebdf85bfbe"}, ] [package.dependencies] typing-extensions = ">=4.12,<5.0" [package.extras] -protobuf = ["grpcio-tools (>=1.83.1)"] +protobuf = ["grpcio-tools (>=1.84.0)"] [[package]] name = "h11" @@ -1792,7 +1803,7 @@ description = "Fast transfer of large files with the Hugging Face Hub." optional = false python-versions = ">=3.8" groups = ["main"] -markers = "platform_machine == \"x86_64\" or platform_machine == \"amd64\" or platform_machine == \"AMD64\" or platform_machine == \"arm64\" or platform_machine == \"aarch64\"" +markers = "(sys_platform == \"linux\" or sys_platform == \"win32\" or platform_machine == \"amd64\" or platform_machine == \"AMD64\" or platform_machine == \"arm64\" or platform_machine == \"aarch64\" or platform_machine == \"x86_64\") and (platform_machine == \"x86_64\" or platform_machine == \"amd64\" or platform_machine == \"AMD64\" or platform_machine == \"arm64\" or platform_machine == \"aarch64\" or sys_platform == \"linux\") and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\" or platform_machine == \"amd64\" or platform_machine == \"AMD64\" or platform_machine == \"arm64\")" files = [ {file = "hf_xet-1.6.0-cp314-cp314t-macosx_10_12_x86_64.whl", hash = "sha256:70cbb9c896901600128cb9b6f06e132954fbede1db30f31f7c6c63f84cb7c31d"}, {file = "hf_xet-1.6.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:23379c2f9ec8696d952b16414a2bae72cad86a52df869b050698ba60f538c675"}, @@ -1840,15 +1851,15 @@ trio = ["trio (>=0.22.0,<1.0)"] [[package]] name = "httpcore2" -version = "2.12.0" +version = "2.13.0" description = "A minimal low-level HTTP client." optional = false python-versions = ">=3.10" groups = ["main"] markers = "sys_platform != \"emscripten\"" files = [ - {file = "httpcore2-2.12.0-py3-none-any.whl", hash = "sha256:7e04258ce01013d7d615e5b910a3b27fac937d7a95038227e79652b4ba3b4ceb"}, - {file = "httpcore2-2.12.0.tar.gz", hash = "sha256:9293522bba0aa7c4c8e9e3f040c16575bd8868e155a77fa30c7a9085a5eae648"}, + {file = "httpcore2-2.13.0-py3-none-any.whl", hash = "sha256:35ae5be347aa40467b4a5dc032ac67ebb6d27189fc97e8cebcf99616f6a1bb9e"}, + {file = "httpcore2-2.13.0.tar.gz", hash = "sha256:2adc8be4fb285fbcd6d894298db3b52c177e74b6674eda3a76bd36be3292a3db"}, ] [package.dependencies] @@ -1859,7 +1870,7 @@ truststore = ">=0.10" asyncio = ["anyio (>=4.5.0,<5.0)"] http2 = ["h2 (>=3,<5)"] socks = ["socksio (==1.*)"] -trio = ["trio (>=0.33.0,<1.0)"] +trio = ["trio (>=0.34.0,<1.0)"] [[package]] name = "httpx" @@ -1900,26 +1911,26 @@ files = [ [[package]] name = "httpx2" -version = "2.12.0" +version = "2.13.0" description = "The next generation HTTP client." optional = false python-versions = ">=3.10" groups = ["main"] files = [ - {file = "httpx2-2.12.0-py3-none-any.whl", hash = "sha256:cc8b6eecb8661c146b8f89a60e97456ee086e91a784ed31ac450c3a9e613dd36"}, - {file = "httpx2-2.12.0.tar.gz", hash = "sha256:7631fe9887a8a2275f4a2540e053aa670fcc50742864a9ae7c66e609fdcf12cf"}, + {file = "httpx2-2.13.0-py3-none-any.whl", hash = "sha256:fc12720cedf72faa26cca6b4ca394e05c894e7d7933fc45cafe767960804e49a"}, + {file = "httpx2-2.13.0.tar.gz", hash = "sha256:81bd07dc67a3701729ef1f777a3c00c915d4539604fdb5afd327f8682f6b7b44"}, ] [package.dependencies] anyio = {version = ">=4.10", markers = "sys_platform != \"emscripten\""} -httpcore2 = {version = "2.12.0", markers = "sys_platform != \"emscripten\""} +httpcore2 = {version = "2.13.0", markers = "sys_platform != \"emscripten\""} httpx2-jsfetch = {version = "*", markers = "sys_platform == \"emscripten\" and python_version >= \"3.12\""} idna = ">=3.18" truststore = {version = ">=0.10", markers = "sys_platform != \"emscripten\""} typing-extensions = {version = ">=4.5.0", markers = "python_version < \"3.13\""} [package.extras] -brotli = ["brotli (>=1.2.0) ; platform_python_implementation == \"CPython\"", "brotlicffi (>=1.2.0.0) ; platform_python_implementation != \"CPython\""] +brotli = ["brotli (>=1.2.0) ; platform_python_implementation == \"CPython\"", "brotlicffi (>=1.2.0.2) ; platform_python_implementation != \"CPython\""] cli = ["click (>=8.4.2)", "pygments (==2.*)", "rich (>=10,<16)"] http2 = ["h2 (>=3,<5)"] socks = ["socksio (==1.*)"] @@ -3306,7 +3317,7 @@ description = "CUBLAS native runtime libraries" optional = false python-versions = ">=3" groups = ["main"] -markers = "platform_system == \"Linux\"" +markers = "(sys_platform == \"linux\" or sys_platform == \"win32\") and (sys_platform == \"linux\" or platform_machine == \"x86_64\") and platform_system == \"Linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\")" files = [ {file = "nvidia_cublas-13.1.1.3-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:b7a210458267ac818974c53038fbec2e969d5c99f305ab15c72522fa9f001dd5"}, {file = "nvidia_cublas-13.1.1.3-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:37936a16db8fe4ac1f065c2139360608a543a09275cb1a1af612e08cfa065436"}, @@ -3316,6 +3327,24 @@ files = [ [package.dependencies] nvidia-cuda-nvrtc = "*" +[[package]] +name = "nvidia-cublas" +version = "13.7.0.27" +description = "CUBLAS native runtime libraries" +optional = false +python-versions = ">=3" +groups = ["main"] +markers = "(sys_platform != \"linux\" or platform_machine != \"aarch64\" and platform_machine != \"x86_64\") and platform_system == \"Linux\" and (sys_platform != \"linux\" and sys_platform != \"win32\" or platform_machine != \"x86_64\")" +files = [ + {file = "nvidia_cublas-13.7.0.27-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:39f240346a8bbc6c4f1e3100dcfa2c6d3fe353db513060329ef28cdd43c27f2f"}, + {file = "nvidia_cublas-13.7.0.27-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:481adf76a6b7585a7a4cd5f53298589abe207415dd584887751b99baa85491d4"}, + {file = "nvidia_cublas-13.7.0.27-py3-none-win_amd64.whl", hash = "sha256:15d3963ae9de7292348ced320c849ff27f76b0f433b81589e9eb87f5407841c9"}, + {file = "nvidia_cublas-13.7.0.27-py3-none-win_arm64.whl", hash = "sha256:4cf58f2c8a4743a2717b0f5c51a90f0d93ada416d71eb3d96b17ca29b7af0106"}, +] + +[package.dependencies] +nvidia-cuda-nvrtc = "*" + [[package]] name = "nvidia-cuda-cupti" version = "13.0.85" @@ -3323,7 +3352,7 @@ description = "CUDA profiling tools runtime libs." optional = false python-versions = ">=3" groups = ["main"] -markers = "(sys_platform == \"win32\" or sys_platform == \"linux\") and platform_system == \"Linux\"" +markers = "(sys_platform == \"linux\" or sys_platform == \"win32\") and (sys_platform == \"linux\" or platform_machine == \"x86_64\") and platform_system == \"Linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\")" files = [ {file = "nvidia_cuda_cupti-13.0.85-py3-none-manylinux_2_25_aarch64.whl", hash = "sha256:796bd679890ee55fb14a94629b698b6db54bcfd833d391d5e94017dd9d7d3151"}, {file = "nvidia_cuda_cupti-13.0.85-py3-none-manylinux_2_25_x86_64.whl", hash = "sha256:4eb01c08e859bf924d222250d2e8f8b8ff6d3db4721288cf35d14252a4d933c8"}, @@ -3351,13 +3380,28 @@ description = "NVRTC native runtime libraries" optional = false python-versions = ">=3" groups = ["main"] -markers = "platform_system == \"Linux\"" +markers = "(sys_platform == \"linux\" or sys_platform == \"win32\") and (sys_platform == \"linux\" or platform_machine == \"x86_64\") and platform_system == \"Linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\")" files = [ {file = "nvidia_cuda_nvrtc-13.0.88-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:ad9b6d2ead2435f11cbb6868809d2adeeee302e9bb94bcf0539c7a40d80e8575"}, {file = "nvidia_cuda_nvrtc-13.0.88-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:d27f20a0ca67a4bb34268a5e951033496c5b74870b868bacd046b1b8e0c3267b"}, {file = "nvidia_cuda_nvrtc-13.0.88-py3-none-win_amd64.whl", hash = "sha256:6bcd4e7f8e205cbe644f5a98f2f799bef9556fefc89dd786e79a16312ce49872"}, ] +[[package]] +name = "nvidia-cuda-nvrtc" +version = "13.4.59" +description = "NVRTC native runtime libraries" +optional = false +python-versions = ">=3" +groups = ["main"] +markers = "(sys_platform != \"linux\" or platform_machine != \"aarch64\" and platform_machine != \"x86_64\") and platform_system == \"Linux\" and (sys_platform != \"linux\" and sys_platform != \"win32\" or platform_machine != \"x86_64\")" +files = [ + {file = "nvidia_cuda_nvrtc-13.4.59-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:f6cb36996797989976bd021d89e263be379668be7a4249cc95a3d11590bda652"}, + {file = "nvidia_cuda_nvrtc-13.4.59-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:da5e1ba65f402b2fdf1c4bfcb16e27df47e8ade27f7936b1524cb10cdd61a1c8"}, + {file = "nvidia_cuda_nvrtc-13.4.59-py3-none-win_amd64.whl", hash = "sha256:93755e0f94e717c61168fb348ab5e2f5c67aa2f918416b1786f55ff813f6923d"}, + {file = "nvidia_cuda_nvrtc-13.4.59-py3-none-win_arm64.whl", hash = "sha256:ec27241be8b65e88d37a25d5e4ec0f4fdb0499439a7d5122cfdc74457752ae81"}, +] + [[package]] name = "nvidia-cuda-runtime" version = "13.0.96" @@ -3365,7 +3409,7 @@ description = "CUDA Runtime native Libraries" optional = false python-versions = ">=3" groups = ["main"] -markers = "(sys_platform == \"win32\" or sys_platform == \"linux\") and platform_system == \"Linux\"" +markers = "(sys_platform == \"linux\" or sys_platform == \"win32\") and (sys_platform == \"linux\" or platform_machine == \"x86_64\") and platform_system == \"Linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\")" files = [ {file = "nvidia_cuda_runtime-13.0.96-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:ef9bcbe90493a2b9d810e43d249adb3d02e98dd30200d86607d8d02687c43f55"}, {file = "nvidia_cuda_runtime-13.0.96-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:7f82250d7782aa23b6cfe765ecc7db554bd3c2870c43f3d1821f1d18aebf0548"}, @@ -3446,7 +3490,7 @@ description = "CUFFT native runtime libraries" optional = false python-versions = ">=3" groups = ["main"] -markers = "(sys_platform == \"win32\" or sys_platform == \"linux\") and platform_system == \"Linux\"" +markers = "(sys_platform == \"linux\" or sys_platform == \"win32\") and (sys_platform == \"linux\" or platform_machine == \"x86_64\") and platform_system == \"Linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\")" files = [ {file = "nvidia_cufft-12.0.0.61-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:2708c852ef8cd89d1d2068bdbece0aa188813a0c934db3779b9b1faa8442e5f5"}, {file = "nvidia_cufft-12.0.0.61-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:6c44f692dce8fd5ffd3e3df134b6cdb9c2f72d99cf40b62c32dde45eea9ddad3"}, @@ -3463,7 +3507,7 @@ description = "cuFile GPUDirect libraries" optional = false python-versions = ">=3" groups = ["main"] -markers = "sys_platform == \"linux\" and platform_system == \"Linux\"" +markers = "sys_platform == \"linux\" and platform_system == \"Linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\")" files = [ {file = "nvidia_cufile-1.15.1.6-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:08a3ecefae5a01c7f5117351c64f17c7c62efa5fffdbe24fc7d298da19cd0b44"}, {file = "nvidia_cufile-1.15.1.6-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:bdc0deedc61f548bddf7733bdc216456c2fdb101d020e1ab4b88d232d5e2f6d1"}, @@ -3476,7 +3520,7 @@ description = "CURAND native runtime libraries" optional = false python-versions = ">=3" groups = ["main"] -markers = "(sys_platform == \"win32\" or sys_platform == \"linux\") and platform_system == \"Linux\"" +markers = "(sys_platform == \"linux\" or sys_platform == \"win32\") and (sys_platform == \"linux\" or platform_machine == \"x86_64\") and platform_system == \"Linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\")" files = [ {file = "nvidia_curand-10.4.0.35-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:133df5a7509c3e292aaa2b477afd0194f06ce4ea24d714d616ff36439cee349a"}, {file = "nvidia_curand-10.4.0.35-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:1aee33a5da6e1db083fe2b90082def8915f30f3248d5896bcec36a579d941bfc"}, @@ -3490,7 +3534,7 @@ description = "CUDA solver native runtime libraries" optional = false python-versions = ">=3" groups = ["main"] -markers = "(sys_platform == \"win32\" or sys_platform == \"linux\") and platform_system == \"Linux\"" +markers = "(sys_platform == \"linux\" or sys_platform == \"win32\") and (sys_platform == \"linux\" or platform_machine == \"x86_64\") and platform_system == \"Linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\")" files = [ {file = "nvidia_cusolver-12.0.4.66-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:02c2457eaa9e39de20f880f4bd8820e6a1cfb9f9a34f820eb12a155aa5bc92d2"}, {file = "nvidia_cusolver-12.0.4.66-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:0a759da5dea5c0ea10fd307de75cdeb59e7ea4fcb8add0924859b944babf1112"}, @@ -3509,7 +3553,7 @@ description = "CUSPARSE native runtime libraries" optional = false python-versions = ">=3" groups = ["main"] -markers = "(sys_platform == \"win32\" or sys_platform == \"linux\") and platform_system == \"Linux\"" +markers = "(sys_platform == \"linux\" or sys_platform == \"win32\") and (sys_platform == \"linux\" or platform_machine == \"x86_64\") and platform_system == \"Linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\")" files = [ {file = "nvidia_cusparse-12.6.3.3-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:80bcc4662f23f1054ee334a15c72b8940402975e0eab63178fc7e670aa59472c"}, {file = "nvidia_cusparse-12.6.3.3-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:2b3c89c88d01ee0e477cb7f82ef60a11a4bcd57b6b87c33f789350b59759360b"}, @@ -3519,6 +3563,24 @@ files = [ [package.dependencies] nvidia-nvjitlink = "*" +[[package]] +name = "nvidia-cusparse" +version = "12.8.6.49" +description = "CUSPARSE native runtime libraries" +optional = false +python-versions = ">=3" +groups = ["main"] +markers = "" +files = [ + {file = "nvidia_cusparse-12.8.6.49-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:3bf4f2e34e8cfe0c691a4afd7bd479bee49f44dc5fa4940d9d54f720cbd6f2f3"}, + {file = "nvidia_cusparse-12.8.6.49-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:f740691172baae1d498f754cfbb565933c7dbcd7b16544d69fa666283d0715d7"}, + {file = "nvidia_cusparse-12.8.6.49-py3-none-win_amd64.whl", hash = "sha256:7578698de43ca4f72371076adcba460182c4ae5b2263a094d6e70bc59db7fe10"}, + {file = "nvidia_cusparse-12.8.6.49-py3-none-win_arm64.whl", hash = "sha256:aeee3f0f53e18345a75c9cb25d34070823d73ae193ab231ee7438fe30abbdfc5"}, +] + +[package.dependencies] +nvidia-nvjitlink = "*" + [[package]] name = "nvidia-cusparselt-cu13" version = "0.8.1" @@ -3702,16 +3764,17 @@ files = [ [[package]] name = "nvidia-nvjitlink" -version = "13.0.88" +version = "13.4.52" description = "Nvidia JIT LTO Library" optional = false python-versions = ">=3" groups = ["main"] -markers = "(sys_platform == \"win32\" or sys_platform == \"linux\") and platform_system == \"Linux\"" +markers = "(sys_platform == \"linux\" or sys_platform == \"win32\") and (sys_platform == \"linux\" or platform_machine == \"x86_64\") and platform_system == \"Linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\")" files = [ - {file = "nvidia_nvjitlink-13.0.88-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:13a74f429e23b921c1109976abefacc69835f2f433ebd323d3946e11d804e47b"}, - {file = "nvidia_nvjitlink-13.0.88-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:e931536ccc7d467a98ba1d8b89ff7fa7f1fa3b13f2b0069118cd7f47bff07d0c"}, - {file = "nvidia_nvjitlink-13.0.88-py3-none-win_amd64.whl", hash = "sha256:634e96e3da9ef845ae744097a1f289238ecf946ce0b82e93cdce14b9782e682f"}, + {file = "nvidia_nvjitlink-13.4.52-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:90401db7e5a580a5067a468b3086e0b65f3b96ab3c42524afd741efc8a0e150a"}, + {file = "nvidia_nvjitlink-13.4.52-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:3a589e3732140839349545efd0db67426523459b12e6952c3c73b6d55900f200"}, + {file = "nvidia_nvjitlink-13.4.52-py3-none-win_amd64.whl", hash = "sha256:4bd7c27adc7832da5e6b8580fa64a42b9eebc94baae3ddcbedf4719f0bab2bb8"}, + {file = "nvidia_nvjitlink-13.4.52-py3-none-win_arm64.whl", hash = "sha256:75f7e1a04b0a63c7d8af9921e1ed45d9ef880d3c6ef976989082482698e01205"}, ] [[package]] @@ -3734,7 +3797,7 @@ description = "NVIDIA Tools Extension" optional = false python-versions = ">=3" groups = ["main"] -markers = "(sys_platform == \"win32\" or sys_platform == \"linux\") and platform_system == \"Linux\"" +markers = "(sys_platform == \"linux\" or sys_platform == \"win32\") and (sys_platform == \"linux\" or platform_machine == \"x86_64\") and platform_system == \"Linux\" and (platform_machine == \"aarch64\" or platform_machine == \"x86_64\")" files = [ {file = "nvidia_nvtx-13.0.85-py3-none-manylinux1_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:4936d1d6780fbe68db454f5e72a42ff64d1fd6397df9f363ae786930fd5c1cd4"}, {file = "nvidia_nvtx-13.0.85-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:cb7780edb6b14107373c835bf8b72e7a178bac7367e23da7acb108f973f157a6"}, @@ -3802,14 +3865,14 @@ PyYAML = ">=5.1.0" [[package]] name = "openai" -version = "3.13.0" +version = "3.14.0" description = "The official Python library for the openai API" optional = false python-versions = ">=3.10" groups = ["main"] files = [ - {file = "openai-3.13.0-py3-none-any.whl", hash = "sha256:e35b1f6fe99245e86e37504d9fad1ad2a363807307c424232c1d849bd0666c8e"}, - {file = "openai-3.13.0.tar.gz", hash = "sha256:a8f87a9b3b9c08eb446d68bd0a80e8ec907c4c35fdea63f4265c7b34b2de3a60"}, + {file = "openai-3.14.0-py3-none-any.whl", hash = "sha256:232a85a1c0ff6820534630fbfdb99313990ed80e66f39892a3abcff4130e2b16"}, + {file = "openai-3.14.0.tar.gz", hash = "sha256:714ba70b91ee1e78e75263694bebfeaca004432032685333a46d0e0206114bf9"}, ] [package.dependencies] @@ -6175,57 +6238,52 @@ files = [ [[package]] name = "torch" -version = "2.12.0" +version = "2.13.0" description = "Tensors and Dynamic neural networks in Python with strong GPU acceleration" optional = false python-versions = ">=3.10" groups = ["main"] files = [ - {file = "torch-2.12.0-cp310-cp310-macosx_14_0_arm64.whl", hash = "sha256:1834bd984f8a2f4f16bdfbeecca9146184b220aa46276bf5756735b5dae12812"}, - {file = "torch-2.12.0-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:d4d029801cb7b6df858804a2a21b00cc2aa0bf0ee5d2ab18d343c9e9e5681f35"}, - {file = "torch-2.12.0-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:d47e7dee68ac4cd7a068b26bcd6b989935427709fae1c8f7bd0019978f829e15"}, - {file = "torch-2.12.0-cp310-cp310-win_amd64.whl", hash = "sha256:cf9839790285dd472e7a16aafcb4a4e6bf58ec1b494045044b0eefb0eb4bd1f2"}, - {file = "torch-2.12.0-cp311-cp311-macosx_14_0_arm64.whl", hash = "sha256:10802fd383bbfed646212e765a72c37d2185205d4f26eb197a254e8ac7ddcb25"}, - {file = "torch-2.12.0-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:c12592630aef72feaf18bd3f197ef587bbfa21131b31c38b23ab2e55fce92e36"}, - {file = "torch-2.12.0-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:415c1b8d0412f67551c8e89a2daca0fb3e56694af0281ba155eaa9da481f58b4"}, - {file = "torch-2.12.0-cp311-cp311-win_amd64.whl", hash = "sha256:dd37188ea325042cb1f6cafa56822b11ada2520c04791a52629b0af25bdfbfd9"}, - {file = "torch-2.12.0-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:b41339df93d491435e790ff8bcbae1c0ce777175889bfd1281d119862793e6a2"}, - {file = "torch-2.12.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:8fbef9f108a863e7722a73740998967e3b074742a834fc5be3a535a2befa7057"}, - {file = "torch-2.12.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:4b4f64c2c2b11f7510d93dd6412b87025ff6eddd6bb61c3b5a3d892ea20c4756"}, - {file = "torch-2.12.0-cp312-cp312-win_amd64.whl", hash = "sha256:8b958caff4a14d3a3b0b2dfc6a378f64dda9728a9dad28c08a0db9ce4dafb549"}, - {file = "torch-2.12.0-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:90dd587a5f61bfe1307148b581e2084fc5bc4a06e2b90a20e9a36b81087ff16b"}, - {file = "torch-2.12.0-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:864392c73b7654f4d2b3ae712f607937d0dbb1101c4555fbb41848106b297f39"}, - {file = "torch-2.12.0-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:5d6b560dfa7d56291c07d615c3bb73e8d9943d9b6d87f76cd0d9d570c4797fa6"}, - {file = "torch-2.12.0-cp313-cp313-win_amd64.whl", hash = "sha256:3fee918902090ade827643e758e98363278815de583c75d111fdd665ebffde9f"}, - {file = "torch-2.12.0-cp313-cp313t-macosx_14_0_arm64.whl", hash = "sha256:10ee1448a9f304d3b987eb4656f664ba6e4d7b410ca7a5a7c642199777a2cf88"}, - {file = "torch-2.12.0-cp313-cp313t-manylinux_2_28_aarch64.whl", hash = "sha256:af68dbf403439cae9ceaeaaf92f8352b460787dcd27b92aa05c40dd4a19c0f1e"}, - {file = "torch-2.12.0-cp313-cp313t-manylinux_2_28_x86_64.whl", hash = "sha256:a6a2eebb237d3b1d9ad3b378e86d9b9e0782afdea8b1e0eba6a13646b9b49c07"}, - {file = "torch-2.12.0-cp313-cp313t-win_amd64.whl", hash = "sha256:2140e373e9a51a3e22ef62e8d14366d0b470d18f0adf19fdc757368077133a34"}, - {file = "torch-2.12.0-cp314-cp314-macosx_14_0_arm64.whl", hash = "sha256:f7dfae4a519197dfa050e98d8e36378a0fb5899625a875c2b54445005a2e404e"}, - {file = "torch-2.12.0-cp314-cp314-manylinux_2_28_aarch64.whl", hash = "sha256:891c769072637c74e9a5a77a3bc782894696d8ffec83b938df8536dee7f0ba78"}, - {file = "torch-2.12.0-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:e2ad3eb85d39c3cab62dfa93ed5a73516e6a53c6713cb97d004004fe089f0f1f"}, - {file = "torch-2.12.0-cp314-cp314-win_amd64.whl", hash = "sha256:c66696857e987efb8bc1777a37357ec4f60ab5e8af6250b83d6034437fa2d8f3"}, - {file = "torch-2.12.0-cp314-cp314t-macosx_14_0_arm64.whl", hash = "sha256:b4556715c8572758625d62b6e0ae3b1f76c440221913a6fb5e100f321fb4fb02"}, - {file = "torch-2.12.0-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:a43ac605a5e13116c72b64c359644cce0229f213dde48d2ae0ae5eb5becf7feb"}, - {file = "torch-2.12.0-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:6a7512adfdd7f6732e40de1c620831e3c75b39b98cef60b11d0c5f0a76473ec5"}, - {file = "torch-2.12.0-cp314-cp314t-win_amd64.whl", hash = "sha256:5f96b63f8287f66a005dd1b5a6abba2920f11156c5e5c4d815f3e2050fd1aa16"}, + {file = "torch-2.13.0-cp310-cp310-macosx_14_0_arm64.whl", hash = "sha256:94f0de129916f77b8dc2c7a8eff644cfeddfe59e39c9f55e9f6e17543410281d"}, + {file = "torch-2.13.0-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:0ab4b69f3ee03a62a002cfbf77b1ca5e88aceb4ea64cb4388bb28f638ddbb045"}, + {file = "torch-2.13.0-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:c78b7b4d04461855a764cf01bae9a462bb88bc93defcfa11235cbc8fdf3e12c4"}, + {file = "torch-2.13.0-cp310-cp310-win_amd64.whl", hash = "sha256:2bd30b6b730d987fa386ce3898933762c5cb8cc82eb0535211d787cc3ce2dfeb"}, + {file = "torch-2.13.0-cp311-cp311-macosx_14_0_arm64.whl", hash = "sha256:e76f9bcecc52b8ff711239a2f7547d5353df95878ab232f0773c1d95928b92f8"}, + {file = "torch-2.13.0-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:092790c696a760c729fd5722835f50b9d81fd7c8f141571f3f3cf4081a8f664c"}, + {file = "torch-2.13.0-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:60fcdcb2f3876e21146cb4524ef06397d727ca9ad5f020818547e25075fe3cb7"}, + {file = "torch-2.13.0-cp311-cp311-win_amd64.whl", hash = "sha256:a0d8b11f16a48d60e2015d8213aa0390744cbebb98e58b62b3514dddc656e330"}, + {file = "torch-2.13.0-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:2fe228aba290d14b9f31b049be550dbd469c3fd3013d7a19705b30454da97027"}, + {file = "torch-2.13.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:572df8be8ffb4599c88cbd6a0726f1f854f4da65d2e3c09f0e2c2283333cd6d4"}, + {file = "torch-2.13.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:796633c4cdf0fe2cdced72d8f88f22e73dbcfce83132763162f6d4bff13b820b"}, + {file = "torch-2.13.0-cp312-cp312-win_amd64.whl", hash = "sha256:024c6cc0c1b085f2f91f20a3dc27b0471d021c31ce84b81be3afdc39f791fd9d"}, + {file = "torch-2.13.0-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:33449899ce5496c1b84b4853179d94fd102028ae1407314d9fb956bb79e70d09"}, + {file = "torch-2.13.0-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:1e09d6a722504957c694faceca843acde562786df1144ebcc5a74075ec7f6005"}, + {file = "torch-2.13.0-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:a3a9a21312872af8a26950b2c15680335a386a1f56ed03e780653d78b9607e9e"}, + {file = "torch-2.13.0-cp313-cp313-win_amd64.whl", hash = "sha256:49b58f1e2c52440abb6f17c28f0335fe6c6d01ad1a7f55b0183b81e4b34d64e6"}, + {file = "torch-2.13.0-cp314-cp314-macosx_14_0_arm64.whl", hash = "sha256:d849b390e07d8d333ce8ecaf91b273c656c598379a19c9acf1318a883f6b391c"}, + {file = "torch-2.13.0-cp314-cp314-manylinux_2_28_aarch64.whl", hash = "sha256:a3893dc2da0a972a8ca5d698c85a9f967559ac5f8ee1797b77408aa8734d073c"}, + {file = "torch-2.13.0-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:49f1ea385c754e54919408a9bb3b5a72b0b755bbe2c916c1d6f70afbec4908a2"}, + {file = "torch-2.13.0-cp314-cp314-win_amd64.whl", hash = "sha256:4f8573e3ce9ebcd53fe922f01077a6085ccdfbe5f12fd215883a9d87d7a744fd"}, + {file = "torch-2.13.0-cp314-cp314t-macosx_14_0_arm64.whl", hash = "sha256:c28def70706c2f9ecc752574766e8ae4da9b810ab6676b611166761a78a9f1e1"}, + {file = "torch-2.13.0-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:31061ff56ed8fbf26c749806905aeb749ebeb819810fd5d52508aa5afd90dddc"}, + {file = "torch-2.13.0-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:cc26eead4cf51d0b544e31e364dcf000846549c273bd148936fe9d24d29acb92"}, + {file = "torch-2.13.0-cp314-cp314t-win_amd64.whl", hash = "sha256:a7de8a313090dc5c7d7ba4bfe5c3be222528f9a4dba1acc83bddb1157360c4b8"}, ] [package.dependencies] -cuda-bindings = {version = ">=13.0.3,<14", markers = "platform_system == \"Linux\""} -cuda-toolkit = {version = "13.0.2", extras = ["cudart", "cufft", "cufile", "cupti", "curand", "cusolver", "cusparse", "nvjitlink", "nvrtc", "nvtx"], markers = "platform_system == \"Linux\""} +cuda-bindings = {version = ">=13.0.3,<14", markers = "platform_system == \"Linux\" and python_version < \"3.15\""} +cuda-toolkit = {version = "13.0.3", extras = ["cublas", "cudart", "cufft", "cufile", "cupti", "curand", "cusolver", "cusparse", "nvjitlink", "nvrtc", "nvtx"], markers = "platform_system == \"Linux\""} filelock = "*" fsspec = ">=0.8.5" jinja2 = "*" networkx = ">=2.5.1" -nvidia-cublas = {version = ">=13.1.0.3,<=13.1.1.3", markers = "platform_system == \"Linux\""} nvidia-cudnn-cu13 = {version = "9.20.0.48", markers = "platform_system == \"Linux\""} nvidia-cusparselt-cu13 = {version = "0.8.1", markers = "platform_system == \"Linux\""} nvidia-nccl-cu13 = {version = "2.29.7", markers = "platform_system == \"Linux\""} nvidia-nvshmem-cu13 = {version = "3.4.5", markers = "platform_system == \"Linux\""} -setuptools = "<82" +setuptools = ">=77.0.3" sympy = ">=1.13.3" -triton = {version = "3.7.0", markers = "platform_system == \"Linux\""} +triton = {version = "3.7.1", markers = "platform_system == \"Linux\" and python_version < \"3.15\""} typing-extensions = ">=4.10.0" [package.extras] @@ -6364,26 +6422,24 @@ vision = ["Pillow (>=10.0.1,<=15.0)", "torchvision"] [[package]] name = "triton" -version = "3.7.0" +version = "3.7.1" description = "A language and compiler for custom Deep Learning operations" optional = false python-versions = "<3.15,>=3.10" groups = ["main"] files = [ - {file = "triton-3.7.0-cp310-cp310-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:223ac302091491436c248a34ee1e6c47a1026486579103c906ffd805be50cb89"}, - {file = "triton-3.7.0-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c631b65668d4951213b948a413c0564184305b77bb45cc9d686d3e1ecc4701a3"}, - {file = "triton-3.7.0-cp311-cp311-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a9e71fc392675fac364e0ecf4ef3f76f85b7f5433a16f4c3c5fe5f05a52c85fe"}, - {file = "triton-3.7.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:22bacffce443f54593dd20f05294d5a40622e0ea9ab632816f87154504356221"}, - {file = "triton-3.7.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a4bf49b00a7a377a68a6da603a876e797614e6455a80e9021669c476a953ad9a"}, - {file = "triton-3.7.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:8f111161d49bf903c0eaedde3962353a3d841c08a836839b7cc1025b8426efcf"}, - {file = "triton-3.7.0-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:abdf6beaa89b1bcfb9a43cd990536ce66091a997841a4814b260b7bee4c88c3c"}, - {file = "triton-3.7.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a35d7afe3f3f058e7ec49fcce09794049e0ffc5c59019ac25ec3413741b8c4e7"}, - {file = "triton-3.7.0-cp313-cp313t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:cc1d61c172d257db80ddf42595131fb196ad2e9bdd751e90fe2ef13531734e8b"}, - {file = "triton-3.7.0-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:70fb9bbdc9f400afc54bbf6eb2670af28829a6ae3996863317964783141daf56"}, - {file = "triton-3.7.0-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c4a44a8476d0d3571eac4e4d1048e1ff75aad81a09ff4602ccfc56c6dea1672e"}, - {file = "triton-3.7.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b9b85e72968a9d8bba5ddb24e9b64aaabaf48affb042f2755cb7cfa92b7531ce"}, - {file = "triton-3.7.0-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:18a160de426fd99f92b0baf509045360afbd3bfaa0b4a5171dde800ec9f09684"}, - {file = "triton-3.7.0-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ce061073102714b725f3660ec6939d94a1da7984b3aa99c921417cae273672f5"}, + {file = "triton-3.7.1-cp310-cp310-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3daf64305d6cea88d3334c65ebc9bcd0c64c9564a977084366aa768d57cbcf64"}, + {file = "triton-3.7.1-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ee89fbf782ec2ad50391dd1cf26cbea4f4467154c37f4773026da8fc31c0f58e"}, + {file = "triton-3.7.1-cp311-cp311-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d4a0e1cd4c4a76370ed74a8432a53cea28716827d19e40ffc732233e35ceb3f6"}, + {file = "triton-3.7.1-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:6744957e9fd610a29680ec2346057d0c86948ed3812468670719f391e94b44a5"}, + {file = "triton-3.7.1-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:9497f2e696ee368862a181a90b2dcc03ca978cc4f602abd67c7d81022a6988e1"}, + {file = "triton-3.7.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:7e40869937a68206ec70d7f25bb7ec6433cb083f9135e1f36dbd318dc449a728"}, + {file = "triton-3.7.1-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:cdbfc09d9ec58bc5e68321525653220de7515c199e7a8097a97c85e62b52cd0a"}, + {file = "triton-3.7.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:58c0e131da05134a2a4788ccbcc0c1105cf0f54c8e98f19e34cd465396dc15eb"}, + {file = "triton-3.7.1-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:fe4ea396a06171f1f1f58cbd39c70b09294398f7dd7c620939bab54ad6f934fa"}, + {file = "triton-3.7.1-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:2020153b08280415ec0da6607834e79166442147e78e144df06b508c75b186d2"}, + {file = "triton-3.7.1-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c58e4c61f0c73b5dba3b5d19b4a7093c32f90dc18b2a7f121a7c16ccd31107b7"}, + {file = "triton-3.7.1-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:10ba85fa2cca4a2fbdeb36bf1cb082f2c252bda55bf9fccd74f65ec5bc647e68"}, ] [package.extras] @@ -6456,7 +6512,7 @@ description = "Provider of IANA time zone data" optional = false python-versions = ">=2" groups = ["main"] -markers = "sys_platform == \"win32\" or sys_platform == \"emscripten\" or python_version == \"3.10\"" +markers = "python_version == \"3.10\" or sys_platform == \"win32\" or sys_platform == \"emscripten\"" files = [ {file = "tzdata-2026.4-py2.py3-none-any.whl", hash = "sha256:c2169a8b0a7a5e9674da5a135ccdfb2b3e671b333ed9fed17b41f73c34476e81"}, {file = "tzdata-2026.4.tar.gz", hash = "sha256:f1b8bd365d8d210c55353f4d7f8d6d8561c0ba50d704b700d195a9424bba0d79"}, @@ -6482,14 +6538,14 @@ zstd = ["backports-zstd (>=1.0.0) ; python_version < \"3.14\""] [[package]] name = "uvicorn" -version = "0.52.4" +version = "0.53.0" description = "The lightning-fast ASGI server." optional = false python-versions = ">=3.10" groups = ["main"] files = [ - {file = "uvicorn-0.52.4-py3-none-any.whl", hash = "sha256:f86e41a149d7d05a9969337e3946a9c171c06a5d42680896daaba624aeac8da1"}, - {file = "uvicorn-0.52.4.tar.gz", hash = "sha256:73acfee47a0b133c5de13d219492d62d8a31e935f4fe6e41a232451a15379f86"}, + {file = "uvicorn-0.53.0-py3-none-any.whl", hash = "sha256:e8dca71ec86dce5f04e333f0d56cdedf942446e6643b9cea1af0d6d3a02cb03e"}, + {file = "uvicorn-0.53.0.tar.gz", hash = "sha256:a9356f0cb89b3b8621529c5d5eebd69bfe154f4c3f68b4cf2de47e45fa855c2e"}, ] [package.dependencies] @@ -7022,4 +7078,4 @@ type = ["pytest-mypy (>=1.0.1) ; platform_python_implementation != \"PyPy\""] [metadata] lock-version = "2.1" python-versions = ">=3.10,<3.13" -content-hash = "7cf2f5f8c434f0e34a542cd206a3cc6af74f77aad39eeec26bcaabe98386a904" +content-hash = "37c9c8fd0819dc9dfdea63196195c64d9a67458bb10a3f979ba90cf87147d339" diff --git a/security_scanning/pyproject.toml b/security_scanning/pyproject.toml index 6aea187b6607..4ac03290e6c4 100644 --- a/security_scanning/pyproject.toml +++ b/security_scanning/pyproject.toml @@ -20,7 +20,7 @@ dependencies = [ "mpi4py (>=4.1.2,<5.0.0)", "numpy (>=2.0.0,<2.4)", "graphviz (>=0.21,<0.22)", - "openai (>=3.13.0,<4.0.0)", + "openai (>=3.14.0,<4.0.0)", "polygraphy (>=0.53.4,<0.54.0)", "psutil (>=7.2.2,<8.0.0)", "nvidia-ml-py (>=13)", @@ -28,8 +28,8 @@ dependencies = [ "h5py (==3.12.1)", "strenum (>=0.4.15,<0.5.0)", "sentencepiece (>=0.1.99)", - "torch (>=2.12.0a0,<=2.13.0a0)", - "nvidia-nccl-cu13 (>=2.29.7,<=2.30.4)", + "torch (>=2.12.0a0,<=2.14.0a0)", + "nvidia-nccl-cu13 (>=2.29.7,<=2.30.7)", "nccl4py (>=0.3.1,<0.4)", "transformers (==5.5.4)", "prometheus-client (>=0.26.0,<0.27.0)", @@ -50,7 +50,7 @@ dependencies = [ "pyzmq (>=27.2.0,<28.0.0)", "fastapi (>=0.136.3)", "starlette (>=1.3.1)", - "uvicorn (>=0.52.4,<0.53.0)", + "uvicorn (>=0.53.0,<0.54.0)", "setuptools (<80)", "packaging (>=24.2)", "ordered-set (>=4.1.0,<5.0.0)", From 6cae275f168a0add00be26f1cda11b8d81ef4a20 Mon Sep 17 00:00:00 2001 From: Tingfeng Xian <289617005+nv-xtf@users.noreply.github.com> Date: Tue, 15 Sep 2026 11:32:07 +0800 Subject: [PATCH 8/8] [TRTLLM-15715][refactor] Extract progress polling and error consensus into DisaggTransferCoordinator (#19128) Signed-off-by: Tingfeng Xian <289617005+nv-xtf@users.noreply.github.com> --- .../orchestration/coordinator.py | 174 +++++++++- .../orchestration/interfaces.py | 10 +- .../_torch/pyexecutor/disagg_adapter.py | 8 + tensorrt_llm/_torch/pyexecutor/py_executor.py | 186 +---------- .../_torch/disaggregation/conftest.py | 20 ++ .../disaggregation/coordinator_harness.py | 114 +++++++ .../_torch/disaggregation/fake_dist.py | 247 ++++++++++++++ .../disaggregation/fake_executor_effects.py | 6 + .../fake_kv_cache_transceiver.py | 5 +- .../disaggregation/test_benchmark_disagg.py | 16 +- .../test_disagg_adapter_contract.py | 39 ++- .../disaggregation/test_disagg_coordinator.py | 5 +- .../test_disagg_coordinator_errors.py | 284 ++++++++++++++++ .../test_disagg_coordinator_progress.py | 256 +++++++++++++++ .../test_disagg_coordinator_transfers.py | 122 +------ .../test_disagg_inflight_cancel_gate.py | 50 --- .../test_disagg_loop_transcript.py | 13 +- .../_torch/disaggregation/test_fake_dist.py | 155 +++++++++ .../_torch/executor/test_py_executor.py | 305 +----------------- .../executor/test_send_kv_async_split.py | 2 +- .../disaggregated/test_chunked_transfer.py | 14 +- .../test_transfer_ownership_regressions.py | 3 +- 22 files changed, 1348 insertions(+), 686 deletions(-) create mode 100644 tests/unittest/_torch/disaggregation/conftest.py create mode 100644 tests/unittest/_torch/disaggregation/coordinator_harness.py create mode 100644 tests/unittest/_torch/disaggregation/fake_dist.py create mode 100644 tests/unittest/_torch/disaggregation/test_disagg_coordinator_errors.py create mode 100644 tests/unittest/_torch/disaggregation/test_disagg_coordinator_progress.py create mode 100644 tests/unittest/_torch/disaggregation/test_fake_dist.py diff --git a/tensorrt_llm/_torch/disaggregation/orchestration/coordinator.py b/tensorrt_llm/_torch/disaggregation/orchestration/coordinator.py index 2a4f69a120d2..05ffb8e895ff 100644 --- a/tensorrt_llm/_torch/disaggregation/orchestration/coordinator.py +++ b/tensorrt_llm/_torch/disaggregation/orchestration/coordinator.py @@ -13,12 +13,14 @@ from dataclasses import dataclass, fields from typing import TYPE_CHECKING, Callable, List, Set, Tuple +from tensorrt_llm._torch.disaggregation.base.transfer import get_unique_rid from tensorrt_llm._torch.disaggregation.kv_cache_transceiver import ( is_disagg_inflight_cancel_enabled, ) from tensorrt_llm._torch.distributed.communicator import ReduceOp from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, LlmRequestState from tensorrt_llm._utils import nvtx_range +from tensorrt_llm.disaggregated_params import DisaggScheduleStyle from tensorrt_llm.logger import logger from .interfaces import ActiveRequestRegistry, ExecutorEffects @@ -52,19 +54,13 @@ class DisaggLoopDelegates: """Executor callables the coordinator still forwards to. Transitional: each field is removed once the corresponding logic moves - into the coordinator (CS-2: progress; CS-3: errors, admission, receive). + into the coordinator (CS-3: admission, receive). """ - handle_errors_synced: Callable[[], None] - prepare_context_schedulable: Callable[[List[LlmRequest]], None] admit: Callable[[List[LlmRequest]], Tuple[List[LlmRequest], bool]] revert_deferred_gen_init: Callable[[List[LlmRequest], List[LlmRequest]], None] receive_gen_init: Callable[[List[LlmRequest]], None] - poll_progress_when_idle: Callable[[], None] prepare_transmission_completed: Callable[["ScheduledRequests"], None] - # Rank-local transfer error handling; reached from the reaps. CS-3. - check_transfer_errors: Callable[[str], None] - requests_in_error_state: Callable[[], List[LlmRequest]] class DisaggTransferCoordinator: @@ -114,13 +110,92 @@ def __init__( # -- loop head ----------------------------------------------------------- + @nvtx_range("handle_errors_synced") def handle_errors_synced(self) -> None: - """Fail requests whose transfer errored; rank-synchronized.""" - self._d.handle_errors_synced() + """Rank-safe disagg cache error and poison handler. + + Called from the top of every executor iteration. Buffer poison is + reduced over the full executor world because one poisoned PP/DP rank + requires the whole distributed executor to stop. ADP TP ranks then + vote on failed request IDs and fail matching local replicas together; + otherwise the downstream ``tp_gather`` in ``_enqueue_responses`` + deadlocks or leaves peer replicas running. + """ + pending_ids = self.take_pending_context_failures() + pending_requests = ( + [req for req in self._registry.active_requests() if get_unique_rid(req) in pending_ids] + if pending_ids + else [] + ) + for request in pending_requests: + request.state = LlmRequestState.DISAGG_TRANS_ERROR + + if self.inflight_cancel_active(): + local_poisoned = self._transceiver.has_poisoned_transfer_buffer() + if self._dist.world_size != 1: + any_poisoned = bool(self._dist.allreduce(int(local_poisoned), op=ReduceOp.MAX)) + else: + any_poisoned = local_poisoned + if any_poisoned: + self._effects.fail_fatal( + "Disagg KV cache transfer buffer is poisoned; process restart is required" + ) + return + if not (self._enable_attention_dp and self._dist.world_size != 1): + if pending_requests: + self.check_transfer_errors("context requests") + return + + local_error_requests = [ + req + for req in self._registry.active_requests() + if req.state == LlmRequestState.DISAGG_TRANS_ERROR + ] + local_vote = { + "error_ids": [self._vote_id(req) for req in local_error_requests], + "blocked_ids": [ + self._vote_id(req) + for req in local_error_requests + if self._is_error_cleanup_blocked(req) + ], + } + all_votes = self._dist.tp_allgather(local_vote) + voted_error_ids = {rid for vote in all_votes for rid in vote["error_ids"]} + blocked_error_ids = {rid for vote in all_votes for rid in vote["blocked_ids"]} + ready_error_ids = voted_error_ids - blocked_error_ids + if not ready_error_ids: + return + local_voted_error_requests = [ + req for req in self._registry.active_requests() if self._vote_id(req) in ready_error_ids + ] + logger.warning( + f"Disagg KV cache transfer error: rank={self._dist.rank} " + f"local_err_count={len(local_error_requests)}, " + f"voted_err_count={len(voted_error_ids)}, " + f"blocked_err_count={len(voted_error_ids & blocked_error_ids)}" + ) + self._effects.fail_requests( + "Disagg KV cache transfer error", local_voted_error_requests, charge_budget=False + ) + + @nvtx_range("prepare_context_schedulable") def prepare_context_schedulable(self, new_requests: List[LlmRequest]) -> None: - """Let the transceiver gate generation-first context requests.""" - self._d.prepare_context_schedulable(new_requests) + """Let the transceiver gate generation-first context requests. + + Context-first context requests are schedulable at once; for + generation-first ones the transceiver decides when the peer is ready. + """ + gen_first_ctx_requests = [ + req + for req in new_requests + if req.is_context_only_request + and req.py_disaggregated_params.schedule_style == DisaggScheduleStyle.GENERATION_FIRST + ] + # Always call prepare_context_requests, with new requests or without, + # so the consensus inside it can promote requests whose peer info has + # arrived on every rank. + self._transceiver.prepare_context_requests(gen_first_ctx_requests) @nvtx_range("poll_gen_transfers") def poll_gen_transfers(self) -> None: @@ -194,8 +269,24 @@ def receive_gen_init(self, admitted: List[LlmRequest]) -> None: self._d.receive_gen_init(admitted) def poll_progress_when_idle(self) -> None: - """Reap completed context sends; rank-symmetric.""" - self._d.poll_progress_when_idle() + """Reap completed context KV transfers so their blocks can be freed. + + A synchronous GEN receive blocks rank-locally, so a multi-rank worker + must not enter the context status collective here. A single-rank + worker cannot diverge and polls only while a send is in flight. + """ + uses_synchronous_gen_transfer = ( + not uses_async_gen_transfer() and not is_gen_only_no_context_benchmark() + ) + should_poll_synchronous_context_status = ( + uses_synchronous_gen_transfer + and self._dist.world_size == 1 + and self._transfers.has_any_inflight_requests() + ) + if uses_synchronous_gen_transfer and not should_poll_synchronous_context_status: + return + + self.reap_context_sends(0) # -- batch execution ----------------------------------------------------- @@ -299,7 +390,7 @@ def reap_context_sends(self, at_least: int = 0) -> None: request.state = LlmRequestState.DISAGG_CONTEXT_COMPLETE self.release_transfer(request) - self._d.check_transfer_errors("context requests") + self.check_transfer_errors("context requests") @nvtx_range("reap_gen_receives") def reap_gen_receives(self, at_least: int = 0) -> None: @@ -312,7 +403,7 @@ def reap_gen_receives(self, at_least: int = 0) -> None: if req_id not in user_canceled_ids: req.state = LlmRequestState.DISAGG_TRANS_ERROR if not self.inflight_cancel_active(): - self._d.check_transfer_errors("generation requests") + self.check_transfer_errors("generation requests") def release_transfer(self, request: LlmRequest) -> None: """Release one transfer claim and terminate once the last owner releases. @@ -500,7 +591,7 @@ def _cancel_timed_out_gen_transfers(self) -> None: def _check_gen_transfer_errors_consensus(self) -> None: """Flush generation transfer errors through a TP-uniform path.""" error_requests = [ - req for req in self._d.requests_in_error_state() if req.is_generation_only_request + req for req in self._requests_in_error_state() if req.is_generation_only_request ] local_needs_flush = bool(error_requests) if self._dist.tp_size > 1: @@ -515,6 +606,45 @@ def _check_gen_transfer_errors_consensus(self) -> None: charge_budget=False, ) + # -- transfer errors ----------------------------------------------------- + + def check_transfer_errors(self, kind: str) -> None: + """Fail requests whose transfer errored, rank-locally. + + Under multi-rank ADP this is a no-op: errors are handled by + ``handle_errors_synced`` at the loop top. Public only because the + executor's synchronous receive path still calls it (CS-3 moves it). + """ + if self._enable_attention_dp and self._dist.world_size != 1: + return + error_requests = self._requests_in_error_state() + if error_requests: + self._effects.fail_requests( + f"Error in kv cache transfer for {kind}", error_requests, charge_budget=False + ) + + def _requests_in_error_state(self) -> List[LlmRequest]: + return [ + req + for req in self._registry.active_requests() + if req.state == LlmRequestState.DISAGG_TRANS_ERROR + and not self._is_error_cleanup_blocked(req) + ] + + def _is_error_cleanup_blocked(self, request: LlmRequest) -> bool: + """Whether a failed request must wait: the cancel path owns it, or a + transfer owner still holds its context blocks.""" + if self._vote_id(request) in self._registry.canceled_request_ids(): + return True + return ( + getattr(request, "is_context_only_request", False) is True + and request.py_request_id in self._transfers.requests_in_transfer() + ) + + @staticmethod + def _vote_id(request: LlmRequest) -> int: + return request.parent_request_id if request.is_child else request.py_request_id + # -- loop tail ----------------------------------------------------------- def pace_idle(self) -> None: @@ -557,12 +687,21 @@ def __init__(self) -> None: delegates=DisaggLoopDelegates(**{f.name: _noop for f in fields(DisaggLoopDelegates)}), ) + def handle_errors_synced(self) -> None: + return None + def admit(self, fitting_gen_init: List[LlmRequest]) -> Tuple[List[LlmRequest], bool]: return fitting_gen_init, False + def prepare_context_schedulable(self, new_requests: List[LlmRequest]) -> None: + return None + def poll_gen_transfers(self) -> None: return None + def poll_progress_when_idle(self) -> None: + return None + def check_transfer_timeouts(self, only_with_context_sends: bool = False) -> None: return None @@ -581,6 +720,9 @@ def release_transfer(self, request: LlmRequest) -> None: "releases are handled by the executor" ) + def check_transfer_errors(self, kind: str) -> None: + return None + def inflight_cancel_active(self) -> bool: return False diff --git a/tensorrt_llm/_torch/disaggregation/orchestration/interfaces.py b/tensorrt_llm/_torch/disaggregation/orchestration/interfaces.py index ec15ce3565b2..8fa04bafd095 100644 --- a/tensorrt_llm/_torch/disaggregation/orchestration/interfaces.py +++ b/tensorrt_llm/_torch/disaggregation/orchestration/interfaces.py @@ -15,7 +15,7 @@ class ExecutorEffects(Protocol): """Executor-owned side effects the coordinator may trigger. - These three are the complete set; adding one is a design decision, not + These four are the complete set; adding one is a design decision, not a convenience. """ @@ -47,6 +47,14 @@ def fail_requests( """Fail requests through the executor's error path.""" ... + def fail_fatal(self, error_msg: str) -> None: + """Mark the executor fatal and fail every active request. + + Called only after a world-wide collective agreed on the failure, so + the executor takes its collective-aligned fatal path on every rank. + """ + ... + class ActiveRequestRegistry(Protocol): """The executor's request bookkeeping, as the coordinator may touch it. diff --git a/tensorrt_llm/_torch/pyexecutor/disagg_adapter.py b/tensorrt_llm/_torch/pyexecutor/disagg_adapter.py index 1c0eacf13ea0..a1b7a84ec476 100644 --- a/tensorrt_llm/_torch/pyexecutor/disagg_adapter.py +++ b/tensorrt_llm/_torch/pyexecutor/disagg_adapter.py @@ -32,6 +32,14 @@ def fail_requests( error_msg=error_msg, requests=requests, charge_budget=charge_budget ) + def fail_fatal(self, error_msg: str) -> None: + executor = self._executor + executor._fatal_error = RuntimeError(f"Fatal error: {error_msg}") + executor.is_shutdown = True + executor._handle_errors( + error_msg, requests=None, charge_budget=False, fatal_is_collective_aligned=True + ) + class PyExecutorRequestRegistry(ActiveRequestRegistry): def __init__(self, executor: "PyExecutor") -> None: diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 48ed6c771255..dc2f60fceb6d 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -3586,19 +3586,12 @@ def _build_disagg_coordinator(self) -> DisaggTransferCoordinator: force_terminate_ctx_for_partial_reuse=getattr( self, "force_terminate_ctx_for_partial_reuse", False), delegates=DisaggLoopDelegates( - handle_errors_synced=self._handle_disagg_cache_errors_synced, - prepare_context_schedulable=self. - _check_disagg_ctx_schedulable_status, admit=self._apply_disagg_transfer_admission, revert_deferred_gen_init=self. _revert_deferred_disagg_gen_init_alloc, receive_gen_init=self._prepare_disagg_gen_init, - poll_progress_when_idle=self. - _check_disagg_transfer_progress_when_idle, prepare_transmission_completed=self. _prepare_disagg_gen_transmission_complete, - check_transfer_errors=self._check_cache_transfer_errors, - requests_in_error_state=self._get_disagg_reqs_in_error_state, )) def _get_disagg_transfer_admission_controller( @@ -3729,26 +3722,6 @@ def _allgather_model_parallel_status( return self.dist.tp_allgather(local_status) return [local_status] - def _check_disagg_transfer_progress_when_idle(self) -> None: - """Reap completed context KV transfers so their blocks can be freed. - - A synchronous GEN receive blocks rank-locally, so a multi-rank worker - must not enter the context status collective here. A single-rank - worker cannot diverge and polls only while a send is in flight. - """ - uses_synchronous_gen_transfer = ( - not self._uses_async_disagg_gen_transfer() - and not self._is_disagg_gen_only_no_context_benchmark()) - should_poll_synchronous_context_status = ( - uses_synchronous_gen_transfer - and self._dist_size(self.dist, "world_size") == 1 - and self.async_transfer_manager.has_any_inflight_requests()) - if (uses_synchronous_gen_transfer - and not should_poll_synchronous_context_status): - return - - self.disagg.reap_context_sends(0) - def _pp_ring_is_drained(self) -> bool: """Return whether no microbatch is queued or awaiting handling.""" return (self.unhandled_batch_counter == 0 @@ -4163,97 +4136,6 @@ def _fail_if_fill_gate_stalled(self, made_progress: bool) -> None: f"completing; the loop would otherwise retry forever without " f"advancing iter_counter.") - @nvtx_range("_handle_disagg_cache_errors_synced") - def _handle_disagg_cache_errors_synced(self): - """Rank-safe disagg cache error and poison handler. - - Called from the top of every executor iteration. Buffer poison is - reduced over the full executor world because one poisoned PP/DP rank - requires the whole distributed executor to stop. ADP TP ranks then - vote on failed request IDs and fail matching local replicas together; - otherwise the downstream ``tp_gather`` in ``_enqueue_responses`` - deadlocks or leaves peer replicas running. - """ - if not self.kv_cache_transceiver: - return - - pending_ids = self.disagg.take_pending_context_failures() - pending_requests = ([ - request for request in self.active_requests - if get_unique_rid(request) in pending_ids - ] if pending_ids else []) - for request in pending_requests: - request.state = LlmRequestState.DISAGG_TRANS_ERROR - - if self.disagg.inflight_cancel_active(): - local_poisoned = self.kv_cache_transceiver.has_poisoned_transfer_buffer( - ) - if self.dist.world_size != 1: - any_poisoned = bool( - self.dist.allreduce(int(local_poisoned), op=ReduceOp.MAX)) - else: - any_poisoned = local_poisoned - if any_poisoned: - error_msg = ( - "Disagg KV cache transfer buffer is poisoned; process " - "restart is required") - self._fatal_error = RuntimeError(f"Fatal error: {error_msg}") - self.is_shutdown = True - self._handle_errors(error_msg, - requests=None, - charge_budget=False, - fatal_is_collective_aligned=True) - return - - if not (self.enable_attention_dp and self.dist.world_size != 1): - if pending_requests: - self._check_cache_transfer_errors("context requests") - return - - local_error_requests = [ - request for request in self.active_requests - if request.state == LlmRequestState.DISAGG_TRANS_ERROR - ] - local_vote = { - "error_ids": [ - self._request_vote_id(request) - for request in local_error_requests - ], - "blocked_ids": [ - self._request_vote_id(request) - for request in local_error_requests - if self._is_disagg_error_cleanup_blocked(request) - ], - } - all_votes = self.dist.tp_allgather(local_vote) - voted_error_ids = { - request_id - for rank_vote in all_votes - for request_id in rank_vote["error_ids"] - } - blocked_error_ids = { - request_id - for rank_vote in all_votes - for request_id in rank_vote["blocked_ids"] - } - ready_error_ids = voted_error_ids - blocked_error_ids - if not ready_error_ids: - return - local_voted_error_requests = [ - request for request in self.active_requests - if self._request_vote_id(request) in ready_error_ids - ] - logger.warning( - f"Disagg KV cache transfer error: rank={self.dist.rank} " - f"local_err_count={len(local_error_requests)}, " - f"voted_err_count={len(voted_error_ids)}, " - f"blocked_err_count={len(voted_error_ids & blocked_error_ids)}") - self._handle_errors( - "Disagg KV cache transfer error", - requests=local_voted_error_requests, - charge_budget=False, - ) - def _emit_initial_stats(self) -> None: """Emit a startup stats snapshot so that cache_config_info is immediately available to external metric scrapers (e.g. the @@ -4379,8 +4261,8 @@ def _executor_loop(self): scheduled_batch, iter_stats = self._prepare_and_schedule_batch() if scheduled_batch is None: - # _handle_disagg_cache_errors_synced() can buffer a - # non-fatal response before scheduling observes shutdown. + # handle_errors_synced() can buffer a non-fatal response + # before scheduling observes shutdown. # Drain it before leaving the loop so the client does not # wait for its own timeout. Scheduling shutdown is # model-parallel synchronized, so every ADP rank reaches @@ -6900,26 +6782,6 @@ def _mark_cross_kv_projection_consumed( req.py_encoder_output = None req.py_skip_cross_kv_projection = True - @nvtx_range("_check_disagg_ctx_schedulable_status") - def _check_disagg_ctx_schedulable_status(self, - new_requests: List[LlmRequest]): - """ - In context-first mode, context requests are schedulable immediately, - otherwise, we need to check if context requests are ready to be scheduled by querying kv cache transceiver - """ - if not self.kv_cache_transceiver: - return - gen_first_ctx_requests = [ - req for req in new_requests - if req.is_context_only_request and req.py_disaggregated_params. - schedule_style == DisaggScheduleStyle.GENERATION_FIRST - ] - # Always call prepare_context_requests when there are new requests - # or previously-waiting requests, so the tp_allgather consensus - # can promote requests whose peer info has arrived on all ranks. - self.kv_cache_transceiver.prepare_context_requests( - gen_first_ctx_requests) - def _count_schedulable_active_requests(self) -> int: """Count active requests that are ready for scheduling. @@ -7550,7 +7412,7 @@ def _recv_disagg_gen_cache(self, new_gen_reqs): # prepared request is left in DISAGG_GENERATION_INIT. for req in new_gen_reqs: self.kv_cache_transceiver.request_and_receive_sync(req) - self._check_cache_transfer_errors("generation requests") + self.disagg.check_transfer_errors("generation requests") return for req in new_gen_reqs: @@ -7600,45 +7462,6 @@ def kv_connector_request_finished(req: LlmRequest): if req.is_finished: kv_connector_request_finished(req) - @staticmethod - def _request_vote_id(request: LlmRequest) -> int: - return (request.parent_request_id - if request.is_child else request.py_request_id) - - def _is_disagg_error_cleanup_blocked(self, request: LlmRequest) -> bool: - request_id = self._request_vote_id(request) - if request_id in getattr(self, "canceled_req_ids", ()): - return True - - async_transfer_manager = getattr(self, "async_transfer_manager", None) - if (getattr(request, "is_context_only_request", False) is True - and async_transfer_manager is not None and request.py_request_id - in async_transfer_manager.requests_in_transfer()): - return True - return False - - def _get_disagg_reqs_in_error_state(self): - return [ - req for req in self.active_requests - if req.state == LlmRequestState.DISAGG_TRANS_ERROR - and not self._is_disagg_error_cleanup_blocked(req) - ] - - def _check_cache_transfer_errors(self, error_msg_prefix: str): - """Check and handle cache transfer errors. - - Under ADP this is a no-op: errors are handled by - ``_handle_disagg_cache_errors_synced`` at the loop top. - """ - if self.enable_attention_dp and self.dist.world_size != 1: - return - error_requests = self._get_disagg_reqs_in_error_state() - if error_requests: - error_msg = f"Error in kv cache transfer for {error_msg_prefix}" - self._handle_errors(error_msg, - requests=error_requests, - charge_budget=False) - def _maybe_prefetch_next_iter_mm_encoders( self, scheduled_batch: ScheduledRequests) -> None: """Best-effort hook for cross-iter MM encoder prefetch. @@ -7967,7 +7790,8 @@ def _handle_errors(self, multi_rank_adp = (self.enable_attention_dp and self.dist.world_size != 1) # ``fatal_is_collective_aligned`` is set only by the synchronized caller - # (_handle_disagg_cache_errors_synced, after its world allreduce), which + # (the coordinator's handle_errors_synced, through fail_fatal after its + # world allreduce), which # guarantees every ADP rank enters the fatal path in the same collective # order. Do NOT infer it from ``self._fatal_error is not None``: a # rank-local setter would then route into a tp_gather while peers are diff --git a/tests/unittest/_torch/disaggregation/conftest.py b/tests/unittest/_torch/disaggregation/conftest.py new file mode 100644 index 000000000000..769c91c0c81a --- /dev/null +++ b/tests/unittest/_torch/disaggregation/conftest.py @@ -0,0 +1,20 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fixtures shared by the coordinator behavior tests in this directory.""" + +import pytest + +from tensorrt_llm._torch.disaggregation.orchestration import coordinator as coordinator_module + + +@pytest.fixture +def inflight_cancel(monkeypatch): + monkeypatch.setattr(coordinator_module, "is_disagg_inflight_cancel_enabled", lambda: True) + + +@pytest.fixture +def clock(monkeypatch): + """Freeze the coordinator's ``time.monotonic``; tests advance ``clock["t"]``.""" + now = {"t": 100.0} + monkeypatch.setattr(coordinator_module.time, "monotonic", lambda: now["t"]) + return now diff --git a/tests/unittest/_torch/disaggregation/coordinator_harness.py b/tests/unittest/_torch/disaggregation/coordinator_harness.py new file mode 100644 index 000000000000..377107ae7257 --- /dev/null +++ b/tests/unittest/_torch/disaggregation/coordinator_harness.py @@ -0,0 +1,114 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Shared harness for ``DisaggTransferCoordinator`` behavior tests. + +The coordinator runs against the contract fake transceiver, a real +``AsyncTransferManager`` and stateful fakes of the executor interfaces, so +tests assert on request state, transfer ownership and what the executor was +asked to do. Multi-rank tests build one harness per rank and share only +``dist``; everything else is rank-local. +""" + +from types import SimpleNamespace +from unittest.mock import Mock + +from fake_executor_effects import FakeExecutorEffects, FakeRequestRegistry +from fake_kv_cache_transceiver import FakeKvCacheTransceiver + +from tensorrt_llm._torch.disaggregation.orchestration.coordinator import DisaggTransferCoordinator +from tensorrt_llm._torch.disaggregation.orchestration.transfer_manager import AsyncTransferManager +from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManagerType +from tensorrt_llm.bindings import LlmRequestState + + +class TransferRequest(SimpleNamespace): + """Request stub with the attributes the transfer paths read.""" + + def __init__(self, rid: int, **overrides) -> None: + defaults = dict( + py_request_id=rid, + request_id=rid, + parent_request_id=None, + is_child=False, + state=LlmRequestState.CONTEXT_INIT, + is_context_only_request=True, + is_context_finished=True, + is_finished_due_to_length=False, + is_finished_due_to_cancellation=False, + is_disagg_generation_init_state=False, + is_disagg_generation_transmission_in_progress=False, + py_kv_transfer_start_time=None, + py_kv_transfer_timed_out=False, + py_disaggregated_params=None, + cached_tokens=0, + response=None, + ) + defaults.update(overrides) + super().__init__(**defaults) + self.state_at_response_creation = None + + def create_response(self, _use_fast_logits, _rank): + self.state_at_response_creation = self.state + return self.response + + @property + def is_generation_only_request(self) -> bool: + return not self.is_context_only_request + + +class CoordinatorHarness: + """One coordinator with its rank-local collaborators. + + ``dist`` defaults to a ``Mock`` carrying ``rank`` / ``tp_size`` / + ``world_size`` for single-rank tests; multi-rank tests pass one FakeDist + rank per harness instead, and the rank object then defines those sizes. + """ + + def __init__( + self, + *, + kv_transfer_timeout_ms=None, + supports_inflight_cancellation=False, + enable_attention_dp=False, + world_size=1, + tp_size=1, + force_terminate_ctx_for_partial_reuse=False, + draft_kv_cache_manager=None, + dist=None, + ) -> None: + self.transceiver = FakeKvCacheTransceiver( + kv_transfer_timeout_ms=kv_transfer_timeout_ms, + supports_inflight_cancellation=supports_inflight_cancellation, + ) + self.transceiver.has_retired_send_session = lambda req: False + self.kv_cache_manager = Mock(spec=["store_blocks_for_reuse", "unpin_blocks_by_id"]) + self.kv_cache_manager.store_blocks_for_reuse.side_effect = lambda req, _: req.py_request_id + resource_manager = SimpleNamespace( + resource_managers={ResourceManagerType.KV_CACHE_MANAGER: self.kv_cache_manager} + ) + self.transfers = AsyncTransferManager(resource_manager) + self.active = [] + self.registry = FakeRequestRegistry(self.active) + self.effects = FakeExecutorEffects() + self.dist = ( + dist if dist is not None else Mock(rank=0, tp_size=tp_size, world_size=world_size) + ) + self.delegates = Mock() + self.coordinator = DisaggTransferCoordinator( + transceiver=self.transceiver, + transfer_manager=self.transfers, + kv_cache_manager=self.kv_cache_manager, + dist=self.dist, + effects=self.effects, + registry=self.registry, + enable_attention_dp=enable_attention_dp, + force_terminate_ctx_for_partial_reuse=force_terminate_ctx_for_partial_reuse, + delegates=self.delegates, + draft_kv_cache_manager=draft_kv_cache_manager, + ) + + def send(self, *requests: TransferRequest) -> None: + self.coordinator.send_completed_context(list(requests)) + + def in_transfer(self, req: TransferRequest) -> bool: + return req.py_request_id in self.transfers.requests_in_transfer() diff --git a/tests/unittest/_torch/disaggregation/fake_dist.py b/tests/unittest/_torch/disaggregation/fake_dist.py new file mode 100644 index 000000000000..d39edac685b9 --- /dev/null +++ b/tests/unittest/_torch/disaggregation/fake_dist.py @@ -0,0 +1,247 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""In-process fake of the executor's ``dist`` object for multi-rank coordinator tests. + +A ``FakeDistGroup`` is a world of ``world_size`` ranks split into TP groups of +``tp_size`` consecutive ranks, one thread per rank. Each collective is a +``threading.Barrier`` rendezvous over its group: the call blocks until every +rank of the group has entered a collective, checks that all of them entered +the same one, then returns the gathered or reduced payloads. The fake +verifies the protocol -- which collectives each rank enters, how often, in +what order and with what payload -- not the blocking semantics of a real +communication backend. + +Barrier timeouts are wall-clock inside ``threading`` and unaffected by tests +that patch ``time.monotonic``. +""" + +import copy +import threading +from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple + +import numpy as np + +from tensorrt_llm._torch.distributed.communicator import ReduceOp + + +class FakeDistTimeout(AssertionError): + """A rank waited for a collective that some peer never entered.""" + + +class FakeDistMismatch(AssertionError): + """Ranks of one group entered different collectives at the same step.""" + + +class _PeerFailed(Exception): + """A pending collective was released because another rank raised.""" + + +_REDUCERS = {ReduceOp.SUM: sum, ReduceOp.MAX: max, ReduceOp.MIN: min} + + +class _Rendezvous: + """Barrier plus per-step payload slots for one group of ranks.""" + + def __init__( + self, + name: str, + ranks: Sequence[int], + timeout_s: float, + last_call: Callable[[int], str], + ) -> None: + self.name = name + self.ranks = tuple(ranks) + self._timeout_s = timeout_s + self._last_call = last_call + self._barrier = threading.Barrier(len(self.ranks), timeout=timeout_s) + self._lock = threading.Lock() + # step -> {rank: (collective, payload)}; a slot is written once, never overwritten. + self._steps: Dict[int, Dict[int, Tuple[str, Any]]] = {} + self._next_step = {rank: 0 for rank in self.ranks} + self._abort_reason: Optional[str] = None + + def exchange(self, rank: int, collective: str, payload: Any) -> List[Any]: + """Enter ``collective`` with ``payload``; return every rank's payload in rank order. + + ``payload`` is stored as is, so callers pass a value they will not + mutate. ``collective`` is the label every rank of the group must match. + """ + with self._lock: + step = self._next_step[rank] + self._next_step[rank] = step + 1 + self._steps.setdefault(step, {})[rank] = (collective, payload) + try: + self._barrier.wait() + except threading.BrokenBarrierError: + raise self._broken(rank, step) from None + with self._lock: + arrivals = dict(self._steps[step]) + entered = {peer: name for peer, (name, _) in arrivals.items()} + if len(set(entered.values())) > 1: + raise FakeDistMismatch( + f"{self.name}, step {step}: ranks entered different collectives: " + + ", ".join(f"rank {peer}: {name}" for peer, name in sorted(entered.items())) + ) + return [arrivals[peer][1] for peer in self.ranks] + + def abort(self, reason: str) -> None: + """Release every waiting rank; they raise ``_PeerFailed`` instead of timing out.""" + with self._lock: + if self._abort_reason is None: + self._abort_reason = reason + self._barrier.abort() + + def _broken(self, rank: int, step: int) -> Exception: + with self._lock: + reason = self._abort_reason + arrivals = dict(self._steps.get(step, {})) + if reason is not None: + return _PeerFailed(reason) + arrived = ", ".join(f"rank {peer} ({name})" for peer, (name, _) in sorted(arrivals.items())) + missing = ", ".join( + f"rank {peer} (last call: {self._last_call(peer)})" + for peer in self.ranks + if peer not in arrivals + ) + return FakeDistTimeout( + f"{self.name}, step {step}: rank {rank} waited {self._timeout_s}s for a collective. " + f"Arrived: {arrived}. Missing: {missing}." + ) + + +class FakeDistRank: + """The ``dist`` object one rank hands to its coordinator. + + Models only what the coordinator uses: a single pipeline / context-parallel + stage, TP collectives over the rank's TP group and ``allreduce`` over the + world. Every call is appended to ``calls`` as ``(collective, payload)`` + with the payload as the caller passed it, including calls a single-rank + group answers locally. + """ + + pp_size = 1 + cp_size = 1 + + def __init__(self, group: "FakeDistGroup", rank: int) -> None: + self._group = group + self.rank = rank + self.world_size = group.world_size + self.tp_size = group.tp_size + self.tp_rank = rank % group.tp_size + self.calls: List[Tuple[str, Any]] = [] + + def tp_allgather(self, obj, *, small_payload: bool = False) -> list: + return self._gather(self._group.tp_rendezvous(self.rank), "tp_allgather", obj) + + def tp_allgather_int64(self, values) -> np.ndarray: + gathered = self._gather(self._group.tp_rendezvous(self.rank), "tp_allgather_int64", values) + rows = [np.asarray(row, dtype=np.int64).reshape(-1) for row in gathered] + if len({row.size for row in rows}) > 1: + raise FakeDistMismatch( + f"tp_allgather_int64: ranks passed vectors of different lengths " + f"{[row.size for row in rows]}" + ) + return np.stack(rows) + + def tp_allreduce(self, obj, op: ReduceOp = ReduceOp.SUM): + return self._reduce(self._group.tp_rendezvous(self.rank), "tp_allreduce", obj, op) + + def allreduce(self, obj, op: ReduceOp = ReduceOp.SUM): + return self._reduce(self._group.world_rendezvous(), "allreduce", obj, op) + + def _gather( + self, rendezvous: _Rendezvous, collective: str, payload: Any, label: Optional[str] = None + ) -> list: + """Record the call and exchange a snapshot of ``payload``; return one copy per rank. + + The snapshot is taken before waiting: a sender may mutate its input as + soon as the call returns, while a peer may not have read the slot yet. + Receivers get their own copies, as they would after unpickling. + ``label`` is what peers must agree on; it defaults to ``collective``. + """ + snapshot = copy.deepcopy(payload) + self.calls.append((collective, snapshot)) + if len(rendezvous.ranks) == 1: + return [copy.deepcopy(snapshot)] + gathered = rendezvous.exchange(self.rank, label or collective, snapshot) + return [copy.deepcopy(item) for item in gathered] + + def _reduce(self, rendezvous: _Rendezvous, collective: str, payload: Any, op: ReduceOp): + # Peers must agree on the operation, not just on entering a reduce. + op = ReduceOp(op) + values = self._gather(rendezvous, collective, payload, label=f"{collective}[{op.name}]") + if len(values) == 1: + return values[0] + reducer = _REDUCERS.get(op) + if reducer is None: + raise NotImplementedError(f"FakeDist does not model {op!r}") + return reducer(values) + + +class FakeDistGroup: + """``world_size`` fake ranks in TP groups of ``tp_size`` consecutive ranks.""" + + def __init__(self, world_size: int, tp_size: int, timeout_s: float = 5.0) -> None: + if world_size % tp_size: + raise ValueError(f"tp_size {tp_size} must divide world_size {world_size}") + self.world_size = world_size + self.tp_size = tp_size + self._ranks = [FakeDistRank(self, rank) for rank in range(world_size)] + self._world = _Rendezvous("world", range(world_size), timeout_s, self._last_call) + self._tp_groups = [ + _Rendezvous( + f"TP group {index}", + range(index * tp_size, (index + 1) * tp_size), + timeout_s, + self._last_call, + ) + for index in range(world_size // tp_size) + ] + + def rank(self, rank: int) -> FakeDistRank: + return self._ranks[rank] + + def tp_rendezvous(self, rank: int) -> _Rendezvous: + return self._tp_groups[rank // self.tp_size] + + def world_rendezvous(self) -> _Rendezvous: + return self._world + + def run(self, fn: Callable[[int], Any]) -> List[Any]: + """Run ``fn(rank)`` on one thread per rank; return the results in rank order. + + The first exception raised on any rank is re-raised here. Peers blocked + in a collective at that moment are released instead of timing out. + """ + results: List[Any] = [None] * self.world_size + failures: List[Tuple[int, Exception]] = [] + lock = threading.Lock() + + def worker(rank: int) -> None: + try: + results[rank] = fn(rank) + except Exception as error: # re-raised on the caller's thread below + with lock: + failures.append((rank, error)) + for rendezvous in (self._world, *self._tp_groups): + rendezvous.abort(f"rank {rank} raised {type(error).__name__}") + + threads = [ + threading.Thread(target=worker, args=(rank,), name=f"fake-dist-rank-{rank}") + for rank in range(self.world_size) + ] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + if failures: + causes = [failure for failure in failures if not isinstance(failure[1], _PeerFailed)] + rank, error = (causes or failures)[0] + if hasattr(error, "add_note"): + error.add_note(f"raised on fake rank {rank}") + raise error + return results + + def _last_call(self, rank: int) -> str: + calls = self._ranks[rank].calls + return calls[-1][0] if calls else "none" diff --git a/tests/unittest/_torch/disaggregation/fake_executor_effects.py b/tests/unittest/_torch/disaggregation/fake_executor_effects.py index ed3f9ed2d1b6..ee6cf87c407a 100644 --- a/tests/unittest/_torch/disaggregation/fake_executor_effects.py +++ b/tests/unittest/_torch/disaggregation/fake_executor_effects.py @@ -23,6 +23,8 @@ def __init__(self) -> None: self.staged_responses: List[Tuple[int, LlmResponse, Optional[LlmRequest]]] = [] # (error_msg, requests, charge_budget) self.failed: List[Tuple[str, List[LlmRequest], bool]] = [] + # Messages of collective-aligned fatal failures. + self.fatal: List[str] = [] # Interleaved history of every effect, for relative-order assertions. self.history: List[Tuple[str, object]] = [] # Optional exception raised from fail_requests, to model a fatal error. @@ -49,6 +51,10 @@ def fail_requests( if self.fail_raises is not None: raise self.fail_raises + def fail_fatal(self, error_msg: str) -> None: + self.fatal.append(error_msg) + self.history.append(("fatal", error_msg)) + class FakeRequestRegistry(ActiveRequestRegistry): """Registry over a caller-owned active list, read live on every call.""" diff --git a/tests/unittest/_torch/disaggregation/fake_kv_cache_transceiver.py b/tests/unittest/_torch/disaggregation/fake_kv_cache_transceiver.py index 529bb6b07d0a..9bfa3643261e 100644 --- a/tests/unittest/_torch/disaggregation/fake_kv_cache_transceiver.py +++ b/tests/unittest/_torch/disaggregation/fake_kv_cache_transceiver.py @@ -196,8 +196,9 @@ def supports_inflight_request_cancellation(self) -> bool: def prepare_context_requests(self, requests: List[LlmRequest]) -> None: # Mirror BindKvCacheTransceiver: a no-op placeholder so the executor - # can invoke it unconditionally. - ... + # can invoke it unconditionally. Logged so tests can pin that it is + # entered every iteration, with or without requests. + self.call_log.append(f"prepare_context_requests:{[req.py_request_id for req in requests]}") def get_disaggregated_params(self) -> Dict[str, object]: return {} diff --git a/tests/unittest/_torch/disaggregation/test_benchmark_disagg.py b/tests/unittest/_torch/disaggregation/test_benchmark_disagg.py index c69a60c2d429..9fd4f8151f40 100644 --- a/tests/unittest/_torch/disaggregation/test_benchmark_disagg.py +++ b/tests/unittest/_torch/disaggregation/test_benchmark_disagg.py @@ -95,7 +95,7 @@ def _stub_transfer_entry_points(ex) -> None: """Mock the coordinator's transfer polls once the executor builds it. The build stays lazy so a test can still swap delegated executor methods - (admission, gen init, idle progress) before the first ``ex.disagg`` use. + (admission, gen init) before the first ``ex.disagg`` use. """ build = ex._build_disagg_coordinator @@ -876,7 +876,6 @@ def test_fetch_called_once_even_in_benchmark_disagg(self): mock_fetch = Mock(return_value=[]) ex._fetch_and_activate_new_requests = mock_fetch - ex._check_disagg_ctx_schedulable_status = Mock() _stub_transfer_entry_points(ex) ex._pad_attention_dp_dummy_request = Mock() ex._schedule = Mock(return_value=(ScheduledRequests(), [], 0)) @@ -1269,7 +1268,6 @@ def _make_executor( ex.active_requests = init_reqs + ready_reqs ex._fetch_and_activate_new_requests = Mock(return_value=[]) - ex._check_disagg_ctx_schedulable_status = Mock() _stub_transfer_entry_points(ex) ex._pad_attention_dp_dummy_request = Mock() ex._prepare_disagg_gen_init = Mock() @@ -1308,20 +1306,20 @@ def test_healthy_fill_phase_does_not_kill(self): ex._handle_errors.assert_not_called() def test_partial_transfer_admission_uses_only_admitted_requests(self) -> None: - """The admitted subset is prepared and passed to the idle check.""" + """Only the admitted subset is prepared for receive; the idle context + poll still runs once regardless of the deferred candidates.""" admitted_req = _make_active_request(in_init=True) deferred_req = _make_active_request(in_init=True) candidates = [admitted_req, deferred_req] ex = self._make_executor(fill_phase_active=True, fitting_init_requests=candidates) ex._apply_disagg_transfer_admission = Mock(return_value=([admitted_req], False)) - ex._check_disagg_transfer_progress_when_idle = Mock() result, _ = ex._prepare_and_schedule_batch() assert result is not None ex._apply_disagg_transfer_admission.assert_called_once_with(candidates) ex._prepare_disagg_gen_init.assert_called_once_with([admitted_req]) - ex._check_disagg_transfer_progress_when_idle.assert_called_once_with() + ex.disagg.reap_context_sends.assert_called_once_with(0) ex._handle_errors.assert_not_called() def test_fill_with_no_init_requests_does_not_kill(self): @@ -1387,7 +1385,6 @@ def test_model_parallel_peer_terminal_no_fit_kills_all_ranks( gather = getattr(ex.dist, gather_name) gather.return_value = all_rank_status ex._apply_disagg_transfer_admission = Mock(return_value=([], True)) - ex._check_disagg_transfer_progress_when_idle = Mock() result, _ = ex._prepare_and_schedule_batch() @@ -1410,7 +1407,6 @@ def test_attention_dp_backpressure_without_terminal_peer_does_not_kill(self): (True, False), ] ex._apply_disagg_transfer_admission = Mock(return_value=([], True)) - ex._check_disagg_transfer_progress_when_idle = Mock() result, _ = ex._prepare_and_schedule_batch() @@ -1427,7 +1423,6 @@ def test_model_parallel_waits_until_all_ranks_have_fetched(self): (True, True), (False, False), ] - ex._check_disagg_transfer_progress_when_idle = Mock() result, _ = ex._prepare_and_schedule_batch() @@ -1452,7 +1447,6 @@ def test_post_fill_skips_fail_fast_vote(self): ex.enable_attention_dp = True ex.dist.tp_size = 2 ex.dist.world_size = 2 - ex._check_disagg_transfer_progress_when_idle = Mock() result, _ = ex._prepare_and_schedule_batch() @@ -1545,7 +1539,6 @@ def _make_executor(self): ex.active_requests = [] ex._fetch_and_activate_new_requests = Mock(return_value=[]) - ex._check_disagg_ctx_schedulable_status = Mock() _stub_transfer_entry_points(ex) ex._pad_attention_dp_dummy_request = Mock() ex._prepare_disagg_gen_init = Mock() @@ -1619,7 +1612,6 @@ def test_full_lifecycle(self): ex._schedule = Mock(return_value=(ScheduledRequests(), [], 0)) ex.active_requests = ready_reqs ex.dist.tp_allgather = Mock() - ex._check_disagg_transfer_progress_when_idle = Mock() result, _ = ex._prepare_and_schedule_batch() assert result is not None diff --git a/tests/unittest/_torch/disaggregation/test_disagg_adapter_contract.py b/tests/unittest/_torch/disaggregation/test_disagg_adapter_contract.py index c15ca039714e..c1d2dfb79fd1 100644 --- a/tests/unittest/_torch/disaggregation/test_disagg_adapter_contract.py +++ b/tests/unittest/_torch/disaggregation/test_disagg_adapter_contract.py @@ -9,7 +9,7 @@ """ from types import SimpleNamespace -from unittest.mock import Mock +from unittest.mock import Mock, call import pytest from fake_executor_effects import FakeExecutorEffects, FakeRequestRegistry @@ -29,6 +29,8 @@ def _executor() -> PyExecutor: executor._pending_response_terminations = [] executor._terminate_request = Mock() executor._handle_errors = Mock() + executor._fatal_error = None + executor.is_shutdown = False executor.active_requests = [] executor.canceled_req_ids = [] return executor @@ -80,6 +82,31 @@ def test_fail_requests_uses_the_executor_error_path() -> None: ) +def test_fail_fatal_marks_the_executor_fatal_before_the_aligned_error_path_runs() -> None: + """The coordinator calls this only after a world-wide collective agreed, + so the executor may enter the collective-aligned fatal path. The fatal + state must already be set when ``_handle_errors`` runs: with + ``charge_budget=False`` it reads ``_fatal_error`` to decide whether to do + the fatal cleanup, so the reverse order would take the plain error path.""" + executor = _executor() + state_on_entry = {} + + def record_state_on_entry(*_args, **_kwargs): + state_on_entry["fatal_error"] = executor._fatal_error + state_on_entry["is_shutdown"] = executor.is_shutdown + + executor._handle_errors.side_effect = record_state_on_entry + + PyExecutorEffects(executor).fail_fatal("poisoned") + + executor._handle_errors.assert_called_once_with( + "poisoned", requests=None, charge_budget=False, fatal_is_collective_aligned=True + ) + assert isinstance(state_on_entry["fatal_error"], RuntimeError) + assert str(state_on_entry["fatal_error"]) == "Fatal error: poisoned" + assert state_on_entry["is_shutdown"] is True + + def test_registry_reads_the_executor_lists_live() -> None: """The executor rebinds active_requests; the registry must not cache.""" executor = _executor() @@ -118,6 +145,7 @@ def test_fake_and_adapter_record_the_same_effects() -> None: ("terminate_request", (request,), {}), ("stage_transfer_response", (7, response, late_request), {}), ("fail_requests", ("boom", [request]), {"charge_budget": False}), + ("fail_fatal", ("poisoned",), {}), ] fake = FakeExecutorEffects() executor = _executor() @@ -132,9 +160,12 @@ def test_fake_and_adapter_record_the_same_effects() -> None: assert executor._pending_transfer_responses == [(7, response)] assert executor._pending_response_terminations == [late_request] assert fake.failed == [("boom", [request], False)] - executor._handle_errors.assert_called_once_with( - error_msg="boom", requests=[request], charge_budget=False - ) + assert fake.fatal == ["poisoned"] + assert executor.is_shutdown is True + assert executor._handle_errors.call_args_list == [ + call(error_msg="boom", requests=[request], charge_budget=False), + call("poisoned", requests=None, charge_budget=False, fatal_is_collective_aligned=True), + ] def test_fake_registry_reads_live_and_removes_like_the_adapter() -> None: diff --git a/tests/unittest/_torch/disaggregation/test_disagg_coordinator.py b/tests/unittest/_torch/disaggregation/test_disagg_coordinator.py index da0624acb678..3e763ab9dd77 100644 --- a/tests/unittest/_torch/disaggregation/test_disagg_coordinator.py +++ b/tests/unittest/_torch/disaggregation/test_disagg_coordinator.py @@ -67,12 +67,13 @@ def test_orchestration_modules_do_not_depend_on_py_executor(module) -> None: def test_executor_facing_surface_is_a_closed_set() -> None: """The executor-owned behavior the coordinator can trigger is a closed set: - three effects plus one registry mutation. Growing it is a design decision, + four effects plus one registry mutation. Growing it is a design decision, not a convenience.""" assert _public_methods(ExecutorEffects) == { "terminate_request", "stage_transfer_response", "fail_requests", + "fail_fatal", } assert _public_methods(ActiveRequestRegistry) == { "active_requests", @@ -87,8 +88,6 @@ def test_delegated_methods_forward_to_their_own_delegate(name: str) -> None: """Each still-delegated entry point must reach exactly its own delegate with the arguments unchanged; a cross-wired or dropped call changes loop behavior and may break rank symmetry for collective-sensitive entry points.""" - if name in ("check_transfer_errors", "requests_in_error_state"): - pytest.skip("reached from inside the reaps, not a coordinator entry point") delegates = DisaggLoopDelegates(**{f.name: Mock() for f in fields(DisaggLoopDelegates)}) coordinator = _delegating_coordinator(delegates) method = getattr(coordinator, name) diff --git a/tests/unittest/_torch/disaggregation/test_disagg_coordinator_errors.py b/tests/unittest/_torch/disaggregation/test_disagg_coordinator_errors.py new file mode 100644 index 000000000000..154fc40ec5a2 --- /dev/null +++ b/tests/unittest/_torch/disaggregation/test_disagg_coordinator_errors.py @@ -0,0 +1,284 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Transfer error handling of the coordinator: the rank-local check, the ADP +vote and the poison consensus. + +Single-rank cases run one ``CoordinatorHarness``; multi-rank cases run one per +rank over a ``FakeDistGroup`` and pin what each rank exchanges and what it +then asks its executor to do. +""" + +from types import SimpleNamespace + +import pytest +from coordinator_harness import CoordinatorHarness, TransferRequest +from fake_dist import FakeDistGroup + +from tensorrt_llm.bindings import LlmRequestState + +pytestmark = pytest.mark.cpu_only + +_POISON_MSG = "Disagg KV cache transfer buffer is poisoned; process restart is required" +_VOTE_MSG = "Disagg KV cache transfer error" + + +def _failed_gen(rid: int, **overrides) -> TransferRequest: + return TransferRequest( + rid, is_context_only_request=False, state=LlmRequestState.DISAGG_TRANS_ERROR, **overrides + ) + + +def _running_gen(rid: int, **overrides) -> TransferRequest: + return TransferRequest( + rid, + is_context_only_request=False, + state=LlmRequestState.GENERATION_IN_PROGRESS, + **overrides, + ) + + +def _single_rank(world_size: int = 1, **kwargs) -> CoordinatorHarness: + """One harness on rank 0 of a world; its dist records but never blocks.""" + return CoordinatorHarness(dist=FakeDistGroup(world_size, world_size).rank(0), **kwargs) + + +def _ranks(group: FakeDistGroup, **kwargs) -> list: + return [CoordinatorHarness(dist=group.rank(rank), **kwargs) for rank in range(group.world_size)] + + +def _vote(error_ids=(), blocked_ids=()) -> tuple: + return "tp_allgather", {"error_ids": list(error_ids), "blocked_ids": list(blocked_ids)} + + +# -- rank-local check -------------------------------------------------------- + + +def test_rank_local_check_fails_failed_requests_and_names_the_kind() -> None: + h = _single_rank() + failed, running = _failed_gen(1), _running_gen(2) + h.active.extend([failed, running]) + + h.coordinator.check_transfer_errors("generation requests") + + assert h.effects.failed == [ + ("Error in kv cache transfer for generation requests", [failed], False) + ] + assert h.dist.calls == [] + + +def test_rank_local_check_defers_to_the_vote_under_multi_rank_adp() -> None: + """Failing here would enter the executor's response collective from one + rank; the loop-top vote fails replicas on every rank together instead.""" + h = _single_rank(world_size=2, enable_attention_dp=True) + h.active.append(_failed_gen(1)) + + h.coordinator.check_transfer_errors("context requests") + + assert h.effects.failed == [] + assert h.dist.calls == [] + + +def test_rank_local_check_handles_errors_on_a_single_adp_rank() -> None: + h = _single_rank(enable_attention_dp=True) + failed = _failed_gen(1) + h.active.append(failed) + + h.coordinator.check_transfer_errors("generation requests") + + assert [requests for _, requests, _ in h.effects.failed] == [[failed]] + + +def test_user_cancelled_failed_requests_are_left_to_the_cancel_path() -> None: + h = _single_rank() + h.active.append(_failed_gen(1)) + h.registry.canceled = [1] + + h.coordinator.check_transfer_errors("generation requests") + + assert h.effects.failed == [] + + +def test_failed_context_send_waits_until_every_transfer_owner_released_it() -> None: + """The KV connector may still hold the request's blocks when its send + fails. The reap releases only the transceiver's claim, and its own error + check must leave the request alone while the connector's claim stands; + the error is applied once the last owner lets go.""" + h = _single_rank() + failed = TransferRequest(7) + h.active.append(failed) + h.send(failed) + h.transfers.start_transfer(failed) # the connector's claim + h.transceiver.finish_send(failed, outcome="error") + + h.coordinator.reap_context_sends(0) + + assert failed.state == LlmRequestState.DISAGG_TRANS_ERROR + assert h.in_transfer(failed) + assert h.active == [failed] + assert h.effects.failed == [] + + h.coordinator.release_transfer(failed) # the connector lets go + h.coordinator.check_transfer_errors("context requests") + + assert not h.in_transfer(failed) + assert h.effects.failed == [ + ("Error in kv cache transfer for context requests", [failed], False) + ] + + +# -- synced handler outside the ADP vote ------------------------------------- + + +def test_synced_handler_leaves_rank_local_errors_to_the_reaps_outside_adp() -> None: + """Without multi-rank ADP the reaps already failed what they could; the + loop-top pass only applies context failures reported after release.""" + h = _single_rank(world_size=2) + h.active.append(_failed_gen(1)) + + h.coordinator.handle_errors_synced() + + assert h.effects.failed == [] + assert h.dist.calls == [] + + +def test_context_failure_reported_after_release_is_applied_at_the_synced_pass() -> None: + """An error id the transfer manager no longer knows is parked by the reap; + the next loop-top pass flips the request to the error state and fails it + as a context request.""" + h = _single_rank() + late = TransferRequest(9) + h.active.append(late) + status = SimpleNamespace(completed_request_ids=[], error_request_ids=[9]) + h.transceiver.check_context_transfer_status = lambda at_least, mark_complete=False: status + + h.coordinator.reap_context_sends(0) + assert h.effects.failed == [] + + h.coordinator.handle_errors_synced() + + assert late.state == LlmRequestState.DISAGG_TRANS_ERROR + assert h.effects.failed == [("Error in kv cache transfer for context requests", [late], False)] + + +# -- ADP vote ---------------------------------------------------------------- + + +def test_peer_error_fails_the_local_replica_and_spares_unrelated_requests() -> None: + group = FakeDistGroup(world_size=2, tp_size=2) + ranks = _ranks(group, enable_attention_dp=True) + replica, unrelated, failed = _running_gen(7), _running_gen(8), _failed_gen(7) + ranks[0].active.extend([replica, unrelated]) + ranks[1].active.append(failed) + + group.run(lambda rank: ranks[rank].coordinator.handle_errors_synced()) + + assert ranks[0].dist.calls == [_vote()] + assert ranks[1].dist.calls == [_vote(error_ids=[7])] + assert ranks[0].effects.failed == [(_VOTE_MSG, [replica], False)] + assert ranks[1].effects.failed == [(_VOTE_MSG, [failed], False)] + + +def test_every_rank_votes_even_when_no_rank_has_errors() -> None: + group = FakeDistGroup(world_size=2, tp_size=2) + ranks = _ranks(group, enable_attention_dp=True) + for h in ranks: + h.active.append(_running_gen(1)) + + group.run(lambda rank: ranks[rank].coordinator.handle_errors_synced()) + + assert [h.dist.calls for h in ranks] == [[_vote()], [_vote()]] + assert [h.effects.failed for h in ranks] == [[], []] + + +def test_peer_error_without_a_local_replica_still_enters_the_error_path() -> None: + """The executor's error path runs a response collective; a rank with no + matching request enters it with an empty list rather than skipping it.""" + group = FakeDistGroup(world_size=2, tp_size=2) + ranks = _ranks(group, enable_attention_dp=True) + ranks[0].active.append(_running_gen(8)) + ranks[1].active.append(_failed_gen(7)) + + group.run(lambda rank: ranks[rank].coordinator.handle_errors_synced()) + + assert ranks[0].effects.failed == [(_VOTE_MSG, [], False)] + assert [requests for _, requests, _ in ranks[1].effects.failed] == [[ranks[1].active[0]]] + + +def test_child_requests_vote_and_fail_by_parent_id() -> None: + group = FakeDistGroup(world_size=2, tp_size=2) + ranks = _ranks(group, enable_attention_dp=True) + child = _running_gen(101, is_child=True, parent_request_id=9) + ranks[0].active.append(child) + ranks[1].active.append(_failed_gen(9)) + + group.run(lambda rank: ranks[rank].coordinator.handle_errors_synced()) + + assert ranks[1].dist.calls == [_vote(error_ids=[9])] + assert ranks[0].effects.failed == [(_VOTE_MSG, [child], False)] + + +@pytest.mark.parametrize("blocker", ["user_cancelled", "context_send_still_owned"]) +def test_a_locally_blocked_request_vetoes_the_vote_on_every_rank(blocker: str) -> None: + """Rank 0 cannot clean request 7 yet: the cancel path owns it, or a + transfer owner still holds its blocks. Its blocked vote keeps every rank, + including the one that reported the error, from failing 7 this round.""" + group = FakeDistGroup(world_size=2, tp_size=2) + ranks = _ranks(group, enable_attention_dp=True) + blocked = TransferRequest(7) + ranks[0].active.append(blocked) + if blocker == "user_cancelled": + ranks[0].registry.canceled = [7] + else: + ranks[0].transfers.start_transfer(blocked) + # The send fails after the transfer started: start_transfer itself moves + # the request to TRANS_IN_PROGRESS, so the error state must come last. + blocked.state = LlmRequestState.DISAGG_TRANS_ERROR + ranks[1].active.append(_failed_gen(7)) + + group.run(lambda rank: ranks[rank].coordinator.handle_errors_synced()) + + assert ranks[0].dist.calls == [_vote(error_ids=[7], blocked_ids=[7])] + assert ranks[1].dist.calls == [_vote(error_ids=[7])] + assert [h.effects.failed for h in ranks] == [[], []] + + +# -- poison consensus -------------------------------------------------------- + + +def test_one_poisoned_rank_takes_the_whole_world_down(inflight_cancel) -> None: + """Poison is reduced over the world, not the TP group: with TP groups of + one, rank 0 learns of rank 1's poison only through the world allreduce. + Both ranks then fail fatally exactly once and skip the vote that would + otherwise have failed their local error requests.""" + group = FakeDistGroup(world_size=2, tp_size=1) + ranks = _ranks(group, enable_attention_dp=True, supports_inflight_cancellation=True) + ranks[1].transceiver.has_poisoned_transfer_buffer = lambda: True + for h in ranks: + h.active.append(_failed_gen(1)) + + group.run(lambda rank: ranks[rank].coordinator.handle_errors_synced()) + + assert [h.dist.calls for h in ranks] == [[("allreduce", 0)], [("allreduce", 1)]] + assert [h.effects.fatal for h in ranks] == [[_POISON_MSG], [_POISON_MSG]] + assert [h.effects.failed for h in ranks] == [[], []] + + +def test_poison_is_only_checked_when_in_flight_cancellation_is_active() -> None: + group = FakeDistGroup(world_size=2, tp_size=1) + ranks = _ranks(group, supports_inflight_cancellation=True) + ranks[1].transceiver.has_poisoned_transfer_buffer = lambda: True + + group.run(lambda rank: ranks[rank].coordinator.handle_errors_synced()) + + assert [h.dist.calls for h in ranks] == [[], []] + assert [h.effects.fatal for h in ranks] == [[], []] + + +def test_single_rank_poison_needs_no_collective(inflight_cancel) -> None: + h = _single_rank(supports_inflight_cancellation=True) + h.transceiver.has_poisoned_transfer_buffer = lambda: True + + h.coordinator.handle_errors_synced() + + assert h.dist.calls == [] + assert h.effects.fatal == [_POISON_MSG] diff --git a/tests/unittest/_torch/disaggregation/test_disagg_coordinator_progress.py b/tests/unittest/_torch/disaggregation/test_disagg_coordinator_progress.py new file mode 100644 index 000000000000..fcb93307ccd2 --- /dev/null +++ b/tests/unittest/_torch/disaggregation/test_disagg_coordinator_progress.py @@ -0,0 +1,256 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Rank-synchronized progress of the coordinator under FakeDist. + +Every test runs one coordinator per rank, each on its own thread, over a +shared ``FakeDistGroup``. The ranks hold different local state; what is +pinned is that each rank still enters the same collectives in the same +order, with the payload the protocol really exchanges, and what each rank +then asks its executor to do. FakeDist checks the protocol only; blocking +semantics of real collectives need multi-process coverage. +""" + +from types import SimpleNamespace + +import pytest +from coordinator_harness import CoordinatorHarness, TransferRequest +from fake_dist import FakeDistGroup + +from tensorrt_llm.bindings import LlmRequestState +from tensorrt_llm.disaggregated_params import DisaggScheduleStyle + +pytestmark = pytest.mark.cpu_only + + +@pytest.fixture(autouse=True) +def _async_transfer_mode(monkeypatch) -> None: + """Asynchronous generation transfers unless a test sets a mode knob itself.""" + monkeypatch.delenv("TRTLLM_DISAGG_BENCHMARK_GEN_ONLY", raising=False) + monkeypatch.delenv("TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP", raising=False) + + +def _ranks(group: FakeDistGroup, **kwargs) -> list: + return [CoordinatorHarness(dist=group.rank(rank), **kwargs) for rank in range(group.world_size)] + + +def _collectives(h: CoordinatorHarness) -> list: + return [name for name, _ in h.dist.calls] + + +def _receiving(h: CoordinatorHarness, rid: int, started_at: float) -> TransferRequest: + req = TransferRequest( + rid, + is_context_only_request=False, + is_disagg_generation_transmission_in_progress=True, + py_kv_transfer_start_time=started_at, + ) + h.active.append(req) + h.transceiver.request_and_receive_async(req) + return req + + +# -- CS-1 collectives under rank skew ---------------------------------------- + + +def test_timeout_drain_exchanges_a_flag_and_each_rank_fails_its_own_requests() -> None: + """Under multi-rank ADP the drain gathers one bool per rank; the rank whose + peer timed out enters the error path with an empty list. Ids are not + exchanged here, so nothing is failed on behalf of a peer.""" + group = FakeDistGroup(world_size=2, tp_size=2) + ranks = _ranks(group, enable_attention_dp=True) + req = TransferRequest(1, is_context_only_request=False) + ranks[0].coordinator.fail_timed_out([req]) + + group.run(lambda rank: ranks[rank].coordinator.handle_timeouts_synced()) + + assert ranks[0].dist.calls == [("tp_allgather_int64", [True])] + assert ranks[1].dist.calls == [("tp_allgather_int64", [False])] + assert ranks[0].effects.failed == [("Request timed out (KV transfer)", [req], False)] + assert ranks[1].effects.failed == [("Request timed out (KV transfer)", [], False)] + + +def test_peer_timeout_is_mirrored_through_the_tp_wide_id_union(inflight_cancel, clock) -> None: + """Only rank 1's copy of the receive has expired. Both ranks enter the same + two collectives -- the flag reduce, then the id gather -- and rank 0 + cancels its copy on the peer's decision. Neither rank cancels twice.""" + group = FakeDistGroup(world_size=2, tp_size=2) + ranks = _ranks(group, kv_transfer_timeout_ms=1000, supports_inflight_cancellation=True) + fresh = _receiving(ranks[0], 1, started_at=clock["t"]) + expired = _receiving(ranks[1], 1, started_at=clock["t"] - 2.0) + + group.run(lambda rank: ranks[rank].coordinator.poll_gen_transfers()) + + assert ranks[0].dist.calls[:2] == [("tp_allreduce", 0), ("tp_allgather", [])] + assert ranks[1].dist.calls[:2] == [("tp_allreduce", 1), ("tp_allgather", [1])] + assert _collectives(ranks[0]) == _collectives(ranks[1]) + assert fresh.py_kv_transfer_timed_out and expired.py_kv_transfer_timed_out + assert [h.transceiver.call_log.count("cancel_request:1") for h in ranks] == [1, 1] + + group.run(lambda rank: ranks[rank].coordinator.poll_gen_transfers()) + + assert [h.transceiver.call_log.count("cancel_request:1") for h in ranks] == [1, 1] + + +def test_generation_error_flush_is_entered_by_the_rank_without_errors_too( + inflight_cancel, +) -> None: + """The flush reduces a flag over TP; the rank without local errors still + enters the executor error path so the response collective stays aligned.""" + group = FakeDistGroup(world_size=2, tp_size=2) + ranks = _ranks(group, supports_inflight_cancellation=True) + error_req = TransferRequest( + 1, is_context_only_request=False, state=LlmRequestState.DISAGG_TRANS_ERROR + ) + ranks[0].active.append(error_req) + + group.run(lambda rank: ranks[rank].coordinator.poll_gen_transfers()) + + assert ranks[0].dist.calls == [("tp_allreduce", 1)] + assert ranks[1].dist.calls == [("tp_allreduce", 0)] + assert ranks[0].effects.failed == [ + ("Error in kv cache transfer for generation requests", [error_req], False) + ] + assert ranks[1].effects.failed == [ + ("Error in kv cache transfer for generation requests", [], False) + ] + + +def test_timeout_check_is_rank_local(clock) -> None: + group = FakeDistGroup(world_size=2, tp_size=2) + ranks = _ranks(group, kv_transfer_timeout_ms=1000) + requests = [_receiving(h, 1, started_at=clock["t"]) for h in ranks] + clock["t"] += 2.0 + + group.run(lambda rank: ranks[rank].coordinator.check_transfer_timeouts()) + + assert all(req.py_kv_transfer_timed_out for req in requests) + assert [h.dist.calls for h in ranks] == [[], []] + + +# -- idle progress poll ------------------------------------------------------ + + +_CONTEXT_POLL = "check_context_transfer_status:0" + + +def test_idle_poll_reaps_context_sends_and_leaves_gen_status_to_the_loop_head() -> None: + """The loop head already polls generation status every iteration. The + context poll is a consensus inside the transceiver and rank-symmetric, so + no dist collective gates it, whatever the model-parallel layout.""" + h = CoordinatorHarness(dist=FakeDistGroup(world_size=16, tp_size=4).rank(0)) + + h.coordinator.poll_progress_when_idle() + + assert h.transceiver.call_log == [_CONTEXT_POLL] + assert h.dist.calls == [] + + +def test_gen_only_benchmark_still_reaps_context_sends_when_idle(monkeypatch) -> None: + monkeypatch.setenv("TRTLLM_DISAGG_BENCHMARK_GEN_ONLY", "1") + h = CoordinatorHarness(dist=FakeDistGroup(world_size=4, tp_size=4).rank(0)) + + h.coordinator.poll_progress_when_idle() + + assert h.transceiver.call_log == [_CONTEXT_POLL] + assert h.dist.calls == [] + + +def test_sync_transfers_skip_the_idle_poll_on_a_multi_rank_worker(monkeypatch) -> None: + """A synchronous GEN receive blocks rank-locally, so a multi-rank worker + must not enter the context status collective from the idle path, even + with a send in flight.""" + monkeypatch.setenv("TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP", "1") + h = CoordinatorHarness(dist=FakeDistGroup(world_size=4, tp_size=4).rank(0)) + h.send(TransferRequest(1)) + sent = list(h.transceiver.call_log) + + h.coordinator.poll_progress_when_idle() + + assert h.transceiver.call_log == sent + assert h.dist.calls == [] + + +@pytest.mark.parametrize("send_in_flight", [False, True]) +def test_sync_transfers_on_a_single_rank_poll_only_while_a_send_is_in_flight( + monkeypatch, send_in_flight: bool +) -> None: + monkeypatch.setenv("TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP", "1") + h = CoordinatorHarness() + if send_in_flight: + h.send(TransferRequest(1)) + before = list(h.transceiver.call_log) + + h.coordinator.poll_progress_when_idle() + + assert h.transceiver.call_log == before + ([_CONTEXT_POLL] if send_in_flight else []) + + +def test_async_idle_poll_is_entered_by_every_rank_regardless_of_local_sends() -> None: + """A rank with nothing in flight enters the context status poll too; the + consensus is inside the transceiver and needs no dist collective.""" + group = FakeDistGroup(world_size=2, tp_size=2) + ranks = _ranks(group) + ranks[1].send(TransferRequest(1)) + + group.run(lambda rank: ranks[rank].coordinator.poll_progress_when_idle()) + + assert [h.transceiver.call_log.count(_CONTEXT_POLL) for h in ranks] == [1, 1] + assert [h.dist.calls for h in ranks] == [[], []] + + +def test_sync_idle_poll_is_skipped_by_every_rank_of_a_multi_rank_worker(monkeypatch) -> None: + monkeypatch.setenv("TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP", "1") + group = FakeDistGroup(world_size=2, tp_size=2) + ranks = _ranks(group) + ranks[1].send(TransferRequest(1)) + logs = [list(h.transceiver.call_log) for h in ranks] + + group.run(lambda rank: ranks[rank].coordinator.poll_progress_when_idle()) + + assert [h.transceiver.call_log for h in ranks] == logs + + +# -- generation-first context gate ------------------------------------------- + + +def _ctx_request(rid: int, schedule_style: DisaggScheduleStyle) -> TransferRequest: + return TransferRequest( + rid, py_disaggregated_params=SimpleNamespace(schedule_style=schedule_style) + ) + + +def test_context_gate_is_entered_even_without_new_requests() -> None: + """The transceiver runs a consensus inside prepare_context_requests, so it + is entered every iteration; skipping it on an empty list would leave a + waiting request unpromoted on the ranks whose peer info arrived late.""" + h = CoordinatorHarness() + + h.coordinator.prepare_context_schedulable([]) + + assert h.transceiver.call_log == ["prepare_context_requests:[]"] + + +def test_context_gate_receives_only_generation_first_context_requests() -> None: + """Disaggregated params are read for context-only requests only.""" + h = CoordinatorHarness() + gen_first = _ctx_request(1, DisaggScheduleStyle.GENERATION_FIRST) + ctx_first = _ctx_request(2, DisaggScheduleStyle.CONTEXT_FIRST) + gen_only = TransferRequest(3, is_context_only_request=False) + + h.coordinator.prepare_context_schedulable([ctx_first, gen_first, gen_only]) + + assert h.transceiver.call_log == ["prepare_context_requests:[1]"] + + +def test_context_gate_is_entered_once_per_rank_whatever_arrived_locally() -> None: + group = FakeDistGroup(world_size=2, tp_size=2) + ranks = _ranks(group) + arrivals = [[], [_ctx_request(1, DisaggScheduleStyle.GENERATION_FIRST)]] + + group.run(lambda rank: ranks[rank].coordinator.prepare_context_schedulable(arrivals[rank])) + + assert [h.transceiver.call_log for h in ranks] == [ + ["prepare_context_requests:[]"], + ["prepare_context_requests:[1]"], + ] + assert [h.dist.calls for h in ranks] == [[], []] diff --git a/tests/unittest/_torch/disaggregation/test_disagg_coordinator_transfers.py b/tests/unittest/_torch/disaggregation/test_disagg_coordinator_transfers.py index 61aa1b141a17..afc32fb31bb0 100644 --- a/tests/unittest/_torch/disaggregation/test_disagg_coordinator_transfers.py +++ b/tests/unittest/_torch/disaggregation/test_disagg_coordinator_transfers.py @@ -12,114 +12,15 @@ from unittest.mock import Mock import pytest -from fake_executor_effects import FakeExecutorEffects, FakeRequestRegistry -from fake_kv_cache_transceiver import FakeKvCacheTransceiver +from coordinator_harness import CoordinatorHarness as _Harness +from coordinator_harness import TransferRequest as _Request from tensorrt_llm._torch.disaggregation.orchestration import coordinator as coordinator_module -from tensorrt_llm._torch.disaggregation.orchestration.coordinator import DisaggTransferCoordinator -from tensorrt_llm._torch.disaggregation.orchestration.transfer_manager import AsyncTransferManager -from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManagerType from tensorrt_llm.bindings import LlmRequestState pytestmark = pytest.mark.cpu_only -class _Request(SimpleNamespace): - """Request stub with the attributes the transfer paths read.""" - - def __init__(self, rid: int, **overrides) -> None: - defaults = dict( - py_request_id=rid, - request_id=rid, - parent_request_id=None, - is_child=False, - state=LlmRequestState.CONTEXT_INIT, - is_context_only_request=True, - is_context_finished=True, - is_finished_due_to_length=False, - is_finished_due_to_cancellation=False, - is_disagg_generation_init_state=False, - is_disagg_generation_transmission_in_progress=False, - py_kv_transfer_start_time=None, - py_kv_transfer_timed_out=False, - py_disaggregated_params=None, - cached_tokens=0, - response=None, - ) - defaults.update(overrides) - super().__init__(**defaults) - self.state_at_response_creation = None - - def create_response(self, _use_fast_logits, _rank): - self.state_at_response_creation = self.state - return self.response - - @property - def is_generation_only_request(self) -> bool: - return not self.is_context_only_request - - -class _Harness: - def __init__( - self, - *, - kv_transfer_timeout_ms=None, - supports_inflight_cancellation=False, - enable_attention_dp=False, - world_size=1, - tp_size=1, - force_terminate_ctx_for_partial_reuse=False, - draft_kv_cache_manager=None, - ) -> None: - self.transceiver = FakeKvCacheTransceiver( - kv_transfer_timeout_ms=kv_transfer_timeout_ms, - supports_inflight_cancellation=supports_inflight_cancellation, - ) - self.transceiver.has_retired_send_session = lambda req: False - self.kv_cache_manager = Mock(spec=["store_blocks_for_reuse", "unpin_blocks_by_id"]) - self.kv_cache_manager.store_blocks_for_reuse.side_effect = lambda req, _: req.py_request_id - resource_manager = SimpleNamespace( - resource_managers={ResourceManagerType.KV_CACHE_MANAGER: self.kv_cache_manager} - ) - self.transfers = AsyncTransferManager(resource_manager) - self.active = [] - self.registry = FakeRequestRegistry(self.active) - self.effects = FakeExecutorEffects() - self.dist = Mock(rank=0, tp_size=tp_size, world_size=world_size) - self.delegates = Mock() - self.delegates.requests_in_error_state.return_value = [] - self.coordinator = DisaggTransferCoordinator( - transceiver=self.transceiver, - transfer_manager=self.transfers, - kv_cache_manager=self.kv_cache_manager, - dist=self.dist, - effects=self.effects, - registry=self.registry, - enable_attention_dp=enable_attention_dp, - force_terminate_ctx_for_partial_reuse=force_terminate_ctx_for_partial_reuse, - delegates=self.delegates, - draft_kv_cache_manager=draft_kv_cache_manager, - ) - - def send(self, *requests: _Request) -> None: - self.coordinator.send_completed_context(list(requests)) - - def in_transfer(self, req: _Request) -> bool: - return req.py_request_id in self.transfers.requests_in_transfer() - - -@pytest.fixture -def inflight_cancel(monkeypatch): - monkeypatch.setattr(coordinator_module, "is_disagg_inflight_cancel_enabled", lambda: True) - - -@pytest.fixture -def clock(monkeypatch): - now = {"t": 100.0} - monkeypatch.setattr(coordinator_module.time, "monotonic", lambda: now["t"]) - return now - - # -- sending ----------------------------------------------------------------- @@ -273,8 +174,9 @@ def test_fast_completion_without_a_response_terminates_immediately() -> None: def test_failed_send_releases_its_claim_but_leaves_the_request_to_the_error_path() -> None: - """The transfer ends (blocks unpinned) but the request stays active so the - rank-synchronized error pass can respond; nothing is terminated here.""" + """The transfer ends (blocks unpinned) and the still-active request is + handed to the executor's error path as a context failure; the reap itself + terminates nothing.""" h = _Harness() req = _Request(1) h.active.append(req) @@ -288,7 +190,7 @@ def test_failed_send_releases_its_claim_but_leaves_the_request_to_the_error_path assert not h.in_transfer(req) assert h.effects.terminated == [] assert h.effects.staged_responses == [] - h.delegates.check_transfer_errors.assert_called_once_with("context requests") + assert h.effects.failed == [("Error in kv cache transfer for context requests", [req], False)] def test_failure_reported_after_release_is_kept_for_the_synced_error_pass() -> None: @@ -373,6 +275,7 @@ def test_remote_cancellation_of_a_receive_fails_the_request_unless_the_user_canc _Request(1, is_context_only_request=False), _Request(2, is_context_only_request=False), ) + h.active.extend([remote, user]) h.registry.canceled = [2] h.transceiver.request_and_receive_async(remote) h.transceiver.request_and_receive_async(user) @@ -383,17 +286,22 @@ def test_remote_cancellation_of_a_receive_fails_the_request_unless_the_user_canc assert remote.state == LlmRequestState.DISAGG_TRANS_ERROR assert user.state == LlmRequestState.DISAGG_GENERATION_TRANS_IN_PROGRESS - h.delegates.check_transfer_errors.assert_called_once_with("generation requests") + assert h.effects.failed == [ + ("Error in kv cache transfer for generation requests", [remote], False) + ] def test_gen_reap_leaves_error_handling_to_the_consensus_path_under_inflight_cancel( inflight_cancel, ) -> None: h = _Harness(supports_inflight_cancellation=True) + h.active.append( + _Request(1, is_context_only_request=False, state=LlmRequestState.DISAGG_TRANS_ERROR) + ) h.coordinator.reap_gen_receives(0) - h.delegates.check_transfer_errors.assert_not_called() + assert h.effects.failed == [] # -- timeouts ---------------------------------------------------------------- @@ -585,7 +493,7 @@ def test_peer_rank_timeout_decision_is_mirrored_locally(inflight_cancel, clock) def test_generation_error_consensus_fails_only_when_some_rank_needs_it(inflight_cancel) -> None: h = _Harness(supports_inflight_cancellation=True, tp_size=2) error_req = _Request(1, is_context_only_request=False, state=LlmRequestState.DISAGG_TRANS_ERROR) - h.delegates.requests_in_error_state.return_value = [error_req] + h.active.append(error_req) h.dist.tp_allgather.return_value = [[], []] # no timed-out receives on any rank h.dist.tp_allreduce.return_value = 0 diff --git a/tests/unittest/_torch/disaggregation/test_disagg_inflight_cancel_gate.py b/tests/unittest/_torch/disaggregation/test_disagg_inflight_cancel_gate.py index 4729f1f18346..fdaeff184a9a 100644 --- a/tests/unittest/_torch/disaggregation/test_disagg_inflight_cancel_gate.py +++ b/tests/unittest/_torch/disaggregation/test_disagg_inflight_cancel_gate.py @@ -29,7 +29,6 @@ from tensorrt_llm._torch.disaggregation.orchestration import coordinator as coordinator_module from tensorrt_llm._torch.disaggregation.orchestration.coordinator import DisaggTransferCoordinator from tensorrt_llm._torch.disaggregation.orchestration.interfaces import ExecutorEffects -from tensorrt_llm._torch.pyexecutor import py_executor as executor_module from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import ( CppMambaHybridCacheManager, MambaHybridCacheManagerV2, @@ -299,27 +298,6 @@ def test_context_transfer_error_keeps_request_active_until_all_owners_release(): effects.terminate_request.assert_not_called() -def test_context_transfer_error_cleanup_waits_for_async_owners(): - request = SimpleNamespace( - state=LlmRequestState.DISAGG_TRANS_ERROR, - py_request_id=7, - is_child=False, - is_context_only_request=True, - ) - executor = object.__new__(PyExecutor) - executor.active_requests = [request] - executor.canceled_req_ids = [] - executor.async_transfer_manager = Mock() - executor.async_transfer_manager.requests_in_transfer.return_value = { - request.py_request_id: request - } - - assert PyExecutor._get_disagg_reqs_in_error_state(executor) == [] - - executor.async_transfer_manager.requests_in_transfer.return_value = {} - assert PyExecutor._get_disagg_reqs_in_error_state(executor) == [request] - - def test_user_cancel_waits_for_context_transfer_owners(monkeypatch): request = SimpleNamespace( state=LlmRequestState.DISAGG_TRANS_ERROR, @@ -372,34 +350,6 @@ def test_flag_unset_generation_driver_skips_cancel_pipeline(): effects.fail_requests.assert_not_called() -def test_peer_buffer_poison_triggers_world_consistent_fatal_cleanup(monkeypatch): - executor = object.__new__(PyExecutor) - executor.kv_cache_transceiver = Mock() - executor.kv_cache_transceiver.supports_inflight_request_cancellation.return_value = True - executor.kv_cache_transceiver.has_poisoned_transfer_buffer.return_value = False - executor.enable_attention_dp = False - executor.dist = SimpleNamespace( - world_size=2, - allreduce=Mock(return_value=1), - ) - executor._fatal_error = None - executor.is_shutdown = False - executor._handle_errors = Mock() - monkeypatch.setattr(coordinator_module, "is_disagg_inflight_cancel_enabled", lambda: True) - - PyExecutor._handle_disagg_cache_errors_synced(executor) - - executor.dist.allreduce.assert_called_once_with(0, op=executor_module.ReduceOp.MAX) - assert isinstance(executor._fatal_error, RuntimeError) - assert executor.is_shutdown - executor._handle_errors.assert_called_once_with( - "Disagg KV cache transfer buffer is poisoned; process restart is required", - requests=None, - charge_budget=False, - fatal_is_collective_aligned=True, - ) - - def test_preclassified_fatal_error_keeps_adp_response_collectives_aligned(): executor = object.__new__(PyExecutor) executor._fatal_error = RuntimeError("already fatal") diff --git a/tests/unittest/_torch/disaggregation/test_disagg_loop_transcript.py b/tests/unittest/_torch/disaggregation/test_disagg_loop_transcript.py index bfa318a793c5..922c79a8b037 100644 --- a/tests/unittest/_torch/disaggregation/test_disagg_loop_transcript.py +++ b/tests/unittest/_torch/disaggregation/test_disagg_loop_transcript.py @@ -10,8 +10,9 @@ Rank symmetry is checked only in the narrow form that fits one process: the first and a non-first PP rank must issue the same collective-sensitive calls in -the same order during an idle iteration. Real multi-rank blocking semantics are -covered elsewhere (Gloo tests; FakeDist arrives with CS-2). Regular disagg PP +the same order during an idle iteration. Rank skew inside the coordinator is +covered by the FakeDist tests (test_disagg_coordinator_progress.py); real +multi-process blocking semantics only by multi-GPU E2E. Regular disagg PP termination advances from executed-batch handling; a recompute-pause fallback can call the same termination handler from an idle iteration. Neither path is covered here (nothing is pending in these iterations); both belong to the @@ -47,9 +48,8 @@ pytestmark = pytest.mark.cpu_only -# Coordinator entry points whose delegates run a rank-consensus collective. -# Derived from the delegate targets in PyExecutor._build_disagg_coordinator; -# update alongside them. +# Coordinator entry points that run a rank-consensus collective, in the +# coordinator itself or in a delegate. Maintained by hand as entry points move. _COLLECTIVE_COORDINATOR_CALLS = { "handle_errors_synced", # dist.allreduce / tp_allgather under ADP "prepare_context_schedulable", # transceiver.prepare_context_requests consensus @@ -113,6 +113,7 @@ def _idle_executor(monkeypatch, calls: list) -> PyExecutor: executor.kv_cache_transceiver = Mock() executor.async_transfer_manager = Mock() executor.async_transfer_manager.has_any_inflight_requests.return_value = False + executor.async_transfer_manager.requests_in_transfer.return_value = {} executor.kv_connector_manager = None executor.device_id = 0 @@ -325,6 +326,8 @@ def _adp_executor(monkeypatch, calls: list, *, rank: int, transceiver) -> PyExec executor.enable_attention_dp = True executor.dist = Mock(rank=rank, tp_size=2, world_size=2) executor.dist.tp_allgather_int64.return_value = Mock(any=lambda: False) + # The real error vote iterates the gathered votes; echo this rank's twice. + executor.dist.tp_allgather.side_effect = lambda obj: [obj, obj] executor.kv_cache_transceiver = transceiver del executor._disagg_coordinator if transceiver is not None: diff --git a/tests/unittest/_torch/disaggregation/test_fake_dist.py b/tests/unittest/_torch/disaggregation/test_fake_dist.py new file mode 100644 index 000000000000..3ba68dbac274 --- /dev/null +++ b/tests/unittest/_torch/disaggregation/test_fake_dist.py @@ -0,0 +1,155 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Contract of the FakeDist test double: grouping, return shapes and diagnostics.""" + +import time + +import numpy as np +import pytest +from coordinator_harness import CoordinatorHarness +from fake_dist import FakeDistGroup, FakeDistMismatch, FakeDistTimeout + +from tensorrt_llm._torch.distributed.communicator import ReduceOp + +pytestmark = pytest.mark.cpu_only + + +def test_tp_collectives_stay_in_the_tp_group_while_allreduce_spans_the_world() -> None: + group = FakeDistGroup(world_size=4, tp_size=2) + + def step(rank): + dist = group.rank(rank) + return dist.tp_allgather(rank), dist.allreduce(rank, op=ReduceOp.SUM) + + results = group.run(step) + + assert [gathered for gathered, _ in results] == [[0, 1], [0, 1], [2, 3], [2, 3]] + assert [total for _, total in results] == [6, 6, 6, 6] + + +def test_world_allreduce_reaches_ranks_in_other_tp_groups() -> None: + """The shape the poison consensus relies on: TP groups of one, MAX over the world.""" + group = FakeDistGroup(world_size=2, tp_size=1) + + flags = group.run(lambda rank: group.rank(rank).allreduce(int(rank == 1), op=ReduceOp.MAX)) + + assert flags == [1, 1] + + +def test_single_rank_tp_group_answers_locally_with_the_real_shapes() -> None: + """A TP group of one never blocks and returns what a one-rank communicator + would: a one-element list, the object itself and a (1, n) matrix. Wrapping + the reduced value would turn a falsy 0 into a truthy [0].""" + dist = FakeDistGroup(world_size=2, tp_size=1).rank(1) + + assert dist.tp_allgather(5) == [5] + reduced = dist.tp_allreduce(0, op=ReduceOp.MAX) + assert reduced == 0 and not reduced + rows = dist.tp_allgather_int64([True, False]) + assert rows.dtype == np.int64 and rows.tolist() == [[1, 0]] + assert dist.calls == [ + ("tp_allgather", 5), + ("tp_allreduce", 0), + ("tp_allgather_int64", [True, False]), + ] + + +def test_int64_allgather_rows_are_ordered_by_tp_rank() -> None: + group = FakeDistGroup(world_size=2, tp_size=2) + + results = group.run(lambda rank: group.rank(rank).tp_allgather_int64([rank, 10 + rank])) + + for rows in results: + assert rows.dtype == np.int64 + assert rows.tolist() == [[0, 10], [1, 11]] + + +def test_gathered_payloads_are_copies() -> None: + """Ranks receive copies, as after pickling; mutating one rank's result must + not leak into another rank's view.""" + group = FakeDistGroup(world_size=2, tp_size=2) + + results = group.run(lambda rank: group.rank(rank).tp_allgather([rank])) + + results[0][1].append("mutated") + assert results[1] == [[0], [1]] + + +class _MutatedAfterCopy(list): + """A payload its sender appends to as soon as the fake has copied it. + + Stands in for a sender that modifies its input right after the call + returns, timed so that the modification always lands before any peer + could read a shared slot; only a snapshot taken before the wait is immune. + """ + + def __deepcopy__(self, memo): + snapshot = list(self) + self.append("late") + return snapshot + + +def test_a_sender_mutating_its_input_after_sending_does_not_reach_its_peers() -> None: + group = FakeDistGroup(world_size=2, tp_size=2) + payload = _MutatedAfterCopy([1]) + + results = group.run(lambda rank: group.rank(rank).tp_allgather(payload if rank == 1 else 0)) + + assert payload == [1, "late"] + assert results == [[0, [1]], [0, [1]]] + assert group.rank(1).calls == [("tp_allgather", [1])] + + +def test_ranks_entering_different_collectives_are_reported() -> None: + group = FakeDistGroup(world_size=2, tp_size=2) + + def step(rank): + dist = group.rank(rank) + return dist.tp_allreduce(1, op=ReduceOp.MAX) if rank == 0 else dist.tp_allgather([1]) + + with pytest.raises( + FakeDistMismatch, match=r"rank 0: tp_allreduce\[MAX\], rank 1: tp_allgather" + ): + group.run(step) + + +def test_ranks_reducing_with_different_operations_are_reported() -> None: + """Entering the same reduce is not enough; the ranks must run the same operation.""" + group = FakeDistGroup(world_size=2, tp_size=2) + + def step(rank): + return group.rank(rank).tp_allreduce(1, op=ReduceOp.MAX if rank == 0 else ReduceOp.MIN) + + with pytest.raises( + FakeDistMismatch, match=r"rank 0: tp_allreduce\[MAX\], rank 1: tp_allreduce\[MIN\]" + ): + group.run(step) + + +def test_a_missing_rank_is_named_when_a_collective_times_out(clock) -> None: + """Only rank 0 drains the timeout consensus. The barrier timeout is + wall-clock inside ``threading``, so the frozen test clock neither stalls + nor shortens it.""" + group = FakeDistGroup(world_size=2, tp_size=2, timeout_s=0.2) + ranks = [CoordinatorHarness(dist=group.rank(i), enable_attention_dp=True) for i in range(2)] + + with pytest.raises(FakeDistTimeout, match=r"Missing: rank 1 \(last call: none\)") as failure: + group.run( + lambda rank: ranks[rank].coordinator.handle_timeouts_synced() if rank == 0 else None + ) + + assert "Arrived: rank 0 (tp_allgather_int64)" in str(failure.value) + + +def test_a_failing_rank_releases_its_peers_and_is_reported_first() -> None: + group = FakeDistGroup(world_size=2, tp_size=2, timeout_s=5.0) + + def step(rank): + if rank == 1: + raise RuntimeError("rank 1 bug") + group.rank(rank).tp_allreduce(1, op=ReduceOp.MAX) + + started = time.perf_counter() + with pytest.raises(RuntimeError, match="rank 1 bug"): + group.run(step) + assert time.perf_counter() - started < 2.0 diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py index 64a6a01c2f41..40e6d5fff0cc 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/tests/unittest/_torch/executor/test_py_executor.py @@ -1150,109 +1150,11 @@ def test_gen_transfer_status_skips_sync_mode(self, monkeypatch): transceiver.check_gen_transfer_status.assert_not_called() - def test_polls_context_transfers_without_blocking(self): - executor = object.__new__(PyExecutor) - executor.dist = Mock(tp_size=1) - executor._disagg_coordinator = Mock() - - PyExecutor._check_disagg_transfer_progress_when_idle(executor) - - executor._disagg_coordinator.reap_context_sends.assert_called_once_with(0) - - def test_does_not_repeat_gen_status_polled_by_loop_head(self): - """The loop head already polls GEN status every iteration.""" - executor = object.__new__(PyExecutor) - executor.dist = Mock(tp_size=1) - executor._disagg_coordinator = Mock() - - PyExecutor._check_disagg_transfer_progress_when_idle(executor) - - executor._disagg_coordinator.reap_gen_receives.assert_not_called() - - def test_idle_poll_enters_no_extra_collective(self): - """The context poll is rank-symmetric, so no gating collective is needed.""" - executor = object.__new__(PyExecutor) - executor.dist = Mock(tp_size=4, cp_size=4, world_size=16) - executor._disagg_coordinator = Mock() - - PyExecutor._check_disagg_transfer_progress_when_idle(executor) - - executor.dist.allreduce.assert_not_called() - executor.dist.tp_allreduce.assert_not_called() - executor.dist.tp_cp_allgather.assert_not_called() - executor._disagg_coordinator.reap_context_sends.assert_called_once_with(0) - - def test_gen_only_no_context_benchmark_polls_context_when_idle( - self, monkeypatch: pytest.MonkeyPatch - ) -> None: - monkeypatch.setenv("TRTLLM_DISAGG_BENCHMARK_GEN_ONLY", "1") - executor = object.__new__(PyExecutor) - executor.dist = Mock(tp_size=4, cp_size=1, world_size=4) - executor._disagg_coordinator = Mock() - - PyExecutor._check_disagg_transfer_progress_when_idle(executor) - - executor.dist.allreduce.assert_not_called() - executor.dist.tp_allreduce.assert_not_called() - executor._disagg_coordinator.reap_context_sends.assert_called_once_with(0) - - def test_sync_transfer_skips_idle_progress_collectives( - self, monkeypatch: pytest.MonkeyPatch - ) -> None: - monkeypatch.setenv("TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP", "1") - executor = object.__new__(PyExecutor) - executor.dist = Mock(tp_size=4, cp_size=1, world_size=4) - executor.async_transfer_manager = Mock() - executor.async_transfer_manager.has_any_inflight_requests.return_value = True - executor._disagg_coordinator = Mock() - - PyExecutor._check_disagg_transfer_progress_when_idle(executor) - - executor.dist.allreduce.assert_not_called() - executor.dist.tp_allreduce.assert_not_called() - executor.dist.tp_cp_allgather.assert_not_called() - executor._disagg_coordinator.reap_gen_receives.assert_not_called() - executor._disagg_coordinator.reap_context_sends.assert_not_called() - - def test_sync_single_rank_ctx_skips_poll_without_inflight_transfer( - self, monkeypatch: pytest.MonkeyPatch - ) -> None: - monkeypatch.setenv("TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP", "1") - executor = object.__new__(PyExecutor) - executor.dist = Mock(tp_size=1, cp_size=1, world_size=1) - executor.async_transfer_manager = Mock() - executor.async_transfer_manager.has_any_inflight_requests.return_value = False - executor._disagg_coordinator = Mock() - - PyExecutor._check_disagg_transfer_progress_when_idle(executor) - - executor._disagg_coordinator.reap_gen_receives.assert_not_called() - executor._disagg_coordinator.reap_context_sends.assert_not_called() - - def test_sync_single_rank_ctx_reaps_idle_transfer( - self, monkeypatch: pytest.MonkeyPatch - ) -> None: - monkeypatch.setenv("TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP", "1") - executor = object.__new__(PyExecutor) - executor.dist = Mock(tp_size=1, cp_size=1, world_size=1) - executor.async_transfer_manager = Mock() - executor.async_transfer_manager.has_any_inflight_requests.return_value = True - executor._disagg_coordinator = Mock() - - PyExecutor._check_disagg_transfer_progress_when_idle(executor) - - executor.dist.allreduce.assert_not_called() - executor.dist.tp_allreduce.assert_not_called() - executor.dist.tp_cp_allgather.assert_not_called() - executor._disagg_coordinator.reap_gen_receives.assert_not_called() - executor._disagg_coordinator.reap_context_sends.assert_called_once_with(0) - def test_sync_receive_does_not_poll_async_status(self, monkeypatch): monkeypatch.setenv("TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP", "1") executor = object.__new__(PyExecutor) executor.kv_cache_transceiver = Mock() executor._disagg_coordinator = Mock() - executor._check_cache_transfer_errors = Mock() requests = [ Mock(state=LlmRequestState.DISAGG_GENERATION_TRANS_COMPLETE), Mock(state=LlmRequestState.DISAGG_GENERATION_TRANS_COMPLETE), @@ -1266,7 +1168,9 @@ def test_sync_receive_does_not_poll_async_status(self, monkeypatch): ] == requests executor.kv_cache_transceiver.request_and_receive_async.assert_not_called() executor._disagg_coordinator.reap_gen_receives.assert_not_called() - executor._check_cache_transfer_errors.assert_called_once_with("generation requests") + executor._disagg_coordinator.check_transfer_errors.assert_called_once_with( + "generation requests" + ) def test_sync_receive_drains_batch_before_rank_aligned_error_vote(self, monkeypatch): monkeypatch.setenv("TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP", "1") @@ -1282,7 +1186,9 @@ def test_sync_receive_drains_batch_before_rank_aligned_error_vote(self, monkeypa {"error_ids": [], "blocked_ids": []}, ] executor._handle_errors = Mock() - executor._check_cache_transfer_errors = Mock() + executor.canceled_req_ids = [] + executor.async_transfer_manager = Mock() + executor.async_transfer_manager.requests_in_transfer.return_value = {} error_request = Mock( py_request_id=1, state=LlmRequestState.DISAGG_GENERATION_INIT, @@ -1303,6 +1209,9 @@ def complete_or_error(req): ) executor.kv_cache_transceiver.request_and_receive_sync.side_effect = complete_or_error + # Build the real coordinator now that its inputs are set; observe the + # rank-local check without replacing it. + executor.disagg.check_transfer_errors = Mock(wraps=executor.disagg.check_transfer_errors) PyExecutor._recv_disagg_gen_cache(executor, [error_request, following_request]) @@ -1314,13 +1223,13 @@ def complete_or_error(req): assert following_request.state == LlmRequestState.DISAGG_GENERATION_TRANS_COMPLETE executor.kv_cache_transceiver.cancel_request.assert_not_called() executor._handle_errors.assert_not_called() - executor._check_cache_transfer_errors.assert_called_once_with("generation requests") + executor.disagg.check_transfer_errors.assert_called_once_with("generation requests") - PyExecutor._handle_disagg_cache_errors_synced(executor) + executor.disagg.handle_errors_synced() executor.dist.tp_allgather.assert_called_once_with(local_vote) executor._handle_errors.assert_called_once_with( - "Disagg KV cache transfer error", + error_msg="Disagg KV cache transfer error", requests=[error_request], charge_budget=False, ) @@ -1488,7 +1397,6 @@ class StopLocalSchedule(RuntimeError): executor._profiler = Mock(return_value=profiler) executor.hang_detector = MagicMock() executor.enable_iter_perf_stats = False - executor._handle_disagg_cache_errors_synced = Mock() executor._fetch_and_activate_new_requests = Mock(return_value=[]) executor.is_shutdown = False executor._handle_control_request = Mock() @@ -1544,12 +1452,10 @@ class StopAfterReconciliation(RuntimeError): executor._is_kv_manager_v2 = True executor._pp_rebalance_drain_iters = None executor._can_pause_for_rebalance = Mock(return_value=False) - executor._handle_disagg_cache_errors_synced = Mock() executor._fetch_and_activate_new_requests = Mock(return_value=[]) executor.is_shutdown = False executor._handle_control_request = Mock() executor.kv_cache_transceiver = Mock() - executor._check_disagg_ctx_schedulable_status = Mock() executor._pad_attention_dp_dummy_request = Mock() executor._pp_retry_until_can_schedule = Mock() executor._mm_encoder_item_scheduling_enabled = False @@ -2898,191 +2804,6 @@ def test_reset_prefix_cache_clears_target_and_draft_reuse_trees(): stub.draft_kv_cache_manager.reset_reuse_state.assert_called_once_with() -# --------------------------------------------------------------------------- -# ADP-safe disagg cache error handling (#13900): all TP ranks enter _handle_errors together. -# --------------------------------------------------------------------------- -def _err_req(request_id=1): - return _make_adp_request(LlmRequestState.DISAGG_TRANS_ERROR, request_id=request_id) - - -def _disagg_error_vote(error_ids=(), blocked_ids=()): - return { - "error_ids": list(error_ids), - "blocked_ids": list(blocked_ids), - } - - -_DEFAULT_KV_CACHE_TRANSCEIVER = object() - - -def _make_disagg_err_stub( - *, - enable_attention_dp=True, - kv_cache_transceiver=_DEFAULT_KV_CACHE_TRANSCEIVER, - world_size=2, - active_requests=None, - tp_allgather_result=None, -): - stub = types.SimpleNamespace() - stub.enable_attention_dp = enable_attention_dp - if kv_cache_transceiver is _DEFAULT_KV_CACHE_TRANSCEIVER: - kv_cache_transceiver = Mock() - kv_cache_transceiver.supports_inflight_request_cancellation.return_value = False - kv_cache_transceiver.has_poisoned_transfer_buffer.return_value = False - stub.kv_cache_transceiver = kv_cache_transceiver - stub.disagg = types.SimpleNamespace( - inflight_cancel_active=lambda: False, - take_pending_context_failures=set, - ) - stub.active_requests = active_requests if active_requests is not None else [] - stub.dist = Mock() - stub.dist.world_size = world_size - stub.dist.rank = 0 - if tp_allgather_result is not None: - stub.dist.tp_allgather = Mock(return_value=tp_allgather_result) - else: - stub.dist.tp_allgather = Mock(side_effect=lambda v: [v]) - stub.handle_errors_calls = [] - - def _rec_handle_errors(error_msg, requests=None, charge_budget=True): - stub.handle_errors_calls.append( - {"error_msg": error_msg, "requests": requests, "charge_budget": charge_budget} - ) - - stub._handle_errors = _rec_handle_errors - stub._request_vote_id = PyExecutor._request_vote_id - for helper in ( - "_handle_disagg_cache_errors_synced", - "_is_disagg_error_cleanup_blocked", - "_get_disagg_reqs_in_error_state", - "_check_cache_transfer_errors", - ): - setattr(stub, helper, types.MethodType(getattr(PyExecutor, helper), stub)) - return stub - - -class TestDisaggCacheErrorsSynced: - def test_guard_short_circuits_without_transceiver(self): - stub = _make_disagg_err_stub(kv_cache_transceiver=None, active_requests=[_err_req()]) - stub._handle_disagg_cache_errors_synced() - stub.dist.tp_allgather.assert_not_called() - assert stub.handle_errors_calls == [] - - def test_guard_short_circuits_without_adp(self): - stub = _make_disagg_err_stub(enable_attention_dp=False, active_requests=[_err_req()]) - stub._handle_disagg_cache_errors_synced() - stub.dist.tp_allgather.assert_not_called() - assert stub.handle_errors_calls == [] - - def test_guard_short_circuits_single_rank(self): - stub = _make_disagg_err_stub(world_size=1, active_requests=[_err_req()]) - stub._handle_disagg_cache_errors_synced() - stub.dist.tp_allgather.assert_not_called() - assert stub.handle_errors_calls == [] - - def test_all_ranks_enter_when_a_peer_has_error(self): - # A peer reports request 7. This rank must fail its matching replica, - # while leaving an unrelated request active. - matching = _make_adp_request(_STATE_GENERATION_IN_PROGRESS, request_id=7) - unrelated = _make_adp_request(_STATE_GENERATION_IN_PROGRESS, request_id=8) - stub = _make_disagg_err_stub( - active_requests=[matching, unrelated], - tp_allgather_result=[_disagg_error_vote(), _disagg_error_vote([7])], - ) - stub._handle_disagg_cache_errors_synced() - assert len(stub.handle_errors_calls) == 1 - assert stub.handle_errors_calls[0]["requests"] == [matching] - assert stub.handle_errors_calls[0]["charge_budget"] is False - - def test_no_handle_when_no_rank_has_error(self): - stub = _make_disagg_err_stub( - active_requests=[], - tp_allgather_result=[_disagg_error_vote(), _disagg_error_vote()], - ) - stub._handle_disagg_cache_errors_synced() - assert stub.handle_errors_calls == [] - - def test_peer_error_without_local_replica_still_enters_handler(self): - unrelated = _make_adp_request(_STATE_GENERATION_IN_PROGRESS, request_id=8) - stub = _make_disagg_err_stub( - active_requests=[unrelated], - tp_allgather_result=[_disagg_error_vote(), _disagg_error_vote([7])], - ) - - stub._handle_disagg_cache_errors_synced() - - assert stub.handle_errors_calls[0]["requests"] == [] - - def test_local_error_req_forwarded_request_scoped(self): - err = _err_req() - ok = _make_adp_request(_STATE_GENERATION_IN_PROGRESS, request_id=2) - stub = _make_disagg_err_stub( - active_requests=[ok, err], tp_allgather_result=[_disagg_error_vote([1])] - ) - stub._handle_disagg_cache_errors_synced() - assert len(stub.handle_errors_calls) == 1 - assert stub.handle_errors_calls[0]["requests"] == [err] - assert stub.handle_errors_calls[0]["charge_budget"] is False - - def test_child_request_votes_by_parent_id(self): - child = _make_adp_request( - _STATE_GENERATION_IN_PROGRESS, - request_id=101, - is_child=True, - parent_request_id=9, - ) - stub = _make_disagg_err_stub( - active_requests=[child], - tp_allgather_result=[_disagg_error_vote(), _disagg_error_vote([9])], - ) - - stub._handle_disagg_cache_errors_synced() - - assert stub.handle_errors_calls[0]["requests"] == [child] - - def test_peer_vote_does_not_clean_up_locally_deferred_request(self): - request = _err_req(request_id=7) - request.is_context_only_request = True - stub = _make_disagg_err_stub( - active_requests=[request], - tp_allgather_result=[ - _disagg_error_vote(blocked_ids=[7]), - _disagg_error_vote([7]), - ], - ) - stub.canceled_req_ids = [] - stub.async_transfer_manager = Mock() - stub.async_transfer_manager.requests_in_transfer.return_value = { - request.py_request_id: request - } - - stub._handle_disagg_cache_errors_synced() - - assert stub.handle_errors_calls == [] - - -class TestCheckCacheTransferErrorsAdpNoop: - def test_noop_under_adp_multirank(self): - # Even with an error req present, ADP+world_size>1 defers to the synced handler. - stub = _make_disagg_err_stub(active_requests=[_err_req()]) - stub._check_cache_transfer_errors("ctx") - assert stub.handle_errors_calls == [] - - def test_handles_error_when_not_adp(self): - err = _err_req() - stub = _make_disagg_err_stub(enable_attention_dp=False, active_requests=[err]) - stub._check_cache_transfer_errors("ctx") - assert len(stub.handle_errors_calls) == 1 - assert stub.handle_errors_calls[0]["requests"] == [err] - assert stub.handle_errors_calls[0]["charge_budget"] is False - - def test_handles_error_on_single_rank(self): - err = _err_req() - stub = _make_disagg_err_stub(world_size=1, active_requests=[err]) - stub._check_cache_transfer_errors("gen") - assert len(stub.handle_errors_calls) == 1 - - class TestPendingTransferResponseFlush: def test_rank_local_fatal_error_does_not_issue_adp_response_gather(self): """A lone fatal rank must fail locally rather than desynchronize TP.""" @@ -3176,7 +2897,6 @@ def _make_executor_loop_stub(): executor._is_kv_manager_v2 = False executor._mm_encoder_item_scheduling_enabled = False executor.is_benchmark_disagg = False - executor._handle_disagg_cache_errors_synced = Mock() executor._flush_pending_transfer_responses = Mock() return executor @@ -3367,7 +3087,6 @@ def _make_one_model_mtp_executor( ex.waiting_queue = [] ex._fetch_and_activate_new_requests = Mock(return_value=[]) - ex._check_disagg_ctx_schedulable_status = Mock() ex._pad_attention_dp_dummy_request = Mock() ex._prefetch_for_context_requests = Mock() ex._prepare_disagg_gen_init = Mock() diff --git a/tests/unittest/_torch/executor/test_send_kv_async_split.py b/tests/unittest/_torch/executor/test_send_kv_async_split.py index 11eb8215c0f2..060e93b5ae4a 100644 --- a/tests/unittest/_torch/executor/test_send_kv_async_split.py +++ b/tests/unittest/_torch/executor/test_send_kv_async_split.py @@ -155,7 +155,7 @@ def _dual_claim_executor() -> PyExecutor: executor.force_terminate_ctx_for_partial_reuse = False executor.dist = SimpleNamespace(rank=0, world_size=2) executor._terminate_request = Mock() - # Make the reap's trailing _check_cache_transfer_errors a no-op. + # Make the reap's trailing check_transfer_errors a no-op (multi-rank ADP). executor.enable_attention_dp = True return executor diff --git a/tests/unittest/disaggregated/test_chunked_transfer.py b/tests/unittest/disaggregated/test_chunked_transfer.py index 56bfea2e6be3..280aee4882c7 100644 --- a/tests/unittest/disaggregated/test_chunked_transfer.py +++ b/tests/unittest/disaggregated/test_chunked_transfer.py @@ -1420,8 +1420,6 @@ def test_send_disagg_ctx_kv_skips_retired_session_without_mutating_state( def test_context_send_failure_is_applied_at_next_loop_boundary(): - from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor - request = _make_send_kv_request(is_last_chunk=False) request.py_kv_transfer_timed_out = False coordinator, transceiver, transfer_manager = _make_send_kv_coordinator( @@ -1439,17 +1437,13 @@ def test_context_send_failure_is_applied_at_next_loop_boundary(): assert request.state == LlmRequestState.CONTEXT_INIT transfer_manager.end_transfer.assert_not_called() - executor = MagicMock() - executor.disagg = coordinator - executor.active_requests = [request] - executor.enable_attention_dp = False - executor.dist.world_size = 1 - - PyExecutor._handle_disagg_cache_errors_synced(executor) + coordinator.handle_errors_synced() assert request.state == LlmRequestState.DISAGG_TRANS_ERROR assert not coordinator.take_pending_context_failures() - executor._check_cache_transfer_errors.assert_called_with("context requests") + coordinator._effects.fail_requests.assert_called_once_with( + "Error in kv cache transfer for context requests", [request], charge_budget=False + ) @pytest.mark.parametrize("is_last_chunk", [False, True]) diff --git a/tests/unittest/disaggregated/test_transfer_ownership_regressions.py b/tests/unittest/disaggregated/test_transfer_ownership_regressions.py index 169e8d3802c9..a92c83ce7de1 100644 --- a/tests/unittest/disaggregated/test_transfer_ownership_regressions.py +++ b/tests/unittest/disaggregated/test_transfer_ownership_regressions.py @@ -1475,7 +1475,8 @@ def test_pre_cancelled_sender_does_not_publish_from_transceiver() -> None: end_transfer=end_transfer, ) executor.active_requests = [request] - executor._check_cache_transfer_errors = Mock() + # The reap's trailing rank-local error check is not under test here. + executor.disagg.check_transfer_errors = Mock() executor.disagg.reap_context_sends(0)