From 1c9f751655bbc1a3428c2670bab1d06baf51dbef Mon Sep 17 00:00:00 2001 From: Pranav Shrestha Date: Fri, 9 Oct 2026 14:24:06 -0700 Subject: [PATCH 01/23] [https://nvbugs/6665906][fix] Cherry-pick #19788 from release/1.3: Restore Mistral Large 3 multimodal inference Forward-port of release/1.3 PR https://github.com/NVIDIA/TensorRT-LLM/pull/19788 (release/1.3 commit 48c796cb6fe53b7f4c80acc41ea003c2fb3e3d25) to main. Adapted to main: adopts main's MultimodalItemMetadata contract from #20391. The shared Pixtral encoder contract builds item metadata in _build_mm_item_metadata and exposes it through get_mm_item_metadata, and the new tests call get_mm_item_metadata and compare against the tuples that MultimodalItemMetadata now stores. Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> Co-authored-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> Signed-off-by: Simeng Liu --- cpp/kernels/fmha_v2/setup.py | 16 +- cpp/tensorrt_llm/kernels/fmhaDispatcher.cpp | 12 +- cpp/tests/unit_tests/kernels/CMakeLists.txt | 1 + .../unit_tests/kernels/fmhaHead104Test.cpp | 254 +++++++++++++ .../_torch/models/modeling_mistral.py | 343 ++++++++++-------- tests/integration/defs/cpp/conftest.py | 35 ++ tests/integration/defs/cpp/test_unit_tests.py | 20 + .../test_lists/test-db/l0_b200.yml | 1 + .../test_lists/test-db/l0_gb202.yml | 2 + .../test_lists/test-db/l0_rtx_pro_6000.yml | 2 + tests/integration/test_lists/waives.txt | 1 - .../test_fmha_v2_head_size_coverage.py | 107 ++++++ .../test_pixtral_head104_attention.py | 65 ++++ .../_torch/modeling/test_modeling_mistral.py | 204 ++++++++++- 14 files changed, 901 insertions(+), 162 deletions(-) create mode 100644 cpp/tests/unit_tests/kernels/fmhaHead104Test.cpp create mode 100644 tests/unittest/_torch/attention/test_fmha_v2_head_size_coverage.py create mode 100644 tests/unittest/_torch/attention/test_pixtral_head104_attention.py diff --git a/cpp/kernels/fmha_v2/setup.py b/cpp/kernels/fmha_v2/setup.py index f15d7f39f080..d524b6db44e3 100644 --- a/cpp/kernels/fmha_v2/setup.py +++ b/cpp/kernels/fmha_v2/setup.py @@ -2199,8 +2199,16 @@ def selected_mask_types(kspec): custom_mask = '0' # encoder models (head_size = 32 / 64 / 128) need packed_qkv input layout + padding mask. elif kspec.input_layout == InputLayout.PACKED_QKV: - # NOTE: 72/80 are added for vision transformer - if kspec.head_size not in [32, 64, 72, 80, 128]: + # NOTE: 72/80 are added for vision transformers. Pixtral's head size + # 104 uses warp specialization on SM90 and non-tiled kernels on SM120. + pixtral_padding = kspec.head_size == 104 and ( + kspec.sm == 100 or + (kspec.sm in (90, 120) and kspec.dtype in + ('fp16', 'bf16', 'fp16_fp32') and not kspec.enable_skip_softmax + and kspec.warp_specialization == (kspec.sm == 90) and + (kspec.sm != 120 or not kspec.tiled))) + if (kspec.head_size not in [32, 64, 72, 80, 128] + and not pixtral_padding): padding_mask = '0' # only cross attention (head_size = 32/64/128) needs contiguous_q_kv input layout + padding mask / custom_mask. elif kspec.input_layout == InputLayout.CONTIGUOUS_Q_KV: @@ -6920,10 +6928,10 @@ def enumerate_kernels(): and kspec.cross_mha == False and kspec.flash_attention == True and kspec.input_layout != InputLayout.SEPARATE_Q_K_V) - # Gemma3 VL support. + # Gemma3 VL and Pixtral VL support. or (kspec.sm == 100 and kspec.dtype in ['fp16', 'bf16', 'fp16_fp32', 'e4m3', 'e4m3_fp32'] - and kspec.head_size == 72 + and kspec.head_size in [72, 104] and kspec.head_size_v == 0 and kspec.sage_block_sizes is None and kspec.version == 2 diff --git a/cpp/tensorrt_llm/kernels/fmhaDispatcher.cpp b/cpp/tensorrt_llm/kernels/fmhaDispatcher.cpp index 5ced0f20c073..071c5a99120e 100644 --- a/cpp/tensorrt_llm/kernels/fmhaDispatcher.cpp +++ b/cpp/tensorrt_llm/kernels/fmhaDispatcher.cpp @@ -63,10 +63,16 @@ QkvLayout AttentionInputLayoutToQkvLayout(AttentionInputLayout layout) FmhaDispatcher::FmhaDispatcher(MHARunnerFixedParams fixedParams) : mFixedParams(fixedParams) - // TRTLLM-GEN only supports power of 2 head sizes (and 80 with padding). - // The exception will fall back to fmha v2. + // FMHA v2 provides head-size 72 and FP16/BF16 full/causal packed-QKV head-size 104 kernels. + // Keep mixed Q/KV precisions on TRTLLM-GEN: fmha v2 requires matching types. // Please update fmha_v2/setup.py if you want to add more supported head sizes. - , mUseTllmGen(tensorrt_llm::common::isSM100Family() && fixedParams.headSize != 72) + , mUseTllmGen(tensorrt_llm::common::isSM100Family() && fixedParams.headSize != 72 + && !(fixedParams.headSize == 104 && fixedParams.headSizeV == 104 + && fixedParams.attentionInputLayout == AttentionInputLayout::PACKED_QKV + && fixedParams.dataType == fixedParams.dataTypeKv + && (fixedParams.dataType == DATA_TYPE_FP16 || fixedParams.dataType == DATA_TYPE_BF16) + && (fixedParams.attentionMaskType == ContextAttentionMaskType::PADDING + || fixedParams.attentionMaskType == ContextAttentionMaskType::CAUSAL))) , mMultiProcessorCount(tensorrt_llm::common::getMultiProcessorCount()) { if (mUseTllmGen) diff --git a/cpp/tests/unit_tests/kernels/CMakeLists.txt b/cpp/tests/unit_tests/kernels/CMakeLists.txt index c0c393a2634a..d99c4b336301 100644 --- a/cpp/tests/unit_tests/kernels/CMakeLists.txt +++ b/cpp/tests/unit_tests/kernels/CMakeLists.txt @@ -101,6 +101,7 @@ endif() add_gtest(sparseKvCacheTest sparseKvCacheTest.cu) add_gtest(prepareCustomMaskTest prepareCustomMaskTest.cpp) add_gtest(fmhaDsv4EpilogueFusionTest fmhaDsv4EpilogueFusionTest.cpp) +add_gtest(fmhaHead104Test fmhaHead104Test.cpp) set(NVFP4_COLD_PAGE_KERNEL_TEST_SRC nvfp4ColdPageKernelsTest.cpp diff --git a/cpp/tests/unit_tests/kernels/fmhaHead104Test.cpp b/cpp/tests/unit_tests/kernels/fmhaHead104Test.cpp new file mode 100644 index 000000000000..126e2977891d --- /dev/null +++ b/cpp/tests/unit_tests/kernels/fmhaHead104Test.cpp @@ -0,0 +1,254 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include + +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/kernels/fmhaDispatcher.h" + +namespace +{ + +namespace tk = tensorrt_llm::kernels; + +struct TestCase +{ + char const* name; + tk::Data_type inputType; + tk::ContextAttentionMaskType mask; + bool forceFixed = false; + bool forceRuntime = false; +}; + +class FmhaHead104Test : public ::testing::TestWithParam +{ +protected: + static constexpr int kHeadDim = 104; + static constexpr int kHeads = 2; + static constexpr float kQScaling = 1.25F; + + void SetUp() override + { + int const sm = tensorrt_llm::common::getSMVersion(); + bool const isSm100 = tensorrt_llm::common::isSM100Family(sm); + if (sm != 90 && sm != 120 && !isSm100) + { + GTEST_SKIP() << "Head-104 packed-QKV kernels require SM90, SM100-family, or SM120."; + } + TLLM_CUDA_CHECK(cudaStreamCreate(&mStream)); + } + + void TearDown() override + { + if (mStream != nullptr) + { + cudaStreamSynchronize(mStream); + } + for (void* allocation : mAllocations) + { + cudaFree(allocation); + } + if (mStream != nullptr) + { + cudaStreamDestroy(mStream); + } + } + + template + T* allocate(size_t count) + { + T* pointer = nullptr; + TLLM_CUDA_CHECK(cudaMalloc(&pointer, count * sizeof(T))); + mAllocations.push_back(pointer); + return pointer; + } + + template + T* upload(std::vector const& values) + { + auto* pointer = allocate(values.size()); + TLLM_CUDA_CHECK(cudaMemcpy(pointer, values.data(), values.size() * sizeof(T), cudaMemcpyHostToDevice)); + return pointer; + } + + static std::vector reference(std::vector const& qkv, std::vector const& lengths, bool causal) + { + int const tokens = std::accumulate(lengths.begin(), lengths.end(), 0); + int const hidden = kHeads * kHeadDim; + std::vector output(tokens * hidden); + int offset = 0; + for (int const length : lengths) + { + std::vector scores(length); + for (int q = 0; q < length; ++q) + { + int const keys = causal ? q + 1 : length; + for (int head = 0; head < kHeads; ++head) + { + for (int k = 0; k < keys; ++k) + { + float dot = 0.0F; + for (int d = 0; d < kHeadDim; ++d) + { + dot += qkv[(offset + q) * 3 * hidden + head * kHeadDim + d] + * qkv[(offset + k) * 3 * hidden + hidden + head * kHeadDim + d]; + } + scores[k] = dot / (std::sqrt(static_cast(kHeadDim)) * kQScaling); + } + float const maximum = *std::max_element(scores.begin(), scores.begin() + keys); + float sum = 0.0F; + for (int k = 0; k < keys; ++k) + { + scores[k] = std::exp(scores[k] - maximum); + sum += scores[k]; + } + for (int d = 0; d < kHeadDim; ++d) + { + float value = 0.0F; + for (int k = 0; k < keys; ++k) + { + value + += scores[k] / sum * qkv[(offset + k) * 3 * hidden + 2 * hidden + head * kHeadDim + d]; + } + output[(offset + q) * hidden + head * kHeadDim + d] = value; + } + } + } + offset += length; + } + return output; + } + + template + void run(std::vector const& lengths) + { + TestCase const& test = GetParam(); + int const tokens = std::accumulate(lengths.begin(), lengths.end(), 0); + int const maximum = *std::max_element(lengths.begin(), lengths.end()); + size_t const outputSize = tokens * kHeads * kHeadDim; + std::vector qkv(3 * outputSize); + std::vector roundedQkv(qkv.size()); + std::mt19937 generator(6665906); + std::normal_distribution normal; + for (size_t i = 0; i < qkv.size(); ++i) + { + qkv[i] = T(normal(generator)); + roundedQkv[i] = static_cast(qkv[i]); + } + std::vector cumulative(lengths.size() + 1, 0); + std::partial_sum(lengths.begin(), lengths.end(), cumulative.begin() + 1); + + tk::MHARunnerFixedParams fixed{}; + fixed.dataType = test.inputType; + fixed.dataTypeKv = test.inputType; + fixed.dataTypeOut = test.inputType; + fixed.forceFp32Acc = test.forceFixed; + fixed.attentionMaskType = test.mask; + fixed.attentionInputLayout = tk::AttentionInputLayout::PACKED_QKV; + fixed.numQHeads = kHeads; + fixed.numKvHeads = kHeads; + fixed.headSize = kHeadDim; + fixed.headSizeV = kHeadDim; + fixed.qScaling = kQScaling; + tk::FmhaDispatcher dispatcher(fixed); + ASSERT_TRUE(dispatcher.isSupported()); + + auto* output = allocate(outputSize); + TLLM_CUDA_CHECK(cudaMemset(output, 0xff, outputSize * sizeof(T))); + tk::MHARunnerParams params{}; + params.b = lengths.size(); + params.qSeqLen = maximum; + params.kvSeqLen = maximum; + params.totalQSeqLen = tokens; + params.totalKvSeqLen = tokens; + params.slidingWindowSize = std::numeric_limits::max(); + params.qkvPtr = upload(qkv); + params.outputPtr = output; + params.cuQSeqLenPtr = upload(cumulative); + params.cuKvSeqLenPtr = params.cuQSeqLenPtr; + params.kvSeqLenPtr = upload(lengths); + params.tileCounterPtr = upload({0}); + params.forceFp32Acc = test.forceRuntime; + params.stream = mStream; + // Direct dispatch must launch a fused kernel; there is no unfused implementation in this test. + dispatcher.run(params); + TLLM_CUDA_CHECK(cudaStreamSynchronize(mStream)); + + std::vector actual(outputSize); + TLLM_CUDA_CHECK(cudaMemcpy(actual.data(), output, outputSize * sizeof(T), cudaMemcpyDeviceToHost)); + auto const expected = reference(roundedQkv, lengths, test.mask == tk::ContextAttentionMaskType::CAUSAL); + double squaredError = 0.0; + double squaredReference = 0.0; + float maxError = 0.0F; + float referencePeak = 0.0F; + for (size_t i = 0; i < outputSize; ++i) + { + float const value = static_cast(actual[i]); + ASSERT_TRUE(std::isfinite(value)) << "Unwritten or nonfinite output at " << i; + float const error = value - expected[i]; + squaredError += static_cast(error) * error; + squaredReference += static_cast(expected[i]) * expected[i]; + maxError = std::max(maxError, std::abs(error)); + referencePeak = std::max(referencePeak, std::abs(expected[i])); + } + bool const bf16 = test.inputType == tk::DATA_TYPE_BF16; + float const relativeRmseLimit = bf16 ? 0.015F : 0.005F; + float const peakErrorLimit = bf16 ? 0.025F : 0.01F; + EXPECT_LE(std::sqrt(squaredError / std::max(squaredReference, 1e-16)), relativeRmseLimit); + EXPECT_LE(maxError, std::max(0.01F, referencePeak * peakErrorLimit)); + } + + cudaStream_t mStream = nullptr; + std::vector mAllocations; +}; + +TEST_P(FmhaHead104Test, PackedRaggedNumerics) +{ + for (auto const& lengths : {std::vector{31, 7}, std::vector{257, 65}}) + { + SCOPED_TRACE("Maximum sequence length: " + std::to_string(lengths.front())); + switch (GetParam().inputType) + { + case tk::DATA_TYPE_FP16: run(lengths); break; + case tk::DATA_TYPE_BF16: run<__nv_bfloat16>(lengths); break; + default: FAIL() << "Unexpected input type"; + } + ASSERT_FALSE(HasFatalFailure()); + } +} + +INSTANTIATE_TEST_SUITE_P(Head104, FmhaHead104Test, + ::testing::Values(TestCase{"Fp16Full", tk::DATA_TYPE_FP16, tk::ContextAttentionMaskType::PADDING}, + TestCase{"Fp16Causal", tk::DATA_TYPE_FP16, tk::ContextAttentionMaskType::CAUSAL}, + TestCase{"Bf16Full", tk::DATA_TYPE_BF16, tk::ContextAttentionMaskType::PADDING}, + TestCase{"Bf16Causal", tk::DATA_TYPE_BF16, tk::ContextAttentionMaskType::CAUSAL}, + TestCase{"Fp16FixedFp32Full", tk::DATA_TYPE_FP16, tk::ContextAttentionMaskType::PADDING, true}, + TestCase{"Fp16RuntimeFp32Full", tk::DATA_TYPE_FP16, tk::ContextAttentionMaskType::PADDING, false, true}), + [](auto const& info) { return info.param.name; }); + +} // namespace diff --git a/tensorrt_llm/_torch/models/modeling_mistral.py b/tensorrt_llm/_torch/models/modeling_mistral.py index 67acf7d93cd9..d8b058f20855 100644 --- a/tensorrt_llm/_torch/models/modeling_mistral.py +++ b/tensorrt_llm/_torch/models/modeling_mistral.py @@ -326,7 +326,11 @@ def get_num_tokens_per_image(self, image_size): Image.new("RGB", (w, h))) return ncols * nrows + nrows - def __call__(self, text, images=None, **kwargs): + def __call__(self, + text: str, + images: list[Image.Image] | None = None, + **kwargs: Any) -> dict[str, torch.Tensor]: + """Tokenize a native request and batch images without losing their sizes.""" mm_items = [{ "type": "image", "base64": encode_base64_image(image) @@ -341,164 +345,55 @@ def __call__(self, text, images=None, **kwargs): }] encoded = self.tokenizer.transformers_tokenizer.apply_chat_template( - conversation, tokenize=True, return_dict=True, return_tensors='pt') + conversation, tokenize=True, return_dict=True, return_tensors=None) processed = { - "input_ids": encoded.input_ids, + "input_ids": torch.tensor([encoded.input_ids]), } # text-only mode for VLM if "pixel_values" in encoded: + # Native images may have different shapes. Preserve their unpadded + # sizes for item scheduling and the vision encoder's patch slicing. + pixel_values, image_sizes = Mistral3VLM.batch_pixel_values([ + torch.as_tensor(image, dtype=self.dtype).unsqueeze(0) + for image in encoded.pixel_values + ], [torch.tensor([size]) for size in encoded.image_sizes]) processed.update({ "pixel_values": - encoded.pixel_values.to(self.dtype), + pixel_values, "attention_mask": - encoded.attention_mask, + torch.tensor([encoded.attention_mask]), "image_sizes": - torch.tensor([encoded.pixel_values.shape[2:]]) + image_sizes, }) return processed -class MistralHFInputProcessor(BaseMultimodalInputProcessor, - BaseMultimodalDummyInputsBuilder): - """Input processor for Mistral VLM checkpoints in HuggingFace format.""" - - def __init__(self, - model_path: str, - config: PretrainedConfig, - tokenizer: AutoTokenizer, - trust_remote_code: bool = True, - **kwargs): - super().__init__(model_path=model_path, - config=config, - tokenizer=tokenizer, - trust_remote_code=trust_remote_code, - **kwargs) - self._config = config - self._dtype = self._config.torch_dtype - self._model_path = model_path - self._tokenizer = (tokenizer if tokenizer is not None else - AutoTokenizer.from_pretrained( - model_path, - config=config, - use_fast=True, - trust_remote_code=True)) - self._processor = AutoProcessor.from_pretrained( - model_path, - use_fast=self.use_fast, - trust_remote_code=trust_remote_code) - logger.info(f"[mistral] HF processor={type(self._processor).__name__} " - f"tokenizer={type(self._tokenizer).__name__}") - - @property - def config(self) -> PretrainedConfig: - return self._config - - @property - def tokenizer(self) -> AutoTokenizer: - return self._tokenizer - - @property - def model_path(self) -> str: - return self._model_path - - @property - def processor(self) -> AutoProcessor: - return self._processor - - @property - def dtype(self) -> torch.dtype: - return self._dtype - - def _build_mm_item_metadata( - self, - _prompt_token_ids: List[int], - multimodal_data: Dict[str, Any], - ) -> Optional[MultimodalItemMetadata]: - """Return Pixtral image items and physical ViT patch counts.""" - image_data = multimodal_data.get("image") - if not isinstance(image_data, dict): - return None - image_sizes = image_data.get("image_sizes") - if image_sizes is None: - return None - patch, merge, _, _ = self._vision_geometry() - encoder_token_lengths = [ - self._vit_tokens(width=int(width), height=int(height), patch=patch) - for height, width in image_sizes - ] - min_tokens_per_image = merge * merge - if any(token_length < min_tokens_per_image or token_length % - min_tokens_per_image for token_length in encoder_token_lengths): - raise ValueError( - "Processed Mistral image geometry must contain a nonempty " - f"multiple of {min_tokens_per_image} encoder tokens") - item_refs = [("image", item_idx) - for item_idx in range(len(encoder_token_lengths))] - output_embedding_lengths = [ - token_length // (merge * merge) - for token_length in encoder_token_lengths - ] - return MultimodalItemMetadata( - item_refs=item_refs, - encoder_token_lengths=encoder_token_lengths, - output_embedding_lengths=output_embedding_lengths) +class MistralPixtralEncoderContract: + """MM encoder item-scheduling contract for Pixtral-based Mistral VLMs. - def get_mm_item_metadata( - self, - prompt_token_ids: List[int], - multimodal_data: Dict[str, Any], - mm_item_order: Optional[List[Dict[str, Any]]] = None, - ) -> Optional[MultimodalItemMetadata]: - metadata = ( - self._build_mm_item_metadata(prompt_token_ids, multimodal_data) - if multimodal_data.get("multimodal_embedding") is None else None) - return metadata if metadata is not None else super( - ).get_mm_item_metadata(prompt_token_ids, multimodal_data, mm_item_order) + `Mistral3VLM` declares `supports_mm_encoder_item_scheduling` on the *model* + class, but the engine requires the contract from the input processor, which + `create_input_processor` picks from `checkpoint_format` alone. Both Mistral + frontends drive the same `PixtralVisionModel` + `Mistral3MultiModalProjector`, + so this geometry belongs to the checkpoint rather than to the preprocessing + frontend; sharing it keeps the capability and the requirement together. - @torch.inference_mode() - def call_with_text_prompt( - self, inputs: TextPrompt, sampling_params: SamplingParams - ) -> Tuple[List[int], ExtraProcessedInputs | None]: - images = inputs.get("multi_modal_data", {}).get("image") - do_rescale = getattr(self.processor.image_processor, "do_rescale", - False) - if images is not None and isinstance(images[0], torch.Tensor): - # The default multimodal input loader will normalize images to [0, 1] when the requested - # format is "pt" (pytorch tensors), but not for "pil" (PIL images). - do_rescale = False + Mixed into processors supplying `config`/`processor`/`dtype`. The + mistral-native frontend serves every native checkpoint, including ones with + no vision encoder at all, so each entry point reports no encoder workload + when the config carries no `vision_config`. + """ - prompt = inputs["prompt"] - processed = self.processor(text=prompt, - images=images, - do_rescale=do_rescale) - input_ids = processed.pop("input_ids").tolist()[0] - # Remaining in `processed`: - # * "attention_mask": [B, num_input_tokens] - # * "pixel_values": [B, C, H, W] - # * "image_sizes": [B, 2] - extra_processed_inputs = None - pixel_values = processed.get("pixel_values") - if pixel_values is not None: - # We have no use for the `attention_mask`. - processed.pop("attention_mask") - # `image_sizes` is a `[B, 2]` tensor indicating the height and width of each image in the - # request. If we keep it as a regular tensor, it would get converted to a CUDA tensor before - # reaching the model forward. Since its values are used to infer the amount of padding - # + slice the patch embeddings, this would incur a D2H copy. We therefore convert it to a - # list here to avoid this. - processed["image_sizes"] = processed["image_sizes"].tolist() - # NOTE: `processed` is a dict-like object, but not actually a dict. - extra_processed_inputs = { - "multimodal_data": { - "image": { - **processed - } - } - } + def _has_vision_encoder(self) -> bool: + """Whether this checkpoint carries a Pixtral vision encoder. - return input_ids, extra_processed_inputs + ``MistralConfigLoader`` synthesizes ``vision_config`` only for vision + checkpoints, so its absence means the encoder geometry below is + meaningless for this checkpoint. + """ + return getattr(self.config, "vision_config", None) is not None # ------------------------------------------------------------------ # Deterministic dummy sizing for KV-cache encoder profiling. @@ -572,11 +467,61 @@ def _dummy_mm_data_for_size( } } + def _build_mm_item_metadata( + self, + _prompt_token_ids: List[int], + multimodal_data: Dict[str, Any], + ) -> Optional[MultimodalItemMetadata]: + """Return Pixtral image items and physical ViT patch counts.""" + if not self._has_vision_encoder(): + return None + image_data = multimodal_data.get("image") + if not isinstance(image_data, dict): + return None + image_sizes = image_data.get("image_sizes") + if image_sizes is None: + return None + patch, merge, _, _ = self._vision_geometry() + encoder_token_lengths = [ + self._vit_tokens(width=int(width), height=int(height), patch=patch) + for height, width in image_sizes + ] + min_tokens_per_image = merge * merge + if any(token_length < min_tokens_per_image or token_length % + min_tokens_per_image for token_length in encoder_token_lengths): + raise ValueError( + "Processed Mistral image geometry must contain a nonempty " + f"multiple of {min_tokens_per_image} encoder tokens") + item_refs = [("image", item_idx) + for item_idx in range(len(encoder_token_lengths))] + output_embedding_lengths = [ + token_length // (merge * merge) + for token_length in encoder_token_lengths + ] + return MultimodalItemMetadata( + item_refs=item_refs, + encoder_token_lengths=encoder_token_lengths, + output_embedding_lengths=output_embedding_lengths) + + def get_mm_item_metadata( + self, + prompt_token_ids: List[int], + multimodal_data: Dict[str, Any], + mm_item_order: Optional[List[Dict[str, Any]]] = None, + ) -> Optional[MultimodalItemMetadata]: + metadata = ( + self._build_mm_item_metadata(prompt_token_ids, multimodal_data) + if multimodal_data.get("multimodal_embedding") is None else None) + return metadata if metadata is not None else super( + ).get_mm_item_metadata(prompt_token_ids, multimodal_data, mm_item_order) + def get_mm_max_tokens_per_item( self, max_num_encoder_tokens: Optional[int] = None, ) -> Dict[str, int]: """Return the largest legal Pixtral image encoder item size.""" + if not self._has_vision_encoder(): + return {} patch, merge, _, max_size = self._vision_geometry() unit = patch * merge edge = max((max_size // unit) * unit, unit) @@ -594,14 +539,18 @@ def get_mm_max_tokens_per_item( return {"image": encoder_tokens} def get_max_mm_encoder_output_embeddings( - self, max_num_encoder_tokens: int) -> int: + self, max_num_encoder_tokens: int) -> Optional[int]: """Bound post-merger embeddings from one Pixtral encoder iteration.""" + if not self._has_vision_encoder(): + return None _, merge, _, _ = self._vision_geometry() return max_num_encoder_tokens // (merge * merge) def get_mm_encoder_attention_metadata_capacity( self, max_num_tokens: int) -> Optional[Dict[str, int]]: """Bound Pixtral contexts by the physical-token budget.""" + if not self._has_vision_encoder(): + return None _, merge, _, _ = self._vision_geometry() min_tokens_per_image = merge * merge return {"attention": max(1, max_num_tokens // min_tokens_per_image)} @@ -616,6 +565,10 @@ def get_dummy_mm_data( """Build processed Pixtral tensors for profiler-selected images.""" if max_num_encoder_tokens <= 0: raise ValueError("max_num_encoder_tokens must be positive") + if not self._has_vision_encoder(): + raise NotImplementedError( + "Checkpoint has no vision encoder; cannot build dummy " + "multimodal data") unsupported_modalities = set(mm_counts) - {"image"} if unsupported_modalities: raise ValueError("Pixtral cannot build dummy data for modalities " @@ -642,6 +595,102 @@ def get_dummy_mm_data( dtype=dtype, ) + +class MistralHFInputProcessor(MistralPixtralEncoderContract, + BaseMultimodalInputProcessor, + BaseMultimodalDummyInputsBuilder): + """Input processor for Mistral VLM checkpoints in HuggingFace format.""" + + def __init__(self, + model_path: str, + config: PretrainedConfig, + tokenizer: AutoTokenizer, + trust_remote_code: bool = True, + **kwargs): + super().__init__(model_path=model_path, + config=config, + tokenizer=tokenizer, + trust_remote_code=trust_remote_code, + **kwargs) + self._config = config + self._dtype = self._config.torch_dtype + self._model_path = model_path + self._tokenizer = (tokenizer if tokenizer is not None else + AutoTokenizer.from_pretrained( + model_path, + config=config, + use_fast=True, + trust_remote_code=True)) + self._processor = AutoProcessor.from_pretrained( + model_path, + use_fast=self.use_fast, + trust_remote_code=trust_remote_code) + logger.info(f"[mistral] HF processor={type(self._processor).__name__} " + f"tokenizer={type(self._tokenizer).__name__}") + + @property + def config(self) -> PretrainedConfig: + return self._config + + @property + def tokenizer(self) -> AutoTokenizer: + return self._tokenizer + + @property + def model_path(self) -> str: + return self._model_path + + @property + def processor(self) -> AutoProcessor: + return self._processor + + @property + def dtype(self) -> torch.dtype: + return self._dtype + + @torch.inference_mode() + def call_with_text_prompt( + self, inputs: TextPrompt, sampling_params: SamplingParams + ) -> Tuple[List[int], ExtraProcessedInputs | None]: + images = inputs.get("multi_modal_data", {}).get("image") + do_rescale = getattr(self.processor.image_processor, "do_rescale", + False) + if images is not None and isinstance(images[0], torch.Tensor): + # The default multimodal input loader will normalize images to [0, 1] when the requested + # format is "pt" (pytorch tensors), but not for "pil" (PIL images). + do_rescale = False + + prompt = inputs["prompt"] + processed = self.processor(text=prompt, + images=images, + do_rescale=do_rescale) + input_ids = processed.pop("input_ids").tolist()[0] + # Remaining in `processed`: + # * "attention_mask": [B, num_input_tokens] + # * "pixel_values": [B, C, H, W] + # * "image_sizes": [B, 2] + extra_processed_inputs = None + pixel_values = processed.get("pixel_values") + if pixel_values is not None: + # We have no use for the `attention_mask`. + processed.pop("attention_mask") + # `image_sizes` is a `[B, 2]` tensor indicating the height and width of each image in the + # request. If we keep it as a regular tensor, it would get converted to a CUDA tensor before + # reaching the model forward. Since its values are used to infer the amount of padding + # + slice the patch embeddings, this would incur a D2H copy. We therefore convert it to a + # list here to avoid this. + processed["image_sizes"] = processed["image_sizes"].tolist() + # NOTE: `processed` is a dict-like object, but not actually a dict. + extra_processed_inputs = { + "multimodal_data": { + "image": { + **processed + } + } + } + + return input_ids, extra_processed_inputs + def get_vocab_size(self) -> int: """Return the vocab size of the model.""" # Unlike some other VLMs, mistral3's vocab size is stored in its `text_config`, not the top-level @@ -666,9 +715,15 @@ def get_mm_special_token_ids(self) -> torch.Tensor: ]) -class MistralNativeInputProcessor(BaseMultimodalInputProcessor, +class MistralNativeInputProcessor(MistralPixtralEncoderContract, + BaseMultimodalInputProcessor, BaseMultimodalDummyInputsBuilder): - """Input processor for Mistral VLM checkpoints in mistral-native format.""" + """Input processor for Mistral VLM checkpoints in mistral-native format. + + Unlike the HF frontend, this one serves *every* mistral-native checkpoint + (text-only, MoE, mamba, audio and vision), which is why the inherited + encoder contract self-gates on the presence of a vision encoder. + """ def __init__( self, diff --git a/tests/integration/defs/cpp/conftest.py b/tests/integration/defs/cpp/conftest.py index 45eacab81501..a9638670d209 100644 --- a/tests/integration/defs/cpp/conftest.py +++ b/tests/integration/defs/cpp/conftest.py @@ -181,6 +181,41 @@ def build_kv_cache_compression_tests(request, build_type): ) +@pytest.fixture(scope="session") +def build_fmha_head104_tests(request: pytest.FixtureRequest, + build_type: str) -> None: + """Build the packed-QKV numerical regression and its native library.""" + cuda_arch = f"{request.param}-real" + _logger.info(f"Using CUDA arch: {cuda_arch}") + build_trt_llm( + build_type=build_type, + cuda_architectures=cuda_arch, + job_count=12, + use_ccache=True, + generator="Ninja", + nixl_root="/opt/nvidia/nvda_nixl", + skip_building_wheel=True, + configure_only=True, + ) + build_dir = _cpp.find_build_dir(build_type) + _cpp.run_command( + [ + "cmake", + "--build", + str(build_dir), + "--config", + build_type, + "--parallel", + "12", + "--target", + "fmhaHead104Test", + ], + cwd=build_dir, + env=_os.environ, + timeout=2700, + ) + + @pytest.fixture(scope="function", autouse=True) def keep_log_files(build_dir): """Backup previous cpp test results when run multiple ctest invocations.""" diff --git a/tests/integration/defs/cpp/test_unit_tests.py b/tests/integration/defs/cpp/test_unit_tests.py index 4a3deb858f90..c6ea44b83f44 100644 --- a/tests/integration/defs/cpp/test_unit_tests.py +++ b/tests/integration/defs/cpp/test_unit_tests.py @@ -1,4 +1,8 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + import os as _os +import pathlib as _pl import defs.cpp.cpp_common as _cpp import pytest @@ -70,3 +74,19 @@ def test_kv_cache_compression_unit_tests(build_kv_cache_compression_tests, env={**_os.environ}, timeout=2700, ) + + +@pytest.mark.parametrize("build_fmha_head104_tests", ["100", "120"], + indirect=True) +def test_fmha_head104_unit_tests(build_fmha_head104_tests: None, + build_dir: _pl.Path) -> None: + """Run the dedicated FMHA target on architectures without the full gtest suite.""" + _cpp.run_command( + [ + f"{build_dir}/tests/unit_tests/kernels/fmhaHead104Test", + f"--gtest_output=xml:{build_dir}/results-unit-tests-fmha-head104.xml", + ], + cwd=build_dir, + env={**_os.environ}, + timeout=2700, + ) diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index b35281ce61c7..c6c97e6cd145 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -18,6 +18,7 @@ l0_b200: tests: # ------------- CPP tests --------------- - cpp/test_unit_tests.py::test_kv_cache_compression_unit_tests[100] + - cpp/test_unit_tests.py::test_fmha_head104_unit_tests[100] - condition: ranges: system_gpu_count: diff --git a/tests/integration/test_lists/test-db/l0_gb202.yml b/tests/integration/test_lists/test-db/l0_gb202.yml index 79bd803a1890..effdf3cb74a7 100644 --- a/tests/integration/test_lists/test-db/l0_gb202.yml +++ b/tests/integration/test_lists/test-db/l0_gb202.yml @@ -15,11 +15,13 @@ l0_gb202: backend: pytorch tests: # ------------- PyTorch tests --------------- + - cpp/test_unit_tests.py::test_fmha_head104_unit_tests[120] - unittest/_torch/modeling -k "modeling_out_of_tree" - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_single_gpu[e8_k1_h512_i512-seq=8-dtype=torch.float16-backend=CUTLASS-quant=NVFP4-routing=Renormalize] - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_single_gpu[e8_k1_h512_i512-seq=8-dtype=torch.bfloat16-backend=CUTLASS-quant=NVFP4-routing=Renormalize] # - unittest/_torch/modeling -k "modeling_qwen3" # https://nvbugs/5234573 - unittest/_torch/attention/test_attention_mla.py + - unittest/_torch/attention/test_pixtral_head104_attention.py - test_e2e.py::test_ptp_quickstart_bert[VANILLA-BertForSequenceClassification-bert/bert-base-uncased-yelp-polarity] - test_e2e.py::test_ptp_quickstart_bert[TRTLLM-BertForSequenceClassification-bert/bert-base-uncased-yelp-polarity] - accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_bf16[latency] diff --git a/tests/integration/test_lists/test-db/l0_rtx_pro_6000.yml b/tests/integration/test_lists/test-db/l0_rtx_pro_6000.yml index bcc8842fd77f..0ecad80b5df8 100644 --- a/tests/integration/test_lists/test-db/l0_rtx_pro_6000.yml +++ b/tests/integration/test_lists/test-db/l0_rtx_pro_6000.yml @@ -14,6 +14,7 @@ l0_rtx_pro_6000: backend: pytorch tests: # ------------- PyTorch tests --------------- + - cpp/test_unit_tests.py::test_fmha_head104_unit_tests[120] - unittest/_torch/modeling -k "modeling_out_of_tree" # NCCL-EP v0.2 owns LL shared-memory sizing; cover the factory regression on SM120. - unittest/_torch/moe/test_communication_factory.py @@ -22,6 +23,7 @@ l0_rtx_pro_6000: - unittest/_torch/visual_gen/test_attention_flashinfer.py # - unittest/_torch/modeling -k "modeling_qwen3" # https://nvbugs/5234573 - unittest/_torch/attention/test_attention_mla.py + - unittest/_torch/attention/test_pixtral_head104_attention.py # SM120 W4A16 / W4A8 mixed-dtype GEMM coverage (paired with FinegrainedMixedDtypeGemm # MAX_SUPPORTED_SM_VERSION_W4A* gates and cutlass_heuristic.cpp SM120 routing). - unittest/_torch/kernels/parallel/test_w4a16_linear.py diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index c0c2e715bfc9..5aaa1d64c671 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -17,7 +17,6 @@ accuracy/test_llm_api_pytorch.py::TestQwen3_235B_A22B::test_fp8[throughput_laten accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_nvfp4[tep4_latency_moe_cutlass-torch_compile=True] SKIP (https://nvbugs/6561558) accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B_Instruct_2507::test_skip_softmax_attention_4gpus[target_sparsity_0.5-fp8kv=True] SKIP (https://nvbugs/6668774) accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_bf16[latency] SKIP (https://nvbugs/6412098) -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-nixl_kvcache-90] SKIP (https://nvbugs/5838199) cpp/test_multi_gpu.py::test_cache_transceiver[8proc-ucx_kvcache-90] SKIP (https://nvbugs/5838199) diff --git a/tests/unittest/_torch/attention/test_fmha_v2_head_size_coverage.py b/tests/unittest/_torch/attention/test_fmha_v2_head_size_coverage.py new file mode 100644 index 000000000000..ce1186ea38bd --- /dev/null +++ b/tests/unittest/_torch/attention/test_fmha_v2_head_size_coverage.py @@ -0,0 +1,107 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Static coverage for the Pixtral head-size-104 FMHA v2 kernel.""" + +import importlib.util +import pathlib +import tempfile +from types import ModuleType +from typing import Any + +import pytest + +pytestmark = pytest.mark.cpu_only + +_SETUP_PY = pathlib.Path(__file__).resolve().parents[4] / "cpp" / "kernels" / "fmha_v2" / "setup.py" + + +@pytest.fixture(scope="module") +def kernel_specs() -> tuple[ModuleType, list[Any]]: + """Return the generator module and specs using the wheel's generation flags.""" + if not _SETUP_PY.is_file(): + pytest.skip(f"fmha_v2 generator not present at {_SETUP_PY}") + + spec = importlib.util.spec_from_file_location("fmha_v2_setup", _SETUP_PY) + module = importlib.util.module_from_spec(spec) + captured = [] + + with pytest.MonkeyPatch.context() as patch, tempfile.TemporaryDirectory() as scratch: + patch.chdir(scratch) + patch.setenv("ENABLE_SM89_QMMA", "1") + patch.setenv("ENABLE_HMMA_FP32", "1") + patch.setenv("ENABLE_SM100", "1") + patch.setenv("ENABLE_SM120", "1") + patch.setenv("GENERATE_CUBIN", "1") + patch.setenv("GENERATE_CU_TRTLLM", "true") + patch.setenv("SCHEDULING_MODE", "1") + patch.delenv("DISABLE_SKIP_SOFTMAX", raising=False) + spec.loader.exec_module(module) + module.generate_files = captured.extend + module.enumerate_kernels() + + assert captured, "enumerate_kernels did not reach generate_files" + return module, [kspec for kspec, *_ in captured] + + +@pytest.mark.parametrize("sm", [90, 100, 120]) +@pytest.mark.parametrize("dtype", ["fp16", "bf16", "fp16_fp32"]) +def test_pixtral_head_size_has_a_padding_mask_kernel( + kernel_specs: tuple[ModuleType, list[Any]], + sm: int, + dtype: str, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Generate head-104 padding kernels for both FP16 accumulation modes and BF16.""" + module, specs = kernel_specs + matching = [ + kspec + for kspec in specs + if kspec.sm == sm + and kspec.head_size == 104 + and kspec.dtype == dtype + and kspec.input_layout == module.InputLayout.PACKED_QKV + and kspec.flash_attention + and kspec.warp_specialization == (sm == 90) + and (sm != 120 or not kspec.tiled) + and not kspec.enable_skip_softmax + ] + assert matching, f"No SM{sm} {dtype} packed-QKV FMHA v2 kernel is generated for head size 104" + + monkeypatch.setenv("GENERATE_CUBIN", "1") + assert any(module.selected_mask_types(kspec)[0] == "1" for kspec in matching), ( + f"Every SM{sm} {dtype} head-size-104 kernel has the padding mask disabled" + ) + + +def test_pixtral_padding_mask_excludes_other_variants( + kernel_specs: tuple[ModuleType, list[Any]], monkeypatch: pytest.MonkeyPatch +) -> None: + """Keep head-104 padding limited to supported architectures and precision modes.""" + module, specs = kernel_specs + packed = [ + kspec + for kspec in specs + if kspec.head_size == 104 + and kspec.input_layout == module.InputLayout.PACKED_QKV + and kspec.flash_attention + ] + added_arches = [kspec for kspec in packed if kspec.sm in (90, 120)] + excluded = { + "fp8": [kspec for kspec in added_arches if kspec.dtype.startswith("e4m3")], + "skip_softmax": [kspec for kspec in added_arches if kspec.enable_skip_softmax], + "sm90_non_warp_specialized": [ + kspec for kspec in added_arches if kspec.sm == 90 and not kspec.warp_specialization + ], + "sm120_tiled": [ + kspec + for kspec in added_arches + if kspec.sm == 120 and kspec.tiled and kspec.dtype in ("fp16", "bf16") + ], + "other_sms": [kspec for kspec in packed if kspec.sm not in (90, 100, 120)], + } + monkeypatch.setenv("GENERATE_CUBIN", "1") + for variant, candidates in excluded.items(): + assert candidates, f"No {variant} specifications were generated to check" + assert all(module.selected_mask_types(kspec)[0] == "0" for kspec in candidates), ( + f"An excluded {variant} head-size-104 kernel has the padding mask enabled" + ) diff --git a/tests/unittest/_torch/attention/test_pixtral_head104_attention.py b/tests/unittest/_torch/attention/test_pixtral_head104_attention.py new file mode 100644 index 000000000000..d042d98c7896 --- /dev/null +++ b/tests/unittest/_torch/attention/test_pixtral_head104_attention.py @@ -0,0 +1,65 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Exercise Pixtral-Large's packed-QKV dispatcher path and workspace bound.""" + +import pytest +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch.attention.backends.interface import ( + AttentionForwardArgs, + PredefinedAttentionMask, +) +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention, TrtllmAttentionMetadata +from tensorrt_llm._utils import get_sm_version + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@torch.inference_mode() +def test_pixtral_head104_attention(dtype: torch.dtype, monkeypatch: pytest.MonkeyPatch) -> None: + """Compare ragged head-104 attention to FP32 SDPA without quadratic workspace.""" + if get_sm_version() not in (90, 100, 103, 120): + pytest.skip("Exercises Pixtral FMHA on SM90, SM100-family, and SM120 GPUs") + # Exercise thop's dispatcher regardless of optional library configuration. + monkeypatch.setenv("TLLM_FMHA_LIBS", "fallback") + generator = torch.Generator(device="cuda").manual_seed(6665906) + lengths = [2048, 32] + num_heads, head_dim = 4, 104 + metadata = TrtllmAttentionMetadata( + max_num_requests=4096, max_num_tokens=16384, kv_cache_manager=None + ) + metadata.num_contexts = len(lengths) + metadata.request_ids = list(range(len(lengths))) + metadata.prompt_lens = lengths + metadata.seq_lens = torch.tensor(lengths, dtype=torch.int) + metadata.max_seq_len = max(lengths) + metadata.prepare() + attention = TrtllmAttention( + layer_idx=0, num_heads=num_heads, num_kv_heads=num_heads, head_dim=head_dim + ) + q, k, v = [ + torch.randn( + sum(lengths), num_heads * head_dim, device="cuda", dtype=dtype, generator=generator + ) + for _ in range(3) + ] + output = attention.forward( + torch.cat((q, k, v), dim=-1), + None, + None, + metadata, + forward_args=AttentionForwardArgs(attention_mask=PredefinedAttentionMask.FULL), + ) + # The unfused path reserves hundreds of GiB for this metadata capacity. + assert metadata.effective_workspace.numel() < 64 * 1024 * 1024 + reference = [] + offset = 0 + for length in lengths: + qi, ki, vi = [ + tensor[offset : offset + length].view(length, num_heads, head_dim).transpose(0, 1) + for tensor in (q, k, v) + ] + result = F.scaled_dot_product_attention(qi.float(), ki.float(), vi.float()) + reference.append(result.transpose(0, 1).reshape(length, num_heads * head_dim)) + offset += length + torch.testing.assert_close(output.float(), torch.cat(reference), atol=1e-2, rtol=1e-2) diff --git a/tests/unittest/_torch/modeling/test_modeling_mistral.py b/tests/unittest/_torch/modeling/test_modeling_mistral.py index 90944222db3e..62f22e2df0f7 100644 --- a/tests/unittest/_torch/modeling/test_modeling_mistral.py +++ b/tests/unittest/_torch/modeling/test_modeling_mistral.py @@ -22,7 +22,10 @@ from tensorrt_llm._torch import model_config as model_config_lib from tensorrt_llm._torch.attention.backends import utils as attention_utils from tensorrt_llm._torch.models import modeling_mistral -from tensorrt_llm._torch.models.modeling_mistral import MistralHFInputProcessor +from tensorrt_llm._torch.models.modeling_mistral import ( + MistralHFInputProcessor, + MistralNativeInputProcessor, +) from tensorrt_llm._torch.models.modeling_utils import MetaInitMode from tensorrt_llm._torch.pyexecutor import resource_manager from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import KeyType @@ -684,25 +687,59 @@ def test_mistral_attention_swa_layer_types(): # ``(h//patch)*(w//patch)`` -- deliberately *not* the hashing path's LLM-side # Pixtral count with framing tokens. # --------------------------------------------------------------------------- -def _make_dummy_processor(*, patch_size=14, spatial_merge_size=2, image_size=1540, num_channels=3): +def _make_dummy_processor( + *, + patch_size: int = 14, + spatial_merge_size: int = 2, + image_size: int = 1540, + num_channels: int = 3, + processor_cls: type[MistralHFInputProcessor] + | type[MistralNativeInputProcessor] = MistralHFInputProcessor, + geometry_from_processor: bool = False, +) -> MistralHFInputProcessor | MistralNativeInputProcessor: """Construct a processor stub with just the geometry the dummy math reads. - Bypasses the real ``__init__`` (tokenizer/processor loading); the empty - ``_processor`` forces ``_vision_geometry`` to fall back to ``vision_config`` - (the HF ``mistral3`` path). + Bypasses the real ``__init__`` (tokenizer/processor loading). By default the + empty ``_processor`` forces ``_vision_geometry`` to fall back to + ``vision_config`` (the HF ``mistral3`` path); ``geometry_from_processor`` + instead exposes ``patch_size``/``image_size`` on the processor -- the + mistral-common source (``mm_config.image_patch_size`` / ``max_image_size``) + -- and blanks them on ``vision_config`` so a regression that ignores the + processor is caught. ``processor_cls`` selects the frontend; both share the + same encoder contract. """ - instance = MistralHFInputProcessor.__new__(MistralHFInputProcessor) + instance = processor_cls.__new__(processor_cls) instance._config = SimpleNamespace( vision_config=SimpleNamespace( - patch_size=patch_size, image_size=image_size, num_channels=num_channels + patch_size=None if geometry_from_processor else patch_size, + image_size=None if geometry_from_processor else image_size, + num_channels=num_channels, ), spatial_merge_size=spatial_merge_size, ) - instance._processor = SimpleNamespace() + instance._processor = ( + SimpleNamespace(patch_size=patch_size, image_size=image_size) + if geometry_from_processor + else SimpleNamespace() + ) instance._dtype = torch.float16 return instance +def _make_text_only_native_processor() -> MistralNativeInputProcessor: + """A mistral-native processor for a checkpoint with no vision encoder. + + ``MistralNativeInputProcessor`` serves every ``checkpoint_format="mistral"`` + checkpoint, so `adapt_config_dict` leaves `vision_config` unset for the + text-only / MoE / mamba / audio shapes. Deliberately stubs *nothing* beyond + that config: the encoder geometry must not be reached at all, so dropping + the guard surfaces as an `AttributeError` rather than a soft assertion. + """ + instance = MistralNativeInputProcessor.__new__(MistralNativeInputProcessor) + instance._config = SimpleNamespace() + return instance + + def test_dummy_mm_max_tokens_per_item_is_image_only(): proc = _make_dummy_processor(patch_size=14, image_size=1540) demand = proc.get_mm_max_tokens_per_item() @@ -719,8 +756,7 @@ def test_attention_metadata_capacity_uses_token_budget(): def test_mistral_item_metadata_separates_patch_and_embedding_units(): - processor = object.__new__(MistralHFInputProcessor) - processor._vision_geometry = lambda: (14, 2, 3, 1024) + processor = _make_dummy_processor() metadata = processor.get_mm_item_metadata([], {"image": {"image_sizes": [[28, 56], [56, 56]]}}) @@ -729,6 +765,79 @@ def test_mistral_item_metadata_separates_patch_and_embedding_units(): assert metadata.output_embedding_lengths == (2, 4) +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "sizes, encoder_lengths, embedding_lengths", + [ + ([[28, 56]], [8], [2]), + ([[28, 56], [28, 56]], [8, 8], [2, 2]), + ([[28, 56], [56, 84]], [8, 24], [2, 6]), + ], +) +def test_native_processed_images_preserve_item_sizes( + sizes: list[list[int]], encoder_lengths: list[int], embedding_lengths: list[int] +) -> None: + """Keep unpadded per-image geometry through native processing and scheduling.""" + proc = _make_dummy_processor(processor_cls=MistralNativeInputProcessor) + pixels = [ + torch.full((3, height, width), float(i + 1)) for i, (height, width) in enumerate(sizes) + ] + encoded = transformers.BatchEncoding( + { + "input_ids": [1, 2, 3], + "attention_mask": [1, 1, 1], + "pixel_values": [pixel.numpy() for pixel in pixels], + "image_sizes": sizes, + } + ) + tokenizer = SimpleNamespace( + transformers_tokenizer=mock.Mock(), + instruct=SimpleNamespace( + mm_encoder=mock.Mock( + spec=modeling_mistral.ImageEncoder, + mm_config=SimpleNamespace(image_patch_size=14, max_image_size=1540), + ) + ), + ) + + def apply_chat_template( + *args: Any, return_tensors: str | None, **kwargs: Any + ) -> transformers.BatchEncoding: + assert return_tensors is None + return encoded + + tokenizer.transformers_tokenizer.apply_chat_template.side_effect = apply_chat_template + proc._processor = modeling_mistral.MistralCommonImageProcessor(tokenizer, proc.dtype) + input_ids, extra = proc.call_with_text_prompt( + { + "prompt": "Describe these images.", + "multi_modal_data": {"image": [Image.new("RGB", (w, h)) for h, w in sizes]}, + }, + tensorrt_llm.SamplingParams(), + ) + + assert input_ids == [1, 2, 3] + image = extra["multimodal_data"]["image"] + assert image["image_sizes"] == sizes + assert image["pixel_values"].shape == ( + len(sizes), + 3, + max(h for h, _ in sizes), + max(w for _, w in sizes), + ) + for i, ((height, width), pixel) in enumerate(zip(sizes, pixels)): + torch.testing.assert_close( + image["pixel_values"][i, :, :height, :width], pixel.to(proc.dtype) + ) + assert torch.count_nonzero(image["pixel_values"][i, :, height:]) == 0 + assert torch.count_nonzero(image["pixel_values"][i, :, :, width:]) == 0 + metadata = proc.get_mm_item_metadata(input_ids, extra["multimodal_data"]) + metadata.validate() + assert metadata.item_refs == tuple(("image", i) for i in range(len(sizes))) + assert metadata.encoder_token_lengths == tuple(encoder_lengths) + assert metadata.output_embedding_lengths == tuple(embedding_lengths) + + @pytest.mark.parametrize("budget", [1024, 4096, 8192]) def test_dummy_get_size_for_max_tokens_fits_and_aligns(budget): proc = _make_dummy_processor() @@ -814,3 +923,78 @@ def test_dummy_mm_data_satisfies_the_encoder_input_contract(): assert batched_pixel_values.shape[0] == num_images assert batched_sizes.shape == (num_images, 2) assert batched_pixel_values.dtype == torch.float16 + + +# --------------------------------------------------------------------------- +# MM encoder item scheduling is enabled from the *model* class +# (`Mistral3VLM.supports_mm_encoder_item_scheduling`), but the contract it +# implies is required of the *input processor*, and `create_input_processor` +# picks the processor from `checkpoint_format` alone. Both Mistral frontends +# therefore have to satisfy the contract. +# --------------------------------------------------------------------------- +@pytest.mark.parametrize( + "processor_cls, geometry_from_processor", + [(MistralHFInputProcessor, False), (MistralNativeInputProcessor, True)], + ids=["hf-config", "native-processor"], +) +@pytest.mark.cpu_only +def test_both_mistral_processors_satisfy_item_scheduling_contract( + processor_cls: type[MistralHFInputProcessor] | type[MistralNativeInputProcessor], + geometry_from_processor: bool, +) -> None: + """Check encoder budgets, request metadata, and dummy inputs for both frontends.""" + proc = _make_dummy_processor( + processor_cls=processor_cls, geometry_from_processor=geometry_from_processor + ) + budget = 8192 + + # Exercise the sequence `ModelEngine.__init__` runs for item scheduling. + max_tokens_per_item = proc.get_mm_max_tokens_per_item() + assert max_tokens_per_item == {"image": 12100} + + capacity = proc.get_mm_encoder_attention_metadata_capacity(budget) + assert capacity == {"attention": 2048} + + max_output_embeddings = proc.get_max_mm_encoder_output_embeddings(budget) + assert max_output_embeddings == 2048 + + # ...and the request-routing hook, which gates the item-scheduled path. + metadata = proc.get_mm_item_metadata([], {"image": {"image_sizes": [[28, 56]]}}) + assert metadata is not None + metadata.validate() + assert metadata.item_refs == (("image", 0),) + assert metadata.encoder_token_lengths == (8,) + assert metadata.output_embedding_lengths == (2,) + + dummy = proc.get_dummy_mm_data(max_num_encoder_tokens=16, mm_counts={"image": 1}) + assert dummy["image"]["pixel_values"].shape == (1, 3, 56, 56) + assert dummy["image"]["pixel_values"].dtype == proc.dtype + assert dummy["image"]["image_sizes"] == [[56, 56]] + dummy_metadata = proc.get_mm_item_metadata([], dummy) + assert dummy_metadata is not None + dummy_metadata.validate() + assert dummy_metadata.item_refs == (("image", 0),) + assert dummy_metadata.encoder_token_lengths == (16,) + assert dummy_metadata.output_embedding_lengths == (4,) + + +# The native processor is shared by every `checkpoint_format="mistral"` +# checkpoint, including the text-only / MoE / mamba / audio shapes that carry no +# `vision_config` at all. Those must keep the base class's neutral behaviour +# rather than raising `AttributeError` out of the encoder geometry -- the engine +# calls `get_mm_max_tokens_per_item` unconditionally. +@pytest.mark.cpu_only +def test_text_only_native_checkpoint_reports_no_encoder_geometry() -> None: + proc = _make_text_only_native_processor() + + assert proc.get_mm_max_tokens_per_item() == {} + assert proc.get_max_mm_encoder_output_embeddings(8192) is None + assert proc.get_mm_encoder_attention_metadata_capacity(8192) is None + # A text-only request carries no image payload, which the routing hook + # already treats as "not item-scheduled". + assert proc.get_mm_item_metadata([], {}) is None + assert proc.get_mm_item_metadata([], {"image": {"image_sizes": [[28, 56]]}}) is None + + # The memory profiler treats `NotImplementedError` as "skip MM profiling". + with pytest.raises(NotImplementedError): + proc.get_dummy_mm_data(max_num_encoder_tokens=8192, mm_counts={"image": 1}) From da07e21397a9cd10e289dac37a1f08e5e722e336 Mon Sep 17 00:00:00 2001 From: Pranav Shrestha Date: Fri, 9 Oct 2026 14:26:09 -0700 Subject: [PATCH 02/23] [https://nvbugs/6562942][fix] Cherry-pick #19815 from release/1.3: Load Qwen3.5 packed INT4 checkpoints Forward-port of release/1.3 PR https://github.com/NVIDIA/TensorRT-LLM/pull/19815 (release/1.3 commit 4450bbc1bbd71503cb3d688ea122887c90c29251) to main. Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> Co-authored-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> Signed-off-by: Simeng Liu --- .../checkpoints/hf/qwen3_5_weight_mapper.py | 108 +++++++- .../checkpoints/hf/qwen4_exp_weight_mapper.py | 43 +--- tensorrt_llm/_torch/models/modeling_utils.py | 39 +++ tensorrt_llm/_torch/modules/linear.py | 30 ++- .../test_lists/test-db/l0_b200.yml | 2 + .../_torch/modeling/test_qwen4_exp_support.py | 6 +- .../hf/test_qwen3_5_weight_mapper.py | 243 ++++++++++++++++++ .../models/test_qwen3_5_int4_loading.py | 131 ++++++++++ .../_torch/modules/test_linear_uneven_tp.py | 60 +++++ 9 files changed, 598 insertions(+), 64 deletions(-) create mode 100644 tests/unittest/_torch/models/test_qwen3_5_int4_loading.py diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/qwen3_5_weight_mapper.py b/tensorrt_llm/_torch/models/checkpoints/hf/qwen3_5_weight_mapper.py index 48a8a1e91eb9..7d94412edcbc 100644 --- a/tensorrt_llm/_torch/models/checkpoints/hf/qwen3_5_weight_mapper.py +++ b/tensorrt_llm/_torch/models/checkpoints/hf/qwen3_5_weight_mapper.py @@ -12,7 +12,10 @@ from tensorrt_llm._torch.models.checkpoints.hf.qwen3_next_weight_mapper import ( Qwen3NextHfWeightMapper, ) -from tensorrt_llm._torch.models.modeling_utils import register_mapper +from tensorrt_llm._torch.models.modeling_utils import ( + concatenate_weights_by_tp_rank, + register_mapper, +) from tensorrt_llm._torch.moe.fused_moe.interface import MoEWeightLoadingMode from tensorrt_llm._torch.moe.fused_moe.weight_owner import is_moe_weight_owner from tensorrt_llm.quantization import QuantAlgo @@ -20,6 +23,19 @@ _FP8_2D_BLOCK_SIZE = 128 +def _check_packed_int4( + tensors: dict[str, torch.Tensor], keys: set[str], rows: tuple[int, ...], prefix: str +) -> tuple[int, ...]: + """Validate packed projection storage and return physical row counts.""" + if not all(tensors[key].dtype in (torch.uint8, torch.int8) for key in keys): + raise ValueError(f"Cannot fuse packed INT4 and unpacked weights for {prefix}") + if len({tensors[key].dtype for key in keys}) != 1: + raise ValueError(f"Packed INT4 projections must have the same storage dtype for {prefix}") + if any(row % 2 != 0 for row in rows): + raise ValueError(f"INT4 projection dimensions must be even for {prefix}") + return tuple(row // 2 for row in rows) + + @register_mapper("HF", "Qwen3_5ForConditionalGeneration") @register_mapper("HF", "Qwen3_5MoeForCausalLM") @register_mapper("HF", "Qwen3_5MoeForConditionalGeneration") @@ -42,7 +58,9 @@ class Qwen3_5MoeHfWeightMapper(Qwen3NextHfWeightMapper): Qwen3Next checkpoints store pre-packed in_proj_qkvz and in_proj_ba tensors. Qwen3.5 checkpoints store them as separate in_proj_qkv + z (or fully split q/k/v/z) and b + a tensors. This mapper packs them - into the grouped-interleaved layout that TRT-LLM expects. + into the grouped-interleaved layout consumed by the Qwen3Next mapper. + INT4 projections are concatenated directly in per-rank dense order, + preserving the two output channels stored in each byte. For FP8_BLOCK_SCALES checkpoints, the packed qkvz tensor is then dequantized to bf16 as a temporary workaround for TP loading (handled in _dequantize_linear_attn_fp8_qkvz). For MIXED_PRECISION @@ -551,7 +569,7 @@ def _dequantize_linear_attn_fp8_per_tensor(self, weights: dict) -> dict: if name not in drop_keys } - def _pack_split_projections(self, weights: dict) -> dict: + def _pack_split_projections(self, weights: dict, dense_layout: bool = False) -> dict: config = self.config.pretrained_config num_k_groups = config.linear_num_key_heads num_v_heads = config.linear_num_value_heads @@ -560,6 +578,13 @@ def _pack_split_projections(self, weights: dict) -> dict: f"linear_num_key_heads ({num_k_groups})" ) + pack_projection_tensor = self._pack_projection_tensor + num_groups = num_k_groups + if dense_layout: + mapping = self.config.mapping + num_groups = 1 if mapping.enable_attention_dp else mapping.tp_size + pack_projection_tensor = concatenate_weights_by_tp_rank + grouped_weights = defaultdict(dict) packed_weights = {} @@ -576,9 +601,43 @@ def _pack_split_projections(self, weights: dict) -> dict: expected_ba = config.linear_num_value_heads for (prefix, suffix), tensors in grouped_weights.items(): + if dense_layout and suffix == "pre_quant_scale": + # Input scales are shared by fused projections, not concatenated by output row. + for projection, keys in ( + ("qkvz", ("qkv", "z") if "qkv" in tensors else ("q", "k", "v", "z")), + ("ba", ("b", "a")), + ): + if not set(keys) & tensors.keys(): + continue + if missing := set(keys) - tensors.keys(): + raise ValueError( + f"Cannot fuse {prefix}.in_proj_{projection}.pre_quant_scale: " + f"missing scales for {sorted(missing)}" + ) + scales = [ + tensors[key] + if isinstance(tensors[key], torch.Tensor) + else tensors[key][...] + for key in keys + ] + if any(scale.shape != (config.hidden_size,) for scale in scales): + raise ValueError( + f"Expected pre_quant_scale shape ({config.hidden_size},) for {prefix}" + ) + if not all(torch.equal(scales[0], scale) for scale in scales[1:]): + raise ValueError( + f"Cannot fuse {prefix}.in_proj_{projection}: pre_quant_scale vectors " + "must be identical; export with shared input scales" + ) + packed_name = f"{prefix}.in_proj_{projection}.{suffix}" + if packed_name in packed_weights: + raise ValueError(f"Packed projection {packed_name} already exists") + packed_weights[packed_name] = scales[0] + continue # `weight_scale_inv` (loaded by FP8BlockScalesLinearMethod) is the # only suffix stored in 2D-block format [ceil(out/block_size), ceil(in/block_size)]; - # weight/bias/weight_scale all keep out_features as their leading dim. + # INT4 weights pack two output channels per byte; their scales and + # biases retain the logical output dimension. if suffix == "weight_scale_inv": row_q = math.ceil(expected_q / _FP8_2D_BLOCK_SIZE) row_v = math.ceil(expected_v / _FP8_2D_BLOCK_SIZE) @@ -590,13 +649,22 @@ def _pack_split_projections(self, weights: dict) -> dict: qkvz_keys = {"qkv", "q", "k", "v", "z"} & tensors.keys() if qkvz_keys: + if ( + dense_layout + and suffix == "weight" + and any(tensors[key].dtype in (torch.uint8, torch.int8) for key in qkvz_keys) + ): + row_q, row_v = _check_packed_int4(tensors, qkvz_keys, (row_q, row_v), prefix) if "qkv" in tensors: missing = {"qkv", "z"} - tensors.keys() assert not missing, ( f"Missing split projections {sorted(missing)} for {prefix}.{suffix}" ) + split_q, split_v = ( + (expected_q, expected_v) if suffix == "weight_scale_inv" else (row_q, row_v) + ) q_tensor, k_tensor, v_tensor = split_packed_qkv( - tensors["qkv"], expected_q, expected_v + tensors["qkv"], split_q, split_v ) else: missing = {"q", "k", "v", "z"} - tensors.keys() @@ -616,8 +684,8 @@ def _pack_split_projections(self, weights: dict) -> dict: assert packed_name not in packed_weights, ( f"Packed projection {packed_name} already exists" ) - packed_weights[packed_name] = self._pack_projection_tensor( - [q_tensor, k_tensor, v_tensor, tensors["z"]], num_k_groups + packed_weights[packed_name] = pack_projection_tensor( + [q_tensor, k_tensor, v_tensor, tensors["z"]], num_groups ) ba_keys = {"b", "a"} & tensors.keys() @@ -628,6 +696,10 @@ def _pack_split_projections(self, weights: dict) -> dict: ) if suffix == "weight": + if dense_layout and any( + tensors[key].dtype in (torch.uint8, torch.int8) for key in ba_keys + ): + (row_ba,) = _check_packed_int4(tensors, ba_keys, (row_ba,), prefix) assert tensors["b"].shape[0] == row_ba assert tensors["a"].shape[0] == row_ba @@ -635,8 +707,8 @@ def _pack_split_projections(self, weights: dict) -> dict: assert packed_name not in packed_weights, ( f"Packed projection {packed_name} already exists" ) - packed_weights[packed_name] = self._pack_projection_tensor( - [tensors["b"], tensors["a"]], num_k_groups + packed_weights[packed_name] = pack_projection_tensor( + [tensors["b"], tensors["a"]], num_groups ) return packed_weights @@ -762,16 +834,26 @@ def preprocess_weights(self, weights: dict, allow_partial_loading: bool = False) if quant_algo in (QuantAlgo.MIXED_PRECISION, QuantAlgo.NVFP4): normalized_weights = self._dequantize_lm_head_nvfp4(normalized_weights) - packed_weights = self._pack_split_projections(normalized_weights) + dense_layout = quant_algo == QuantAlgo.W4A16_AWQ + if dense_layout: + # Apply MTP renaming and non-projection transforms while in_proj + # tensors are still split. The inherited mapper passes those + # through, so the directly fused result is not permuted again. + normalized_weights = super().preprocess_weights( + normalized_weights, allow_partial_loading=allow_partial_loading + ) + packed_weights = self._pack_split_projections(normalized_weights, dense_layout=dense_layout) if quant_algo == QuantAlgo.FP8_BLOCK_SCALES and not is_modelopt_pb_wo: packed_weights = self._dequantize_linear_attn_fp8_qkvz(packed_weights) if not getattr(self.config.pretrained_config, "num_experts", 0): packed_weights = self._remap_dense_mlp_weights(packed_weights) - processed_weights = super().preprocess_weights( - packed_weights, allow_partial_loading=allow_partial_loading - ) + processed_weights = packed_weights + if not dense_layout: + processed_weights = super().preprocess_weights( + packed_weights, allow_partial_loading=allow_partial_loading + ) if is_consumable and not isinstance(processed_weights, ConsumableWeightsDict): return ConsumableWeightsDict(processed_weights) return processed_weights diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/qwen4_exp_weight_mapper.py b/tensorrt_llm/_torch/models/checkpoints/hf/qwen4_exp_weight_mapper.py index f26a112f92b5..def3e78adffb 100644 --- a/tensorrt_llm/_torch/models/checkpoints/hf/qwen4_exp_weight_mapper.py +++ b/tensorrt_llm/_torch/models/checkpoints/hf/qwen4_exp_weight_mapper.py @@ -76,7 +76,10 @@ from torch import nn from tensorrt_llm._torch.models.checkpoints.hf.qwen2_moe_weight_mapper import Qwen2MoeHfWeightMapper -from tensorrt_llm._torch.models.modeling_utils import register_mapper +from tensorrt_llm._torch.models.modeling_utils import ( + concatenate_weights_by_tp_rank, + register_mapper, +) from tensorrt_llm._torch.modules.qwen4_exp.ple import Qwen4ExpPinnedHostEmbedding from tensorrt_llm._torch.moe.fused_moe.interface import MoEWeightLoadingMode from tensorrt_llm._torch.moe.fused_moe.weight_owner import is_moe_weight_owner @@ -89,36 +92,6 @@ _PER_EXPERT_PROJECTION_PATTERN = re.compile(r"^\d+\.(?:gate_proj|up_proj|down_proj)\.") -def _rank_block(components: list[torch.Tensor], tp_size: int) -> torch.Tensor: - """Concat `components` (each `[out_i, ...]`) into the per-rank column- - parallel row order a contiguous `TensorParallelMode.COLUMN` split recovers. - - At `tp_size == 1` this is a plain `[c0 | c1 | ...]` concat. At - `tp_size > 1` each component is split across ranks and the rank-`r` slices - are grouped together (`[c0_r0 c1_r0 ... | c0_r1 c1_r1 ... | ...]`), so a - later contiguous per-rank row split hands each rank its own head slice of - every component. - """ - if not components: - raise ValueError("rank blocking requires at least one projection") - if tp_size <= 0: - raise ValueError(f"tp_size must be positive, got {tp_size}") - trailing_shape = components[0].shape[1:] - for component in components: - if component.ndim == 0 or component.shape[1:] != trailing_shape: - raise ValueError("rank-blocked projections must have matching trailing dimensions") - if component.shape[0] % tp_size: - raise ValueError( - f"projection rows {component.shape[0]} are not divisible by tp_size={tp_size}" - ) - if tp_size == 1: - return torch.cat(components, dim=0) - rows: list[torch.Tensor] = [] - for rank in range(tp_size): - rows.extend(split(c, tp_size, rank) for c in components) - return torch.cat(rows, dim=0) - - def _normalize_moe_module_weights( module_weights: dict[str, torch.Tensor], config ) -> tuple[dict[str, torch.Tensor], MoEWeightLoadingMode]: @@ -374,7 +347,7 @@ def preprocess_weights(self, weights: dict, allow_partial_loading: bool = False) f"{key} convolution width is {w.shape[1]}, expected {conv_kernel}" ) conv_q, conv_k, conv_v = torch.split(w, [key_dim, key_dim, value_dim], dim=0) - new_weights[key] = _rank_block([conv_q, conv_k, conv_v], tp_size) + new_weights[key] = concatenate_weights_by_tp_rank([conv_q, conv_k, conv_v], tp_size) else: new_weights[key] = tensor @@ -387,7 +360,7 @@ def preprocess_weights(self, weights: dict, allow_partial_loading: bool = False) ) # Split the pre-fused in_proj_qkv [Q|K|V] into its Q/K/V sub-tensors # BEFORE rank-blocking. Passing the concatenated [Q|K|V] as one - # component makes _rank_block cut it at contiguous (key+key+value)/tp + # component cuts it at contiguous (key+key+value)/tp # boundaries, which scrambles Q/K/V across ranks at tp>1 (rank r would # get [Q[..], K[..]] instead of [Q_r|K_r|V_r]) -> garbage GDN output. # Each of Q, K, V, and Z is sharded independently. Mirror the @@ -413,8 +386,8 @@ def preprocess_weights(self, weights: dict, allow_partial_loading: bool = False) f"{tuple(parts[projection].shape)}, expected {expected_ba_shape}" ) q_t, k_t, v_t = torch.split(qkv_t, [key_dim, key_dim, value_dim], dim=0) - qkvz = _rank_block([q_t, k_t, v_t, parts["z"][:]], tp_size) - ba = _rank_block([parts["b"][:], parts["a"][:]], tp_size) + qkvz = concatenate_weights_by_tp_rank([q_t, k_t, v_t, parts["z"][:]], tp_size) + ba = concatenate_weights_by_tp_rank([parts["b"][:], parts["a"][:]], tp_size) expected_qkvz_rows = 2 * key_dim + 2 * value_dim if qkvz.shape[0] != expected_qkvz_rows: raise ValueError( diff --git a/tensorrt_llm/_torch/models/modeling_utils.py b/tensorrt_llm/_torch/models/modeling_utils.py index 661db4f9cb1f..1e283dda7b1a 100755 --- a/tensorrt_llm/_torch/models/modeling_utils.py +++ b/tensorrt_llm/_torch/models/modeling_utils.py @@ -33,6 +33,7 @@ from ..modules.rms_norm import RMSNorm from ..moe.fused_moe import MoE, VanillaMoE, is_moe_weight_owner from ..speculative import SpecMetadata +from ..utils import split from ._arch_index import (MODEL_ARCH_TO_MODULE, SPEC_MODE_TO_MODULE, is_builtin_zoo_module) @@ -108,6 +109,44 @@ def __torch_dispatch__(self, func, types, args=(), kwargs=None): return func(*args, **kwargs) +def concatenate_weights_by_tp_rank(components: list[torch.Tensor], + tp_size: int) -> torch.Tensor: + """Concatenate projection components in rank-major row order. + + Args: + components: Tensors shaped ``[rows_i, ...]`` with matching trailing + dimensions. Each physical row count must be divisible by + ``tp_size``. Packed weights remain packed; callers must ensure + logical channel boundaries align with physical row boundaries. + tp_size: Number of tensor-parallel ranks. + + Returns: + A tensor ordered as ``[c0_r0 | c1_r0 | ... | c0_r1 | c1_r1 | ...]``. + A contiguous row split gives each rank its slice of every component. + With one rank, this is a plain concatenation along dimension zero. + """ + if not components: + raise ValueError("rank blocking requires at least one projection") + if tp_size <= 0: + raise ValueError(f"tp_size must be positive, got {tp_size}") + trailing_shape = components[0].shape[1:] + for component in components: + if component.ndim == 0 or component.shape[1:] != trailing_shape: + raise ValueError( + "rank-blocked projections must have matching trailing dimensions" + ) + if component.shape[0] % tp_size: + raise ValueError( + f"projection rows {component.shape[0]} are not divisible by tp_size={tp_size}" + ) + if tp_size == 1: + return torch.cat(components, dim=0) + rows: list[torch.Tensor] = [] + for rank in range(tp_size): + rows.extend(split(c, tp_size, rank) for c in components) + return torch.cat(rows, dim=0) + + def duplicate_kv_weight(weight: torch.Tensor, num_kv_heads: int, tensor_parallel_size: int): diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py index 655ae90476cc..732b52c3da5c 100644 --- a/tensorrt_llm/_torch/modules/linear.py +++ b/tensorrt_llm/_torch/modules/linear.py @@ -3013,23 +3013,27 @@ def load_weights_vanilla(self, module: Linear, weights: List[Dict]) -> None: elm_packing = 2 if module.tp_mode == TensorParallelMode.COLUMN else 1 load_weights_vanilla_helper(module, weights, elm_packing=elm_packing) - # Use the same device as the weight tensor - # as we register pre_quant_scale after sharded model weights are moved to respective gpus - device = module.weight.device - pre_quant_scale = module.load_shard( - weights[0]["pre_quant_scale"], - # pre_quant_scale applies to activation as opposed to weight, so flip tp_mode the other way around - flip_tp=True, - device=device, - ) + pre_quant_scale = weights[0].get("pre_quant_scale") + if pre_quant_scale is not None: + # Use the device of the loaded weight for the activation scale. + device = module.weight.device + pre_quant_scale = module.load_shard( + pre_quant_scale, + # Activation scales shard along the input dimension. + flip_tp=True, + device=device, + ) - module.pre_quant_scale = Parameter( - torch.ones((module.in_features, ), dtype=pre_quant_scale.dtype), - requires_grad=False).to(device=device) + module.pre_quant_scale = Parameter( + torch.ones((module.in_features, ), dtype=pre_quant_scale.dtype), + requires_grad=False).to(device=device) + + copy_weight(module.pre_quant_scale, pre_quant_scale) + else: + module.pre_quant_scale = None weight_scale = self.load_weight_scales(module, weights)[0] - copy_weight(module.pre_quant_scale, pre_quant_scale) copy_weight(module.weight_scale, weight_scale.T.contiguous()) def load_weights_fused_qkv_linear(self, module: Linear, diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index c6c97e6cd145..7a32677263ca 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -34,6 +34,8 @@ l0_b200: backend: pytorch tests: # ------------- PyTorch tests --------------- + - unittest/_torch/models/test_qwen3_5_int4_loading.py + - unittest/_torch/modules/test_linear_uneven_tp.py::TestW4A16AWQMLP::test_optional_pre_quant_scale - llmapi/test_llm_api_pytorch_t5.py::test_t5_v2_decode_beyond_encoder_length[False] - llmapi/test_llm_api_pytorch_t5.py::test_t5_v2_decode_beyond_encoder_length[True] - llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_block_reuse_is_disabled_for_v2 diff --git a/tests/unittest/_torch/modeling/test_qwen4_exp_support.py b/tests/unittest/_torch/modeling/test_qwen4_exp_support.py index 7dac22c46df8..36447728bbbc 100644 --- a/tests/unittest/_torch/modeling/test_qwen4_exp_support.py +++ b/tests/unittest/_torch/modeling/test_qwen4_exp_support.py @@ -693,17 +693,17 @@ def test_runtime_config_rejects_unsupported_semantics( def test_mapper_normalizes_bf16_and_per_expert_fp8_weights() -> None: from tensorrt_llm._torch.models.checkpoints.hf.qwen4_exp_weight_mapper import ( _normalize_moe_module_weights, - _rank_block, ) + from tensorrt_llm._torch.models.modeling_utils import concatenate_weights_by_tp_rank from tensorrt_llm._torch.moe.fused_moe.interface import MoEWeightLoadingMode q = torch.arange(8, dtype=torch.float32).reshape(4, 2) z = torch.arange(8, 16, dtype=torch.float32).reshape(4, 2) - blocked = _rank_block([q, z], tp_size=2) + blocked = concatenate_weights_by_tp_rank([q, z], tp_size=2) expected = torch.cat((q[:2], z[:2], q[2:], z[2:])) torch.testing.assert_close(blocked, expected) with pytest.raises(ValueError, match="not divisible"): - _rank_block([torch.empty(3, 2)], tp_size=2) + concatenate_weights_by_tp_rank([torch.empty(3, 2)], tp_size=2) config = SimpleNamespace(hidden_size=4, moe_intermediate_size=3, num_experts=2) fused, mode = _normalize_moe_module_weights( diff --git a/tests/unittest/_torch/models/checkpoints/hf/test_qwen3_5_weight_mapper.py b/tests/unittest/_torch/models/checkpoints/hf/test_qwen3_5_weight_mapper.py index 4768b957deed..016fa73ecdfd 100644 --- a/tests/unittest/_torch/models/checkpoints/hf/test_qwen3_5_weight_mapper.py +++ b/tests/unittest/_torch/models/checkpoints/hf/test_qwen3_5_weight_mapper.py @@ -144,6 +144,249 @@ def test_fp8_rowwise_full() -> None: assert out[f"{_ATTN_PREFIX}.in_proj_ba.weight_scale"].shape == (_PACKED_BA_ROWS, 1) +def _pack_int4_rows(weight: torch.Tensor) -> torch.Tensor: + return ((weight[0::2] & 15) | ((weight[1::2] & 15) << 4)).to(torch.uint8) + + +def _unpack_int4_rows(weight: torch.Tensor) -> torch.Tensor: + packed = weight.view(torch.uint8) + values = torch.stack((packed & 15, packed >> 4), dim=1).flatten(0, 1).to(torch.int8) + return torch.where(values >= 8, values - 16, values) + + +def _make_int4_weights( + split_qkv: bool = False, quantize_ba: bool = True, head_dim: int = 4 +) -> tuple[Qwen3_5MoeHfWeightMapper, dict, dict, dict]: + mapper = _make_mapper(QuantAlgo.W4A16_AWQ) + config = mapper.config.pretrained_config + config.linear_num_key_heads = 16 + config.linear_num_value_heads = 48 + config.linear_key_head_dim = head_dim + config.linear_value_head_dim = head_dim + generator = torch.Generator().manual_seed(6562942) + logical_weights, scales, weights = {}, {}, {} + for name, rows in ( + ("q", 16 * head_dim), + ("k", 16 * head_dim), + ("v", 48 * head_dim), + ("z", 48 * head_dim), + ("b", 48), + ("a", 48), + ): + logical_weights[name] = torch.randint( + -8, 8, (rows, 256), generator=generator, dtype=torch.int8 + ) + scales[name] = torch.rand(rows, 2, generator=generator) + 0.1 + if name in ("b", "a") and not quantize_ba: + logical_weights[name] = logical_weights[name].to(torch.bfloat16) + weights[f"{_ATTN_PREFIX}.in_proj_{name}.weight"] = logical_weights[name] + else: + weights[f"{_ATTN_PREFIX}.in_proj_{name}.weight"] = _pack_int4_rows( + logical_weights[name] + ) + weights[f"{_ATTN_PREFIX}.in_proj_{name}.weight_scale"] = scales[name] + if not split_qkv: + for suffix in ("weight", "weight_scale"): + weights[f"{_ATTN_PREFIX}.in_proj_qkv.{suffix}"] = torch.cat( + [weights.pop(f"{_ATTN_PREFIX}.in_proj_{name}.{suffix}") for name in ("q", "k", "v")] + ) + return mapper, weights, logical_weights, scales + + +@pytest.mark.parametrize( + "tp_size,split_qkv,quantize_ba,storage_dtype,head_dim,has_pre_quant_scale", + [ + # Cover both checkpoint layouts and gate formats in the TP1 and TP2 paths. + (1, False, False, torch.uint8, 4, False), + pytest.param(1, False, True, torch.uint8, 128, False, id="qwen35-27b"), + (1, True, False, torch.uint8, 4, False), + (1, True, True, torch.uint8, 4, False), + (2, False, False, torch.uint8, 4, False), + (2, False, True, torch.uint8, 4, False), + (2, True, False, torch.uint8, 4, False), + (2, True, True, torch.uint8, 4, False), + # Shared activation scales in both checkpoint layouts. + (2, False, True, torch.uint8, 4, True), + (2, True, True, torch.uint8, 4, True), + # Signed storage in both layouts; TP8 gives three packed gate rows per rank. + (2, True, True, torch.int8, 4, False), + (8, False, True, torch.int8, 4, False), + ], +) +def test_int4_projections_preserve_values_and_scales( + tp_size: int, + split_qkv: bool, + quantize_ba: bool, + storage_dtype: torch.dtype, + head_dim: int, + has_pre_quant_scale: bool, +) -> None: + mapper, weights, logical_weights, scales = _make_int4_weights(split_qkv, quantize_ba, head_dim) + if head_dim == 128: + assert weights[f"{_ATTN_PREFIX}.in_proj_qkv.weight"].shape[0] == 5120 + weights = { + name: value.view(storage_dtype) if value.dtype == torch.uint8 else value + for name, value in weights.items() + } + mapper.config.mapping.tp_size = tp_size + input_scales = {"qkvz": torch.linspace(0.5, 1.5, 256), "ba": torch.linspace(1.0, 2.0, 256)} + if has_pre_quant_scale: + mapper.config.pretrained_config.hidden_size = 256 + for name in list(weights): + if name.endswith(".weight"): + projection = "ba" if ".in_proj_b." in name or ".in_proj_a." in name else "qkvz" + weights[name.removesuffix("weight") + "pre_quant_scale"] = input_scales[ + projection + ].clone() + out = mapper.preprocess_weights(weights) + activation = torch.randn(3, 256, generator=torch.Generator().manual_seed(17)) + + for projection, components in (("qkvz", ("q", "k", "v", "z")), ("ba", ("b", "a"))): + quantized = projection == "qkvz" or quantize_ba + scale_key = f"{_ATTN_PREFIX}.in_proj_{projection}.pre_quant_scale" + if has_pre_quant_scale: + torch.testing.assert_close(out[scale_key], input_scales[projection]) + fused_input = activation * out[scale_key] + reference_input = activation * input_scales[projection] + else: + assert scale_key not in out + fused_input = reference_input = activation + packed_weight = out[f"{_ATTN_PREFIX}.in_proj_{projection}.weight"] + if quantized: + assert packed_weight.dtype == storage_dtype + assert packed_weight.shape[0] * 2 == sum( + logical_weights[name].shape[0] for name in components + ) + fused_weight = _unpack_int4_rows(packed_weight) + fused_scale = out[f"{_ATTN_PREFIX}.in_proj_{projection}.weight_scale"] + assert fused_scale.shape == (fused_weight.shape[0], 2) + fused_dequantized = fused_weight.float() * fused_scale.repeat_interleave(128, dim=1) + else: + assert packed_weight.dtype == torch.bfloat16 + fused_weight = packed_weight + fused_dequantized = fused_weight.float() + + # Each rank's fused projection must compute the same outputs as the + # individual checkpoint projections, with every scale attached to its row. + for rank, shard in enumerate(fused_dequantized.chunk(tp_size, dim=0)): + references = [] + row_offset = rank * (fused_weight.shape[0] // tp_size) + for name in components: + original = logical_weights[name].chunk(tp_size, dim=0)[rank] + torch.testing.assert_close( + fused_weight[row_offset : row_offset + original.shape[0]], original + ) + row_offset += original.shape[0] + dequantized = logical_weights[name].float() + if quantized: + dequantized = dequantized * scales[name].repeat_interleave(128, dim=1) + references.append((reference_input @ dequantized.T).chunk(tp_size, dim=1)[rank]) + torch.testing.assert_close(fused_input @ shard.T, torch.cat(references, dim=1)) + + +@pytest.mark.parametrize("projection", ["z", "a"]) +@pytest.mark.parametrize("invalid", ["different", "missing", "shape"]) +def test_int4_rejects_incompatible_pre_quant_scales(projection: str, invalid: str) -> None: + """A fused Linear cannot apply different input scales to its components.""" + mapper, weights, _, _ = _make_int4_weights() + mapper.config.pretrained_config.hidden_size = 256 + for name in ("qkv", "z", "b", "a"): + weights[f"{_ATTN_PREFIX}.in_proj_{name}.pre_quant_scale"] = torch.ones(256) + key = f"{_ATTN_PREFIX}.in_proj_{projection}.pre_quant_scale" + if invalid == "missing": + del weights[key] + elif invalid == "shape": + weights[key] = torch.ones(1, 256) + else: + weights[key][0] = 2 + with pytest.raises(ValueError, match="pre_quant_scale"): + mapper.preprocess_weights(weights) + + +def test_int4_rejects_non_byte_aligned_tp_shards() -> None: + mapper, weights, _, _ = _make_int4_weights() + mapper.config.mapping.tp_size = 16 + with pytest.raises(ValueError, match="projection rows 24 are not divisible by tp_size=16"): + mapper.preprocess_weights(weights) + + +@pytest.mark.parametrize("projection", ["z", "a"]) +@pytest.mark.parametrize("other_dtype", [torch.int8, torch.bfloat16]) +def test_int4_rejects_mixed_storage_dtypes(projection: str, other_dtype: torch.dtype) -> None: + mapper, weights, _, _ = _make_int4_weights() + name = f"{_ATTN_PREFIX}.in_proj_{projection}.weight" + weights[name] = weights[name].to(other_dtype) + message = "same storage dtype" if other_dtype == torch.int8 else "packed INT4 and unpacked" + with pytest.raises(ValueError, match=message): + mapper.preprocess_weights(weights) + + +@pytest.mark.parametrize("projection", ["q", "v", "ba"]) +def test_int4_rejects_odd_logical_dimensions(projection: str) -> None: + mapper, weights, _, _ = _make_int4_weights() + config = mapper.config.pretrained_config + config.linear_num_key_heads = 1 + if projection == "ba": + config.linear_num_value_heads = 47 + weights = { + key: value + for key, value in weights.items() + if ".in_proj_b." in key or ".in_proj_a." in key + } + elif projection == "q": + config.linear_key_head_dim = 3 + else: + config.linear_num_value_heads = 1 + config.linear_value_head_dim = 3 + with pytest.raises(ValueError, match="projection dimensions must be even"): + mapper.preprocess_weights(weights) + + +def test_int4_excluded_qkvz_keeps_bf16_layout() -> None: + mapper, weights, logical_weights, _ = _make_int4_weights(split_qkv=True) + mapper.config.mapping.tp_size = 2 + for name in ("q", "k", "v", "z"): + weights[f"{_ATTN_PREFIX}.in_proj_{name}.weight"] = logical_weights[name].to(torch.bfloat16) + del weights[f"{_ATTN_PREFIX}.in_proj_{name}.weight_scale"] + out = mapper.preprocess_weights(weights) + fused = out[f"{_ATTN_PREFIX}.in_proj_qkvz.weight"] + assert fused.dtype == torch.bfloat16 + assert f"{_ATTN_PREFIX}.in_proj_qkvz.weight_scale" not in out + for rank, shard in enumerate(fused.chunk(2)): + torch.testing.assert_close( + shard, + torch.cat([logical_weights[name].chunk(2)[rank] for name in ("q", "k", "v", "z")]).to( + torch.bfloat16 + ), + ) + + +def test_int4_attention_dp_uses_unsplit_projections() -> None: + mapper, weights, _, _ = _make_int4_weights() + reference = mapper.preprocess_weights(weights) + mapper.config.mapping.tp_size = 16 + mapper.config.mapping.enable_attention_dp = True + out = mapper.preprocess_weights(weights) + for name in reference: + torch.testing.assert_close(out[name], reference[name]) + + +def test_int4_partial_loading_and_mtp_names() -> None: + mapper, weights, _, _ = _make_int4_weights() + weights = { + name.replace("model.layers.0", "mtp.layers.0"): value for name, value in weights.items() + } + reference = mapper.preprocess_weights(weights) + assert "model.layers.1.linear_attn.in_proj_qkvz.weight" in reference + out = {} + for name, value in weights.items(): + out.update(mapper.preprocess_weights({name: value}, allow_partial_loading=True)) + mapper.finalize_update_weights() + assert out.keys() == reference.keys() + for name in reference: + torch.testing.assert_close(out[name], reference[name]) + + def test_modelopt_fp8_per_tensor_linear_attention() -> None: # ModelOpt stores QKV and Z as separate per-tensor FP8 projections with independent scales. checkpoint_prefix = "model.language_model.layers.0.linear_attn" diff --git a/tests/unittest/_torch/models/test_qwen3_5_int4_loading.py b/tests/unittest/_torch/models/test_qwen3_5_int4_loading.py new file mode 100644 index 000000000000..ac5705e124eb --- /dev/null +++ b/tests/unittest/_torch/models/test_qwen3_5_int4_loading.py @@ -0,0 +1,131 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +import torch +from transformers import Qwen3NextConfig + +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.models.checkpoints.hf.qwen3_5_weight_mapper import Qwen3_5MoeHfWeightMapper +from tensorrt_llm._torch.models.modeling_auto import AutoModelForCausalLM +from tensorrt_llm._utils import get_sm_version +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.models.modeling_utils import QuantConfig +from tensorrt_llm.quantization import QuantAlgo + + +@pytest.mark.skipif(get_sm_version() < 80, reason="W4A16 AWQ requires Ampere or newer") +@pytest.mark.parametrize("has_pre_quant_scale", [False, True]) +def test_int4_mapper_to_gdn_tp2(monkeypatch: pytest.MonkeyPatch, has_pre_quant_scale: bool) -> None: + """Load mapper output through real GDN projections and check each TP rank's outputs.""" + hidden_size, tp_size = 256, 2 + # The unused embedding/MLP collectives must not allocate multi-GPU workspaces. + monkeypatch.setattr( + "tensorrt_llm._torch.distributed.ops.get_allreduce_workspace", lambda _: None + ) + prefix = "model.layers.0.linear_attn" + generator = torch.Generator().manual_seed(6562942) + weights, reference_weights = {}, {} + for name, rows in (("q", 2048), ("k", 2048), ("v", 6144), ("z", 6144), ("b", 48), ("a", 48)): + # Two nonzero inputs per row keep the independent reference exactly representable. + logical = torch.zeros(rows, hidden_size, dtype=torch.int8) + row_ids = torch.arange(rows) + for group in range(2): + columns = group * 128 + torch.randint(128, (rows,), generator=generator) + logical[row_ids, columns] = torch.randint( + -8, 8, (rows,), generator=generator, dtype=torch.int8 + ) + if name in ("a", "b"): + weights[f"{prefix}.in_proj_{name}.weight"] = logical.to(torch.bfloat16) + reference_weights[name] = logical.float() + else: + scales = 2.0 ** torch.randint(-3, 1, (rows, 2), generator=generator) + weights[f"{prefix}.in_proj_{name}.weight"] = ( + (logical[0::2] & 15) | ((logical[1::2] & 15) << 4) + ).to(torch.uint8) + weights[f"{prefix}.in_proj_{name}.weight_scale"] = scales + reference_weights[name] = logical.float() * scales.repeat_interleave(128, dim=1) + for suffix in ("weight", "weight_scale"): + weights[f"{prefix}.in_proj_qkv.{suffix}"] = torch.cat( + [weights.pop(f"{prefix}.in_proj_{name}.{suffix}") for name in ("q", "k", "v")] + ) + input_scale = torch.tensor([0.5, 2.0], dtype=torch.bfloat16).repeat_interleave(128) + if has_pre_quant_scale: + for name in ("qkv", "z"): + weights[f"{prefix}.in_proj_{name}.pre_quant_scale"] = input_scale.clone() + x = torch.randint(-2, 3, (3, hidden_size), generator=generator).to("cuda", torch.bfloat16) + + # Projection forwards have no collectives, so both real TP mappings run on one GPU. + for rank in range(tp_size): + config = Qwen3NextConfig( + architectures=["Qwen3_5ForCausalLM"], + hidden_size=hidden_size, + intermediate_size=512, + num_hidden_layers=1, + num_attention_heads=8, + num_key_value_heads=2, + num_experts=0, + vocab_size=256, + linear_num_key_heads=16, + linear_num_value_heads=48, + linear_key_head_dim=128, + linear_value_head_dim=128, + layer_types=["linear_attention"], + torch_dtype=torch.bfloat16, + ) + model_config = ModelConfig( + pretrained_config=config, + mapping=Mapping(world_size=tp_size, tp_size=tp_size, rank=rank), + quant_config=QuantConfig( + quant_algo=QuantAlgo.W4A16_AWQ, + group_size=128, + pre_quant_scale=True, + exclude_modules=[f"{prefix}.in_proj_a", f"{prefix}.in_proj_b"], + ), + ) + model = AutoModelForCausalLM.from_config(model_config) + mapper = Qwen3_5MoeHfWeightMapper() + mapper.init_model_and_config(model, model_config) + mapped = mapper.preprocess_weights(weights) + gdn = model.model.layers[0].linear_attn + assert model_config.skip_create_weights_in_init + assert gdn.in_proj_ba.quant_config.quant_algo is None + for projection, names in (("qkvz", ("q", "k", "v", "z")), ("ba", ("b", "a"))): + linear = getattr(gdn, f"in_proj_{projection}") + key_prefix = f"{prefix}.in_proj_{projection}." + linear.load_weights( + [ + { + key.removeprefix(key_prefix): value + for key, value in mapped.items() + if key.startswith(key_prefix) + } + ] + ) + linear.post_load_weights() + linear.cuda() + reference_input = x + if has_pre_quant_scale and projection == "qkvz": + reference_input = x * input_scale.cuda() + expected = torch.cat( + [ + torch.nn.functional.linear( + reference_input.float(), reference_weights[name].cuda() + ).chunk(tp_size, dim=-1)[rank] + for name in names + ], + dim=-1, + ).to(torch.bfloat16) + torch.testing.assert_close(linear(x), expected, rtol=0, atol=0) diff --git a/tests/unittest/_torch/modules/test_linear_uneven_tp.py b/tests/unittest/_torch/modules/test_linear_uneven_tp.py index a183dfd50260..43ddd1b47728 100644 --- a/tests/unittest/_torch/modules/test_linear_uneven_tp.py +++ b/tests/unittest/_torch/modules/test_linear_uneven_tp.py @@ -1157,6 +1157,66 @@ class TestW4A16AWQMLP(_AWQMLPMixin): quant_algo = QuantAlgo.W4A16_AWQ + @pytest.mark.parametrize( + "has_scale,reload_without_scale", [(False, False), (True, False), (True, True)] + ) + @pytest.mark.parametrize("tp_mode", [TensorParallelMode.COLUMN, TensorParallelMode.ROW]) + def test_optional_pre_quant_scale( + self, has_scale: bool, reload_without_scale: bool, tp_mode: TensorParallelMode + ) -> None: + dim, tp_size = 256, 2 + dtype = torch.float16 + weights = build_weights(dim, dim, self.quant_algo, bias=False) + pre_quant_scale = torch.linspace(0.5, 1.5, dim, dtype=dtype, device="cuda") + if has_scale: + weights[0]["pre_quant_scale"] = pre_quant_scale + else: + weights[0].pop("pre_quant_scale") + + x = torch.randn(3, dim, dtype=dtype, device="cuda") + outputs = [] + for rank in range(tp_size): + linear = Linear( + dim, + dim, + bias=False, + dtype=dtype, + mapping=FakeMapping(tp_size, rank), + quant_config=QuantConfig( + quant_algo=self.quant_algo, group_size=128, has_zero_point=False + ), + tensor_parallel_mode=tp_mode, + reduce_output=False, + ) + linear.load_weights(weights) + linear.post_load_weights() + linear.cuda() + if has_scale: + expected_scale = pre_quant_scale + if tp_mode == TensorParallelMode.ROW: + expected_scale = pre_quant_scale.chunk(tp_size)[rank] + torch.testing.assert_close(linear.pre_quant_scale, expected_scale) + else: + assert linear.pre_quant_scale is None + if reload_without_scale: + unscaled_weights = [ + {key: value for key, value in weights[0].items() if key != "pre_quant_scale"} + ] + linear.load_weights(unscaled_weights) + linear.post_load_weights() + assert linear.pre_quant_scale is None + shard_input = x.chunk(tp_size, dim=-1)[rank] if tp_mode == TensorParallelMode.ROW else x + outputs.append(linear(shard_input.contiguous())) + + result = ( + torch.cat(outputs, dim=-1) if tp_mode == TensorParallelMode.COLUMN else sum(outputs) + ) + # The synthetic checkpoint maps input i to output i + 1 with unit weight. + reference_weight = torch.diag(torch.ones(dim - 1, dtype=dtype, device="cuda"), diagonal=-1) + scaled_input = x * pre_quant_scale if has_scale and not reload_without_scale else x + expected = torch.nn.functional.linear(scaled_input, reference_weight) + torch.testing.assert_close(result, expected, rtol=1e-2, atol=1e-2) + @pytest.mark.skipif( not (get_sm_version() in (89, 90) or is_sm_100f()), From f8117136b164bd6ccb1f98e6d4cd06b6b7e365fe Mon Sep 17 00:00:00 2001 From: TensorRT LLM AI Agent <296075020+trtllm-agent@users.noreply.github.com> Date: Mon, 12 Oct 2026 10:00:07 +0800 Subject: [PATCH 03/23] [https://nvbugs/6877539][fix] Cherry-pick #20029 from release/1.3: Fold GQA group into SDPA query axis in MiniMax-M3 decode attention Forward-port of release/1.3 PR https://github.com/NVIDIA/TensorRT-LLM/pull/20029 (release/1.3 commit 9f79700d057e9eeb479d1f9251299111ad4660eb) to main. Also updates a main-only comment above the GQA check that still described decode as one SDPA call per KV head. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> Signed-off-by: chenfeiz0326 Co-authored-by: chenfeiz0326 Co-authored-by: Claude Signed-off-by: Simeng Liu --- .../_torch/models/modeling_minimaxm3.py | 49 +++++----- .../unittest/_torch/models/test_minimax_m3.py | 97 +++++++++++++++++++ 2 files changed, 122 insertions(+), 24 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_minimaxm3.py b/tensorrt_llm/_torch/models/modeling_minimaxm3.py index 311600ab67ca..f39afc84a1db 100644 --- a/tensorrt_llm/_torch/models/modeling_minimaxm3.py +++ b/tensorrt_llm/_torch/models/modeling_minimaxm3.py @@ -1941,8 +1941,8 @@ def _sdpa_dense_attention_core( v_cache_view, m3_meta.req_to_token, m3_meta.slot_ids, max_k ) - # SDPA expands KV heads to query heads for prefill. Decode processes - # one KV head at a time to bound the temporary GQA expansion. + # SDPA expands KV heads to query heads for prefill. Decode folds each + # GQA group into the query axis and reads K/V unexpanded. if self.num_heads % max(self.num_key_value_heads, 1) != 0: raise RuntimeError( f"Dense GQA expects num_heads ({self.num_heads}) divisible " @@ -2007,30 +2007,31 @@ def _sdpa_dense_attention_core( else: # Decode: one query token per request at position seq_lens - 1. valid = kv_positions < seq_lens_dev.unsqueeze(-1) # [batch, max_k] - q_b = q_view.view(batch, 1, self.num_heads, self.head_dim).transpose( - 1, 2 - ) # [batch, H, 1, d] mask_b = valid.unsqueeze(1).unsqueeze(1) # [batch, 1, 1, k] - # Expand K/V one KV head at a time: all heads at once needs an - # O(batch * max_k * num_heads) temporary that overflows the CUDA-graph - # pool under attention DP. - out_b = q.new_empty(batch, self.num_heads, 1, self.head_dim) + # Fold the GQA group into SDPA's query-length axis: the ``group`` Q + # heads sharing KV head h become ``group`` query rows of one SDPA + # head that reads KV head h unexpanded. Exact for decode (every + # query row sits at seq_lens - 1, so one key mask serves them all), + # and it needs neither an O(batch * max_k * num_heads) + # repeat_interleave temporary in the CUDA-graph pool nor one SDPA + # launch per KV head. + num_kv_heads = max(self.num_key_value_heads, 1) + q_b = q_view.view(batch, num_kv_heads, group, self.head_dim) # [batch, KV, group, d] + k_b = k_padded.transpose(1, 2) # [batch, KV, k, d] + v_b = v_padded.transpose(1, 2) # [batch, KV, k, d] with sdpa_kernel(_DENSE_SDPA_BACKENDS): - for h in range(max(self.num_key_value_heads, 1)): - qh = slice(h * group, (h + 1) * group) - k_h = k_padded[:, :, h : h + 1].repeat_interleave(group, dim=2) - v_h = v_padded[:, :, h : h + 1].repeat_interleave(group, dim=2) - out_b[:, qh] = torch.nn.functional.scaled_dot_product_attention( - q_b[:, qh].to(q.dtype), - k_h.transpose(1, 2).to(q.dtype), - v_h.transpose(1, 2).to(q.dtype), - attn_mask=mask_b, - dropout_p=0.0, - is_causal=False, - ) # [batch, group, 1, d] - # Drop the singleton query axis; a transpose(1, 2).reshape would - # scramble (head, head_dim) when H != head_dim. - output.view(batch, self.num_heads, self.head_dim).copy_(out_b.squeeze(2)) + out_b = torch.nn.functional.scaled_dot_product_attention( + q_b.to(q.dtype), + k_b.to(q.dtype), + v_b.to(q.dtype), + attn_mask=mask_b, + dropout_p=0.0, + is_causal=False, + ) # [batch, KV, group, d] + # Q head h * group + g is query row g of SDPA head h, so viewing the + # output as [batch, KV, group, d] keeps (head, head_dim) ordering; a + # transpose(1, 2).reshape would scramble it when H != head_dim. + output.view(batch, num_kv_heads, group, self.head_dim).copy_(out_b) return output diff --git a/tests/unittest/_torch/models/test_minimax_m3.py b/tests/unittest/_torch/models/test_minimax_m3.py index 5e3ce1d1292b..310df8921cac 100644 --- a/tests/unittest/_torch/models/test_minimax_m3.py +++ b/tests/unittest/_torch/models/test_minimax_m3.py @@ -3009,3 +3009,100 @@ def check_forward(packed, unused_k, unused_v, meta, args): ) kernel.forward.assert_called_once() torch.testing.assert_close(output, q) + + +def _dense_sdpa_decode_reference(q, k_ctx, v_ctx, seq_lens, num_kv_heads, head_dim): + """Per-request GQA decode via explicit ``repeat_interleave`` (pre-#18872 math). + + ``q`` is ``[batch, num_heads * head_dim]``; ``k_ctx`` / ``v_ctx`` are + ``[batch, max_len, num_kv_heads, head_dim]`` holding each request's full + context, valid for positions ``< seq_lens[b]``. + """ + batch = q.shape[0] + num_heads = q.shape[1] // head_dim + group = num_heads // num_kv_heads + out = torch.empty_like(q, dtype=torch.float32) + for b in range(batch): + n = int(seq_lens[b]) + q_b = q[b].view(num_heads, 1, head_dim).float() # [H, 1, d] + k_b = k_ctx[b, :n].float().repeat_interleave(group, dim=1).transpose(0, 1) # [H, n, d] + v_b = v_ctx[b, :n].float().repeat_interleave(group, dim=1).transpose(0, 1) + o_b = torch.nn.functional.scaled_dot_product_attention(q_b, k_b, v_b) # [H, 1, d] + out[b] = o_b.reshape(-1) + return out + + +@pytest.mark.skipif(not _has_cuda(), reason="dense SDPA decode runs on the CUDA KV pool") +@pytest.mark.parametrize("num_kv_heads", [1, 2, 4]) +@pytest.mark.parametrize("group", [1, 4]) +def test_minimax_m3_sdpa_dense_decode_matches_repeat_interleave_reference(num_kv_heads, group): + """The decode branch of ``_sdpa_dense_attention_core`` matches plain GQA math. + + Covers ``group > 1``, ``num_kv_heads > 1`` and a batch with mixed + ``seq_lens`` (shorter rows padded out to ``max_k``), so a head-order or + mask-broadcast mistake in the GQA-into-query-axis fold shows up here even + though it would not move throughput (NVBug 6877539). + """ + torch.manual_seed(0) + device, dtype = torch.device("cuda"), torch.bfloat16 + head_dim, tokens_per_block = 64, 4 + num_heads = num_kv_heads * group + seq_lens = torch.tensor([5, 1, 11, 8], dtype=torch.int32) # mixed; max 11 -> 3 blocks + batch, max_len = int(seq_lens.numel()), int(seq_lens.max()) + blocks_per_req = -(-max_len // tokens_per_block) + + # Requests own scattered, non-contiguous pages so a gather or slot mix-up + # cannot pass by accident. + num_pages = batch * blocks_per_req + 3 + pages = torch.randperm(num_pages)[: batch * blocks_per_req].view(batch, blocks_per_req) + pos = torch.arange(blocks_per_req * tokens_per_block) + req_to_token = ( + pages[:, pos // tokens_per_block] * tokens_per_block + pos % tokens_per_block + ).to(device=device, dtype=torch.int32) + + k_ctx = torch.randn(batch, max_len, num_kv_heads, head_dim, device=device, dtype=dtype) + v_ctx = torch.randn_like(k_ctx) + # Pool layout [num_pages, kv_factor, tokens_per_block, num_kv_heads, head_dim]. + # Fill every unused slot with garbage: the seq_lens mask must hide it. + kv_pool = ( + torch.randn( + num_pages, 2, tokens_per_block, num_kv_heads, head_dim, device=device, dtype=dtype + ) + * 50 + ) + for b in range(batch): + n = int(seq_lens[b]) + slots = req_to_token[b, : n - 1].long() # context before this decode step + page, within = slots // tokens_per_block, slots % tokens_per_block + kv_pool[page, 0, within] = k_ctx[b, : n - 1] + kv_pool[page, 1, within] = v_ctx[b, : n - 1] + # The new token's K/V arrives through ``k`` / ``v`` and is written by the method. + out_cache_loc = torch.stack( + [req_to_token[b, int(seq_lens[b]) - 1] for b in range(batch)] + ).long() + k_new = k_ctx[torch.arange(batch), seq_lens.long() - 1].reshape(batch, -1) + v_new = v_ctx[torch.arange(batch), seq_lens.long() - 1].reshape(batch, -1) + q = torch.randn(batch, num_heads * head_dim, device=device, dtype=dtype) + + layer = SimpleNamespace( + layer_idx=0, num_heads=num_heads, num_key_value_heads=num_kv_heads, head_dim=head_dim + ) + m3_meta = SimpleNamespace( + slot_ids=torch.arange(batch, device=device, dtype=torch.int32), + req_to_token=req_to_token, + seq_lens=seq_lens.to(device), + max_seqlen_k=max_len, + is_prefill=False, + ) + attn_metadata = SimpleNamespace( + kv_cache_manager=SimpleNamespace(get_buffers=lambda layer_idx: kv_pool), + minimax_m3={"metadata": m3_meta, "out_cache_loc": out_cache_loc}, + is_cuda_graph=False, + max_seq_len=max_len, + ) + output = torch.empty(batch, num_heads * head_dim, device=device, dtype=dtype) + + MiniMaxM3Attention._sdpa_dense_attention_core(layer, q, k_new, v_new, attn_metadata, output) + + ref = _dense_sdpa_decode_reference(q, k_ctx, v_ctx, seq_lens, num_kv_heads, head_dim) + torch.testing.assert_close(output.float(), ref, atol=2e-2, rtol=2e-2) From 682c90e6a849ac98fa1794f5f5cbd5305df04a26 Mon Sep 17 00:00:00 2001 From: Ivy Zhang <25222398+crazydemo@users.noreply.github.com> Date: Mon, 12 Oct 2026 18:27:40 +0800 Subject: [PATCH 04/23] [None][test] Cherry-pick #20132 from release/1.3: Move cases to multinode test list Forward-port of release/1.3 PR https://github.com/NVIDIA/TensorRT-LLM/pull/20132 (release/1.3 commit 0ea0e59b012a581fc1940c268cb7f0809d10554f) to main. Signed-off-by: Ivy Zhang <25222398+crazydemo@users.noreply.github.com> Signed-off-by: Simeng Liu --- tests/integration/test_lists/qa/llm_function_core.txt | 1 - tests/integration/test_lists/qa/llm_function_multinode.txt | 1 + 2 files changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/integration/test_lists/qa/llm_function_core.txt b/tests/integration/test_lists/qa/llm_function_core.txt index 1fbcfa32289f..e02e1bfb7144 100644 --- a/tests/integration/test_lists/qa/llm_function_core.txt +++ b/tests/integration/test_lists/qa/llm_function_core.txt @@ -604,7 +604,6 @@ accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_prims_ts_bfloat16 accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_w4a8_mxfp4[fp8-latency] accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_w4a8_mxfp4[mxfp8-latency] accuracy/test_llm_api_pytorch.py::TestQwen3_8_27B::test_dflash2[VANILLA] -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 diff --git a/tests/integration/test_lists/qa/llm_function_multinode.txt b/tests/integration/test_lists/qa/llm_function_multinode.txt index 5a94176558fb..75512ab5e2fa 100644 --- a/tests/integration/test_lists/qa/llm_function_multinode.txt +++ b/tests/integration/test_lists/qa/llm_function_multinode.txt @@ -1,6 +1,7 @@ accuracy/test_llm_api_pytorch.py::TestQwen3_8_2_4T_A95B::test_fp8_tp16_mtp3_trtllm accuracy/test_llm_api_pytorch.py::TestQwen3_8_2_4T_A95B::test_fp8_adp16_deepgemm accuracy/test_llm_api_pytorch.py::TestQwen3_8_2_4T_A95B::test_nvfp4_adp16_cutedsl +accuracy/test_llm_api_pytorch.py::TestQwen3_8_2_4T_A95B::test_nvfp4_tp8_mtp3_trtllm test_e2e.py::test_multi_nodes_eval[DeepSeek-R1/DeepSeek-R1-0528-FP4-tp16-mmlu] test_e2e.py::test_multi_nodes_eval[MiniMax-M3-tp16-mmlu] test_e2e.py::test_openai_disagg_multi_nodes_completion[ctx_tp2pp1-gen_tp2pp1] From eaeba343db88394226d39fc3579e5d0df5e02b06 Mon Sep 17 00:00:00 2001 From: chenfeiz0326 Date: Tue, 13 Oct 2026 12:52:46 +0800 Subject: [PATCH 05/23] [https://nvbugs/6463987][fix] Cherry-pick #20136 from release/1.3: Pin GB300 GLM-5 ctx workers to their GPU's NUMA node in perf-sanity Forward-port of release/1.3 PR https://github.com/NVIDIA/TensorRT-LLM/pull/20136 (release/1.3 commit 47c047945ef3b2e596ff19092bfa8461379e05e8) to main. The release PR cherry-picked main PR #20103, which was closed without merging, so this change was not on main. Signed-off-by: chenfeiz0326 <203214996+chenfeiz0326@users.noreply.github.com> Co-authored-by: chenfeiz0326 <203214996+chenfeiz0326@users.noreply.github.com> Co-authored-by: Claude Signed-off-by: Simeng Liu --- ...-5-fp4_1k1k_con1_ctx1_dep2_gen1_tep4_eplb0_mtp3_ccb-NIXL.yaml | 1 + ...4_1k1k_con4096_ctx1_dep2_gen1_dep8_eplb256_mtp1_ccb-NIXL.yaml | 1 + ...fp4_1k1k_con512_ctx1_dep2_gen1_dep32_eplb0_mtp3_ccb-NIXL.yaml | 1 + ...4_8k1k_con1024_ctx1_dep2_gen1_dep8_eplb256_mtp1_ccb-NIXL.yaml | 1 + ...-5-fp4_8k1k_con1_ctx1_dep2_gen1_tep8_eplb0_mtp3_ccb-NIXL.yaml | 1 + ...fp4_8k1k_con512_ctx1_dep2_gen1_dep32_eplb0_mtp3_ccb-NIXL.yaml | 1 + 6 files changed, 6 insertions(+) diff --git a/tests/scripts/perf-sanity/disaggregated/gb300_glm-5-fp4_1k1k_con1_ctx1_dep2_gen1_tep4_eplb0_mtp3_ccb-NIXL.yaml b/tests/scripts/perf-sanity/disaggregated/gb300_glm-5-fp4_1k1k_con1_ctx1_dep2_gen1_tep4_eplb0_mtp3_ccb-NIXL.yaml index b40c86dd6cec..072a16fc482d 100644 --- a/tests/scripts/perf-sanity/disaggregated/gb300_glm-5-fp4_1k1k_con1_ctx1_dep2_gen1_tep4_eplb0_mtp3_ccb-NIXL.yaml +++ b/tests/scripts/perf-sanity/disaggregated/gb300_glm-5-fp4_1k1k_con1_ctx1_dep2_gen1_tep4_eplb0_mtp3_ccb-NIXL.yaml @@ -36,6 +36,7 @@ environment: build_wheel: false work_dir: worker_env_var: TLLM_LOG_LEVEL=INFO TRTLLM_SERVER_DISABLE_GC=1 TRTLLM_WORKER_DISABLE_GC=1 TRTLLM_ENABLE_PDL=1 ENROOT_ALLOW_DEV=yes TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS=3 + ctx_worker_env_var: TLLM_NUMA_AWARE_WORKER_AFFINITY=1 server_env_var: TRTLLM_SERVER_DISABLE_GC=1 profiling: nsys_on: false diff --git a/tests/scripts/perf-sanity/disaggregated/gb300_glm-5-fp4_1k1k_con4096_ctx1_dep2_gen1_dep8_eplb256_mtp1_ccb-NIXL.yaml b/tests/scripts/perf-sanity/disaggregated/gb300_glm-5-fp4_1k1k_con4096_ctx1_dep2_gen1_dep8_eplb256_mtp1_ccb-NIXL.yaml index f2333ddcb500..b584ba200982 100644 --- a/tests/scripts/perf-sanity/disaggregated/gb300_glm-5-fp4_1k1k_con4096_ctx1_dep2_gen1_dep8_eplb256_mtp1_ccb-NIXL.yaml +++ b/tests/scripts/perf-sanity/disaggregated/gb300_glm-5-fp4_1k1k_con4096_ctx1_dep2_gen1_dep8_eplb256_mtp1_ccb-NIXL.yaml @@ -36,6 +36,7 @@ environment: build_wheel: false work_dir: worker_env_var: TLLM_LOG_LEVEL=INFO TRTLLM_SERVER_DISABLE_GC=1 TRTLLM_WORKER_DISABLE_GC=1 TRTLLM_ENABLE_PDL=1 ENROOT_ALLOW_DEV=yes TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS=1 + ctx_worker_env_var: TLLM_NUMA_AWARE_WORKER_AFFINITY=1 server_env_var: TRTLLM_SERVER_DISABLE_GC=1 profiling: nsys_on: false diff --git a/tests/scripts/perf-sanity/disaggregated/gb300_glm-5-fp4_1k1k_con512_ctx1_dep2_gen1_dep32_eplb0_mtp3_ccb-NIXL.yaml b/tests/scripts/perf-sanity/disaggregated/gb300_glm-5-fp4_1k1k_con512_ctx1_dep2_gen1_dep32_eplb0_mtp3_ccb-NIXL.yaml index eb90301b362c..9503cfce509e 100644 --- a/tests/scripts/perf-sanity/disaggregated/gb300_glm-5-fp4_1k1k_con512_ctx1_dep2_gen1_dep32_eplb0_mtp3_ccb-NIXL.yaml +++ b/tests/scripts/perf-sanity/disaggregated/gb300_glm-5-fp4_1k1k_con512_ctx1_dep2_gen1_dep32_eplb0_mtp3_ccb-NIXL.yaml @@ -36,6 +36,7 @@ environment: build_wheel: false work_dir: worker_env_var: TLLM_LOG_LEVEL=INFO TRTLLM_SERVER_DISABLE_GC=1 TRTLLM_WORKER_DISABLE_GC=1 TRTLLM_ENABLE_PDL=1 ENROOT_ALLOW_DEV=yes TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS=3 + ctx_worker_env_var: TLLM_NUMA_AWARE_WORKER_AFFINITY=1 server_env_var: TRTLLM_SERVER_DISABLE_GC=1 profiling: nsys_on: false diff --git a/tests/scripts/perf-sanity/disaggregated/gb300_glm-5-fp4_8k1k_con1024_ctx1_dep2_gen1_dep8_eplb256_mtp1_ccb-NIXL.yaml b/tests/scripts/perf-sanity/disaggregated/gb300_glm-5-fp4_8k1k_con1024_ctx1_dep2_gen1_dep8_eplb256_mtp1_ccb-NIXL.yaml index 27f35f458bcf..b8876eb5ad8e 100644 --- a/tests/scripts/perf-sanity/disaggregated/gb300_glm-5-fp4_8k1k_con1024_ctx1_dep2_gen1_dep8_eplb256_mtp1_ccb-NIXL.yaml +++ b/tests/scripts/perf-sanity/disaggregated/gb300_glm-5-fp4_8k1k_con1024_ctx1_dep2_gen1_dep8_eplb256_mtp1_ccb-NIXL.yaml @@ -36,6 +36,7 @@ environment: build_wheel: false work_dir: worker_env_var: TLLM_LOG_LEVEL=INFO TRTLLM_SERVER_DISABLE_GC=1 TRTLLM_WORKER_DISABLE_GC=1 TRTLLM_ENABLE_PDL=1 ENROOT_ALLOW_DEV=yes TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS=1 + ctx_worker_env_var: TLLM_NUMA_AWARE_WORKER_AFFINITY=1 server_env_var: TRTLLM_SERVER_DISABLE_GC=1 profiling: nsys_on: false diff --git a/tests/scripts/perf-sanity/disaggregated/gb300_glm-5-fp4_8k1k_con1_ctx1_dep2_gen1_tep8_eplb0_mtp3_ccb-NIXL.yaml b/tests/scripts/perf-sanity/disaggregated/gb300_glm-5-fp4_8k1k_con1_ctx1_dep2_gen1_tep8_eplb0_mtp3_ccb-NIXL.yaml index 54e6601eb2fe..4a36bb288c10 100644 --- a/tests/scripts/perf-sanity/disaggregated/gb300_glm-5-fp4_8k1k_con1_ctx1_dep2_gen1_tep8_eplb0_mtp3_ccb-NIXL.yaml +++ b/tests/scripts/perf-sanity/disaggregated/gb300_glm-5-fp4_8k1k_con1_ctx1_dep2_gen1_tep8_eplb0_mtp3_ccb-NIXL.yaml @@ -36,6 +36,7 @@ environment: build_wheel: false work_dir: worker_env_var: TLLM_LOG_LEVEL=INFO TRTLLM_SERVER_DISABLE_GC=1 TRTLLM_WORKER_DISABLE_GC=1 TRTLLM_ENABLE_PDL=1 ENROOT_ALLOW_DEV=yes TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS=3 + ctx_worker_env_var: TLLM_NUMA_AWARE_WORKER_AFFINITY=1 server_env_var: TRTLLM_SERVER_DISABLE_GC=1 profiling: nsys_on: false diff --git a/tests/scripts/perf-sanity/disaggregated/gb300_glm-5-fp4_8k1k_con512_ctx1_dep2_gen1_dep32_eplb0_mtp3_ccb-NIXL.yaml b/tests/scripts/perf-sanity/disaggregated/gb300_glm-5-fp4_8k1k_con512_ctx1_dep2_gen1_dep32_eplb0_mtp3_ccb-NIXL.yaml index 6dd306a299ad..75af7e2ab134 100644 --- a/tests/scripts/perf-sanity/disaggregated/gb300_glm-5-fp4_8k1k_con512_ctx1_dep2_gen1_dep32_eplb0_mtp3_ccb-NIXL.yaml +++ b/tests/scripts/perf-sanity/disaggregated/gb300_glm-5-fp4_8k1k_con512_ctx1_dep2_gen1_dep32_eplb0_mtp3_ccb-NIXL.yaml @@ -36,6 +36,7 @@ environment: build_wheel: false work_dir: worker_env_var: TLLM_LOG_LEVEL=INFO TRTLLM_SERVER_DISABLE_GC=1 TRTLLM_WORKER_DISABLE_GC=1 TRTLLM_ENABLE_PDL=1 ENROOT_ALLOW_DEV=yes TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS=3 + ctx_worker_env_var: TLLM_NUMA_AWARE_WORKER_AFFINITY=1 server_env_var: TRTLLM_SERVER_DISABLE_GC=1 profiling: nsys_on: false From ca66314afaab62c1a780151fa0b5800377b8a06c Mon Sep 17 00:00:00 2001 From: TensorRT LLM AI Agent <296075020+trtllm-agent@users.noreply.github.com> Date: Wed, 14 Oct 2026 08:01:16 +0800 Subject: [PATCH 06/23] =?UTF-8?q?[https://nvbugs/6859934][fix]=20Cherry-pi?= =?UTF-8?q?ck=20#19847=20from=20release/1.3:=20Process-wide=20workspace=20?= =?UTF-8?q?caches=20keyed=20by=20`Mapping`=20with=20a=20per-mapping=20lock?= =?UTF-8?q?=E2=80=A6?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Forward-port of release/1.3 PR https://github.com/NVIDIA/TensorRT-LLM/pull/19847 (release/1.3 commit 84461c627ccf9cf6dc2be38b1b5988e9e5bfd440) to main. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> Signed-off-by: Mike Iovine <6158008+mikeiovine@users.noreply.github.com> Co-authored-by: Claude Co-authored-by: Mike Iovine Signed-off-by: Simeng Liu --- tensorrt_llm/_torch/distributed/ops.py | 63 +++++----- .../test_allreduce_workspace_cache.py | 109 ++++++++++++++++++ 2 files changed, 145 insertions(+), 27 deletions(-) create mode 100644 tests/unittest/_torch/distributed/test_allreduce_workspace_cache.py diff --git a/tensorrt_llm/_torch/distributed/ops.py b/tensorrt_llm/_torch/distributed/ops.py index 0f63e2bb5e7e..aee22893da35 100644 --- a/tensorrt_llm/_torch/distributed/ops.py +++ b/tensorrt_llm/_torch/distributed/ops.py @@ -57,7 +57,22 @@ "TLLM_MNNVL_ONESHOT_THRESHOLD_BYTES must be a non-negative integer, got " f"{_MNNVL_ONE_SHOT_THRESHOLD_BYTES}") -_thread_local = threading.local() +# Allocating these workspaces is a TP-wide collective, so whether a rank +# allocates must not depend on which thread builds its model: RpcWorker builds +# rank 0 on an RPC executor thread and the other ranks on their main thread. +# The caches are therefore shared by all threads of the process, and a +# per-mapping lock serializes each allocation. +_allreduce_workspaces: Dict[Mapping, Tuple[list, torch.Tensor]] = {} +_lowprecision_allreduce_workspaces: Dict[Mapping, Tuple[list, + torch.Tensor]] = {} +_allreduce_workspace_locks: Dict[Mapping, threading.Lock] = {} +_allreduce_workspace_locks_guard = threading.Lock() + + +def _get_allreduce_workspace_lock(mapping: Mapping) -> threading.Lock: + with _allreduce_workspace_locks_guard: + return _allreduce_workspace_locks.setdefault(mapping, threading.Lock()) + # Mirrored by the autotune context so Dynamo can guard this branch without # tracing the stateful AutoTuner singleton. @@ -156,35 +171,29 @@ def warmup_named_allreduce_topologies(model: nn.Module) -> None: def get_allreduce_workspace(mapping: Mapping) -> torch.LongTensor: - if not hasattr(_thread_local, f'allreduce_workspaces_{mapping.pp_rank}'): - setattr(_thread_local, f'allreduce_workspaces_{mapping.pp_rank}', {}) - - allreduce_workspaces = getattr(_thread_local, - f'allreduce_workspaces_{mapping.pp_rank}') - if mapping not in allreduce_workspaces: - ipc_buffers, workspace = CustomAllReduceHelper.allocate_allreduce_fusion_workspace( - mapping, - CustomAllReduceHelper.max_workspace_size_auto( - mapping.tp_size, support_deterministic=False), - ) - allreduce_workspaces[mapping] = (ipc_buffers, workspace) - return allreduce_workspaces[mapping][1] + with _get_allreduce_workspace_lock(mapping): + if mapping not in _allreduce_workspaces: + ipc_buffers, workspace = CustomAllReduceHelper.allocate_allreduce_fusion_workspace( + mapping, + CustomAllReduceHelper.max_workspace_size_auto( + mapping.tp_size, support_deterministic=False), + ) + _allreduce_workspaces[mapping] = (ipc_buffers, workspace) + return _allreduce_workspaces[mapping][1] def allocate_low_presicion_allreduce_workspace(mapping: Mapping) -> None: - if not hasattr(_thread_local, 'lowprecision_allreduce_workspaces'): - _thread_local.lowprecision_allreduce_workspaces = {} - lowprecision_allreduce_workspaces = _thread_local.lowprecision_allreduce_workspaces - if mapping not in lowprecision_allreduce_workspaces: - ipc_buffers, workspace = CustomAllReduceHelper.allocate_lowprecision_workspace( - mapping, - CustomAllReduceHelper.max_workspace_size_lowprecision( - mapping.tp_size), - ) - lowprecision_allreduce_workspaces[mapping] = (ipc_buffers, workspace) - CustomAllReduceHelper.initialize_lowprecision_buffers( - workspace, mapping.tp_size) - return + with _get_allreduce_workspace_lock(mapping): + if mapping not in _lowprecision_allreduce_workspaces: + ipc_buffers, workspace = CustomAllReduceHelper.allocate_lowprecision_workspace( + mapping, + CustomAllReduceHelper.max_workspace_size_lowprecision( + mapping.tp_size), + ) + _lowprecision_allreduce_workspaces[mapping] = (ipc_buffers, + workspace) + CustomAllReduceHelper.initialize_lowprecision_buffers( + workspace, mapping.tp_size) def _initialize_allreduce_mnnvl_protocol(workspace: _MnnvlWorkspace, diff --git a/tests/unittest/_torch/distributed/test_allreduce_workspace_cache.py b/tests/unittest/_torch/distributed/test_allreduce_workspace_cache.py new file mode 100644 index 000000000000..53f2bbcd1ce8 --- /dev/null +++ b/tests/unittest/_torch/distributed/test_allreduce_workspace_cache.py @@ -0,0 +1,109 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Process-wide caching of the custom AllReduce IPC workspaces. + +Allocating a workspace runs a TP-wide host collective, so every rank must take +the same allocate-or-reuse decision. A rank cannot know which thread its peers +build their model on (RpcWorker builds rank 0 on an RPC executor thread and the +other ranks on their main thread), so a cache keyed per thread lets a reused +worker skip the collective while a peer enters it, and both ranks deadlock. +""" + +import threading +from concurrent.futures import ThreadPoolExecutor + +import pytest + +from tensorrt_llm._torch.distributed import ops +from tensorrt_llm.mapping import Mapping + +pytestmark = pytest.mark.cpu_only + + +@pytest.fixture +def allocations(monkeypatch: pytest.MonkeyPatch) -> list: + """Replace the IPC allocators with recorders and start from empty caches.""" + calls = [] + + def allocate_fusion(mapping, size): + calls.append(("fusion", mapping)) + return [], object() + + def allocate_lowprecision(mapping, size): + calls.append(("lowprecision", mapping)) + return [], object() + + helper = ops.CustomAllReduceHelper + monkeypatch.setattr(helper, "allocate_allreduce_fusion_workspace", allocate_fusion) + monkeypatch.setattr(helper, "allocate_lowprecision_workspace", allocate_lowprecision) + monkeypatch.setattr(helper, "max_workspace_size_auto", lambda *args, **kwargs: 1) + monkeypatch.setattr(helper, "max_workspace_size_lowprecision", lambda *args: 1) + monkeypatch.setattr(helper, "initialize_lowprecision_buffers", lambda *args: None) + monkeypatch.setattr(ops, "_allreduce_workspaces", {}) + monkeypatch.setattr(ops, "_lowprecision_allreduce_workspaces", {}) + monkeypatch.setattr(ops, "_allreduce_workspace_locks", {}) + return calls + + +def _run_on_new_thread(fn, *args): + with ThreadPoolExecutor(max_workers=1) as executor: + return executor.submit(fn, *args).result() + + +def _tp2_mapping() -> Mapping: + return Mapping(world_size=2, rank=0, gpus_per_node=2, tp_size=2) + + +def test_workspace_is_reused_from_another_thread(allocations: list) -> None: + mapping = _tp2_mapping() + + first = ops.get_allreduce_workspace(mapping) + second = _run_on_new_thread(ops.get_allreduce_workspace, mapping) + + assert second is first + assert allocations == [("fusion", mapping)] + + +def test_lowprecision_workspace_is_reused_from_another_thread(allocations: list) -> None: + mapping = _tp2_mapping() + + ops.allocate_low_presicion_allreduce_workspace(mapping) + _run_on_new_thread(ops.allocate_low_presicion_allreduce_workspace, mapping) + + assert allocations == [("lowprecision", mapping)] + + +def test_concurrent_first_use_allocates_once(allocations: list) -> None: + mapping = _tp2_mapping() + num_threads = 8 + barrier = threading.Barrier(num_threads) + + def get_after_barrier(): + barrier.wait() + return ops.get_allreduce_workspace(mapping) + + with ThreadPoolExecutor(max_workers=num_threads) as executor: + workspaces = list(executor.map(lambda _: get_after_barrier(), range(num_threads))) + + assert all(workspace is workspaces[0] for workspace in workspaces) + assert allocations == [("fusion", mapping)] + + +def test_distinct_mappings_get_distinct_workspaces(allocations: list) -> None: + tp2 = _tp2_mapping() + tp2_pp2 = Mapping(world_size=4, rank=0, gpus_per_node=4, tp_size=2, pp_size=2) + + assert ops.get_allreduce_workspace(tp2) is not ops.get_allreduce_workspace(tp2_pp2) + assert allocations == [("fusion", tp2), ("fusion", tp2_pp2)] From e6ec4c51dbc1437d883e8350a4b5c2cb2867a35d Mon Sep 17 00:00:00 2001 From: sunnyqgg <159101675+sunnyqgg@users.noreply.github.com> Date: Wed, 14 Oct 2026 09:43:56 +0800 Subject: [PATCH 07/23] [https://nvbugs/6792615][fix] Cherry-pick #20152 from release/1.3: Keep session-prefetch output off sys.stdout Forward-port of release/1.3 PR https://github.com/NVIDIA/TensorRT-LLM/pull/20152 (release/1.3 commit ca9c5b62c54cb650e877c3ba3341704fac9a7d56) to main. Signed-off-by: qgai Signed-off-by: Simeng Liu --- tests/test_common/session_prefetcher.py | 61 ++++++++++++++----------- 1 file changed, 35 insertions(+), 26 deletions(-) diff --git a/tests/test_common/session_prefetcher.py b/tests/test_common/session_prefetcher.py index 3442a030b8ac..33a3097445e5 100644 --- a/tests/test_common/session_prefetcher.py +++ b/tests/test_common/session_prefetcher.py @@ -78,6 +78,27 @@ _FALLBACK_IDENTITY_TIMEOUT = 300.0 +def _log(message: str) -> None: + """Write one diagnostic line straight to file descriptor 1. + + Most messages come from the build, warm, and discard threads. ``print()`` + would resolve ``sys.stdout``, which pytest swaps on every test phase and + frees at teardown of each test's ``capfd``/``capsys`` capture. CPython + 3.12.3, the CI interpreter, keeps only a borrowed reference to it inside + ``print()``, so a print racing that teardown can write to a freed stream + and crash the pytest process. Writing the descriptor reaches the same + destination, since pytest's fd capture redirects fd 1, without touching a + Python stream object. + + Best-effort: a failed write is dropped so it can never skip the cleanup + that follows a log call. + """ + try: + os.write(1, f"{message}\n".encode("utf-8", "replace")) + except OSError: + pass + + def _fallback_identity_timeout() -> float: """Mirror the lower-level env contract while its module is still loading.""" raw = os.environ.get("TRTLLM_MPI_IDENTITY_TIMEOUT") @@ -223,10 +244,9 @@ def warm_page_cache(model_dir: str) -> float: if available is not None and total_bytes > available: # Larger than RAM: pages would be evicted before the test loads them — # pure filer traffic with zero benefit (e.g. multi-hundred-GB models). - print( + _log( f"[session-prefetch] skipping warm of {model_dir}: {total_bytes >> 30} GiB " - f"exceeds available host memory ({available >> 30} GiB)", - flush=True, + f"exceeds available host memory ({available >> 30} GiB)" ) return 0.0 t0 = time.monotonic() @@ -246,10 +266,9 @@ def _read(path): with ThreadPoolExecutor(max_workers=4, thread_name_prefix="session-prefetch-io") as ex: total = sum(ex.map(_read, files)) gib = total / (1 << 30) - print( + _log( f"[session-prefetch] warmed page cache: {gib:.1f} GiB from " - f"{model_dir} in {time.monotonic() - t0:.1f}s", - flush=True, + f"{model_dir} in {time.monotonic() - t0:.1f}s" ) return gib @@ -316,10 +335,7 @@ def _reap_dead_pool(session: object) -> None: # not enumerate. Swallowing is required -- the SIGKILL sweep below is # what actually reclaims the survivors' GPU memory, and it must run # whatever abandon() did. Surface it instead of passing silently. - print( - f"[session-prefetch] WARNING: abandon() failed on a part-dead pool: {exc!r}", - flush=True, - ) + _log(f"[session-prefetch] WARNING: abandon() failed on a part-dead pool: {exc!r}") mpi_session = sys.modules.get("tensorrt_llm.llmapi.mpi_session") process_start_time = getattr(mpi_session, "_process_start_time", None) if process_start_time is None: @@ -419,7 +435,7 @@ def _warm(self, model_dir: str) -> None: else: self.stats["warm_noops"] += 1 # no local weights / RAM guard except Exception as e: # warming must never break the tests - print(f"[session-prefetch] page-cache warm failed (harmless): {e}", flush=True) + _log(f"[session-prefetch] page-cache warm failed (harmless): {e}") def schedule_shadow(self, spec: int, env_overlay=None) -> None: """Start building a spare ``spec``-worker pool in the background. @@ -464,17 +480,13 @@ def _build(self, spec: int, gen: int, env_overlay=None) -> None: # out blocks its shutdown on actual worker exit. session = MpiPoolSession(n_workers=spec, wait_shutdown=True, env_overrides=env_overlay) if any(session.submit_sync(_worker_import_report_cuda)): - print( + _log( "[session-prefetch] note: tensorrt_llm import initialized an idle " - "CUDA context in the prefetched workers (library version behavior)", - flush=True, + "CUDA context in the prefetched workers (library version behavior)" ) self._publish(spec, session, snapshot, gen) except Exception as e: # prefetch must never break the tests - print( - f"[session-prefetch] background build failed (falling back to synchronous): {e}", - flush=True, - ) + _log(f"[session-prefetch] background build failed (falling back to synchronous): {e}") def _publish(self, spec, session, snapshot, gen: int) -> None: """Publish a finished background build, unless it was abandoned. @@ -491,7 +503,7 @@ def _publish(self, spec, session, snapshot, gen: int) -> None: self.stats["pools_built"] += 1 return self.stats["pools_discarded_superseded"] += 1 - print("[session-prefetch] discarding superseded background build", flush=True) + _log("[session-prefetch] discarding superseded background build") session.shutdown() def _drain(self, timeout: float | None = None) -> _Built | None: @@ -561,10 +573,7 @@ def take(self, spec: int) -> object | None: # outer pytest timeout. The dead world cannot be joined, so # abandon it before the synchronous fallback. self.stats["pools_discarded_dead"] += 1 - print( - "[session-prefetch] discarding prefetched pool with dead worker", - flush=True, - ) + _log("[session-prefetch] discarding prefetched pool with dead worker") threading.Thread( target=_reap_dead_pool, args=(built.session,), @@ -578,7 +587,7 @@ def take(self, spec: int) -> object | None: # wait_shutdown=True, so its shutdown() blocked until the workers # actually exited (and released their memory). self.stats["pools_handed_over"] += 1 - print(f"[session-prefetch] handing over prefetched {spec}-worker pool", flush=True) + _log(f"[session-prefetch] handing over prefetched {spec}-worker pool") return built.session # Spec/env/sys.path mismatch (test skipped, reordered, or changed # state the frozen workers would not see): discard. @@ -666,7 +675,7 @@ def dispose(self) -> None: # build remains registered, preventing another shadow from being # launched; the daemon thread is bounded by the lower-level # identity timeout unless the MPI runtime itself is wedged. - print(f"[session-prefetch] cleanup timed out: {e}", flush=True) + _log(f"[session-prefetch] cleanup timed out: {e}") built = None if built is not None: built.session.shutdown() @@ -678,7 +687,7 @@ def dispose(self) -> None: parts = ", ".join(f"{k}={v}" for k, v in sorted(self.stats.items())) if self._warmed_gib: parts += f", warmed_gib={self._warmed_gib:.1f}" - print(f"[session-prefetch] session summary: {parts}", flush=True) + _log(f"[session-prefetch] session summary: {parts}") PREFETCHER = SessionPrefetcher() From 455c15f3d4774bae3156c9afe153f82538091abe Mon Sep 17 00:00:00 2001 From: Stanley Sun <190317771+StanleySun639@users.noreply.github.com> Date: Wed, 14 Oct 2026 14:04:44 +0800 Subject: [PATCH 08/23] [https://nvbugs/6918015][doc] Cherry-pick #20193 from release/1.3: Fix spelling and capitalization in release/1.3 docs Forward-port of release/1.3 PR https://github.com/NVIDIA/TensorRT-LLM/pull/20193 (release/1.3 commit f92ba89f71c6260e95192213db9bbc05250b83cc) to main. Signed-off-by: Stanley Sun Signed-off-by: Simeng Liu --- ...3_Inference_Time_Compute_Implementation_in_TensorRT-LLM.md | 4 ++-- docs/source/overview.md | 2 +- docs/source/release-notes.md | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/docs/source/blogs/tech_blog/blog13_Inference_Time_Compute_Implementation_in_TensorRT-LLM.md b/docs/source/blogs/tech_blog/blog13_Inference_Time_Compute_Implementation_in_TensorRT-LLM.md index e5e2750468fa..36e249f6ba70 100644 --- a/docs/source/blogs/tech_blog/blog13_Inference_Time_Compute_Implementation_in_TensorRT-LLM.md +++ b/docs/source/blogs/tech_blog/blog13_Inference_Time_Compute_Implementation_in_TensorRT-LLM.md @@ -200,12 +200,12 @@ is a certainty-based, training-free approach to accelerate Chain-of-Thought (CoT #### Motivation of Dynasor-CoT LLM reasoning is highly token-inefficient, often requiring far more tokens to achieve the same accuracy as non-reasoning models. A major source of this inefficiency is that reasoning models tend to **self-doubt**; they often reach the correct answer early but then engage in extended verification behaviors like double-checking and reassessment. -For instance, Figure 2 compares a traditional Qwen-7B model with a reasoning-focused, Deepseek-distilled Qwen-7B model on a simple question. While the traditional model reaches its answer in 180 tokens, the reasoning model expends 1,000 tokens on iterative verification, despite having already found the correct answer at token 340. This represents a significant waste of tokens for diminishing returns on accuracy. +For instance, Figure 2 compares a traditional Qwen-7B model with a reasoning-focused, DeepSeek-distilled Qwen-7B model on a simple question. While the traditional model reaches its answer in 180 tokens, the reasoning model expends 1,000 tokens on iterative verification, despite having already found the correct answer at token 340. This represents a significant waste of tokens for diminishing returns on accuracy.
Motivation
-

Figure 2. An example answer from reasoning model (Deepseek-distilled Qwen-2.5 7B) vs traditional model (Qwen-2.5 7B) on one of the problem in MATH500 dataset.

+

Figure 2. An example answer from reasoning model (DeepSeek-distilled Qwen-2.5 7B) vs traditional model (Qwen-2.5 7B) on one of the problem in MATH500 dataset.

#### The "Probe" technique Dynasor-CoT uses a **"Probe-In-The-Middle"** (or "probe" for short) technique, which prompts reasoning models to output early-stage results during intermediate steps of reasoning. Imagine you're in a math exam working on a hard problem. When time is up, you're forced to write down your final answer, regardless of how confident you are. diff --git a/docs/source/overview.md b/docs/source/overview.md index c993f2fcb6a6..0b116cbc1efe 100644 --- a/docs/source/overview.md +++ b/docs/source/overview.md @@ -8,7 +8,7 @@ ## Key Capabilities -### 🔥 **Architected on Pytorch** +### 🔥 **Architected on PyTorch** TensorRT LLM provides a high-level Python [LLM API](./quick-start-guide.md#run-offline-inference-with-llm-api) that supports a wide range of inference setups - from single-GPU to multi-GPU or multi-node deployments. It includes built-in support for various parallelism strategies and advanced features. The LLM API integrates seamlessly with the broader inference ecosystem, including NVIDIA [Dynamo](https://github.com/ai-dynamo/dynamo) and the [Triton Inference Server](https://github.com/triton-inference-server/server). diff --git a/docs/source/release-notes.md b/docs/source/release-notes.md index d83bc8d1fac8..c533f4613c92 100644 --- a/docs/source/release-notes.md +++ b/docs/source/release-notes.md @@ -1084,7 +1084,7 @@ TensorRT LLM 1.0 brings 2 major changes: the PyTorch-based architecture is now s - Supported phi 3 vision multimodal. ### Fixed Issues -- Fixed brokens outputs for the cases when batch size is larger than 1. (#1539) +- Fixed broken outputs for the cases when batch size is larger than 1. (#1539) - Fixed `top_k` type in `executor.py`, thanks to the contribution from @vonjackustc in #1329. - Fixed stop and bad word list pointer offset in Python runtime, thanks to the contribution from @fjosw in #1486. - Fixed some typos for Whisper model, thanks to the contribution from @Pzzzzz5142 in #1328. From 50dbedc1670f48fe7a317e05a1ea7ae0d2e61cd1 Mon Sep 17 00:00:00 2001 From: sunnyqgg <159101675+sunnyqgg@users.noreply.github.com> Date: Wed, 14 Oct 2026 21:13:03 +0800 Subject: [PATCH 09/23] [https://nvbugs/6379406][fix] Cherry-pick #20156 from release/1.3: Use coherent managed memory for EPLB when GDRCopy is unavailable Forward-port of release/1.3 PR https://github.com/NVIDIA/TensorRT-LLM/pull/20156 (release/1.3 commit 248825ed123474bd95302f716b95a2ea1db1cca1) to main. Signed-off-by: qgai Signed-off-by: Simeng Liu --- .../hostAccessibleDeviceAllocator.cpp | 49 +++++++++++++--- .../hostAccessibleDeviceAllocator.h | 9 ++- .../hostAccessibleDeviceAllocatorTest.cu | 58 ++++++++++++++++++- examples/wide_ep/README.md | 2 +- 4 files changed, 104 insertions(+), 14 deletions(-) diff --git a/cpp/tensorrt_llm/runtime/moeLoadBalancer/hostAccessibleDeviceAllocator.cpp b/cpp/tensorrt_llm/runtime/moeLoadBalancer/hostAccessibleDeviceAllocator.cpp index 446d20191aa5..804d2916170c 100644 --- a/cpp/tensorrt_llm/runtime/moeLoadBalancer/hostAccessibleDeviceAllocator.cpp +++ b/cpp/tensorrt_llm/runtime/moeLoadBalancer/hostAccessibleDeviceAllocator.cpp @@ -1,5 +1,5 @@ /* - * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION. All rights reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -166,20 +166,37 @@ void NumaHugePagePoolAllocator::shutdown() munmap(mMmapBasePtr, kReservedVirtualMemorySize + kHugePageSize); } +namespace +{ + +//! Whether the host can access device-resident managed memory directly, without page faults or migration. +//! This holds on hardware-coherent platforms such as Grace-based GB200/GB300, including when GPU memory is not +//! exposed as a NUMA node (CDMM mode). +bool isHostDirectManagedAccessSupported() +{ + int deviceId = -1; + TLLM_CUDA_CHECK(cudaGetDevice(&deviceId)); + int directAccess = 0; + TLLM_CUDA_CHECK(cudaDeviceGetAttribute(&directAccess, cudaDevAttrDirectManagedMemAccessFromHost, deviceId)); + return directAccess != 0; +} + +} // namespace + bool HostAccessibleDeviceAllocator::mAllowManagedFallback = false; bool HostAccessibleDeviceAllocator::isSupported() { - if (!tensorrt_llm::common::getEnvEplbForceGdrcopy() - && TopologyDetector::getInstance().getCurrentGpuMemoryNumaId() >= 0) + bool const forceGdrcopy = tensorrt_llm::common::getEnvEplbForceGdrcopy(); + if (!forceGdrcopy && TopologyDetector::getInstance().getCurrentGpuMemoryNumaId() >= 0) { // we are on systems that GPU memory is also a NUMA node. return true; } if (!tensorrt_llm::runtime::gdrcopy::isInitialized() && !tensorrt_llm::runtime::gdrcopy::initialize()) { - // system don't support GDRCopy. - return mAllowManagedFallback; + // System doesn't support GDRCopy, managed memory still works if the host can access it directly. + return mAllowManagedFallback || (!forceGdrcopy && isHostDirectManagedAccessSupported()); } return true; } @@ -207,6 +224,7 @@ void HostAccessibleDeviceAllocator::init() mGpuMemNumaId = TopologyDetector::getInstance().getCurrentGpuMemoryNumaId(); } + mHostDirectManagedAccess = false; if (mGpuMemNumaId < 0) { // We only use GDRCopy when there is no NUMA node for GPU memory. @@ -219,6 +237,13 @@ void HostAccessibleDeviceAllocator::init() { mGdrHandle = tensorrt_llm::runtime::gdrcopy::open(); } + else if (!tensorrt_llm::common::getEnvEplbForceGdrcopy() && isHostDirectManagedAccessSupported()) + { + mHostDirectManagedAccess = true; + TLLM_LOG_WARNING( + "GPU memory is not a NUMA node and GDRCopy is not available, using managed memory with direct host " + "access as host accessible device memory."); + } } mIsInited = true; } @@ -372,11 +397,17 @@ void* HostAccessibleDeviceAllocator::allocate(size_t memorySize) } else { - TLLM_CHECK_WITH_INFO( - mAllowManagedFallback, "HostAccessibleDeviceAllocator is not supported on the current system."); + TLLM_CHECK_WITH_INFO(mAllowManagedFallback || mHostDirectManagedAccess, + "HostAccessibleDeviceAllocator is not supported on the current system."); TLLM_CUDA_CHECK(cudaMallocManaged(&devPtr, memorySize)); TLLM_CUDA_CHECK(cudaMemAdvise(devPtr, memorySize, cudaMemAdviseSetPreferredLocation, cudaMemLocation{cudaMemLocationTypeDevice, {currentDevId}})); + if (mHostDirectManagedAccess) + { + // Keep the device-resident pages mapped for the host, so host accesses neither fault nor migrate them. + TLLM_CUDA_CHECK(cudaMemAdvise( + devPtr, memorySize, cudaMemAdviseSetAccessedBy, cudaMemLocation{cudaMemLocationTypeHost, {0}})); + } hostPtr = devPtr; } recordAllocation(devPtr, memorySize, hostPtr, memDesc); @@ -401,8 +432,8 @@ void HostAccessibleDeviceAllocator::free(void* ptr) } else { - TLLM_CHECK_WITH_INFO( - mAllowManagedFallback, "HostAccessibleDeviceAllocator is not supported on the current system."); + TLLM_CHECK_WITH_INFO(mAllowManagedFallback || mHostDirectManagedAccess, + "HostAccessibleDeviceAllocator is not supported on the current system."); TLLM_CUDA_CHECK(cudaFree(ptr)); } void* hostPtr = it->second.hostPtr; diff --git a/cpp/tensorrt_llm/runtime/moeLoadBalancer/hostAccessibleDeviceAllocator.h b/cpp/tensorrt_llm/runtime/moeLoadBalancer/hostAccessibleDeviceAllocator.h index 0cff7de95681..be921d4fb82d 100644 --- a/cpp/tensorrt_llm/runtime/moeLoadBalancer/hostAccessibleDeviceAllocator.h +++ b/cpp/tensorrt_llm/runtime/moeLoadBalancer/hostAccessibleDeviceAllocator.h @@ -1,5 +1,5 @@ /* - * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION. All rights reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -157,10 +157,13 @@ class HostAccessibleDeviceAllocator */ AllocationInfo getAllocationInfoFromDevPtr(void const* devPtr); - // if GPU memory has NUMA id, then CPU can direct access that. We should use this. + // Host access to device memory, in order of preference: + // 1. GPU memory is a NUMA node (mGpuMemNumaId >= 0), so the CPU can access it directly. int mGpuMemNumaId = -1; - // if Not, we should use GDRCopy + // 2. Otherwise, GDRCopy maps device memory for the host. gdr_t mGdrHandle = nullptr; + // 3. Without GDRCopy, managed memory that the host can access directly (hardware-coherent platforms). + bool mHostDirectManagedAccess = false; int mDevId = -1; diff --git a/cpp/tests/unit_tests/runtime/hostAccessibleDeviceAllocatorTest.cu b/cpp/tests/unit_tests/runtime/hostAccessibleDeviceAllocatorTest.cu index d35f8a21c348..60c1c8757b8d 100644 --- a/cpp/tests/unit_tests/runtime/hostAccessibleDeviceAllocatorTest.cu +++ b/cpp/tests/unit_tests/runtime/hostAccessibleDeviceAllocatorTest.cu @@ -1,5 +1,5 @@ /* - * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION. All rights reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -15,6 +15,7 @@ */ #include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/common/envUtils.h" #include "tensorrt_llm/runtime/moeLoadBalancer/hostAccessibleDeviceAllocator.h" #include #include @@ -233,4 +234,59 @@ TEST_F(HostAccessibleDeviceAllocatorTest, GetHostPtrForUnallocated) EXPECT_EQ(allocator->getHostPtr(devPtr), nullptr); } +// Without a NUMA node for GPU memory (e.g. CDMM mode) and without GDRCopy, host accessible memory falls back to +// managed memory on hardware-coherent platforms (e.g. GB200/GB300), where the host can access device-resident managed +// memory directly. The test only runs in that configuration, e.g. on a CDMM node where /dev/gdrdrv is absent. +TEST_F(HostAccessibleDeviceAllocatorTest, ManagedMemoryFallbackWithoutGdrCopy) +{ + int deviceId = -1; + TLLM_CUDA_CHECK(cudaGetDevice(&deviceId)); + int directAccess = 0; + TLLM_CUDA_CHECK(cudaDeviceGetAttribute(&directAccess, cudaDevAttrDirectManagedMemAccessFromHost, deviceId)); + if (directAccess == 0) + { + GTEST_SKIP() << "Host cannot directly access device-resident managed memory on this system."; + } + if (tensorrt_llm::common::getEnvEplbForceGdrcopy()) + { + GTEST_SKIP() << "GDRCopy is forced by TRTLLM_EPLB_FORCE_GDRCOPY."; + } + if (TopologyDetector::getInstance().getCurrentGpuMemoryNumaId() >= 0 || gdrcopy::initialize()) + { + GTEST_SKIP() << "GPU memory is a NUMA node or GDRCopy is available, so the managed memory fallback is unused."; + } + ASSERT_TRUE(allocator->isSupported()); + + constexpr size_t numInts = 256; + constexpr size_t allocSize = numInts * sizeof(int); + void* devPtr = allocator->allocate(allocSize); + ASSERT_NE(devPtr, nullptr); + void* hostPtr = allocator->getHostPtr(devPtr); + ASSERT_NE(hostPtr, nullptr); + + // 1. Copy from the host with memcpyToDevice, as the online EPLB weight updates do, then verify on the device. + std::vector pattern(numInts); + std::iota(pattern.begin(), pattern.end(), 0); + allocator->memcpyToDevice(hostPtr, pattern.data(), allocSize); + bool deviceVerificationResult = false; + bool* dResult = nullptr; + TLLM_CUDA_CHECK(cudaMalloc(&dResult, sizeof(bool))); + verifyDataOnDevice<<<1, 1>>>(static_cast(devPtr), numInts, dResult); + TLLM_CUDA_CHECK(cudaDeviceSynchronize()); + TLLM_CUDA_CHECK(cudaMemcpy(&deviceVerificationResult, dResult, sizeof(bool), cudaMemcpyDeviceToHost)); + TLLM_CUDA_CHECK(cudaFree(dResult)); + EXPECT_TRUE(deviceVerificationResult); + + // 2. Write from the device, then read and verify from the host. + writeDataOnDevice<<<1, 1>>>(static_cast(devPtr), numInts); + TLLM_CUDA_CHECK(cudaDeviceSynchronize()); + int const* hostIntPtr = static_cast(hostPtr); + for (size_t i = 0; i < numInts; ++i) + { + EXPECT_EQ(hostIntPtr[i], static_cast(numInts - i)); + } + + allocator->free(devPtr); +} + } // namespace tensorrt_llm::runtime::unit_tests diff --git a/examples/wide_ep/README.md b/examples/wide_ep/README.md index 4cac840c86d7..94564fbdb800 100644 --- a/examples/wide_ep/README.md +++ b/examples/wide_ep/README.md @@ -37,7 +37,7 @@ Starting from R580 Driver, [Coherent Driver-Based Memory Management (CDMM)](http When CDMM mode is off, GPU memory is exposed as NUMA nodes, so no additional prerequisites are required. -When CDMM mode is on, GPU memory doesn't exist in NUMA nodes. In that case, if online EPLB is needed, [GDRCopy](https://github.com/NVIDIA/gdrcopy?tab=readme-ov-file#build-and-installation) needs to be installed. +When CDMM mode is on, GPU memory doesn't exist in NUMA nodes. In that case, online EPLB uses [GDRCopy](https://github.com/NVIDIA/gdrcopy?tab=readme-ov-file#build-and-installation) when it is installed. Without GDRCopy, online EPLB falls back to CUDA managed memory that the host accesses directly over the coherent NVLink-C2C link; this works on hardware-coherent platforms such as GB200/GB300, while other platforms still require GDRCopy. When GDRCopy is installed and the kernel module is loaded, you should be able to see the device file `/dev/gdrdrv` and kernel module `gdrdrv` by `lsmod`. The device file needs to be mapped into the container. From 5c078b528f1747063592c7c7601a3344db2385ad Mon Sep 17 00:00:00 2001 From: TensorRT LLM AI Agent <296075020+trtllm-agent@users.noreply.github.com> Date: Thu, 15 Oct 2026 03:06:36 +0800 Subject: [PATCH 10/23] =?UTF-8?q?[https://nvbugs/6911393][fix]=20Cherry-pi?= =?UTF-8?q?ck=20#20195=20from=20release/1.3:=20Run=20the=20launcher=20test?= =?UTF-8?q?'s=20python3=20stand-in=20with=20-S=20so=20start-up=20hooks=20c?= =?UTF-8?q?annot=E2=80=A6?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Forward-port of release/1.3 PR https://github.com/NVIDIA/TensorRT-LLM/pull/20195 (release/1.3 commit cd4cc1778f21c7ec8e009b8c3df089c468393a48) to main. Adapted to main: main never had the release-only waives for nvbugs/6911393, so only the launcher test change is ported. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> Signed-off-by: Mike Iovine <6158008+mikeiovine@users.noreply.github.com> Co-authored-by: Claude Co-authored-by: Mike Iovine <6158008+mikeiovine@users.noreply.github.com> Signed-off-by: Simeng Liu --- tests/unittest/llmapi/test_mpi_launcher_shutdown.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/unittest/llmapi/test_mpi_launcher_shutdown.py b/tests/unittest/llmapi/test_mpi_launcher_shutdown.py index 3c5cd54c0e94..861b46669c3b 100644 --- a/tests/unittest/llmapi/test_mpi_launcher_shutdown.py +++ b/tests/unittest/llmapi/test_mpi_launcher_shutdown.py @@ -231,7 +231,10 @@ def _launcher_env(tmp_path: Path, mode: str) -> dict[str, str]: stub_bin = tmp_path / "bin" stub_bin.mkdir() python_stub = stub_bin / "python3" - python_stub.write_text(f"#!{sys.executable}\n{_STUB}") + # Like the launcher's own `python3 -S` helpers, skip site so interpreter + # startup hooks (e.g. a CI coverage sitecustomize) cannot import mpi4py and + # initialize MPI in every stand-in process. + python_stub.write_text(f"#!{sys.executable} -S\n{_STUB}") python_stub.chmod(0o755) for name in ("server_ready", "task_ready", "stop_requested"): os.mkfifo(tmp_path / name) From cbecf90da72a6e1b5cdae1c5e3833cb85521f380 Mon Sep 17 00:00:00 2001 From: chenfeiz0326 Date: Thu, 15 Oct 2026 15:42:09 +0800 Subject: [PATCH 11/23] [https://nvbugs/6825983][perf] Cherry-pick #20264 from release/1.3: Share TRTLLM-Gen FMHA cubins across isolated MPI ranks Forward-port of release/1.3 PR https://github.com/NVIDIA/TensorRT-LLM/pull/20264 (release/1.3 commit 74173d41b9dfbe5504a67f52f5d96b277af649b1) to main. The release PR cherry-picked main PR #20257, which was closed without merging, so this change was not on main. Signed-off-by: chenfeiz0326 <203214996+chenfeiz0326@users.noreply.github.com> Co-authored-by: chenfeiz0326 <203214996+chenfeiz0326@users.noreply.github.com> Signed-off-by: Simeng Liu --- docs/source/llm-api/index.md | 9 + tensorrt_llm/_torch/flashinfer_cubin_pool.py | 203 +++++++++++++++ tensorrt_llm/_torch/flashinfer_utils.py | 5 + .../integration/test_lists/test-db/l0_cpu.yml | 1 + .../flashinfer/test_flashinfer_cubin_pool.py | 237 ++++++++++++++++++ 5 files changed, 455 insertions(+) create mode 100644 tensorrt_llm/_torch/flashinfer_cubin_pool.py create mode 100644 tests/unittest/_torch/flashinfer/test_flashinfer_cubin_pool.py diff --git a/docs/source/llm-api/index.md b/docs/source/llm-api/index.md index fc6222cfac1f..8cb069d9093f 100644 --- a/docs/source/llm-api/index.md +++ b/docs/source/llm-api/index.md @@ -345,6 +345,15 @@ artifacts stay warm across launches. Set `TRTLLM_FLASHINFER_WORKSPACE_PER_PROCESS=0` before invoking the launcher or creating the LLM instance to disable this behavior. +Downloaded TRTLLM-Gen FMHA kernel cubins are the exception: FlashInfer only +loads them and never compiles them, so ranks with persistent slots share them +through a `cubins` directory next to the slots. A rank looks a kernel cubin up +in its own slot, then in the shared directory, and downloads it only when +neither has it, publishing the download to the shared directory with an atomic +rename. Every cubin read from the shared directory is checked against its +SHA-256. Headers and other compiler inputs are never shared. Set +`TRTLLM_FLASHINFER_SHARED_CUBINS=0` to keep these cubins per-rank as well. + Persistent slots are not pruned automatically, so their count can grow with peak job concurrency. When no TensorRT-LLM processes are using the cache, the slots may be deleted safely; subsequent launches rebuild the removed JIT diff --git a/tensorrt_llm/_torch/flashinfer_cubin_pool.py b/tensorrt_llm/_torch/flashinfer_cubin_pool.py new file mode 100644 index 000000000000..45e37dbb62ae --- /dev/null +++ b/tensorrt_llm/_torch/flashinfer_cubin_pool.py @@ -0,0 +1,203 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Share FlashInfer runtime cubins between isolated MPI ranks. + +``trtllm-llmapi-launch`` and ``MpiPoolSession`` give every rank its own +FlashInfer workspace, and with ``FLASHINFER_CUBIN_DIR`` unset FlashInfer keeps +its downloaded artifacts inside that workspace too. That isolation is required +for compiler inputs: FlashInfer downloads trtllm-gen export headers into the +same directory and compiles against them, and a header replaced by a peer +rank's download mid-compilation breaks the build. + +Runtime cubins are different. A cubin is never handed to a compiler; FlashInfer +only reads its bytes and passes them to ``cuModuleLoadData``. Keeping them +per-rank means every rank downloads every kernel it uses separately, and a +TRTLLM-Gen FMHA kernel that warmup does not reach is downloaded on the serving +path the first time a request shape selects it. Under attention DP only the +rank that served a request has already fetched its kernels, so the other ranks +stall on the download in lockstep and the whole step waits for them. + +This module installs a read-through pool for the TRTLLM-Gen FMHA kernel cubins +only, the ``*.cubin`` files directly under FlashInfer's ``TRTLLM_GEN_FMHA`` +artifact directory. A lookup first tries the rank-local cache, then a pool +directory that sits next to the per-rank workspace slots, and only then +downloads. A rank that downloads a cubin publishes it to the pool with an +atomic rename, so peers see either a complete file or none. Peers read pooled +cubins in place, and every byte read from the pool is checked against the +artifact's SHA-256 before use. Headers, checksum files and every other +artifact take FlashInfer's own path unchanged and stay per-rank. +""" + +import hashlib +import os +import uuid +from pathlib import Path +from typing import Callable, Optional + +from ..logger import logger + +_CUBIN_SUFFIX = ".cubin" +_POOL_DIR_NAME = "cubins" +_MANAGED_ENV = "TRTLLM_FLASHINFER_WORKSPACE_MANAGED" +_WORKSPACE_ENV = "FLASHINFER_WORKSPACE_BASE" +_CUBIN_DIR_ENV = "FLASHINFER_CUBIN_DIR" +_DISABLE_ENV = "TRTLLM_FLASHINFER_SHARED_CUBINS" +_INSTALLED_MARKER = "_trtllm_shared_cubin_pool" + + +def resolve_shared_cubin_dir(environ: Optional[dict] = None) -> Optional[Path]: + """Return the shared cubin pool for this process, or None when not applicable. + + The pool exists only for TensorRT-LLM-managed per-rank workspaces: the + workspace is a ``rank-N`` slot and the pool is the ``cubins`` directory + beside it. A user-configured ``FLASHINFER_CUBIN_DIR`` is already shared and + is left alone, as is an unmanaged ``FLASHINFER_WORKSPACE_BASE``. + ``TRTLLM_FLASHINFER_SHARED_CUBINS=0`` disables the pool. + + Args: + environ: Environment to read; defaults to ``os.environ``. + """ + env = os.environ if environ is None else environ + if env.get(_DISABLE_ENV, "1") == "0": + return None + if env.get(_CUBIN_DIR_ENV): + return None + if env.get(_MANAGED_ENV) != "1": + return None + workspace = env.get(_WORKSPACE_ENV) + if not workspace: + return None + workspace_path = Path(workspace) + if not workspace_path.name.startswith("rank-"): + # Temporary fallback workspaces are process-unique and removed at exit, + # so there is no stable directory to share with peers. + return None + return workspace_path.parent / _POOL_DIR_NAME + + +def _read_verified(path: Path, sha256: str) -> bytes: + try: + data = path.read_bytes() + except OSError: + return b"" + if hashlib.sha256(data).hexdigest() != sha256: + return b"" + return data + + +def _is_pooled_artifact(file_name: str, pooled_prefix: str) -> bool: + if not file_name.endswith(_CUBIN_SUFFIX): + return False + prefix = pooled_prefix.rstrip("/") + "/" + if not file_name.startswith(prefix): + return False + # Only the kernel cubins directly under the artifact directory are pooled; + # FlashInfer joins the directory and the kernel name with a doubled slash. + return "/" not in file_name[len(prefix) :].lstrip("/") + + +def _write_atomically(path: Path, data: bytes) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + temp_path = path.with_name(f"{path.name}.{uuid.uuid4().hex}.tmp") + try: + temp_path.write_bytes(data) + os.replace(temp_path, path) + finally: + if temp_path.exists(): + temp_path.unlink() + + +def make_pooled_get_artifact( + get_artifact: Callable[..., bytes], + local_cubin_dir: Path, + shared_cubin_dir: Path, + pooled_prefix: str, +) -> Callable[..., bytes]: + """Wrap FlashInfer's ``get_artifact`` with a shared pool for runtime cubins. + + Args: + get_artifact: FlashInfer's ``flashinfer.jit.cubin_loader.get_artifact``. + local_cubin_dir: This rank's FlashInfer cubin directory. + shared_cubin_dir: Pool directory shared by the job's ranks. + pooled_prefix: Artifact directory whose kernel cubins are pooled. + """ + + def pooled_get_artifact(file_name: str, sha256: str, *args, **kwargs) -> bytes: + if not _is_pooled_artifact(file_name, pooled_prefix): + return get_artifact(file_name, sha256, *args, **kwargs) + + data = _read_verified(local_cubin_dir / file_name, sha256) + if data: + return data + + shared_path = shared_cubin_dir / file_name + data = _read_verified(shared_path, sha256) + if data: + return data + + data = get_artifact(file_name, sha256, *args, **kwargs) + if data: + try: + _write_atomically(shared_path, data) + except OSError as error: + logger.debug(f"Could not share FlashInfer cubin {file_name}: {error}") + return data + + setattr(pooled_get_artifact, _INSTALLED_MARKER, True) + return pooled_get_artifact + + +def install_shared_cubin_pool() -> bool: + """Route FlashInfer's runtime cubin lookups through the shared pool. + + Returns True when the pool was installed. A no-op when the process does not + use a TensorRT-LLM-managed per-rank workspace, when the pool is disabled, or + when FlashInfer does not expose the expected loader. + """ + shared_cubin_dir = resolve_shared_cubin_dir() + if shared_cubin_dir is None: + return False + try: + from flashinfer.jit import cubin_loader + except ImportError: + return False + + try: + from flashinfer.artifacts import ArtifactPath + except ImportError: + return False + + get_artifact = getattr(cubin_loader, "get_artifact", None) + local_cubin_dir = getattr(cubin_loader, "FLASHINFER_CUBIN_DIR", None) + pooled_prefix = getattr(ArtifactPath, "TRTLLM_GEN_FMHA", None) + if not callable(get_artifact) or local_cubin_dir is None or not pooled_prefix: + return False + if getattr(get_artifact, _INSTALLED_MARKER, False): + return True + local_cubin_dir = Path(local_cubin_dir) + if local_cubin_dir.resolve() == shared_cubin_dir.resolve(): + return False + + pooled = make_pooled_get_artifact( + get_artifact, local_cubin_dir, shared_cubin_dir, pooled_prefix + ) + # The C++ cubin callback installed by setup_cubin_loader resolves + # get_artifact from this module's globals on every call, so rebinding the + # module attribute reaches the TRTLLM-Gen FMHA runtime loader. + cubin_loader.get_artifact = pooled + if getattr(cubin_loader, "get_cubin", None) is get_artifact: + cubin_loader.get_cubin = pooled + logger.info(f"FlashInfer runtime cubins are shared across ranks via {shared_cubin_dir}") + return True diff --git a/tensorrt_llm/_torch/flashinfer_utils.py b/tensorrt_llm/_torch/flashinfer_utils.py index a7eafaaf9282..7c478ff936aa 100644 --- a/tensorrt_llm/_torch/flashinfer_utils.py +++ b/tensorrt_llm/_torch/flashinfer_utils.py @@ -25,3 +25,8 @@ def get_env_enable_pdl() -> bool: print( "flashinfer is not installed properly, please try pip install or building from source codes" ) + +if IS_FLASHINFER_AVAILABLE: + from .flashinfer_cubin_pool import install_shared_cubin_pool + + install_shared_cubin_pool() diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 86d013d5461b..dfd9abaf7971 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -118,6 +118,7 @@ l0_cpu: - unittest/_torch/speculative/hw_agnostic # l0_h100 selects with -m "not cpu_only", so this entry is what runs it. - unittest/_torch/speculative/test_spec_slot_pool_sizing.py + - unittest/_torch/flashinfer/test_flashinfer_cubin_pool.py - unittest/_torch/test_custom_config_registration.py - unittest/_torch/test_hf_quant_config.py - unittest/_torch/test_mmap_utils.py diff --git a/tests/unittest/_torch/flashinfer/test_flashinfer_cubin_pool.py b/tests/unittest/_torch/flashinfer/test_flashinfer_cubin_pool.py new file mode 100644 index 000000000000..d4f1d5f5966d --- /dev/null +++ b/tests/unittest/_torch/flashinfer/test_flashinfer_cubin_pool.py @@ -0,0 +1,237 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Tests for sharing FlashInfer runtime cubins across isolated MPI ranks.""" + +import hashlib +import sys +import threading +import types +from pathlib import Path + +import pytest + +from tensorrt_llm._torch import flashinfer_cubin_pool as pool + +pytestmark = pytest.mark.cpu_only + +_FMHA = "deadbeef/fmha/trtllm-gen/" +# FlashInfer joins the artifact directory and the kernel name with a doubled slash. +_CUBIN = _FMHA + "/fmhaSm100fKernel_Test.cubin" +_HEADER = _FMHA + "/include/flashInferMetaInfo.h" +_GEMM_CUBIN = "deadbeef/gemm-0000000-0000000/Gemm_Test.cubin" + + +def _sha(data: bytes) -> str: + return hashlib.sha256(data).hexdigest() + + +class _Rank: + """One rank's view of FlashInfer: a private cubin dir and a recording downloader.""" + + def __init__(self, workspace_root: Path, rank: int, downloads: list): + self.workspace = workspace_root / f"rank-{rank}" + self.cubin_dir = self.workspace / ".cache" / "flashinfer" / "cubins" + self.rank = rank + self.downloads = downloads + self.env = { + "TRTLLM_FLASHINFER_WORKSPACE_MANAGED": "1", + "FLASHINFER_WORKSPACE_BASE": str(self.workspace), + } + + def download(self, file_name: str, sha256: str, session=None) -> bytes: + self.downloads.append((self.rank, file_name)) + data = file_name.encode() + path = self.cubin_dir / file_name + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(data) + return data + + def pooled(self): + shared = pool.resolve_shared_cubin_dir(self.env) + assert shared is not None + return pool.make_pooled_get_artifact(self.download, self.cubin_dir, shared, _FMHA) + + +@pytest.fixture +def ranks(tmp_path: Path): + downloads: list = [] + root = tmp_path / "flashinfer" + return root, downloads, [_Rank(root, rank, downloads) for rank in range(4)] + + +def test_resolve_shared_dir_sits_beside_rank_slots(tmp_path: Path) -> None: + env = { + "TRTLLM_FLASHINFER_WORKSPACE_MANAGED": "1", + "FLASHINFER_WORKSPACE_BASE": str(tmp_path / "rank-3"), + } + assert pool.resolve_shared_cubin_dir(env) == tmp_path / "cubins" + + +@pytest.mark.parametrize( + "env", + [ + pytest.param({"FLASHINFER_WORKSPACE_BASE": "/x/rank-0"}, id="unmanaged-workspace"), + pytest.param({"TRTLLM_FLASHINFER_WORKSPACE_MANAGED": "1"}, id="no-workspace"), + pytest.param( + { + "TRTLLM_FLASHINFER_WORKSPACE_MANAGED": "1", + "FLASHINFER_WORKSPACE_BASE": "/tmp/trtllm-flashinfer-rank-0.abcd1234", + }, + id="temporary-fallback-workspace", + ), + pytest.param( + { + "TRTLLM_FLASHINFER_WORKSPACE_MANAGED": "1", + "FLASHINFER_WORKSPACE_BASE": "/x/rank-0", + "FLASHINFER_CUBIN_DIR": "/shared/cubins", + }, + id="explicit-cubin-dir", + ), + pytest.param( + { + "TRTLLM_FLASHINFER_WORKSPACE_MANAGED": "1", + "FLASHINFER_WORKSPACE_BASE": "/x/rank-0", + "TRTLLM_FLASHINFER_SHARED_CUBINS": "0", + }, + id="disabled", + ), + ], +) +def test_resolve_shared_dir_is_off_outside_managed_rank_slots(env: dict) -> None: + assert pool.resolve_shared_cubin_dir(env) is None + + +def test_cubin_downloaded_once_is_reused_by_peer_ranks(ranks) -> None: + root, downloads, (r0, r1, r2, _) = ranks + data = _CUBIN.encode() + + assert r0.pooled()(_CUBIN, _sha(data)) == data + assert r1.pooled()(_CUBIN, _sha(data)) == data + assert r2.pooled()(_CUBIN, _sha(data)) == data + + assert downloads == [(0, _CUBIN)] + assert (root / "cubins" / _CUBIN).read_bytes() == data + # Peers read the pooled cubin in place; nothing is written to their caches. + assert not (r1.cubin_dir / _CUBIN).exists() + assert not (r2.cubin_dir / _CUBIN).exists() + + +@pytest.mark.parametrize( + "file_name", + [ + pytest.param(_HEADER, id="fmha-header"), + pytest.param(_FMHA + "/checksums.txt", id="fmha-checksums"), + pytest.param(_FMHA + "/include/Test.cubin", id="fmha-subdirectory"), + pytest.param(_GEMM_CUBIN, id="other-artifact-cubin"), + ], +) +def test_only_fmha_kernel_cubins_are_shared(ranks, file_name: str) -> None: + root, downloads, (r0, r1, _, _) = ranks + data = file_name.encode() + + assert r0.pooled()(file_name, _sha(data)) == data + assert r1.pooled()(file_name, _sha(data)) == data + + assert downloads == [(0, file_name), (1, file_name)] + assert not (root / "cubins").exists() + + +def test_corrupt_pool_entry_is_ignored_and_replaced(ranks) -> None: + root, downloads, (r0, r1, _, _) = ranks + data = _CUBIN.encode() + shared = root / "cubins" / _CUBIN + shared.parent.mkdir(parents=True) + shared.write_bytes(b"truncated") + + assert r0.pooled()(_CUBIN, _sha(data)) == data + assert downloads == [(0, _CUBIN)] + assert shared.read_bytes() == data + + assert r1.pooled()(_CUBIN, _sha(data)) == data + assert downloads == [(0, _CUBIN)] + + +def test_failed_download_is_not_published(ranks) -> None: + root, _, (r0, _, _, _) = ranks + + def failing_download(file_name: str, sha256: str, session=None) -> bytes: + return b"" + + shared = pool.resolve_shared_cubin_dir(r0.env) + pooled = pool.make_pooled_get_artifact(failing_download, r0.cubin_dir, shared, _FMHA) + assert pooled(_CUBIN, _sha(_CUBIN.encode())) == b"" + assert not (root / "cubins" / _CUBIN).exists() + + +def test_concurrent_ranks_never_read_a_partial_cubin(tmp_path: Path) -> None: + downloads: list = [] + root = tmp_path / "flashinfer" + names = [_CUBIN.replace("Test", f"Test{i}") for i in range(4)] + errors: list = [] + + def run(rank: int) -> None: + pooled = _Rank(root, rank, downloads).pooled() + try: + for _ in range(20): + for name in names: + assert pooled(name, _sha(name.encode())) == name.encode() + except AssertionError as error: + errors.append(error) + + threads = [threading.Thread(target=run, args=(rank,)) for rank in range(16)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + assert not errors + assert not list((root / "cubins").rglob("*.tmp")) + + +def test_install_wraps_flashinfer_loader_once( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + workspace = tmp_path / "rank-0" + monkeypatch.setenv("TRTLLM_FLASHINFER_WORKSPACE_MANAGED", "1") + monkeypatch.setenv("FLASHINFER_WORKSPACE_BASE", str(workspace)) + monkeypatch.delenv("FLASHINFER_CUBIN_DIR", raising=False) + monkeypatch.delenv("TRTLLM_FLASHINFER_SHARED_CUBINS", raising=False) + + def get_artifact(file_name: str, sha256: str, session=None) -> bytes: + return b"" + + cubin_loader = types.SimpleNamespace( + get_artifact=get_artifact, + get_cubin=get_artifact, + FLASHINFER_CUBIN_DIR=workspace / ".cache" / "flashinfer" / "cubins", + ) + artifacts = types.SimpleNamespace(ArtifactPath=types.SimpleNamespace(TRTLLM_GEN_FMHA=_FMHA)) + jit = types.ModuleType("flashinfer.jit") + jit.cubin_loader = cubin_loader + flashinfer = types.ModuleType("flashinfer") + flashinfer.jit = jit + flashinfer.artifacts = artifacts + monkeypatch.setitem(sys.modules, "flashinfer", flashinfer) + monkeypatch.setitem(sys.modules, "flashinfer.jit", jit) + monkeypatch.setitem(sys.modules, "flashinfer.jit.cubin_loader", cubin_loader) + monkeypatch.setitem(sys.modules, "flashinfer.artifacts", artifacts) + + assert pool.install_shared_cubin_pool() + wrapped = cubin_loader.get_artifact + assert wrapped is not get_artifact + assert cubin_loader.get_cubin is wrapped + + assert pool.install_shared_cubin_pool() + assert cubin_loader.get_artifact is wrapped From b7cd5ebe600bf68094344f9179c052442a9b64ea Mon Sep 17 00:00:00 2001 From: asfiyab-nvidia <117682710+asfiyab-nvidia@users.noreply.github.com> Date: Thu, 15 Oct 2026 09:08:38 -0700 Subject: [PATCH 12/23] [https://nvbugs/6912133][fix] Cherry-pick #20248 from release/1.3: Detect dead MPI pool manager thread in MpiPoolSession Forward-port of release/1.3 PR https://github.com/NVIDIA/TensorRT-LLM/pull/20248 (release/1.3 commit 9ec0b503824c2886edec29f0ca99f85dbf0f327b) to main. Signed-off-by: Asfiya Baig Signed-off-by: Simeng Liu --- tensorrt_llm/llmapi/mpi_session.py | 12 ++++++++++++ tests/unittest/executor/test_proxy_fast_death.py | 10 ++++++++++ 2 files changed, 22 insertions(+) diff --git a/tensorrt_llm/llmapi/mpi_session.py b/tensorrt_llm/llmapi/mpi_session.py index 8de04bd8f8ee..3509d70940d3 100644 --- a/tensorrt_llm/llmapi/mpi_session.py +++ b/tensorrt_llm/llmapi/mpi_session.py @@ -588,6 +588,18 @@ def _wait_workers_exit(self, timeout: float = 30.0) -> None: return time.sleep(0.05) + def check_worker_error(self) -> Optional[BaseException]: + """Detect a dead pool manager thread, e.g. after a failed spawn. + + mpi4py leaves the futures pending forever in that case. + """ + pool = self.mpi_pool + thread = getattr(getattr(pool, '_pool', None), 'thread', None) + if thread is None or thread.is_alive() or pool._shutdown: + return None + return RuntimeError( + "MPI pool manager thread died (worker spawn failed?)") + def release_exit_joins(self): if self.mpi_pool is not None: _abandon_mpi_pool_threads(self.mpi_pool) diff --git a/tests/unittest/executor/test_proxy_fast_death.py b/tests/unittest/executor/test_proxy_fast_death.py index 11804de01bb0..6a138c17eea3 100644 --- a/tests/unittest/executor/test_proxy_fast_death.py +++ b/tests/unittest/executor/test_proxy_fast_death.py @@ -689,6 +689,16 @@ def test_abandon_mpi_pool_threads_tolerates_missing_pool(): _abandon_mpi_pool_threads(object()) # no _pool attribute +def test_pool_session_reports_dead_manager_thread(): + from tensorrt_llm.llmapi.mpi_session import MpiPoolSession + + session = object.__new__(MpiPoolSession) + session.mpi_pool = _Mock(_shutdown=False) + session.mpi_pool._pool.thread.is_alive.return_value = False + assert isinstance(session.check_worker_error(), RuntimeError) + session.mpi_pool = None # keep __del__ from shutting down the fake pool + + def test_mark_engine_dead_releases_exit_joins_immediately(): """Exit-join release must happen at detection time, not at teardown. From ab40fa8e28fd97716f5dbe8786efd771e0e60cd1 Mon Sep 17 00:00:00 2001 From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Date: Mon, 19 Oct 2026 13:02:01 +0800 Subject: [PATCH 13/23] [https://nvbugs/6862662][fix] Cherry-pick #20415 from release/1.3: Bound FP8 MLA page-table loads for partial tiles Forward-port of release/1.3 PR https://github.com/NVIDIA/TensorRT-LLM/pull/20415 (release/1.3 commit c98395b58e049731d6f9b8dfb5ab50b26e9fc766) to main. Adapted to main: the fix is applied to tensorrt_llm/_torch/attention/kernels/blackwell/mla/mla_decode_fp8.py, where #19275 moved the CuTe DSL MLA decode kernels, and the new test passes the FP8 softmax/output scales as device tensors at inputs[10:12], the runner input contract main adopted in #19960. Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Signed-off-by: Simeng Liu --- .../kernels/blackwell/mla/mla_decode_fp8.py | 19 +++- .../attention/test_cute_dsl_mla_page_table.py | 106 ++++++++++++++++++ 2 files changed, 119 insertions(+), 6 deletions(-) create mode 100644 tests/unittest/_torch/attention/test_cute_dsl_mla_page_table.py diff --git a/tensorrt_llm/_torch/attention/kernels/blackwell/mla/mla_decode_fp8.py b/tensorrt_llm/_torch/attention/kernels/blackwell/mla/mla_decode_fp8.py index 7a750bd38846..ec03795d625b 100644 --- a/tensorrt_llm/_torch/attention/kernels/blackwell/mla/mla_decode_fp8.py +++ b/tensorrt_llm/_torch/attention/kernels/blackwell/mla/mla_decode_fp8.py @@ -2626,10 +2626,15 @@ def load_tma_qk_one_k_tile( k_idx = cute.make_rmem_tensor(cute.make_layout(page_per_tile), cutlass.Int32) for i in cutlass.range_constexpr(page_per_tile): - k_idx[i] = (common_params.mPT[k_index] if self.mma_qk_tiler[1] // - self.page_size == 1 else common_params.mPT[ - (k_index * qk_params.tiled_mma_qk.thr_id.shape + - common_params.blk_coord[0]) * page_per_tile + i]) + page_idx = (k_index if self.mma_qk_tiler[1] // + self.page_size == 1 else + (k_index * qk_params.tiled_mma_qk.thr_id.shape + + common_params.blk_coord[0]) * page_per_tile + i) + # A partial K tile can extend beyond the allocated page table. + # Its masked lanes still issue TMA loads, so use a valid page. + k_idx[i] = 0 + if page_idx < common_params.mPT.shape[0]: + k_idx[i] = common_params.mPT[page_idx] # load q once at first iteration load_q_pipeline = common_params.load_q_pipeline if load_q: @@ -2705,8 +2710,10 @@ def load_tma_v_one_k_tile( k_idx = cute.make_rmem_tensor(cute.make_layout(page_per_tile), cutlass.Int32) for i in cutlass.range_constexpr(page_per_tile): - k_idx[i] = (common_params.mPT[k_index] if page_per_tile == 1 else - common_params.mPT[k_index * page_per_tile + i]) + page_idx = k_index * page_per_tile + i + k_idx[i] = 0 + if page_idx < common_params.mPT.shape[0]: + k_idx[i] = common_params.mPT[page_idx] # get the mbar ptr from pipeline. tma_bar_ptr = common_params.load_v_pipeline.producer_get_barrier( load_v_producer_state) diff --git a/tests/unittest/_torch/attention/test_cute_dsl_mla_page_table.py b/tests/unittest/_torch/attention/test_cute_dsl_mla_page_table.py new file mode 100644 index 000000000000..f66d18188e07 --- /dev/null +++ b/tests/unittest/_torch/attention/test_cute_dsl_mla_page_table.py @@ -0,0 +1,106 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Exercise partial MLA tiles with page tables that have no tile padding. + +Run under compute-sanitizer memcheck with PYTORCH_NO_CUDA_MEMORY_CACHING=1 +so reads past a tensor allocation are reported even if masked by attention. +""" + +import math + +import pytest +import torch + +from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE +from tensorrt_llm._utils import get_sm_version + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() + or not IS_CUTLASS_DSL_AVAILABLE + or get_sm_version() not in (100, 103), + reason="CuTe DSL MLA decode requires SM100/SM103", +) + + +@pytest.mark.parametrize("page_size", [32, 64]) +@pytest.mark.parametrize("is_persistent", [False, True]) +@pytest.mark.parametrize("split_kv", [1, 2]) +def test_cute_dsl_mla_partial_page_table( + page_size: int, is_persistent: bool, split_kv: int +) -> None: + import cutlass + + from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import CuteDSLNVMlaDecodeBlackwellRunner + + torch.manual_seed(17) + batch_size, num_heads, latent_dim, rope_dim = 64, 128, 512, 64 + device = torch.device("cuda") + # Uniform attention keeps the unnormalized probabilities exact in FP8, + # isolating page selection and masking from softmax quantization error. + q_storage = torch.zeros( + batch_size, 1, num_heads, latent_dim + rope_dim, dtype=torch.float8_e4m3fn, device=device + ) + q_latent = q_storage[..., :latent_dim].permute(2, 3, 1, 0) + q_rope = q_storage[..., latent_dim:].permute(2, 3, 1, 0) + output_storage = torch.empty( + batch_size, 1, num_heads, latent_dim, dtype=torch.bfloat16, device=device + ) + output = output_storage.permute(2, 3, 1, 0) + runner = CuteDSLNVMlaDecodeBlackwellRunner( + in_dtype=cutlass.Float8E4M3FN, + num_heads=num_heads, + seq_len_q=1, + page_size=page_size, + max_batch_size=batch_size, + ) + workspace = torch.empty( + runner.get_max_padded_workspace_size(num_heads, 1, latent_dim, batch_size, cutlass.Float32), + dtype=torch.uint8, + device=device, + ) + scale = 1.0 / math.sqrt(latent_dim + rope_dim) + # FP8 runners read both scales from device tensors at inputs[10:12]. + scale_inputs = [ + None, # kv_bounds + torch.tensor([scale], dtype=torch.float32, device=device), + torch.tensor([1.0], dtype=torch.float32, device=device), + ] + + # Reuse the compiled kernel across padded and unpadded table widths. + # The interleaved K/V slots match KVCacheManagerV1's page-table strides. + for num_blocks in (4, 1, 2, 3, 5): + kv_len = num_blocks * page_size - 1 + num_pages = batch_size * num_blocks + cache = ( + torch.randn( + num_pages, page_size, latent_dim + rope_dim, dtype=torch.bfloat16, device=device + ) + * 0.1 + ).to(torch.float8_e4m3fn) + table_storage = torch.zeros(batch_size, 2, num_blocks, dtype=torch.int32, device=device) + table_storage[:, 0, :] = torch.arange(num_pages, dtype=torch.int32, device=device).view( + batch_size, num_blocks + ) + page_table = table_storage[:, 0, :].transpose(0, 1) + cache_seqs = torch.full((batch_size,), kv_len, dtype=torch.int32, device=device) + runner.forward( + [ + q_latent, + q_rope, + cache[..., :latent_dim].permute(1, 2, 0), + cache[..., latent_dim:].permute(1, 2, 0), + page_table, + cache_seqs, + output, + workspace, + None, + ] + + scale_inputs, + tactic=((128, 128), (128, 256), split_kv, is_persistent), + ) + torch.cuda.synchronize() + + values = cache.float().view(batch_size, num_blocks * page_size, -1)[:, :kv_len, :latent_dim] + expected = values.mean(dim=1, keepdim=True).expand(-1, num_heads, -1) + torch.testing.assert_close(output_storage[:, 0].float(), expected, atol=1e-4, rtol=5e-3) From ca51b44a79d690eba94cf6e0b19db845ca282732 Mon Sep 17 00:00:00 2001 From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Date: Mon, 19 Oct 2026 14:36:26 +0800 Subject: [PATCH 14/23] [None][fix] Cherry-pick #20381 from release/1.3: Follow up PrimTS refresh backport review tests and docs Forward-port of release/1.3 PR https://github.com/NVIDIA/TensorRT-LLM/pull/20381 (release/1.3 commit 4317a841aaa29e8f25a1f4c9981b4624dbb40f9c) to main. Adapted to main: main keeps PrimTS block-sparse plan caches on the attention metadata (TrtllmAttentionMetadata.fmha_plan_caches) and documents SageAttention for the contiguous block-sparse path. The summary-contract paragraph goes right after the routing-payload paragraph in ATTENTION_DEVELOPER_GUIDE.md, ahead of both. With the module-level cpu_only marker removed, the per-test cpu_only marker also goes on main's test_block_sparse_plan_cache_lives_on_the_metadata and on the three CPU-only SageAttention tests from main's #20224, so they keep running on the CPU stage. Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Signed-off-by: Simeng Liu --- .../sparse-attention-development-guide.md | 4 ++ docs/source/features/sparse-attention.md | 5 ++ .../attention/ATTENTION_DEVELOPER_GUIDE.md | 5 ++ .../sparse/test_prims_ts_block_sparse.py | 65 +++++++++++++++---- .../test_prims_ts_attention_backend.py | 62 ++++++++++++++++++ .../_torch/attention/test_prims_ts_fmha.py | 23 +++++-- 6 files changed, 145 insertions(+), 19 deletions(-) diff --git a/docs/source/developer-guide/sparse-attention-development-guide.md b/docs/source/developer-guide/sparse-attention-development-guide.md index 18c189cd5555..d95ce89497f5 100644 --- a/docs/source/developer-guide/sparse-attention-development-guide.md +++ b/docs/source/developer-guide/sparse-attention-development-guide.md @@ -358,6 +358,10 @@ three prediction methods. A `VanillaAttention` implementation instead overrides (BSR `block_indptr`/`block_indices` or a packed `exact_block_bits` bitmask), optional K/V summaries for proxy routes, and optional `kv_valid_bits` masking ragged KV tails. +- **Summary contract**: both `k_summary` and `v_summary` are arithmetic means + over the tokens present in each KV block, excluding padding in a partial + final block. The kernel applies the represented token count to each proxy; + do not pass V sums or pre-weight the summaries. - **Default**: hands through `SparseBackendForwardArgs.block_sparse_inputs`, so modules that predict before the core forward do not override it. Override it to predict inside the backend from the flattened Q/K/V, diff --git a/docs/source/features/sparse-attention.md b/docs/source/features/sparse-attention.md index 5cef178da442..1820f64d4137 100644 --- a/docs/source/features/sparse-attention.md +++ b/docs/source/features/sparse-attention.md @@ -163,6 +163,11 @@ routes is never served by a dense kernel. Both the contiguous prefill path, used by diffusion models that keep no KV cache, and the paged generation path are provided by the vendored PrimTS kernels. +Both `k_summary` and `v_summary` are arithmetic means over the tokens present +in each KV block. A partial final block excludes padding from its mean. The +kernel weights each proxy by its represented token count; do not pass V sums +or pre-weight the summaries. + This is an attention capability, not a standalone public `SparseAttentionConfig` algorithm. A user-facing algorithm must also provide the selector, metadata, and backend integration. diff --git a/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md b/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md index 17cc33280f11..c1b22569d485 100644 --- a/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md @@ -233,6 +233,11 @@ routes, and optional token-validity bits mask ragged KV tails. Plans contain only static format, proxy, geometry, and capacity choices; every run receives the live routes, summaries, validity bits, page tables, and sequence lengths. +Both `k_summary` and `v_summary` are arithmetic means over the tokens present +in each KV block, excluding padding in a partial final block. The kernel +weights each proxy by its represented token count; do not pass V sums or +pre-weight the summaries. + The contiguous block-sparse path also consumes the SageAttention recipe in `AttentionForwardArgs`. Prediction sees the original Q/K/V. The attention core then quantizes the Q/K/V of a block-sparse call before FMHA selection, with diff --git a/tests/unittest/_torch/attention/sparse/test_prims_ts_block_sparse.py b/tests/unittest/_torch/attention/sparse/test_prims_ts_block_sparse.py index f8c682f65a6f..e472ecd15102 100644 --- a/tests/unittest/_torch/attention/sparse/test_prims_ts_block_sparse.py +++ b/tests/unittest/_torch/attention/sparse/test_prims_ts_block_sparse.py @@ -40,8 +40,6 @@ from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager from tensorrt_llm.functional import PositionEmbeddingType -pytestmark = pytest.mark.cpu_only - _REQUIRES_PRIMTS_GPU = pytest.mark.skipif( getSMVersion() not in (100, 103), reason="PrimTS block-sparse attention requires SM100 or SM103", @@ -104,23 +102,47 @@ def _proxy_reference( block_size = 64 exact_tokens = torch.arange( exact_block * block_size, - (exact_block + 1) * block_size, + min((exact_block + 1) * block_size, k.shape[1]), device=q.device, ) proxy_blocks = [block for block in range(k_summary.shape[1]) if block != exact_block] + proxy_token_counts = torch.tensor( + [min(block_size, k.shape[1] - block * block_size) for block in proxy_blocks], + device=q.device, + dtype=torch.float32, + ) q_rows = q[0, :, 0].float() exact_logits = q_rows @ k[0, exact_tokens, 0].float().T proxy_logits = q_rows @ k_summary[0, proxy_blocks, 0].float().T logits = torch.cat((exact_logits, proxy_logits), dim=1) / math.sqrt(q.shape[-1]) weights = torch.exp(logits - logits.amax(dim=1, keepdim=True)) - exact_weights, proxy_weights = weights.split((block_size, len(proxy_blocks)), dim=1) + exact_weights, proxy_weights = weights.split((exact_tokens.numel(), len(proxy_blocks)), dim=1) + proxy_weights = proxy_weights * proxy_token_counts numerator = exact_weights @ v[0, exact_tokens, 0].float() - numerator += (proxy_weights * block_size) @ v_summary[0, proxy_blocks, 0].float() + numerator += proxy_weights @ v_summary[0, proxy_blocks, 0].float() denominator = exact_weights.sum(dim=1, keepdim=True) - denominator += proxy_weights.sum(dim=1, keepdim=True) * block_size + denominator += proxy_weights.sum(dim=1, keepdim=True) return (numerator / denominator).to(q.dtype)[None, :, None] +@pytest.mark.cpu_only +@pytest.mark.parametrize("kv_seq_len", [160, 192]) +@pytest.mark.parametrize("exact_block", [0, 2]) +def test_proxy_reference_weights_present_tokens(kv_seq_len: int, exact_block: int) -> None: + q = torch.zeros((1, 64, 1, 128)) + k = torch.zeros((1, kv_seq_len, 1, 128)) + v = torch.zeros_like(k) + v[:, 128:] = 1 + k_summary = torch.stack([block.mean(dim=1) for block in k.split(64, dim=1)], dim=1) + v_summary = torch.stack([block.mean(dim=1) for block in v.split(64, dim=1)], dim=1) + + actual = _proxy_reference(q, k, v, k_summary, v_summary, exact_block) + + # Uniform logits give every present token equal weight, including the tail. + expected = v.mean(dim=1, keepdim=True).expand_as(q) + torch.testing.assert_close(actual, expected) + + class _Attention: def __init__(self) -> None: self.sparse_params = None @@ -214,6 +236,7 @@ def _paged_case(): return attention, fmha, q, metadata, args +@pytest.mark.cpu_only def test_block_sparse_route_mode_is_derived_from_payload() -> None: bsr = _bsr_inputs() exact = _bitmask_inputs(proxy=False) @@ -236,6 +259,7 @@ def test_block_sparse_route_mode_is_derived_from_payload() -> None: ({"k_summary": torch.empty(0)}, "k_summary and v_summary"), ], ) +@pytest.mark.cpu_only def test_block_sparse_payload_rejects_ambiguous_combinations(overrides, message) -> None: kwargs = { "q_block_size": 64, @@ -250,6 +274,7 @@ def test_block_sparse_payload_rejects_ambiguous_combinations(overrides, message) BlockSparseForwardInputs(**kwargs) +@pytest.mark.cpu_only def test_block_sparse_support_is_phase_specific_and_paged_proxy_is_rejected( monkeypatch, ) -> None: @@ -267,6 +292,7 @@ def test_block_sparse_support_is_phase_specific_and_paged_proxy_is_rejected( assert reason == "paged block-sparse attention only supports BSR exact routes" +@pytest.mark.cpu_only def test_contiguous_proxy_routes_reject_causal_mask_before_planning(monkeypatch) -> None: _attention, fmha, q, k, v, metadata, args = _contiguous_case() _set_block_sparse_inputs(args, _bitmask_inputs(proxy=True)) @@ -287,6 +313,7 @@ def test_contiguous_proxy_routes_reject_causal_mask_before_planning(monkeypatch) @pytest.mark.parametrize("paged", [False, True]) +@pytest.mark.cpu_only def test_block_sparse_support_rejects_invalid_static_kernel_profile( monkeypatch, paged, @@ -326,6 +353,7 @@ def test_block_sparse_support_rejects_invalid_static_kernel_profile( assert reason == "block-sparse requires head_dim=128" +@pytest.mark.cpu_only def test_contiguous_wrappers_cache_static_profile_and_keep_routes_live(monkeypatch) -> None: _attention, fmha, q, k, v, metadata, args = _contiguous_case() wrapper = Mock() @@ -369,6 +397,7 @@ def test_contiguous_wrappers_cache_static_profile_and_keep_routes_live(monkeypat assert call.kwargs["v_summary"] is inputs.v_summary +@pytest.mark.cpu_only def test_block_sparse_plan_key_includes_attention_head_topology() -> None: inputs = _bitmask_inputs(proxy=True) q = torch.empty((128, 256), dtype=torch.bfloat16) @@ -392,6 +421,7 @@ def _key(fmha): assert _key(first) != _key(second) +@pytest.mark.cpu_only def test_block_sparse_plan_cache_lives_on_the_metadata(monkeypatch) -> None: """Layers that run with one metadata object share its planned wrappers.""" wrapper = Mock() @@ -419,6 +449,7 @@ def _set_sage_recipe(args: AttentionForwardArgs) -> None: args.sage_attn_num_elts_per_blk_v = 1 +@pytest.mark.cpu_only def test_attention_core_quantizes_sage_inputs_ahead_of_the_fmha(monkeypatch) -> None: attention, _fmha, q, k, v, metadata, args = _contiguous_case() # The core quantizes whichever libraries the layer has; each library decides whether it @@ -485,6 +516,7 @@ def quantize(): sage_quantize.assert_not_called() +@pytest.mark.cpu_only def test_contiguous_sage_requires_inputs_quantized_by_the_attention_core() -> None: _attention, fmha, q, k, v, metadata, args = _contiguous_case() _set_sage_recipe(args) @@ -494,6 +526,7 @@ def test_contiguous_sage_requires_inputs_quantized_by_the_attention_core() -> No assert reason == "SageAttention requires inputs quantized by the attention core" +@pytest.mark.cpu_only def test_dense_sage_with_v_smoothing_runs_on_a_dense_plan() -> None: _attention, fmha, q, k, v, metadata, args = _contiguous_case() _set_block_sparse_inputs(args, None) @@ -526,6 +559,7 @@ def test_dense_sage_with_v_smoothing_runs_on_a_dense_plan() -> None: assert key.unsupported_reason() is None +@pytest.mark.cpu_only def test_paged_wrapper_uses_zero_copy_padded_row_stride_block_tables(monkeypatch) -> None: attention, fmha, q, metadata, args = _paged_case() wrapper = Mock() @@ -602,6 +636,7 @@ def snapshot(*_args, **kwargs): assert snapshots[1][3] is _get_block_sparse_inputs(args).block_indptr +@pytest.mark.cpu_only def test_paged_block_tables_remain_live_across_graph_replay(monkeypatch) -> None: attention, fmha, q, metadata, args = _paged_case() wrapper = Mock() @@ -649,6 +684,7 @@ def snapshot(*_args, **kwargs): torch.testing.assert_close(seen[1][1], torch.arange(8, dtype=torch.int32).view(2, 4) + 10) +@pytest.mark.cpu_only def test_prepare_workspace_checks_capture_before_resize(monkeypatch) -> None: _attention, fmha, _q, _metadata, _args = _paged_case() query_device = torch.device("cuda:1") @@ -683,6 +719,7 @@ def test_prepare_workspace_checks_capture_before_resize(monkeypatch) -> None: assert workspace.numel() == 0 +@pytest.mark.cpu_only def test_prepare_workspace_skips_generation_layout_for_contiguous_requests(monkeypatch) -> None: _attention, fmha, q, _k, _v, metadata, args = _contiguous_case() layout = Mock() @@ -741,16 +778,19 @@ def test_real_gpu_raw_routes_and_token_mask_match_reference() -> None: @_REQUIRES_PRIMTS_GPU +@pytest.mark.parametrize("kv_seq_len", [160, 192]) @torch.no_grad() -def test_real_gpu_proxy_adapter_replays_live_routes_and_summaries() -> None: +def test_real_gpu_proxy_adapter_replays_live_routes_and_summaries(kv_seq_len: int) -> None: torch.manual_seed(20260901) q = torch.randn((1, 64, 1, 128), device="cuda", dtype=torch.bfloat16) - k = torch.randn((1, 192, 1, 128), device="cuda", dtype=torch.bfloat16) + k = torch.randn((1, kv_seq_len, 1, 128), device="cuda", dtype=torch.bfloat16) v = torch.randn_like(k) - k_blocks = k.float().view(1, 3, 64, 1, 128) - v_blocks = v.float().view(1, 3, 64, 1, 128) - initial_k_summary = k_blocks.mean(dim=2).to(k.dtype) - initial_v_summary = v_blocks.mean(dim=2).to(v.dtype) + initial_k_summary = torch.stack( + [block.float().mean(dim=1) for block in k.split(64, dim=1)], dim=1 + ).to(k.dtype) + initial_v_summary = torch.stack( + [block.float().mean(dim=1) for block in v.split(64, dim=1)], dim=1 + ).to(v.dtype) live_k_summary = initial_k_summary.clone() live_v_summary = initial_v_summary.clone() live_exact_bits = torch.tensor([[[[1]]]], device="cuda", dtype=torch.uint32) @@ -1009,6 +1049,7 @@ def test_real_gpu_paged_routes_match_reference_across_topologies( torch.testing.assert_close(actual, expected, rtol=2e-2, atol=2e-2) +@pytest.mark.cpu_only def test_paged_support_rejects_pages_smaller_than_a_route_fragment() -> None: page_size = 32 _attention, fmha, q, metadata, args = _paged_case() diff --git a/tests/unittest/_torch/attention/test_prims_ts_attention_backend.py b/tests/unittest/_torch/attention/test_prims_ts_attention_backend.py index d0e61dfb5e54..2d2fccc1ee91 100644 --- a/tests/unittest/_torch/attention/test_prims_ts_attention_backend.py +++ b/tests/unittest/_torch/attention/test_prims_ts_attention_backend.py @@ -571,6 +571,68 @@ def check_output() -> None: check_output() +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16], ids=["fp16", "bf16"]) +@pytest.mark.parametrize("log2_jump", [7.0, 9.0], ids=["freeze-row-max", "rescale-row-max"]) +@torch.no_grad() +def test_prims_ts_context_lazy_rescale_with_late_key_block( + dtype: torch.dtype, log2_jump: float +) -> None: + from tensorrt_llm._torch.attention.backends.prims_ts import BatchPrefillPagedTSWrapper + + device = torch.device("cuda") + q_len, kv_len, page_size, head_dim = 64, 512, 64, 128 + sm_scale = head_dim**-0.5 + query = torch.zeros((q_len, 1, head_dim), device=device, dtype=dtype) + query[..., 0] = 1 + k_cache = torch.zeros((kv_len // page_size, 1, page_size, head_dim), device=device, dtype=dtype) + v_cache = torch.ones_like(k_cache) + # Only the last page raises the row maximum, after earlier KV tiles have + # contributed to the accumulator. The skip threshold is 8 in log2 logits. + k_cache[-1, :, :, 0] = log2_jump / (sm_scale * math.log2(math.e)) + v_cache[-1] = -1 + qo_indptr = torch.tensor([0, q_len], device=device, dtype=torch.int32) + block_tables = torch.arange(kv_len // page_size, device=device, dtype=torch.int32)[None] + seq_lens = torch.tensor([kv_len], device=device, dtype=torch.int32) + output = torch.empty_like(query) + + logical_k = k_cache[:, 0].flatten(0, 1).float() + logical_v = v_cache[:, 0].flatten(0, 1).float() + scores = (query[:, 0].float() @ logical_k.T) * sm_scale + actual_log2_jump = (scores[0, -1] - scores[0, 0]).item() * math.log2(math.e) + assert (actual_log2_jump > 8.0) == (log2_jump > 8.0) + expected = (torch.softmax(scores, dim=-1) @ logical_v).to(dtype)[:, None] + + wrapper = BatchPrefillPagedTSWrapper(kv_layout="HND") + wrapper.plan( + device=device, + batch_size=1, + max_seq_len_q=q_len, + max_kv_len=kv_len, + num_qo_heads=1, + num_kv_heads=1, + head_dim=head_dim, + q_dtype=dtype, + k_dtype=dtype, + v_dtype=dtype, + out_dtype=dtype, + page_size=page_size, + mask_type="dense", + sm_scale=sm_scale, + ) + wrapper.run(query, k_cache, v_cache, qo_indptr, block_tables, seq_lens, out=output) + torch.testing.assert_close(output, expected, rtol=1e-2, atol=1e-2) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + wrapper.run( + query, k_cache, v_cache, qo_indptr, block_tables, seq_lens, out=output, validate=False + ) + output.zero_() + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(output, expected, rtol=1e-2, atol=1e-2) + + def test_prims_ts_context_wrapper_cuda_graph_replay_with_updated_metadata( monkeypatch: pytest.MonkeyPatch, ) -> None: diff --git a/tests/unittest/_torch/attention/test_prims_ts_fmha.py b/tests/unittest/_torch/attention/test_prims_ts_fmha.py index 8d2f771b88b2..283b053fe967 100644 --- a/tests/unittest/_torch/attention/test_prims_ts_fmha.py +++ b/tests/unittest/_torch/attention/test_prims_ts_fmha.py @@ -1156,7 +1156,7 @@ def test_generation_wrapper_plans_once_and_reads_live_fixed_metadata( False, ) ) - wrapper = Mock() + wrapper = create_autospec(prims_decode_module.BatchDecodePagedTSWrapper, instance=True) wrapper._plan_state = SimpleNamespace( policy=( ("use_split_kv", use_split_kv), @@ -1324,7 +1324,10 @@ def test_decode_layer_adapters_bind_the_same_shared_workspace( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) - wrappers = [Mock(), Mock()] + wrappers = [ + create_autospec(prims_decode_module.BatchDecodePagedTSWrapper, instance=True) + for _ in range(2) + ] wrapper_factory = Mock(side_effect=wrappers) monkeypatch.setattr( prims_decode_module, @@ -1410,7 +1413,10 @@ def test_decode_wrapper_cache_plans_each_batch_once_and_reuses_a_b_a( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) - wrappers = [Mock(), Mock()] + wrappers = [ + create_autospec(prims_decode_module.BatchDecodePagedTSWrapper, instance=True) + for _ in range(2) + ] wrapper_factory = Mock(side_effect=wrappers) monkeypatch.setattr( prims_decode_module, @@ -1510,7 +1516,7 @@ def test_mla_eager_wrapper_plans_once_and_reads_live_fixed_metadata( dtype=torch.int32, ) build_metadata = Mock(return_value=(kv_cache, block_tables, None)) - wrapper = Mock() + wrapper = create_autospec(prims_mla_module.BatchMLADecodePagedTSWrapper, instance=True) wrapper_factory = Mock(return_value=wrapper) monkeypatch.setattr( prims_ts_module.thop, @@ -1704,7 +1710,10 @@ def test_mla_wrapper_cache_plans_each_batch_once_and_reuses_a_b_a( monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) attn = _Attention(head_dim=576, is_mla=True, num_heads=4) fmha = PrimsTSFmha(attn) - wrappers = [Mock(), Mock()] + wrappers = [ + create_autospec(prims_mla_module.BatchMLADecodePagedTSWrapper, instance=True) + for _ in range(2) + ] wrapper_factory = Mock(side_effect=wrappers) monkeypatch.setattr( prims_mla_module, @@ -1770,7 +1779,7 @@ def test_mla_wrapper_capture_uses_cached_plan_and_rejects_plan_miss( "is_current_stream_capturing", lambda: capturing, ) - wrapper = Mock() + wrapper = create_autospec(prims_mla_module.BatchMLADecodePagedTSWrapper, instance=True) wrapper_factory = Mock(return_value=wrapper) monkeypatch.setattr( prims_mla_module, @@ -1816,7 +1825,7 @@ def test_mla_wrapper_receives_v2_bound_and_shared_workspace( dtype=torch.int32, ) build_metadata = Mock(return_value=(kv_cache, block_tables, None)) - wrapper = Mock() + wrapper = create_autospec(prims_mla_module.BatchMLADecodePagedTSWrapper, instance=True) wrapper_factory = Mock(return_value=wrapper) monkeypatch.setattr( prims_ts_module.thop, From e16a1568b2124cf3b39dc5a638ed3bea95929f53 Mon Sep 17 00:00:00 2001 From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Date: Mon, 19 Oct 2026 18:15:57 +0800 Subject: [PATCH 15/23] [https://nvbugs/6932241][fix] Cherry-pick #20588 from release/1.3: Scope NCCL window leases in helix attention tests Forward-port of release/1.3 PR https://github.com/NVIDIA/TensorRT-LLM/pull/20588 (release/1.3 commit 356b1a2ff5cf376b18d43f1849fa685b62b79681) to main. Adapted to main: main never had the release-only waives for nvbugs/6932241 and 6932265, so only the helix test changes are ported. Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Signed-off-by: Simeng Liu --- .../_torch/attention/multi_gpu/test_mha_helix.py | 7 +++++-- .../_torch/attention/multi_gpu/test_mla_helix.py | 9 +++++++-- 2 files changed, 12 insertions(+), 4 deletions(-) diff --git a/tests/unittest/_torch/attention/multi_gpu/test_mha_helix.py b/tests/unittest/_torch/attention/multi_gpu/test_mha_helix.py index 2a0443c3daef..f1ccc1e4e099 100644 --- a/tests/unittest/_torch/attention/multi_gpu/test_mha_helix.py +++ b/tests/unittest/_torch/attention/multi_gpu/test_mha_helix.py @@ -45,6 +45,7 @@ from tensorrt_llm._torch.attention.backends.utils import get_attention_backend from tensorrt_llm._torch.distributed.ops import cp_allgather from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.nccl_window_tensor_scope import nccl_window_tensor_scope from tensorrt_llm._torch.utils import model_extra_attrs from tensorrt_llm.functional import PositionEmbeddingType from tensorrt_llm.mapping import CpType, Mapping @@ -209,8 +210,10 @@ def _run_attention_distributed( scenario.batch, ctx_len_per_gpu, kv_cache_manager, helix_is_inactive_rank, position_ids_gen ) extra_attrs["attention_metadata"] = weakref.ref(attn_metadata) - with model_extra_attrs(extra_attrs): - output = attn(position_ids_gen, input_gen, attn_metadata) + # The o_proj all-reduce returns NCCL-window tensors, whose leases must be released by a + # scope rather than a tensor destructor. The clone outlives the scope for the accuracy check. + with model_extra_attrs(extra_attrs), nccl_window_tensor_scope(input_gen): + output = attn(position_ids_gen, input_gen, attn_metadata).clone() print(f"Rank {rank} {world_size}-GPU: result: {output[0, :8]} / {output[-1, -8:]}") kv_cache_manager.shutdown() diff --git a/tests/unittest/_torch/attention/multi_gpu/test_mla_helix.py b/tests/unittest/_torch/attention/multi_gpu/test_mla_helix.py index 741b2511bbb2..7b378a6428b1 100644 --- a/tests/unittest/_torch/attention/multi_gpu/test_mla_helix.py +++ b/tests/unittest/_torch/attention/multi_gpu/test_mla_helix.py @@ -47,6 +47,7 @@ from tensorrt_llm._torch.attention.mla import MLA from tensorrt_llm._torch.distributed.ops import cp_allgather from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.nccl_window_tensor_scope import nccl_window_tensor_scope from tensorrt_llm._torch.utils import model_extra_attrs from tensorrt_llm.functional import PositionEmbeddingType from tensorrt_llm.mapping import CpType, Mapping @@ -406,8 +407,12 @@ def _run_mla_distributed( enable_context_mla_with_cached_kv=True, ) extra_attrs["attention_metadata"] = weakref.ref(attn_metadata) - with model_extra_attrs(extra_attrs): - output = mla(position_ids_gen, input_gen, attn_metadata, latent_cache_gen=latent_cache_gen) + # The o_proj all-reduce returns NCCL-window tensors, whose leases must be released by a + # scope rather than a tensor destructor. The clone outlives the scope for the accuracy check. + with model_extra_attrs(extra_attrs), nccl_window_tensor_scope(input_gen): + output = mla( + position_ids_gen, input_gen, attn_metadata, latent_cache_gen=latent_cache_gen + ).clone() print(f"Rank {rank} {world_size}-GPU: result: {output[0, :8]} / {output[-1, -8:]}") kv_cache_manager.shutdown() From 57b691f506fabc2e3af39ecbfdec5055a48e10fb Mon Sep 17 00:00:00 2001 From: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com> Date: Mon, 19 Oct 2026 20:11:20 -0700 Subject: [PATCH 16/23] [https://nvbugs/6553508][fix] Cherry-pick #20344 from release/1.3: Reject conflicting LoRA adapter identities Forward-port of release/1.3 PR https://github.com/NVIDIA/TensorRT-LLM/pull/20344 (release/1.3 commit d74787637463c78eeb1d43dfd7711937345b69d4) to main. Adapted to main: #20043 made BaseWorker._load_lora_adapter return None and load with _lora_model_config only, so the identity checks wrap main's call and the new tests check load calls and bindings instead of a return value. Main's test_lora_ownership.py (#20605) builds BaseWorker without __init__, so its _make_runtime also sets the adapter identity map and locks. Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com> Signed-off-by: Simeng Liu --- tensorrt_llm/executor/base_worker.py | 52 ++++- .../_torch/peft/test_lora_ownership.py | 4 + tests/unittest/executor/test_base_worker.py | 207 ++++++++++++++++++ 3 files changed, 259 insertions(+), 4 deletions(-) diff --git a/tensorrt_llm/executor/base_worker.py b/tensorrt_llm/executor/base_worker.py index 4cf9d669ef72..924a5615c5dc 100644 --- a/tensorrt_llm/executor/base_worker.py +++ b/tensorrt_llm/executor/base_worker.py @@ -22,6 +22,7 @@ import weakref from pathlib import Path from queue import Empty, Queue +from threading import Lock from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Tuple, Union import torch @@ -130,6 +131,12 @@ def __init__( self._await_response_helper = AwaitResponseHelper(weakref.proxy(self)) self._backend = None if llm_args is None else llm_args.backend self._lora_config = llm_args.lora_config if llm_args is not None else None + # Maps adapter ID to (canonical path, checkpoint source). + self._lora_adapter_identities: dict[str, tuple[str, str]] = {} + # Weak values avoid retaining locks for failed, one-off adapter IDs. + self._lora_adapter_id_locks = weakref.WeakValueDictionary() + # Ensures concurrent requests for one adapter ID share the same lock. + self._lora_adapter_id_locks_guard = Lock() self._resource_governor_queue = None if global_mpi_size() > 1: @@ -316,10 +323,47 @@ def start(self): def _load_lora_adapter(self, lora_request: LoRARequest) -> None: """Ensure the adapter is loaded into the local cache.""" adapter_id = str(lora_request.adapter_id) - self._lora_manager.load_from_ckpt([lora_request.path], - model_config=self._lora_model_config, - uids=[adapter_id], - ckpt_source=lora_request.ckpt_source) + with self._lora_adapter_id_locks_guard: + adapter_lock = self._lora_adapter_id_locks.get(adapter_id) + if adapter_lock is None: + adapter_lock = Lock() + self._lora_adapter_id_locks[adapter_id] = adapter_lock + + with adapter_lock: + previous = self._lora_adapter_identities.get(adapter_id) + identity = (lora_request.path, lora_request.ckpt_source) + # An empty path retains the trusted API's ID-only cache lookup. + if lora_request.path: + same_source = (previous is not None + and previous[1] == lora_request.ckpt_source) + if previous is None or (same_source and previous != identity): + # Preserve path-alias support while avoiding filesystem + # access for the common exact-identity reuse case. + identity = (str(Path(lora_request.path).resolve()), + lora_request.ckpt_source) + if previous is not None and previous != identity: + logger.warning( + "Rejecting conflicting LoRA adapter identity for ID %s: " + "existing=%r, requested=%r", adapter_id, previous, + identity) + raise ValueError( + f"LoRA adapter ID {adapter_id} is already bound to a " + "different path or checkpoint source; use a new ID.") + loaded = False + try: + self._lora_manager.load_from_ckpt( + [lora_request.path], + model_config=self._lora_model_config, + uids=[adapter_id], + ckpt_source=lora_request.ckpt_source) + loaded = True + finally: + # Failed lookups are retryable unless the loader published UID + # state. Keep bindings for partial loads and across eviction. + if (loaded or adapter_id in self._lora_manager.cpp_lora_weights + or adapter_id in self._lora_manager.cpp_lora_config): + self._lora_adapter_identities.setdefault( + adapter_id, identity) def _load_prompt_adapter(self, prompt_adapter_request: PromptAdapterRequest): diff --git a/tests/unittest/_torch/peft/test_lora_ownership.py b/tests/unittest/_torch/peft/test_lora_ownership.py index 7ff67517b423..64299f7db1b9 100644 --- a/tests/unittest/_torch/peft/test_lora_ownership.py +++ b/tests/unittest/_torch/peft/test_lora_ownership.py @@ -7,6 +7,7 @@ from dataclasses import dataclass from importlib import import_module from pathlib import Path +from threading import Lock from types import SimpleNamespace from unittest.mock import Mock @@ -154,6 +155,9 @@ def _make_runtime( worker._runtime_model_config = None worker._lora_model_config = model_config worker._lora_manager = manager + worker._lora_adapter_identities = {} + worker._lora_adapter_id_locks = {} + worker._lora_adapter_id_locks_guard = Lock() return _Runtime(config, model_config, manager, worker, cache, transfer) diff --git a/tests/unittest/executor/test_base_worker.py b/tests/unittest/executor/test_base_worker.py index 7776a39902ca..6cade02c2e35 100644 --- a/tests/unittest/executor/test_base_worker.py +++ b/tests/unittest/executor/test_base_worker.py @@ -1,7 +1,11 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 import time +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path +from threading import Event, Lock from types import SimpleNamespace +from unittest.mock import Mock, patch import pytest import torch @@ -194,6 +198,209 @@ def test_lora_request_does_not_probe_filesystem_on_init(tmp_path): assert request.path == missing_path +@pytest.fixture +def lora_identity_worker() -> BaseWorker: + worker = object.__new__(BaseWorker) + worker.doing_shutdown = False + worker._lora_model_config = None + worker._lora_adapter_identities = {} + worker._lora_adapter_id_locks = {} + worker._lora_adapter_id_locks_guard = Lock() + worker._lora_manager = Mock() + worker._lora_manager.cpp_lora_weights = {} + worker._lora_manager.cpp_lora_config = {} + worker._lora_manager.load_from_ckpt.return_value = ["7"] + worker.engine = Mock() + return worker + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("change_source", [False, True]) +def test_lora_identity_rejects_conflicting_request( + lora_identity_worker: BaseWorker, tmp_path: Path, + change_source: bool) -> None: + worker = lora_identity_worker + path = str(tmp_path / "adapter") + worker._load_lora_adapter(LoRARequest("first", 7, path)) + conflict = LoRARequest("second", 7, + path if change_source else str(tmp_path / "other"), + "nemo" if change_source else "hf") + with patch("tensorrt_llm.executor.base_worker.logger.warning") as warning: + with pytest.raises(RequestError, match="ID 7 is already bound"): + worker._enqueue_request(SimpleNamespace(id=1, + lora_request=conflict)) + assert warning.call_args.args[1:3] == ("7", (path, "hf")) + assert worker._lora_manager.load_from_ckpt.call_count == 1 + worker.engine.enqueue_request.assert_not_called() + + +@pytest.mark.cpu_only +def test_lora_identity_preserves_reuse(lora_identity_worker: BaseWorker, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch) -> None: + worker = lora_identity_worker + path = tmp_path / "adapter" + path.mkdir() + alias = tmp_path / "alias" + alias.symlink_to(path, target_is_directory=True) + monkeypatch.chdir(tmp_path) + worker._load_lora_adapter(LoRARequest("first", 7, str(path))) + worker._lora_manager.load_from_ckpt.return_value = [] + with patch.object(Path, + "resolve", + side_effect=AssertionError("unexpected resolve")): + worker._load_lora_adapter(LoRARequest("renamed", 7, str(path))) + for alternate in ("adapter", str(alias), ""): + worker._load_lora_adapter(LoRARequest("renamed", 7, alternate)) + assert worker._lora_adapter_identities == {"7": (str(path), "hf")} + worker._lora_manager.load_from_ckpt.assert_called_with([""], + model_config=None, + uids=["7"], + ckpt_source="hf") + worker._lora_manager.load_from_ckpt.return_value = ["8"] + worker._load_lora_adapter( + LoRARequest("independent", 8, str(tmp_path / "other"))) + + +@pytest.mark.cpu_only +def test_lora_identity_preserves_id_only_first_use( + lora_identity_worker: BaseWorker, tmp_path: Path) -> None: + worker = lora_identity_worker + worker._lora_manager.load_from_ckpt.return_value = [] + worker._load_lora_adapter(LoRARequest("cached", 7)) + assert worker._lora_adapter_identities == {"7": ("", "hf")} + with pytest.raises(ValueError, match="already bound"): + worker._load_lora_adapter(LoRARequest("new", 7, str(tmp_path))) + assert worker._lora_manager.load_from_ckpt.call_count == 1 + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("empty_path", [False, True]) +def test_lora_identity_failed_lookup_is_retryable( + lora_identity_worker: BaseWorker, tmp_path: Path, + empty_path: bool) -> None: + worker = lora_identity_worker + worker._lora_manager.load_from_ckpt.side_effect = FileNotFoundError( + "missing") + with pytest.raises(FileNotFoundError, match="missing"): + worker._load_lora_adapter( + LoRARequest("missing", 7, + "" if empty_path else str(tmp_path / "typo"))) + assert worker._lora_adapter_identities == {} + worker._lora_manager.load_from_ckpt.side_effect = None + worker._load_lora_adapter( + LoRARequest("corrected", 7, str(tmp_path / "adapter"))) + assert worker._lora_adapter_identities == { + "7": (str(tmp_path / "adapter"), "hf") + } + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("published_cache", + ["cpp_lora_weights", "cpp_lora_config"]) +def test_lora_identity_retains_partial_binding(lora_identity_worker: BaseWorker, + tmp_path: Path, + published_cache: str) -> None: + worker = lora_identity_worker + + def fail_after_publish(*args, **kwargs) -> None: + getattr(worker._lora_manager, published_cache)["7"] = [] + raise RuntimeError("load failed") + + worker._lora_manager.load_from_ckpt.side_effect = fail_after_publish + with pytest.raises(RuntimeError, match="load failed"): + worker._load_lora_adapter(LoRARequest("first", 7, str(tmp_path / "a"))) + with pytest.raises(ValueError, match="already bound"): + worker._load_lora_adapter(LoRARequest("second", 7, str(tmp_path / "b"))) + assert worker._lora_manager.load_from_ckpt.call_count == 1 + + +@pytest.mark.cpu_only +def test_lora_identity_serializes_registration(lora_identity_worker: BaseWorker, + tmp_path: Path) -> None: + worker = lora_identity_worker + load_started = Event() + release_load = Event() + + class SignalingLock: + + def __init__(self) -> None: + self._lock = Lock() + self._attempts_lock = Lock() + self._attempts = 0 + self.second_attempted = Event() + + def __enter__(self): + with self._attempts_lock: + self._attempts += 1 + if self._attempts == 2: + self.second_attempted.set() + self._lock.acquire() + return self + + def __exit__(self, exc_type, exc_value, traceback) -> None: + self._lock.release() + + def load_adapter(*args, **kwargs): + load_started.set() + assert release_load.wait(timeout=5) + return ["7"] + + adapter_lock = SignalingLock() + worker._lora_adapter_id_locks["7"] = adapter_lock + worker._lora_manager.load_from_ckpt.side_effect = load_adapter + requests = [ + LoRARequest("adapter", 7, str(tmp_path / name)) for name in ("a", "b") + ] + with ThreadPoolExecutor(max_workers=2) as pool: + first = pool.submit(worker._load_lora_adapter, requests[0]) + assert load_started.wait(timeout=5) + second = pool.submit(worker._load_lora_adapter, requests[1]) + try: + assert adapter_lock.second_attempted.wait(timeout=5) + finally: + release_load.set() + futures = [first, second] + assert sum(future.exception() is None for future in futures) == 1 + errors = [future.exception() for future in futures if future.exception()] + assert isinstance(errors[0], ValueError) + assert "already bound" in str(errors[0]) + assert worker._lora_manager.load_from_ckpt.call_count == 1 + + +@pytest.mark.cpu_only +def test_lora_identity_allows_independent_registration( + lora_identity_worker: BaseWorker, tmp_path: Path) -> None: + worker = lora_identity_worker + first_started = Event() + release_first = Event() + second_loaded = Event() + + def load_adapter(*args, uids, **kwargs): + adapter_id = uids[0] + if adapter_id == "7": + first_started.set() + assert release_first.wait(timeout=5) + else: + second_loaded.set() + return [adapter_id] + + worker._lora_manager.load_from_ckpt.side_effect = load_adapter + with ThreadPoolExecutor(max_workers=2) as pool: + first = pool.submit(worker._load_lora_adapter, + LoRARequest("first", 7, str(tmp_path / "a"))) + assert first_started.wait(timeout=5) + second = pool.submit(worker._load_lora_adapter, + LoRARequest("second", 8, str(tmp_path / "b"))) + try: + assert second_loaded.wait(timeout=5) + assert second.result(timeout=5) is None + finally: + release_first.set() + assert first.result(timeout=5) is None + assert set(worker._lora_adapter_identities) == {"7", "8"} + + def create_fake_llm_args(engine_path, tp_size: int = 1): """Create TorchLlmArgs for testing. From a9b9c772a4ede4f5e7e5b139a7164a18408193aa Mon Sep 17 00:00:00 2001 From: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com> Date: Tue, 20 Oct 2026 21:12:22 -0700 Subject: [PATCH 17/23] [https://nvbugs/6857649][fix] Cherry-pick #20346 from release/1.3: Validate native transfer destinations before submission Forward-port of release/1.3 PR https://github.com/NVIDIA/TensorRT-LLM/pull/20346 (release/1.3 commit 2331e1f45dac1b3d820cf34567bed73d8556204f) to main. Adapted to main: main-only disaggregated tests that drive Sender._submit_transfer (test_retirement_deadline.py, test_kv_transfer_fail_stop.py, test_transfer_late_settlement.py) now pass a WRITE request whose remote descriptors pass the new pre-submission check. Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com> Signed-off-by: Simeng Liu --- .../nixl_utils/transferAgent.cpp | 6 +- .../unit_tests/executor/transferAgentTest.cpp | 69 +++++++++++ .../_torch/disaggregation/native/transfer.py | 29 ++++- .../_torch/disaggregation/nixl/_agent_py.py | 17 ++- tests/unittest/disaggregated/test_agent.py | 41 +++++++ .../test_kv_transfer_fail_stop.py | 11 +- .../disaggregated/test_retirement_deadline.py | 23 ++-- .../test_transfer_late_settlement.py | 8 +- .../test_transfer_ownership_regressions.py | 110 ++++++++++++++++++ 9 files changed, 296 insertions(+), 18 deletions(-) diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp index 61633d141244..85f015d3e458 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp @@ -1234,7 +1234,11 @@ bool NixlTransferAgent::checkRemoteDescs(std::string const& name, MemoryDescs co { std::shared_lock lock(mLock); TLLM_CHECK_WITH_INFO(!mShutdown.load(), "NixlTransferAgent::checkRemoteDescs called after shutdown"); - auto status = mRawAgent->checkRemoteMD(name, NixlHelper::convertXferDist(memoryDescs)); + static VramRegionMap const kEmptyMap; + auto const remoteIt = mRemoteVramRegionInfo.find(name); + auto const& remoteRegionMap = (remoteIt != mRemoteVramRegionInfo.end()) ? remoteIt->second : kEmptyMap; + auto const splitDescs = VmmDescSplitter::splitDescsWithRegionMap(memoryDescs, remoteRegionMap); + auto status = mRawAgent->checkRemoteMD(name, NixlHelper::convertXferDist(splitDescs)); TLLM_CHECK_WITH_INFO(status == NIXL_SUCCESS || status == NIXL_ERR_NOT_FOUND, "checkRemoteMD failed with status: %s", nixlEnumStrings::statusStr(status).c_str()); return status == NIXL_SUCCESS; diff --git a/cpp/tests/unit_tests/executor/transferAgentTest.cpp b/cpp/tests/unit_tests/executor/transferAgentTest.cpp index dcf6dd5a6e1e..c63c20dc5b75 100644 --- a/cpp/tests/unit_tests/executor/transferAgentTest.cpp +++ b/cpp/tests/unit_tests/executor/transferAgentTest.cpp @@ -21,6 +21,7 @@ #include #include +#include #include #include #include @@ -125,6 +126,74 @@ TEST_P(TransferAgentTest, Basic) xferAgent0->invalidateRemoteAgent(agent1); } +TEST_P(TransferAgentTest, CheckRemoteDescriptorRanges) +{ + if (backend != "nixl") + { + GTEST_SKIP() << "NIXL metadata validation"; + } + std::string const agent0{"agent0"}, agent1{"agent1"}; + BaseAgentConfig config0{agent0, true, false, true}, config1{agent1, true, false, true}; + auto xferAgent0 = makeTransferAgent(config0); + auto xferAgent1 = makeTransferAgent(config1); + std::vector source(64, 10); + std::vector destination(128, 1); + RegisteredHostMemory regSource(MemoryDescs{MemoryType::kDRAM, {MemoryDesc{source}}}, xferAgent0.get()); + RegisteredHostMemory regDestination( + MemoryDescs{MemoryType::kDRAM, {MemoryDesc{destination.data(), source.size(), 0}}}, xferAgent1.get()); + xferAgent0->loadRemoteAgent(agent1, xferAgent1->getLocalAgentDesc()); + ASSERT_TRUE(xferAgent0->checkRemoteDescs(agent1, regDestination.getDescs())); + + // The backing allocation is large enough, but the requested range exceeds its registration. + MemoryDescs const overrun{MemoryType::kDRAM, {MemoryDesc{destination.data() + 32, source.size(), 0}}}; + EXPECT_FALSE(xferAgent0->checkRemoteDescs(agent1, overrun)); + EXPECT_EQ(destination, std::vector(128, 1)); + + EXPECT_FALSE(xferAgent0->checkRemoteDescs("unknown-agent", regDestination.getDescs())); + + MemoryDescs const subSource{MemoryType::kDRAM, {MemoryDesc{source.data(), 16, 0}}}; + MemoryDescs const subDestination{MemoryType::kDRAM, {MemoryDesc{destination.data() + 16, 16, 0}}}; + EXPECT_TRUE(xferAgent0->checkRemoteDescs(agent1, subDestination)); + TransferRequest const validWrite{TransferOp::kWRITE, subSource, subDestination, agent1}; + auto status = xferAgent0->submitTransferRequests(validWrite); + ASSERT_EQ(status->wait(), TransferState::kSUCCESS); + auto expected = std::vector(128, 1); + std::fill(expected.begin() + 16, expected.begin() + 32, 10); + EXPECT_EQ(destination, expected); + xferAgent0->invalidateRemoteAgent(agent1); +} + +TEST_P(TransferAgentTest, CheckRemoteDescriptorsAcrossVmmChunks) +{ + if (backend != "nixl") + { + GTEST_SKIP() << "NIXL metadata validation"; + } + constexpr size_t kChunkSize = 64; + constexpr size_t kTotalSize = 2 * kChunkSize; + char* destination = nullptr; + TLLM_CUDA_CHECK(cudaMalloc(&destination, kTotalSize)); + std::unique_ptr allocation{destination, &cudaFree}; + std::string const agent0{"agent0"}, agent1{"agent1"}; + BaseAgentConfig config0{agent0, true, false, true}, config1{agent1, true, false, true}; + auto xferAgent0 = makeTransferAgent(config0); + auto xferAgent1 = makeTransferAgent(config1); + RegisteredHostMemory registered( + MemoryDescs{MemoryType::kVRAM, + {MemoryDesc{destination, kChunkSize, 0}, MemoryDesc{destination + kChunkSize, kChunkSize, 0}}}, + xferAgent1.get()); + // Advertise a chunked region over two real registrations without requiring a VMM allocator. + auto const metadata = xferAgent1->getLocalAgentDesc(); + AgentDesc const chunked{ + metadata.getBackendAgentDesc(), {{reinterpret_cast(destination), kTotalSize, kChunkSize}}}; + xferAgent0->loadRemoteAgent(agent1, chunked); + EXPECT_TRUE( + xferAgent0->checkRemoteDescs(agent1, MemoryDescs{MemoryType::kVRAM, {MemoryDesc{destination, kTotalSize, 0}}})); + EXPECT_FALSE(xferAgent0->checkRemoteDescs( + agent1, MemoryDescs{MemoryType::kVRAM, {MemoryDesc{destination, kTotalSize + 1, 0}}})); + xferAgent0->invalidateRemoteAgent(agent1); +} + TEST_P(TransferAgentTest, Basic2) { diff --git a/tensorrt_llm/_torch/disaggregation/native/transfer.py b/tensorrt_llm/_torch/disaggregation/native/transfer.py index 29243078fb1c..ae5bacc4e985 100644 --- a/tensorrt_llm/_torch/disaggregation/native/transfer.py +++ b/tensorrt_llm/_torch/disaggregation/native/transfer.py @@ -782,6 +782,18 @@ def _submit_transfer( peer_rank: int, request: TransferRequest, ) -> tuple[bool, Optional[str]]: + try: + # This query cannot submit a transfer. Reject before acquiring backend ownership. + if request.op == TransferOp.WRITE and not self._agent.check_remote_descs( + request.remote_name, request.dst_descs + ): + raise RuntimeError("NIXL WRITE remote descriptors unavailable") + except Exception as error: + if self._enforce_physical_ownership: + task.retire_unsubmitted_physical_operation(peer_rank) + raise _TransferNotSubmittedError( + "NIXL sender rejected transfer before backend submission" + ) from error if not self._enforce_physical_ownership: status = self._agent.submit_transfer_requests(request) completed = status.wait() @@ -1063,6 +1075,19 @@ def _deliver_kv_to_agent(self, write_meta: WriteMeta): submitted_and_completed = False if write_meta.src_ptrs.size > 0: try: + # Bounce rewrites the submitted destination to its registered arena, + # but the receiver later scatters into these original KV addresses. + if write_meta.bounce_dst_base is not None: + if write_meta.dst_device_id is None: + raise RuntimeError("bounce destination validation requires dst_device_id") + final_dst_descs = MemoryDescs.from_arrays_uniform_device( + MemoryType.VRAM, + write_meta.dst_ptrs, + write_meta.sizes, + write_meta.dst_device_id, + ) + if not self._agent.check_remote_descs(write_meta.peer_name, final_dst_descs): + raise RuntimeError("NIXL WRITE final remote descriptors unavailable") request, send_slot_id = build_send_request( self._bounce, write_meta, @@ -1072,12 +1097,12 @@ def _deliver_kv_to_agent(self, write_meta: WriteMeta): # Don't let a gather fault escape: without a result the receiver would hang and its # region leak. Tell the receiver it failed and fail the local task instead. logger.error( - f"_deliver_kv_to_agent: failed to build the KV send request for " + f"_deliver_kv_to_agent: failed to validate or build the KV send request for " f"{write_meta.unique_rid} slice={write_meta.slice_id}: {e}" ) if owned: task.retire_unsubmitted_physical_operation(write_meta.peer_rank) - task.fail(RuntimeError(f"build_send_request failed: {e}")) + task.fail(RuntimeError(f"validate/build send request failed: {e}")) self._get_result_dealer(write_meta.peer_endpoint, legacy_shared=True).send( _make_kv_result_msg( self._instance_rank, diff --git a/tensorrt_llm/_torch/disaggregation/nixl/_agent_py.py b/tensorrt_llm/_torch/disaggregation/nixl/_agent_py.py index b8c4f096b715..55c2187f2ba6 100644 --- a/tensorrt_llm/_torch/disaggregation/nixl/_agent_py.py +++ b/tensorrt_llm/_torch/disaggregation/nixl/_agent_py.py @@ -21,7 +21,13 @@ from tensorrt_llm.logger import logger # Import base classes for type compatibility -from ..base.agent import BaseTransferAgent, RegMemoryDescs, TransferRequest, TransferStatus +from ..base.agent import ( + BaseTransferAgent, + MemoryDescs, + RegMemoryDescs, + TransferRequest, + TransferStatus, +) # The PyPI ``nixl`` meta-package provides a top-level ``nixl`` shim that # redirects to ``nixl_cu13`` or ``nixl_cu12``. When only ``nixl-cu13`` (or @@ -168,8 +174,13 @@ def get_local_agent_desc(self): def invalidate_remote_agent(self, name: str): self.agent.remove_remote_agent(name) - def check_remote_descs(self, name: str, memory_descs: list[int]) -> bool: - raise NotImplementedError + def check_remote_descs(self, name: str, memory_descs: MemoryDescs) -> bool: + if len(memory_descs.descs) == 0: + return self.agent.check_remote_metadata(name) + descs = self.agent.get_xfer_descs(memory_descs.descs, memory_descs.type) + if descs is None: + raise RuntimeError("nixl get_xfer_descs returned None for remote descriptors") + return self.agent.check_remote_metadata(name, descs) def notify_sync_message(self, name: str, sync_message: str): raise NotImplementedError diff --git a/tests/unittest/disaggregated/test_agent.py b/tests/unittest/disaggregated/test_agent.py index 5d7a8a3d62de..030789f1dd49 100644 --- a/tests/unittest/disaggregated/test_agent.py +++ b/tests/unittest/disaggregated/test_agent.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + import os from dataclasses import dataclass, field from unittest import TestCase @@ -68,6 +71,44 @@ def test_python_agent_accepts_topology_without_forwarding_it_to_nixl(): nixl_agent.assert_called_once_with("testAgent", agent_config.return_value) +@pytest.mark.cpu_only +@pytest.mark.parametrize("outcome", [True, False, RuntimeError("metadata lookup failed")]) +def test_python_check_remote_descs( + outcome: bool | RuntimeError, +) -> None: + agent_py = pytest.importorskip(_AGENT_PY_MODULE, exc_type=ImportError) + with patch.object(agent_py, "nixl_agent_config"), patch.object(agent_py, "nixl_agent"): + agent = agent_py.NixlTransferAgent("testAgent") + backend = agent.agent + memory_descs = Mock(descs=[(1024, 64, 0)], type="VRAM") + converted = backend.get_xfer_descs.return_value + if isinstance(outcome, Exception): + backend.check_remote_metadata.side_effect = outcome + else: + backend.check_remote_metadata.return_value = outcome + + if isinstance(outcome, Exception): + with pytest.raises(RuntimeError): + agent.check_remote_descs("peer", memory_descs) + else: + assert agent.check_remote_descs("peer", memory_descs) is outcome + backend.get_xfer_descs.assert_called_once_with(memory_descs.descs, "VRAM") + backend.check_remote_metadata.assert_called_once_with("peer", converted) + backend.initialize_xfer.assert_not_called() + backend.transfer.assert_not_called() + + +@pytest.mark.cpu_only +def test_python_check_remote_descs_rejects_failed_conversion() -> None: + agent_py = pytest.importorskip(_AGENT_PY_MODULE, exc_type=ImportError) + with patch.object(agent_py, "nixl_agent_config"), patch.object(agent_py, "nixl_agent"): + agent = agent_py.NixlTransferAgent("testAgent") + agent.agent.get_xfer_descs.return_value = None + with pytest.raises(RuntimeError, match="returned None"): + agent.check_remote_descs("peer", Mock(descs=[(1024, 64, 0)], type="VRAM")) + agent.agent.check_remote_metadata.assert_not_called() + + def _convert_to_memory_descs(reg_descs: RegMemoryDescs) -> MemoryDescs: tuples = [(ptr, size, device_id) for (ptr, size, device_id, _) in reg_descs.descs] diff --git a/tests/unittest/disaggregated/test_kv_transfer_fail_stop.py b/tests/unittest/disaggregated/test_kv_transfer_fail_stop.py index 29d38828ccc8..3efba725f0e2 100644 --- a/tests/unittest/disaggregated/test_kv_transfer_fail_stop.py +++ b/tests/unittest/disaggregated/test_kv_transfer_fail_stop.py @@ -338,7 +338,7 @@ def test_failed_sender_write_retains_roots_until_fatal_deadline() -> None: task = transfer_mod.SendTaskBase(DisaggregatedParams(disagg_request_id=12)) task.bind_logical_outcomes(transfer_mod._LogicalOutcomes(owner)) assert task.begin_physical_operation(7) - request = object() + request = Mock(op=transfer_mod.TransferOp.WRITE, remote_name="peer") status = Mock(wait=Mock(return_value=False), is_completed=Mock(return_value=False)) status.last_status_str.return_value = "ERROR" sender = object.__new__(transfer_mod.Sender) @@ -347,7 +347,10 @@ def test_failed_sender_write_retains_roots_until_fatal_deadline() -> None: sender._enforce_physical_ownership = True sender._ownership_poisoned = None sender._ownership_poison_lock = watchdog.lock - sender._agent = Mock(submit_transfer_requests=Mock(return_value=status)) + sender._agent = Mock( + submit_transfer_requests=Mock(return_value=status), + check_remote_descs=Mock(return_value=True), + ) assert sender._submit_transfer(task, 7, request) == (False, "ERROR") operation = task._physical_operations[7] @@ -360,7 +363,9 @@ def test_failed_sender_write_retains_roots_until_fatal_deadline() -> None: later = transfer_mod.SendTaskBase(DisaggregatedParams(disagg_request_id=13)) assert later.begin_physical_operation(8) with pytest.raises(transfer_mod._TransferNotSubmittedError, match="before backend submission"): - sender._submit_transfer(later, 8, object()) + sender._submit_transfer( + later, 8, Mock(op=transfer_mod.TransferOp.WRITE, remote_name="peer") + ) sender._agent.submit_transfer_requests.assert_called_once_with(request) assert later._physical_operations[8].state is transfer_mod._PhysicalOperationState.NOT_SUBMITTED assert later.resources_drained diff --git a/tests/unittest/disaggregated/test_retirement_deadline.py b/tests/unittest/disaggregated/test_retirement_deadline.py index 44375e8db8b2..61625a97caba 100644 --- a/tests/unittest/disaggregated/test_retirement_deadline.py +++ b/tests/unittest/disaggregated/test_retirement_deadline.py @@ -429,7 +429,7 @@ def test_blocked_backend_cannot_delay_fatal_deadline(blocked_call: str) -> None: task = transfer_mod.SendTaskBase(DisaggregatedParams(disagg_request_id=4)) task.bind_logical_outcomes(transfer_mod._LogicalOutcomes(retirement)) assert task.begin_physical_operation(7) - request = object() + request = _write_request() status = Mock(wait=Mock(return_value=True), is_completed=Mock(return_value=True)) def block() -> object: @@ -444,7 +444,10 @@ def block() -> object: sender._enforce_physical_ownership = True sender._ownership_poisoned = None sender._ownership_poison_lock = watchdog.lock - sender._agent = Mock(submit_transfer_requests=Mock(return_value=status)) + sender._agent = Mock( + submit_transfer_requests=Mock(return_value=status), + check_remote_descs=Mock(return_value=True), + ) if blocked_call == "query": task.begin_backend_submission(7, request) task.record_backend_submission(7, status) @@ -608,6 +611,11 @@ def _rx_session(receiver: transfer_mod.Receiver) -> transfer_mod.RxSession: ) +def _write_request() -> Mock: + """A NIXL WRITE request whose remote descriptors pass the sender preflight.""" + return Mock(op=transfer_mod.TransferOp.WRITE, remote_name="peer") + + def _tx_session(watchdog: RetirementWatchdog, *, need_aux: bool = True) -> transfer_mod.TxSession: """Use real session/admission logic with only queue and backend boundary doubles.""" sender = object.__new__(transfer_mod.Sender) @@ -625,6 +633,7 @@ def _tx_session(watchdog: RetirementWatchdog, *, need_aux: bool = True) -> trans sender._peer_requests_timestamps = {} sender.dispatch_task = Mock() sender._agent = Mock() + sender._agent.check_remote_descs.return_value = True sender._agent.submit_transfer_requests.return_value.wait.return_value = True params = DisaggregatedParams(disagg_request_id=4) if need_aux: @@ -637,7 +646,7 @@ def _tx_session(watchdog: RetirementWatchdog, *, need_aux: bool = True) -> trans def _submit_piece(session: transfer_mod.TxSession, task: transfer_mod.SendTaskBase) -> None: """Drive the real sender's submission, backend-DONE and delivery transitions.""" assert session._sender._begin_task_operation(task, 7) - assert session._sender._submit_transfer(task, 7, object()) == (True, None) + assert session._sender._submit_transfer(task, 7, _write_request()) == (True, None) task.complete() @@ -664,12 +673,12 @@ def test_idle_session_keeps_original_deadline_for_later_piece(piece: str, submit assert session._sender._begin_task_operation(later, 7) clock.return_value = submit_at if submit_at < 15.0: - assert session._sender._submit_transfer(later, 7, object()) == (True, None) + assert session._sender._submit_transfer(later, 7, _write_request()) == (True, None) later.complete() assert session.is_completed() else: with pytest.raises(transfer_mod._TransferNotSubmittedError): - session._sender._submit_transfer(later, 7, object()) + session._sender._submit_transfer(later, 7, _write_request()) assert ( later._physical_operations[7].state is transfer_mod._PhysicalOperationState.NOT_SUBMITTED @@ -889,7 +898,7 @@ def test_completion_at_transfer_deadline_cannot_replace_timeout() -> None: _submit_piece(session, session.kv_tasks[0]) task = session.send_aux() assert session._sender._begin_task_operation(task, 7) - assert session._sender._submit_transfer(task, 7, object()) == (True, None) + assert session._sender._submit_transfer(task, 7, _write_request()) == (True, None) clock.return_value = 15.0 task.complete() assert task.logical_outcome.status is SessionStatus.ERROR @@ -912,7 +921,7 @@ def test_expired_fanout_rejects_queued_peer_but_retains_submitted_sibling() -> N assert session._sender._begin_task_operation(task, 8) clock.return_value = 15.0 with pytest.raises(transfer_mod._TransferNotSubmittedError): - session._sender._submit_transfer(task, 8, object()) + session._sender._submit_transfer(task, 8, _write_request()) session._sender._agent.submit_transfer_requests.assert_not_called() assert task._physical_operations[8].state is transfer_mod._PhysicalOperationState.NOT_SUBMITTED assert list(session._retirement._claims.values()) == [task._physical_operations[7]] diff --git a/tests/unittest/disaggregated/test_transfer_late_settlement.py b/tests/unittest/disaggregated/test_transfer_late_settlement.py index a1321ae6cfab..ab788fb12dc5 100644 --- a/tests/unittest/disaggregated/test_transfer_late_settlement.py +++ b/tests/unittest/disaggregated/test_transfer_late_settlement.py @@ -243,8 +243,12 @@ def test_sender_late_done_reports_physical_failure_only( task.expected_transfers = 1 status = Mock(wait=Mock(return_value=False), is_completed=Mock(return_value=False)) status.last_status_str.return_value = "ERROR" - sender._agent = SimpleNamespace(name="nixl", submit_transfer_requests=lambda request: status) - request = SimpleNamespace(op="WRITE", remote_name="gen") + sender._agent = SimpleNamespace( + name="nixl", + submit_transfer_requests=lambda request: status, + check_remote_descs=lambda _name, _descs: True, + ) + request = SimpleNamespace(op="WRITE", remote_name="gen", dst_descs=None) monkeypatch.setattr(transfer_mod.Sender, "_make_agent_request", Mock(return_value=request)) dealer = Mock() sender._get_result_dealer = Mock(return_value=dealer) diff --git a/tests/unittest/disaggregated/test_transfer_ownership_regressions.py b/tests/unittest/disaggregated/test_transfer_ownership_regressions.py index 8fecebd650a1..3d52376b9d14 100644 --- a/tests/unittest/disaggregated/test_transfer_ownership_regressions.py +++ b/tests/unittest/disaggregated/test_transfer_ownership_regressions.py @@ -1790,6 +1790,62 @@ def test_kv_build_failure_reports_safe_pre_submission_failure(monkeypatch) -> No assert transfer_mod._AGENT_RESULT_BY_CODE[status_code] == AgentResult.FAILED +@pytest.mark.cpu_only +def test_bounce_rejects_unregistered_final_destination_before_build() -> None: + rid = 104 + peer_rank = 2 + sender = _make_owned_sender() + sender._device_id = 0 + sender._agent = Mock() + sender._agent.check_remote_descs.return_value = False + sender._bounce = Mock() + dealer = Mock() + sender._get_result_dealer = Mock(return_value=dealer) + task = transfer_mod.KVSendTask( + _sole_piece(), + DisaggregatedParams(disagg_request_id=rid), + slice_id=0, + ) + assert task.begin_physical_operation(peer_rank) + sender._sessions = { + rid: SimpleNamespace( + kv_tasks=[task], + lock=threading.Lock(), + status=SessionStatus.READY, + ) + } + write_meta = transfer_mod.WriteMeta( + task=task, + expected_transfers=1, + peer_name="gen2", + peer_rank=peer_rank, + peer_endpoint="receiver", + unique_rid=rid, + src_ptrs=np.array([0x1000], dtype=np.int64), + dst_ptrs=np.array([0x2000], dtype=np.int64), + sizes=np.array([0x100], dtype=np.int64), + dst_device_id=0, + slice_id=0, + receiver_slice_id=7, + is_last_slice=False, + bounce_dst_base=0x3000, + ) + + sender._deliver_kv_to_agent(write_meta) + + peer_name, checked_descs = sender._agent.check_remote_descs.call_args.args + assert peer_name == "gen2" + assert [(desc.addr, desc.len, desc.device_id) for desc in checked_descs.descs] == [ + (0x2000, 0x100, 0) + ] + sender._bounce.build_request.assert_not_called() + sender._agent.submit_transfer_requests.assert_not_called() + assert task.resources_drained + assert task.status == transfer_mod.TaskStatus.ERROR + status_code = transfer_mod._KV_RESULT_PREFIX.unpack(dealer.send.call_args.args[0][1])[4] + assert transfer_mod._AGENT_RESULT_BY_CODE[status_code] == AgentResult.FAILED + + @pytest.mark.cpu_only def test_owned_kv_build_failure_settles_every_published_peer(monkeypatch) -> None: rid = 100 @@ -1901,6 +1957,7 @@ def submit(_request): sender._agent = SimpleNamespace( name="nixl", submit_transfer_requests=submit, + check_remote_descs=lambda _name, _descs: True, ) dealer = Mock() sender._get_result_dealer = Mock(return_value=dealer) @@ -2245,6 +2302,59 @@ def test_sender_physical_operation_rejects_unproven_retirement(submitted: bool) task.retire_backend_done_physical_operation(peer_rank) +@pytest.mark.cpu_only +@pytest.mark.parametrize("lookup_raises", [False, True]) +def test_rejected_descriptor_does_not_poison_later_write(lookup_raises: bool) -> None: + rejected = transfer_mod.SendTaskBase(DisaggregatedParams(disagg_request_id=12)) + assert rejected.begin_physical_operation(7) + sender = _make_owned_sender() + sender._shutdown = True + status = Mock(wait=Mock(return_value=True)) + sender._agent = Mock(submit_transfer_requests=Mock(return_value=status)) + sender._agent.check_remote_descs.side_effect = [ + RuntimeError("metadata lookup failed") if lookup_raises else False, + True, + ] + request = Mock(op=transfer_mod.TransferOp.WRITE, remote_name="peer") + with pytest.raises(transfer_mod._TransferNotSubmittedError): + sender._submit_transfer(rejected, 7, request) + sender._agent.check_remote_descs.assert_called_once_with("peer", request.dst_descs) + sender._agent.submit_transfer_requests.assert_not_called() + operation = rejected._physical_operations[7] + assert operation.state is transfer_mod._PhysicalOperationState.NOT_SUBMITTED + assert operation.request is None + assert operation.status is None + assert rejected.resources_drained + assert sender._ownership_poisoned is None + + later = transfer_mod.SendTaskBase(DisaggregatedParams(disagg_request_id=13)) + assert later.begin_physical_operation(8) + valid = Mock(op=transfer_mod.TransferOp.WRITE, remote_name="peer") + assert sender._submit_transfer(later, 8, valid) == (True, None) + sender._agent.submit_transfer_requests.assert_called_once_with(valid) + assert later._physical_operations[8].state is transfer_mod._PhysicalOperationState.BACKEND_DONE + assert later.resources_drained + assert sender._ownership_poisoned is None + + +@pytest.mark.cpu_only +def test_submission_exception_after_preflight_remains_in_doubt() -> None: + task = transfer_mod.SendTaskBase(DisaggregatedParams(disagg_request_id=12)) + assert task.begin_physical_operation(7) + sender = _make_owned_sender() + sender._shutdown = True + error = RuntimeError("submission outcome unknown") + sender._agent = Mock( + check_remote_descs=Mock(return_value=True), + submit_transfer_requests=Mock(side_effect=error), + ) + request = Mock(op=transfer_mod.TransferOp.WRITE, remote_name="peer") + assert sender._submit_transfer(task, 7, request) == (False, str(error)) + assert sender._ownership_poisoned is error + assert task._physical_operations[7].state is transfer_mod._PhysicalOperationState.IN_DOUBT + assert not task.resources_drained + + @pytest.mark.cpu_only def test_sender_retires_only_after_backend_done() -> None: sender = _make_owned_sender() From 11905318811fe44f059c46016b364d5e90c34f79 Mon Sep 17 00:00:00 2001 From: xingfei xi <95731198+xxi-nv@users.noreply.github.com> Date: Wed, 21 Oct 2026 13:17:36 +0800 Subject: [PATCH 18/23] [https://nvbugs/6384589][fix] Cherry-pick #20531 from release/1.3: Mount a 640G /dev/shm for online-EPLB perf cases Forward-port of release/1.3 PR https://github.com/NVIDIA/TensorRT-LLM/pull/20531 (release/1.3 commit 240195683da269a78d0957e3b2fafafd4d9c6e13) to main. Adapted to main: the new tests reuse the submit-module fixtures and the _fake_sinfo helper that main already defines in test_perf_submit.py. Signed-off-by: xxi Signed-off-by: Simeng Liu --- jenkins/scripts/perf/benchmark_utils.py | 112 ++++++++++ jenkins/scripts/perf/local/submit.py | 10 +- jenkins/scripts/perf/submit.py | 8 + tests/unittest/scripts/test_perf_submit.py | 247 +++++++++++++++++++++ 4 files changed, 375 insertions(+), 2 deletions(-) diff --git a/jenkins/scripts/perf/benchmark_utils.py b/jenkins/scripts/perf/benchmark_utils.py index b7ffa13c5aa5..06395bdbe334 100644 --- a/jenkins/scripts/perf/benchmark_utils.py +++ b/jenkins/scripts/perf/benchmark_utils.py @@ -14,6 +14,11 @@ # See the License for the specific language governing permissions and # limitations under the License. +import os +from typing import Iterator, Optional + +import yaml + GEN_ONLY_NO_CONTEXT_MODE = "gen_only_no_context" DISAGG_BENCHMARK_MODES = ("e2e", "gen_only") @@ -194,3 +199,110 @@ def parse_positive_concurrency(value: object) -> int: if concurrency <= 0: raise ValueError(f"benchmark.concurrency_list must be a positive integer, got {value!r}") return concurrency + + +# Online EPLB stages a host copy of every MoE layer's expert weights in POSIX +# shared memory on each node, so the container's /dev/shm has to hold all of +# them: ~532 GiB for Kimi-K2.5 NVFP4, ~380 GiB for GLM-5 NVFP4. /dev/shm +# defaults to half of host RAM (480 GiB on a GB300 node). +EPLB_SHM_TMPFS_MOUNT = "tmpfs:/dev/shm:size=640G" + + +def _iter_worker_llm_configs(config: Optional[dict]) -> Iterator[dict]: + """Yield every per-server LLM API config dict in a perf-sanity YAML. + + Disaggregated YAMLs keep them under ``worker_config.{ctx,gen}``; aggregated + YAMLs keep one per ``server_configs`` entry. + """ + for role_config in ((config or {}).get("worker_config") or {}).values(): + if isinstance(role_config, dict): + yield role_config + for server_config in (config or {}).get("server_configs") or []: + if isinstance(server_config, dict): + yield server_config + + +def _load_yaml(path: str, llm_src: str) -> object: + """Load a YAML file; a relative ``path`` is resolved against ``llm_src``.""" + if not os.path.isabs(path): + path = os.path.join(llm_src, path) + with open(path, "r") as f: + return yaml.safe_load(f) + + +def _merged_llm_api_config(server_config: dict, llm_src: str) -> dict: + """Return the LLM API config one server is launched with. + + The file at ``extra_llm_api_config_path`` is the base and the inline fields + replace its top-level keys, matching ``ServerConfig.merged_llm_api_config_data`` + in ``tests/integration/defs/perf/test_perf_sanity.py``. + """ + extra_config_path = server_config.get("extra_llm_api_config_path") + if not extra_config_path: + return server_config + extra_config = _load_yaml(extra_config_path, llm_src) or {} + if not isinstance(extra_config, dict): + raise ValueError( + f"extra_llm_api_config_path {extra_config_path} does not contain a YAML mapping" + ) + return {**extra_config, **server_config} + + +def uses_online_eplb(config: Optional[dict], llm_src: str = "") -> bool: + """Return True when any server in the YAML enables online EPLB. + + Online EPLB (``moe_config.load_balancer.layer_updates_per_iter > 0``) is the + only EPLB mode that keeps expert weights in host shared memory. + ``moe_config`` is read from the config each server is launched with, so a + load balancer set only in its ``extra_llm_api_config_path`` file counts. + ``load_balancer`` is either an inline dict or a path to an EPLB YAML. + Relative paths are resolved against ``llm_src``. + + Args: + config: The parsed perf-sanity config YAML. + llm_src: TensorRT-LLM source root used to resolve relative paths. + + Returns: + bool + """ + for server_config in _iter_worker_llm_configs(config): + llm_config = _merged_llm_api_config(server_config, llm_src) + load_balancer = (llm_config.get("moe_config") or {}).get("load_balancer") + if isinstance(load_balancer, str) and load_balancer: + load_balancer = _load_yaml(load_balancer, llm_src) + if ( + isinstance(load_balancer, dict) + and (load_balancer.get("layer_updates_per_iter") or 0) > 0 + ): + return True + return False + + +def _mount_target(entry: str) -> str: + """Return the container path of one ``--container-mounts`` entry. + + pyxis mounts a source-only entry such as ``/dev/shm`` on the same path. + """ + parts = entry.split(":") + target = parts[1] if len(parts) > 1 else parts[0] + return target.rstrip("/") or "/" + + +def add_eplb_shm_mount(mounts: str, config: Optional[dict], llm_src: str = "") -> str: + """Append the large /dev/shm tmpfs mount when the YAML uses online EPLB. + + Args: + mounts: Comma-separated ``--container-mounts`` value (may be empty). + config: The parsed perf-sanity config YAML. + llm_src: TensorRT-LLM source root used to resolve relative paths. + + Returns: + str: ``mounts`` unchanged, or with ``EPLB_SHM_TMPFS_MOUNT`` appended. A + caller-provided mount on ``/dev/shm`` is kept as is. + """ + if not uses_online_eplb(config, llm_src): + return mounts + entries = [entry for entry in mounts.split(",") if entry] + if any(_mount_target(entry) == "/dev/shm" for entry in entries): + return mounts + return ",".join(entries + [EPLB_SHM_TMPFS_MOUNT]) diff --git a/jenkins/scripts/perf/local/submit.py b/jenkins/scripts/perf/local/submit.py index 3ab46aab3a88..2d496adc8e4f 100755 --- a/jenkins/scripts/perf/local/submit.py +++ b/jenkins/scripts/perf/local/submit.py @@ -32,6 +32,7 @@ DISAGG_BENCHMARK_MODES, DISAGG_CONFIG_MODES, GEN_ONLY_NO_CONTEXT_MODE, + add_eplb_shm_mount, ctx_only_world_size, dwdp_ctx_layout, dwdp_ctx_size, @@ -573,7 +574,9 @@ def generate_sbatch_params(args, hardware_config, work_dir): return lines -def generate_srun_args(args, runtime_mode, timestamp, llm_src="", hardware_config=None): +def generate_srun_args( + args, runtime_mode, timestamp, llm_src="", hardware_config=None, config=None +): """Generate srun arguments.""" is_aggr = runtime_mode == "aggregated" container_name = f"{'aggr' if is_aggr else 'disagg'}_test-{timestamp}" @@ -591,6 +594,7 @@ def generate_srun_args(args, runtime_mode, timestamp, llm_src="", hardware_confi if llm_src and llm_src not in mounts: llm_src_mount = f"{llm_src}:{llm_src}" mounts = f"{mounts},{llm_src_mount}" if mounts else llm_src_mount + mounts = add_eplb_shm_mount(mounts, config, llm_src) if mounts: lines.append(f"--container-mounts={mounts}") @@ -961,7 +965,9 @@ def main(): sbatch_lines = generate_sbatch_params(args, hardware_config, work_dir) # Generate srun args - srun_args_lines = generate_srun_args(args, runtime_mode, timestamp, llm_src, hardware_config) + srun_args_lines = generate_srun_args( + args, runtime_mode, timestamp, llm_src, hardware_config, config + ) # Generate pytest command pytest_command, test_list_content, test_list_path = generate_pytest_command( diff --git a/jenkins/scripts/perf/submit.py b/jenkins/scripts/perf/submit.py index 9005a24bef54..ed58bfc28668 100755 --- a/jenkins/scripts/perf/submit.py +++ b/jenkins/scripts/perf/submit.py @@ -53,6 +53,7 @@ AGGREGATED_DISAGG_YAML_MODES, DISAGG_BENCHMARK_MODES, GEN_ONLY_NO_CONTEXT_MODE, + add_eplb_shm_mount, ctx_only_world_size, dwdp_ctx_layout, dwdp_ctx_size, @@ -949,6 +950,13 @@ def main(): with open(args.srun_args, "r") as f: srun_args_content = f.read() srun_args_lines = srun_args_content.split() + mounts_prefix = "--container-mounts=" + srun_args_lines = [ + mounts_prefix + add_eplb_shm_mount(line[len(mounts_prefix) :], config, args.llm_src) + if line.startswith(mounts_prefix) + else line + for line in srun_args_lines + ] ( agg_pytest_command, diff --git a/tests/unittest/scripts/test_perf_submit.py b/tests/unittest/scripts/test_perf_submit.py index a524301a7155..4b56aff18ba7 100644 --- a/tests/unittest/scripts/test_perf_submit.py +++ b/tests/unittest/scripts/test_perf_submit.py @@ -585,3 +585,250 @@ def test_ctx_only_without_dwdp_sizes_from_ctx_parallelism(submit_module: ModuleT hardware = _aggregated_hardware(submit_module, config, "ctx_only") assert hardware["gpus_per_server"] == 8 assert hardware["total_nodes"] == 2 + + +# ------------------------- online-EPLB /dev/shm mount ---------------------------- +BENCHMARK_UTILS_PATH = REPO_ROOT / "jenkins" / "scripts" / "perf" / "benchmark_utils.py" +KIMI_K25_EPLB416_CONFIG = ( + "tests/scripts/perf-sanity/disaggregated/" + "gb300_kimi-k25-thinking-fp4_8k1k_con1024_ctx1_dep4_gen1_dep32_eplb416_mtp3_ccb-NIXL.yaml" +) +DSV4_STATIC_EPLB_CONFIG = ( + "tests/scripts/perf-sanity/disaggregated/" + "gb300_deepseek-v4-pro-fp4_8k1k_con180_ctx3_dep4_gen1_dep32_eplb384_mtp3_ccb-NIXL.yaml" +) + + +@pytest.fixture +def benchmark_utils_module(monkeypatch: pytest.MonkeyPatch) -> ModuleType: + return _load_module(BENCHMARK_UTILS_PATH, monkeypatch) + + +def _load_repo_yaml(relative_path: str) -> dict: + import yaml + + with open(REPO_ROOT / relative_path) as f: + return yaml.safe_load(f) + + +def _lb_config(load_balancer, role="gen") -> dict: + return {"worker_config": {role: {"moe_config": {"load_balancer": load_balancer}}}} + + +@pytest.mark.parametrize( + "config, expected", + [ + (_lb_config({"num_slots": 416, "layer_updates_per_iter": 1}), True), + (_lb_config({"num_slots": 416, "layer_updates_per_iter": 1}, role="ctx"), True), + (_lb_config({"num_slots": 384, "layer_updates_per_iter": 0}), False), + (_lb_config({"num_slots": 384}), False), + (_lb_config(None), False), + ({"worker_config": {"gen": {"moe_config": None}}}, False), + ({"worker_config": {"gen": None}}, False), + ( + { + "server_configs": [ + {"moe_config": {"load_balancer": {"layer_updates_per_iter": 0}}}, + {"moe_config": {"load_balancer": {"layer_updates_per_iter": 2}}}, + ] + }, + True, + ), + ({}, False), + (None, False), + ], +) +def test_uses_online_eplb(benchmark_utils_module: ModuleType, config, expected): + assert benchmark_utils_module.uses_online_eplb(config) is expected + + +def test_uses_online_eplb_reads_load_balancer_file( + benchmark_utils_module: ModuleType, tmp_path: Path +): + (tmp_path / "online.yaml").write_text("num_slots: 288\nlayer_updates_per_iter: 1\n") + (tmp_path / "static.yaml").write_text("num_slots: 288\nlayer_updates_per_iter: 0\n") + assert benchmark_utils_module.uses_online_eplb(_lb_config("online.yaml"), str(tmp_path)) + assert not benchmark_utils_module.uses_online_eplb(_lb_config("static.yaml"), str(tmp_path)) + assert benchmark_utils_module.uses_online_eplb( + _lb_config(str(tmp_path / "online.yaml")), "/nonexistent" + ) + + +@pytest.mark.parametrize( + "server_config, expected", + [ + ({"extra_llm_api_config_path": "online_extra.yaml"}, True), + ({"extra_llm_api_config_path": "lb_file_extra.yaml"}, True), + ({"extra_llm_api_config_path": "static_extra.yaml"}, False), + ({"extra_llm_api_config_path": "empty_extra.yaml"}, False), + # Inline fields replace the file's top-level keys, so this moe_config wins. + ( + { + "extra_llm_api_config_path": "online_extra.yaml", + "moe_config": {"backend": "CUTEDSL"}, + }, + False, + ), + ], +) +def test_uses_online_eplb_reads_extra_llm_api_config( + benchmark_utils_module: ModuleType, tmp_path: Path, server_config, expected +): + (tmp_path / "online_extra.yaml").write_text( + "moe_config:\n load_balancer:\n num_slots: 288\n layer_updates_per_iter: 1\n" + ) + (tmp_path / "static_extra.yaml").write_text( + "moe_config:\n load_balancer:\n num_slots: 288\n layer_updates_per_iter: 0\n" + ) + (tmp_path / "lb_file_extra.yaml").write_text("moe_config:\n load_balancer: online.yaml\n") + (tmp_path / "online.yaml").write_text("num_slots: 288\nlayer_updates_per_iter: 1\n") + (tmp_path / "empty_extra.yaml").write_text("") + disagg_config = {"worker_config": {"gen": server_config}} + aggr_config = {"server_configs": [server_config]} + assert benchmark_utils_module.uses_online_eplb(disagg_config, str(tmp_path)) is expected + assert benchmark_utils_module.uses_online_eplb(aggr_config, str(tmp_path)) is expected + + +def test_uses_online_eplb_rejects_non_mapping_extra_llm_api_config( + benchmark_utils_module: ModuleType, tmp_path: Path +): + (tmp_path / "list_extra.yaml").write_text("- moe_config\n") + config = {"worker_config": {"gen": {"extra_llm_api_config_path": "list_extra.yaml"}}} + with pytest.raises(ValueError, match="does not contain a YAML mapping"): + benchmark_utils_module.uses_online_eplb(config, str(tmp_path)) + + +def test_uses_online_eplb_on_checked_in_configs(benchmark_utils_module: ModuleType): + assert benchmark_utils_module.uses_online_eplb( + _load_repo_yaml(KIMI_K25_EPLB416_CONFIG), str(REPO_ROOT) + ) + assert not benchmark_utils_module.uses_online_eplb( + _load_repo_yaml(DSV4_STATIC_EPLB_CONFIG), str(REPO_ROOT) + ) + + +@pytest.mark.parametrize( + "mounts, expected_suffix", + [ + ("", "tmpfs:/dev/shm:size=640G"), + ("/a:/a", "/a:/a,tmpfs:/dev/shm:size=640G"), + ("/a:/a,", "/a:/a,tmpfs:/dev/shm:size=640G"), + ("/x:/dev/shm/sub", "/x:/dev/shm/sub,tmpfs:/dev/shm:size=640G"), + ("/dev/shm:/host_shm", "/dev/shm:/host_shm,tmpfs:/dev/shm:size=640G"), + ], +) +def test_add_eplb_shm_mount_appends_for_online_eplb( + benchmark_utils_module: ModuleType, mounts, expected_suffix +): + config = _lb_config({"num_slots": 416, "layer_updates_per_iter": 1}) + assert benchmark_utils_module.add_eplb_shm_mount(mounts, config) == expected_suffix + + +@pytest.mark.parametrize( + "mounts", + [ + "/a:/a,tmpfs:/dev/shm:size=900G", + "/dev/shm", + "/a:/a,/dev/shm", + "/host/shm:/dev/shm/", + ], +) +def test_add_eplb_shm_mount_keeps_caller_dev_shm(benchmark_utils_module: ModuleType, mounts): + config = _lb_config({"num_slots": 416, "layer_updates_per_iter": 1}) + assert benchmark_utils_module.add_eplb_shm_mount(mounts, config) == mounts + + +def test_add_eplb_shm_mount_noop_without_online_eplb(benchmark_utils_module: ModuleType): + config = _lb_config({"num_slots": 384, "layer_updates_per_iter": 0}) + assert benchmark_utils_module.add_eplb_shm_mount("/a:/a", config) == "/a:/a" + + +def test_ci_submit_adds_eplb_shm_mount_to_launch_script( + ci_submit_module: ModuleType, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +): + test_id = ( + "perf/test_perf_sanity.py::test_e2e[disagg-gen_only-" + "gb300_kimi-k25-thinking-fp4_8k1k_con1024_ctx1_dep4_gen1_dep32_eplb416_mtp3_ccb-NIXL]" + ) + test_list = tmp_path / "test_list.txt" + test_list.write_text(test_id + "\n") + prefix = tmp_path / "prefix.sh" + prefix.write_text( + "#!/bin/bash\n" + "export jobWorkspace=/ws\n" + 'export pytestCommand="LLM_ROOT=/src LLM_MODELS_ROOT=/models ' + "/src/tensorrt_llm/llmapi/trtllm-llmapi-launch pytest perf/test_perf_sanity.py " + '--test-prefix=/src/tests/integration/defs --test-list=/ws/test_list.txt --output-dir=/ws"\n' + ) + srun_args = tmp_path / "srun_args.txt" + srun_args.write_text("--container-image=/img.sqsh --container-mounts=/a:/a,/b:/b\n") + launch = tmp_path / "launch.sh" + draft = REPO_ROOT / "jenkins" / "scripts" / "perf" / "disaggregated" / "slurm_launch_draft.sh" + monkeypatch.setattr( + "sys.argv", + [ + "submit.py", + "--draft-launch-sh", + str(draft), + "--launch-sh", + str(launch), + "--run-sh", + "/run.sh", + "--install-sh", + "/install.sh", + "--llm-src", + str(REPO_ROOT), + "--test-list", + str(test_list), + "--script-prefix", + str(prefix), + "--srun-args", + str(srun_args), + "--stage-name", + "GB300-8_GPUs-2_Nodes-PyTorch-Disagg-PerfSanity-Post-Merge-1", + ], + ) + ci_submit_module.main() + assert '"--container-mounts=/a:/a,/b:/b,tmpfs:/dev/shm:size=640G"' in launch.read_text() + + +def test_local_submit_adds_eplb_shm_mount_to_launch_script( + local_submit_module: ModuleType, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +): + test_id = ( + "perf/test_perf_sanity.py::test_e2e[disagg-e2e-" + "gb300_kimi-k25-thinking-fp4_8k1k_con1024_ctx1_dep4_gen1_dep32_eplb416_mtp3_ccb-NIXL]" + ) + _fake_sinfo(monkeypatch, local_submit_module, "gpu:4\n") + launch = tmp_path / "launch.sh" + monkeypatch.setattr( + "sys.argv", + [ + "submit.py", + "--test-list", + test_id, + "--partition", + "batch", + "--account", + "acct", + "--job-name", + "job", + "--image", + "/img.sqsh", + "--mounts", + "/a:/a", + "--work-dir", + str(tmp_path), + "--launch-sh", + str(launch), + "--llm-src", + str(REPO_ROOT), + "--llm-models-root", + "/models", + "--cluster-name", + "oci-aga", + ], + ) + local_submit_module.main() + content = launch.read_text() + assert f'"--container-mounts=/a:/a,{REPO_ROOT}:{REPO_ROOT},tmpfs:/dev/shm:size=640G"' in content From 5020a21ea5383175631b6d861515075c7db52423 Mon Sep 17 00:00:00 2001 From: Lori Ren Date: Wed, 21 Oct 2026 13:18:48 +0800 Subject: [PATCH 19/23] [https://nvbugs/6944374][fix] Cherry-pick #20657 from release/1.3: Fix rank-asymmetric NCCL window allocation hang during PARALLEL autotuning Forward-port of release/1.3 PR https://github.com/NVIDIA/TensorRT-LLM/pull/20657 (release/1.3 commit 0c88e434d593fd05f82569a16671d92bb80d02b2) to main. Signed-off-by: Lori Ren Signed-off-by: Simeng Liu --- cpp/tensorrt_llm/common/ncclUtils.cpp | 13 ++++ cpp/tensorrt_llm/common/ncclUtils.h | 5 ++ cpp/tensorrt_llm/thop/allreduceOp.cpp | 8 +++ tensorrt_llm/_torch/autotuner.py | 13 +++- tensorrt_llm/_torch/nccl_window_graph.py | 15 ++++ .../test_allocate_output_buffer_kinds.py | 72 +++++++++++++++++++ 6 files changed, 123 insertions(+), 3 deletions(-) diff --git a/cpp/tensorrt_llm/common/ncclUtils.cpp b/cpp/tensorrt_llm/common/ncclUtils.cpp index c908be4fb56d..6718d56007ce 100644 --- a/cpp/tensorrt_llm/common/ncclUtils.cpp +++ b/cpp/tensorrt_llm/common/ncclUtils.cpp @@ -24,6 +24,7 @@ #include #include #include +#include namespace { @@ -398,6 +399,7 @@ size_t NcclCommResourceManager::getResourceCount(ncclComm_t comm) const noexcept namespace { thread_local std::optional gGraphPoolOwner; +thread_local bool gWindowAllocationSuppressed = false; } // namespace NCCLWindowAllocator& NCCLWindowAllocator::getInstance() @@ -417,6 +419,11 @@ void NCCLWindowAllocator::setGraphPoolOwner(int64_t owner) gGraphPoolOwner = owner < 0 ? std::nullopt : std::optional{static_cast(owner)}; } +bool NCCLWindowAllocator::setWindowAllocationSuppressed(bool suppressed) +{ + return std::exchange(gWindowAllocationSuppressed, suppressed); +} + void NCCLWindowAllocator::releaseGraphPoolOwner(int64_t owner) { TLLM_CHECK_WITH_INFO(owner >= 0, "NCCL window graph-pool owner must be non-negative"); @@ -469,6 +476,12 @@ NCCLWindowBuffer NCCLWindowAllocator::requestBuffer(ncclComm_t comm, size_t size { return NCCLWindowBuffer(); } + if (gWindowAllocationSuppressed) + { + TLLM_LOG_TRACE("[NCCLUtil] NCCL window allocation suppressed on this thread for comm %p (requested: %zu)", + static_cast(comm), size); + return NCCLWindowBuffer(); + } TLLM_CHECK_WITH_INFO(comm != nullptr, "NCCL communicator cannot be null"); TLLM_CHECK_WITH_INFO(size > 0, "Buffer size must be greater than 0"); diff --git a/cpp/tensorrt_llm/common/ncclUtils.h b/cpp/tensorrt_llm/common/ncclUtils.h index 76d6b66d4a49..ff96a987030d 100644 --- a/cpp/tensorrt_llm/common/ncclUtils.h +++ b/cpp/tensorrt_llm/common/ncclUtils.h @@ -247,6 +247,11 @@ class NCCLWindowAllocator // still in use are returned when their deterministic tensor scope ends. void releaseGraphPoolOwner(int64_t owner); + // While set on this host thread, requestBuffer returns an invalid buffer without touching the + // pool, so callers fall back to plain memory. Used where ranks are not SPMD, such as PARALLEL + // autotuning. Returns the previous setting. + bool setWindowAllocationSuppressed(bool suppressed); + // Search for a buffer by pointer. Returns an invalid buffer if not found. // This matches the UBManager.search_buffer() interface. NCCLWindowBuffer searchBuffer(ncclComm_t comm, void* ptr) const; diff --git a/cpp/tensorrt_llm/thop/allreduceOp.cpp b/cpp/tensorrt_llm/thop/allreduceOp.cpp index 85fbd59793d5..157a3cf46179 100644 --- a/cpp/tensorrt_llm/thop/allreduceOp.cpp +++ b/cpp/tensorrt_llm/thop/allreduceOp.cpp @@ -2609,6 +2609,7 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m) "end_nccl_window_tensor_scope.tracked(Tensor[] inputs, Tensor[] outputs, bool failed, int entry_depth) -> ()"); m.def("set_nccl_window_graph_owner(int owner) -> ()"); m.def("release_nccl_window_graph_owner(int owner) -> ()"); + m.def("set_nccl_window_allocation_suppressed(bool suppressed) -> bool"); m.def("get_nccl_window_buffer_release_epoch() -> int"); m.def("promote_nccl_window_buffer_releases(int device, int release_epoch) -> ()"); m.def( @@ -2671,6 +2672,12 @@ TORCH_LIBRARY_IMPL(trtllm, CatchAll, m) m.impl("release_nccl_window_graph_owner", [](int64_t owner) { tensorrt_llm::common::nccl_util::NCCLWindowAllocator::getInstance().releaseGraphPoolOwner(owner); }); + m.impl("set_nccl_window_allocation_suppressed", + [](bool suppressed) + { + return tensorrt_llm::common::nccl_util::NCCLWindowAllocator::getInstance().setWindowAllocationSuppressed( + suppressed); + }); m.impl("get_nccl_window_buffer_release_epoch", []() { @@ -2687,6 +2694,7 @@ TORCH_LIBRARY_IMPL(trtllm, CatchAll, m) #else m.impl("set_nccl_window_graph_owner", [](int64_t) {}); m.impl("release_nccl_window_graph_owner", [](int64_t) {}); + m.impl("set_nccl_window_allocation_suppressed", [](bool) { return false; }); m.impl("get_nccl_window_buffer_release_epoch", []() { return int64_t{0}; }); m.impl("promote_nccl_window_buffer_releases", [](int64_t, int64_t) {}); #endif diff --git a/tensorrt_llm/_torch/autotuner.py b/tensorrt_llm/_torch/autotuner.py index f55006af4213..dd1bc070a97d 100644 --- a/tensorrt_llm/_torch/autotuner.py +++ b/tensorrt_llm/_torch/autotuner.py @@ -34,8 +34,8 @@ import tensorrt_llm from tensorrt_llm._torch.distributed import Distributed from tensorrt_llm._torch.nccl_window_graph import ( - abandon_nccl_window_graph_owner, nccl_window_graph_capture, - release_nccl_window_graph_owner) + abandon_nccl_window_graph_owner, nccl_window_allocation_suppressed, + nccl_window_graph_capture, release_nccl_window_graph_owner) from tensorrt_llm._torch.nccl_window_tensor_scope import \ nccl_window_tensor_scope from tensorrt_llm._utils import confidential_compute_enabled, nvtx_range @@ -1332,7 +1332,14 @@ def _profile_runners( tuning_config: TuningConfig, **kwargs, ) -> float: - with nccl_window_tensor_scope(input_tensors) as scope: + # PARALLEL ranks profile different tactics, so window requests here would diverge the + # per-rank pools and let a later request register a window collectively on some ranks only. + window_guard = (nccl_window_allocation_suppressed() + if tuning_config.distributed_tuning_strategy + == DistributedTuningStrategy.PARALLEL + and self._is_distributed() else + contextlib.nullcontext()) + with window_guard, nccl_window_tensor_scope(input_tensors) as scope: scope.escape(input_tensors) if tuning_config.inputs_pre_hook is not None: input_tensors = tuning_config.inputs_pre_hook(input_tensors) diff --git a/tensorrt_llm/_torch/nccl_window_graph.py b/tensorrt_llm/_torch/nccl_window_graph.py index a5f97ad524bf..9fc416dc7686 100644 --- a/tensorrt_llm/_torch/nccl_window_graph.py +++ b/tensorrt_llm/_torch/nccl_window_graph.py @@ -131,3 +131,18 @@ def release_nccl_window_graph_owner(pool: Any) -> None: if owner is not None: torch.ops.trtllm.release_nccl_window_graph_owner(owner) del _pool_owners[key] + + +@contextlib.contextmanager +def nccl_window_allocation_suppressed() -> Iterator[None]: + """Serve NCCL-window requests on this thread from plain memory. + + The window pool is left untouched, so work that does not run in SPMD order + across ranks cannot make a later request miss on only some ranks and enter + the collective registration alone. + """ + previous = torch.ops.trtllm.set_nccl_window_allocation_suppressed(True) + try: + yield + finally: + torch.ops.trtllm.set_nccl_window_allocation_suppressed(previous) diff --git a/tests/unittest/_torch/multi_gpu/test_allocate_output_buffer_kinds.py b/tests/unittest/_torch/multi_gpu/test_allocate_output_buffer_kinds.py index 5eb80ee2aba5..9baecff37331 100644 --- a/tests/unittest/_torch/multi_gpu/test_allocate_output_buffer_kinds.py +++ b/tests/unittest/_torch/multi_gpu/test_allocate_output_buffer_kinds.py @@ -29,6 +29,13 @@ import tensorrt_llm import tensorrt_llm.bindings.internal.userbuffers as ub +from tensorrt_llm._torch.autotuner import ( + AutoTuner, + DistributedTuningStrategy, + TunableRunner, + TuningConfig, + autotune, +) from tensorrt_llm._torch.distributed import ( AllReduce, AllReduceFusionOp, @@ -222,6 +229,56 @@ def _allocate(ref, kind, tp_group): return True +def _run_parallel_tuning_then_window_allreduce(tp_size, tp_rank, seq_len, hidden_size): + """PARALLEL-tune a window-output runner, then all-reduce in a fresh window on every rank.""" + from tensorrt_llm.bindings.internal.thop import BufferKind + + def _allocate(ref, tp_group): + return torch.ops.trtllm.allocate_output(ref, int(BufferKind.NCCL_WINDOW), tp_group) + + class WindowOutputRunner(TunableRunner): + def __init__(self, tp_group): + self.tp_group = tp_group + + def get_valid_tactics(self, inputs, profile, **kwargs): + return [0] + + def forward(self, inputs, *, tactic=-1, **kwargs): + out, _ = _allocate(inputs[0], self.tp_group) + return out.copy_(inputs[0]) + + mapping = Mapping(world_size=tp_size, tp_size=tp_size, rank=tp_rank) + ref = torch.ones(seq_len, hidden_size, dtype=torch.bfloat16, device="cuda") + with discard_nccl_window_tensor_outputs(ref): + _, actual_kind_int = _allocate(ref, mapping.tp_group) + if actual_kind_int != int(BufferKind.NCCL_WINDOW): + return _SKIP + + tuner = AutoTuner.get() + tuner.clear_cache() + tuner.setup_distributed_state(mapping) + try: + # A single tactic is profiled by rank 0 only. + with autotune(): + tuner.choose_one( + "test_parallel_tuning_then_window_allreduce", + [WindowOutputRunner(mapping.tp_group)], + TuningConfig(distributed_tuning_strategy=DistributedTuningStrategy.PARALLEL), + [ref], + ) + finally: + tuner.setup_distributed_state(Mapping(), None) + + allreduce = AllReduce(mapping=mapping, strategy=AllReduceStrategy.NCCL_SYMMETRIC) + with discard_nccl_window_tensor_outputs(ref): + out, actual_kind_int = _allocate(ref, mapping.tp_group) + assert actual_kind_int == int(BufferKind.NCCL_WINDOW) + out.copy_(ref) + result = allreduce(out, all_reduce_params=AllReduceParams(fusion_op=AllReduceFusionOp.NONE)) + torch.testing.assert_close(result, torch.full_like(result, tp_size), atol=1e-2, rtol=1e-2) + return True + + @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Requires at least 2 GPUs") @pytest.mark.parametrize( "seq_len,hidden_size", @@ -297,3 +354,18 @@ def test_allreduce_nccl_window_buffer(seq_len, hidden_size, mpi_pool_executor): ) for r in results: assert r is True or r == _SKIP, f"Unexpected result: {r}" + + +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Requires at least 2 GPUs") +@pytest.mark.timeout(300) +@pytest.mark.parametrize("mpi_pool_executor", [2], indirect=True) +def test_parallel_tuning_keeps_nccl_windows_symmetric(mpi_pool_executor): + tp_size = mpi_pool_executor.num_workers + results = list( + mpi_pool_executor.map( + run_single_rank, + *zip(*[(tp_size, _run_parallel_tuning_then_window_allreduce, 256, 4096)] * tp_size), + ) + ) + for r in results: + assert r is True or r == _SKIP, f"Unexpected result: {r}" From 11f6764a42f5ded2cd32b7106f8453a606e4d440 Mon Sep 17 00:00:00 2001 From: dominicshanshan <30051912+dominicshanshan@users.noreply.github.com> Date: Wed, 21 Oct 2026 13:43:17 +0800 Subject: [PATCH 20/23] [https://nvbugs/6885370][fix] Cherry-pick #20284 from release/1.3: Keep Nemotron-H SSM cache dtype unset Forward-port of release/1.3 PR https://github.com/NVIDIA/TensorRT-LLM/pull/20284 (release/1.3 commit 0312481b64563cde81c149980987959acf36bd0f) to main. Signed-off-by: Wangshanshan <30051912+dominicshanshan@users.noreply.github.com> Signed-off-by: Simeng Liu --- tensorrt_llm/_torch/configs/nemotron_h.py | 2 +- tensorrt_llm/bench/tuning/heuristics.py | 7 +-- .../_torch/test_custom_config_registration.py | 28 +++++++++++ tests/unittest/others/test_bench.py | 47 +++++++++++++++++++ 4 files changed, 80 insertions(+), 4 deletions(-) diff --git a/tensorrt_llm/_torch/configs/nemotron_h.py b/tensorrt_llm/_torch/configs/nemotron_h.py index 0ac1286aead8..200372a9dc06 100644 --- a/tensorrt_llm/_torch/configs/nemotron_h.py +++ b/tensorrt_llm/_torch/configs/nemotron_h.py @@ -135,7 +135,7 @@ def __init__( mamba_head_dim: int = 64, mamba_hidden_act: str = "silu", mamba_proj_bias: bool = False, - mamba_ssm_cache_dtype: str = "float32", + mamba_ssm_cache_dtype: str | None = None, n_groups: int = 8, conv_kernel: int = 4, expand: int = 2, diff --git a/tensorrt_llm/bench/tuning/heuristics.py b/tensorrt_llm/bench/tuning/heuristics.py index 5d5391679ae5..967bb3b71b7d 100755 --- a/tensorrt_llm/bench/tuning/heuristics.py +++ b/tensorrt_llm/bench/tuning/heuristics.py @@ -122,9 +122,10 @@ def calc_engine_setting( bytes_per_elem = BYTES_PER_ELEM.get(QuantAlgo.NO_QUANT) if isinstance(model_config, (NemotronHybridConfig, Qwen3HybridConfig)): mamba_ssm_cache_dtype = model_config.mamba_ssm_cache_dtype - if mamba_ssm_cache_dtype != "auto": - if str_dtype_to_torch(mamba_ssm_cache_dtype) == torch.float32: - bytes_per_elem = 4.0 + if mamba_ssm_cache_dtype in (None, "auto"): + mamba_ssm_cache_dtype = model_config.dtype + if str_dtype_to_torch(mamba_ssm_cache_dtype) == torch.float32: + bytes_per_elem = 4.0 gb_per_extra_cache = model_config.extra_model_cache_in_gb(bytes_per_elem, target_seq_len) kv_cache_max_requests = cache_memory / (gb_per_token * target_seq_len + gb_per_extra_cache) diff --git a/tests/unittest/_torch/test_custom_config_registration.py b/tests/unittest/_torch/test_custom_config_registration.py index 63563cc6bd5e..19f3567b2692 100644 --- a/tests/unittest/_torch/test_custom_config_registration.py +++ b/tests/unittest/_torch/test_custom_config_registration.py @@ -377,6 +377,34 @@ def test_nemotron_h_dense_pattern_survives_autoconfig(tmp_path): assert cfg.num_nextn_predict_layers == 0 +@pytest.mark.parametrize( + "declared_dtype, expected_dtype", + [ + (None, torch.bfloat16), + ("float32", torch.float32), + ], +) +def test_nemotron_h_dense_ssm_cache_dtype_follows_checkpoint( + tmp_path, declared_dtype, expected_dtype +): + # NemotronHConfig must not invent an SSM cache dtype: without a declared one + # the recurrent-state cache is sized in the model dtype, as with the + # checkpoint's own remote-code config. + from tensorrt_llm._torch.configs import NemotronHConfig + from tensorrt_llm._torch.pyexecutor.config_utils import resolve_auto_ssm_cache_dtype + + config_dict = _nemotron_h_min_config(_NEMOTRON_H_DENSE_PATTERN) + config_dict["torch_dtype"] = "bfloat16" + if declared_dtype is not None: + config_dict["mamba_ssm_cache_dtype"] = declared_dtype + + cfg = AutoConfig.from_pretrained(_write_config(tmp_path, "nemotron_h_ssm", config_dict)) + + assert isinstance(cfg, NemotronHConfig) + assert cfg.mamba_ssm_cache_dtype == declared_dtype + assert resolve_auto_ssm_cache_dtype(cfg, torch.float16) == expected_dtype + + def test_nemotron_h_representable_pattern_uses_native_config(tmp_path): # Patterns the installed transformers can express must keep loading through # the native class, so this shim cannot regress checkpoints that work today. diff --git a/tests/unittest/others/test_bench.py b/tests/unittest/others/test_bench.py index a81652ee535a..090bafa7dbb7 100644 --- a/tests/unittest/others/test_bench.py +++ b/tests/unittest/others/test_bench.py @@ -21,12 +21,15 @@ from tensorrt_llm.bench.dataclasses.configuration import RuntimeConfig from tensorrt_llm.bench.dataclasses.reporting import PerfItemTuple, ReportUtility, StatsKeeper from tensorrt_llm.bench.dataclasses.statistics import PercentileStats +from tensorrt_llm.bench.tuning import heuristics +from tensorrt_llm.bench.tuning.dataclasses import NemotronHybridConfig from tensorrt_llm.bench.utils import VALID_QUANT_ALGOS from tensorrt_llm.bench.utils.data import ( DatasetFormatError, create_dataset_from_stream, initialize_tokenizer, ) +from tensorrt_llm.llmapi.llm_utils import QuantConfig from tensorrt_llm.quantization.mode import QuantAlgo pytestmark = pytest.mark.cpu_only @@ -236,3 +239,47 @@ def test_acceptance_length_weighted_by_iterations(): # (= sum(output_tokens) / total_decoding_iterations). This would fail under # output-token weighting. assert stats.acceptance_length_percentiles.average == pytest.approx(stats.acceptance_length) + + +def _nemotron_nano_12b_v2_bench_config(mamba_ssm_cache_dtype, dtype): + """trtllm-bench model config with the NVIDIA-Nemotron-Nano-12B-v2 geometry.""" + return NemotronHybridConfig( + name="nvidia/NVIDIA-Nemotron-Nano-12B-v2", + model_type="nemotron_h", + param_count=12_310_001_152, + hybrid_override_pattern="M-M-M-M*-M-M-M-M*-M-M-M-M*-M-M-M-M*-M-M-M-M*-M-M-M-M*-M-M-M-M-", + num_hidden_layers=62, + num_attention_heads=40, + num_key_value_heads=8, + hidden_size=5120, + head_dim=128, + ssm_state_size=128, + conv_kernel=4, + mamba_num_heads=128, + mamba_head_dim=80, + n_groups=8, + dtype=dtype, + mamba_ssm_cache_dtype=mamba_ssm_cache_dtype, + ) + + +def test_calc_engine_setting_sizes_unset_ssm_cache_dtype_from_model_dtype(monkeypatch): + """An unset mamba_ssm_cache_dtype is sized like the model dtype.""" + monkeypatch.setattr(heuristics, "get_device_memory", lambda: 80.0) + + def engine_setting(mamba_ssm_cache_dtype, dtype): + return heuristics.calc_engine_setting( + _nemotron_nano_12b_v2_bench_config(mamba_ssm_cache_dtype, dtype), + QuantConfig(), + tp_size=1, + pp_size=1, + target_input_len=128, + target_output_len=128, + ) + + bf16_setting = engine_setting("bfloat16", "bfloat16") + fp32_setting = engine_setting("float32", "float32") + assert fp32_setting[0] < bf16_setting[0] + for unset in (None, "auto"): + assert engine_setting(unset, "bfloat16") == bf16_setting + assert engine_setting(unset, "float32") == fp32_setting From dc6104d10ddf6dc5cdb43b31ea43ce2dbdf2f720 Mon Sep 17 00:00:00 2001 From: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com> Date: Wed, 21 Oct 2026 15:35:11 -0700 Subject: [PATCH 21/23] [https://nvbugs/6887489][fix] Cherry-pick #20651 from release/1.3: Restrict coordinator HTTP cluster storage to loopback Forward-port of release/1.3 PR https://github.com/NVIDIA/TensorRT-LLM/pull/20651 (release/1.3 commit 135d95ad5077ad93381a57e245e42b0d06a7ca05) to main. Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com> Signed-off-by: Simeng Liu --- tensorrt_llm/serve/disagg_coordinator.py | 5 +++++ tests/unittest/disaggregated/test_coordinator_worker.py | 9 +++++++++ 2 files changed, 14 insertions(+) diff --git a/tensorrt_llm/serve/disagg_coordinator.py b/tensorrt_llm/serve/disagg_coordinator.py index b134f616daf3..0b5a006fdec5 100644 --- a/tensorrt_llm/serve/disagg_coordinator.py +++ b/tensorrt_llm/serve/disagg_coordinator.py @@ -57,6 +57,7 @@ ClusterStorage, WatchEventType, create_cluster_storage, + validate_http_cluster_storage_scope, ) from tensorrt_llm.serve.disagg_auto_scaling import DisaggClusterManager, WorkerInfo from tensorrt_llm.serve.metadata_server import create_metadata_server @@ -142,6 +143,10 @@ def __init__( self._config = config self._client_factory = client_factory self._metadata_config = metadata_config + if config.disagg_cluster_config: + validate_http_cluster_storage_scope( + config.disagg_cluster_config.cluster_uri, config.hostname + ) # The coordinator owns routing state, so it builds the owner routers here # (is_delegating_client=False): one shared namespace-aware core + a single # ZMQ ingest server. Sole place owner routers are created (fleet workers diff --git a/tests/unittest/disaggregated/test_coordinator_worker.py b/tests/unittest/disaggregated/test_coordinator_worker.py index 9e19e5c004e9..848fcfa41a7b 100644 --- a/tests/unittest/disaggregated/test_coordinator_worker.py +++ b/tests/unittest/disaggregated/test_coordinator_worker.py @@ -223,6 +223,15 @@ def test_coordinator_rejects_unknown_role(): coordinator._router_for_role("typo") +def test_coordinator_rejects_non_loopback_http_cluster_storage(): + config = _make_config([], [], "round_robin", "round_robin") + config.hostname = "0.0.0.0" + config.disagg_cluster_config = DisaggClusterConfig(cluster_uri="http://127.0.0.1:18000") + + with pytest.raises(ValueError, match="loopback-only"): + DisaggCoordinatorService(config, _client_factory) + + def test_kv_router_rejects_mixed_prepared_hash_algorithms(): router = KvCacheAwareRouter(server_role=ServerRole.CONTEXT, servers=["server-a", "server-b"]) router._prepared_ready_servers.update(router.servers) From ab6d70ed2810969613b6c572010cd3599b889725 Mon Sep 17 00:00:00 2001 From: ruodil <200874449+ruodil@users.noreply.github.com> Date: Thu, 22 Oct 2026 16:28:29 +0800 Subject: [PATCH 22/23] [https://nvbugs/6964933][test] Cherry-pick #20794 from release/1.3: Use FlashInfer for Gemma4 NVFP4 perf on SM107 Forward-port of release/1.3 PR https://github.com/NVIDIA/TensorRT-LLM/pull/20794 (release/1.3 commit 025f9ae470071934a4a87b31e53970c7a3b082f9) to main. Also updates the comment on main's SM121 Gemma 4 override (#20210) to mention the new SM107 override. Signed-off-by: Ruodi Lu Co-authored-by: Ruodi Lu Signed-off-by: Simeng Liu --- .../defs/perf/pytorch_model_config.py | 17 ++++++++++++----- 1 file changed, 12 insertions(+), 5 deletions(-) diff --git a/tests/integration/defs/perf/pytorch_model_config.py b/tests/integration/defs/perf/pytorch_model_config.py index a3763e8dab15..c3295421b3b8 100644 --- a/tests/integration/defs/perf/pytorch_model_config.py +++ b/tests/integration/defs/perf/pytorch_model_config.py @@ -524,9 +524,8 @@ def get_model_yaml_config(model_label: str, } }, # Gemma 4 NVFP4: VSWA (1024-token sliding window) plus per-layer - # head_dim 256/512. The FLASHINFER backend is selected by the model's - # own get_model_defaults(), so it is not repeated here. KV-cache reuse - # is disabled to match the validated accuracy configuration. + # head_dim 256/512. KV-cache reuse is disabled to match the validated + # accuracy configuration. { 'patterns': ['gemma_4_'], 'config': { @@ -947,6 +946,13 @@ def get_model_yaml_config(model_label: str, recursive_update(base_config, pattern_config['config']) break # Stop checking other patterns for this config once we find a match + # Gemma4 NVFP4 needs FlashInfer for its 512-wide attention heads on SM107. + if (model_label.lower().startswith( + ('gemma_4_26b_a4b_nvfp4-', 'gemma_4_31b_it_nvfp4-')) + and '-pytorch-' in model_label.lower() + and _get_sm_version_safe() == 107): + base_config['attn_backend'] = 'FLASHINFER' + # DeepSeek FP8 (block-scale) models on Blackwell (SM100+) must use the # DEEPGEMM MoE backend; the default CUTLASS backend's FP8 block-scale path # JIT-compiles DeepGEMM kernels that only target Hopper (SM90). @@ -957,8 +963,9 @@ def get_model_yaml_config(model_label: str, moe_config = base_config.setdefault('moe_config', {}) moe_config.setdefault('backend', 'DEEPGEMM') - # Gemma 4 on SM121 (GB10/Spark) needs FLASHINFER set explicitly; other - # GPUs keep the model's own get_model_defaults() choice. + # Gemma 4 on SM121 (GB10/Spark) needs FLASHINFER set explicitly. Gemma 4 + # NVFP4 on SM107 is forced to FLASHINFER above; other GPUs keep the + # model's own get_model_defaults() choice. if 'gemma_4_' in model_label.lower() and _get_sm_version_safe() == 121: base_config['attn_backend'] = 'FLASHINFER' From 034f5ddf5d554203ddcddcd23b0b4baf89d59ed3 Mon Sep 17 00:00:00 2001 From: Mgluhovskoi Date: Thu, 22 Oct 2026 19:06:52 +0200 Subject: [PATCH 23/23] [None][fix] Cherry-pick #20777 from release/1.3: Validate telemetry deployment IDs as UUIDs Forward-port of release/1.3 PR https://github.com/NVIDIA/TensorRT-LLM/pull/20777 (release/1.3 commit 6c172f3aa574d951babe274797ffd44ed77e823f) to main. Adapted to main: the new payload test patches _HEARTBEAT_STOP, which main renamed from _REPORTER_STOP. Signed-off-by: Maxim Gluhovskoi Signed-off-by: Simeng Liu --- tensorrt_llm/usage/schemas/README.md | 14 +++-- tensorrt_llm/usage/usage_lib.py | 14 +++-- tests/unittest/usage/test_reporter.py | 74 +++++++++++++++++++-------- 3 files changed, 72 insertions(+), 30 deletions(-) diff --git a/tensorrt_llm/usage/schemas/README.md b/tensorrt_llm/usage/schemas/README.md index c8a939877c53..715da3f9eb12 100644 --- a/tensorrt_llm/usage/schemas/README.md +++ b/tensorrt_llm/usage/schemas/README.md @@ -109,7 +109,7 @@ exclude these sentinel values from parallelism averages and distributions. | `llmApiConfigJson` | string | JSON-serialized sanitized, type-driven effective LLM API configuration. See [LLM API config capture](#llm-api-config-capture). | `'{"tensor_parallel_size":2,...}'` | | `llmApiConfigMetaJson` | string | JSON-serialized metadata for LLM API configuration capture. | `'{"capture_succeeded":true,...}'` | | `disaggRole` | ShortString | Disaggregated serving role. Empty if not disaggregated. | `""`, `"context"`, `"generation"`, `"coordinator"`, `"server_coordinator"`, `"ctx0"`, `"gen0"` | -| `deploymentId` | ShortString | Shared ID across disaggregated workers. Empty if not disaggregated. | `""`, `"dep-abc123"` | +| `deploymentId` | ShortString | Shared disaggregated deployment UUID, normalized to lowercase hex. Empty if unset or invalid. | `""`, `"7cd46f928e324879b82a7cabd3c0ecde"` | ##### Architecture privacy policy @@ -163,7 +163,7 @@ exit, a terminal report, or telemetry opt-out, with no heartbeat-count limit. | `seq` | PositiveInt | Zero-based heartbeat sequence number, saturating at uint32 max. | `0`, `1`, `42` | | `ingressPoint` | ShortString | Invocation boundary for the process session. | `"cli_serve"` | | `disaggRole` | ShortString | Disaggregated role: `context`, `generation`, `coordinator`, `server_coordinator`, or compatible legacy values such as `ctx0`/`gen0`. | `"context"` | -| `deploymentId` | ShortString | Optional shared disaggregated deployment ID. | `"dep-abc123"` | +| `deploymentId` | ShortString | Optional shared disaggregated deployment UUID. | `"7cd46f928e324879b82a7cabd3c0ecde"` | Every heartbeat also contains the five aggregate LLM lifecycle counters above. @@ -183,7 +183,7 @@ outcome. Missing terminal events remain unknown; they are not confirmed crashes. | `reportingSource` | enum | `self`, `supervisor`, or `executor_proxy`. | | `ingressPoint` | ShortString | Entry point copied onto the terminal event so terminal-only early failures remain attributable. | | `disaggRole` | ShortString | Disaggregated role (`context`, `generation`, `coordinator`, `server_coordinator`, or compatible legacy `ctx0`/`gen0`), or empty when unavailable/not applicable. | -| `deploymentId` | ShortString | Optional shared disaggregated deployment ID. | +| `deploymentId` | ShortString | Optional shared disaggregated deployment UUID. | Every terminal report also contains the five aggregate LLM lifecycle counters above. Delivery is best-effort and waits no more than 0.5 seconds; the local @@ -310,7 +310,13 @@ over time. | `TRTLLM_USAGE_HEARTBEAT_INTERVAL` | `600` | Heartbeat interval in seconds. | | `TRTLLM_USAGE_FORCE_ENABLED` | `0` | Set to `1` to force-enable telemetry in CI/test environments. | | `TRTLLM_DISAGG_ROLE` | unset | Disaggregated serving role (`context`, `generation`, `coordinator`, `server_coordinator`, or compatible legacy values such as `ctx0`/`gen0`). | -| `TRTLLM_DISAGG_DEPLOYMENT_ID` | unset | Shared deployment ID across disaggregated workers. | +| `TRTLLM_DISAGG_DEPLOYMENT_ID` | unset | Shared deployment UUID across disaggregated workers. Invalid values are reported as an empty ID. | + +The disaggregated launcher generates a random deployment UUID and shares it +through the environment. Custom launchers can also supply a shared UUID via +`TRTLLM_DISAGG_DEPLOYMENT_ID`. Telemetry validates it with `uuid.UUID` and emits +32 lowercase hexadecimal characters, or an empty string if invalid. This checks +format, not who generated the UUID or whether it contains encoded identifying data. ## For Developers: Adding a New Field diff --git a/tensorrt_llm/usage/usage_lib.py b/tensorrt_llm/usage/usage_lib.py index 64b11c636b6b..28004d27a11e 100644 --- a/tensorrt_llm/usage/usage_lib.py +++ b/tensorrt_llm/usage/usage_lib.py @@ -85,6 +85,14 @@ _T = TypeVar("_T") +def _get_disagg_deployment_id() -> str: + """Return the environment's deployment UUID as lowercase hex, or empty if invalid.""" + try: + return uuid.UUID(os.environ.get(_DISAGG_DEPLOYMENT_ID_ENV, "")).hex + except ValueError: + return "" + + @dataclass(frozen=True) class TerminalOutcome: """One typed process-boundary or previously observed terminal outcome.""" @@ -794,7 +802,7 @@ def _refresh_metadata_unlocked(self) -> None: if not self.disagg_role: self.disagg_role = os.environ.get(_DISAGG_ROLE_ENV, "") if not self.deployment_id: - self.deployment_id = os.environ.get(_DISAGG_DEPLOYMENT_ID_ENV, "") + self.deployment_id = _get_disagg_deployment_id() normalized_role = self.disagg_role.lower() if self.component == "server" and ( @@ -1071,9 +1079,7 @@ def _empty_event_snapshot(usage_context: str = "") -> dict[str, Any]: return { "ingressPoint": _clamp_str(usage_context, schema._SHORT_STR), "disaggRole": _clamp_str(os.environ.get(_DISAGG_ROLE_ENV, ""), schema._SHORT_STR), - "deploymentId": _clamp_str( - os.environ.get(_DISAGG_DEPLOYMENT_ID_ENV, ""), schema._SHORT_STR - ), + "deploymentId": _get_disagg_deployment_id(), "llmInitializationAttempts": 0, "llmInstancesCreated": 0, "activeLlmInstances": 0, diff --git a/tests/unittest/usage/test_reporter.py b/tests/unittest/usage/test_reporter.py index 43c8f43668a2..896dc511868b 100644 --- a/tests/unittest/usage/test_reporter.py +++ b/tests/unittest/usage/test_reporter.py @@ -595,31 +595,54 @@ def fake_send(payload): class TestDisaggMetadata: - """Verify _background_reporter() reads disagg env vars into initial report.""" + """Verify disagg correlation accepts UUIDs but excludes arbitrary identifying text.""" - def test_disagg_env_vars_appear_in_payload(self, monkeypatch, reporter_session): - """Disagg env vars appear as disaggRole and deploymentId in payload.""" + @pytest.mark.parametrize( + ("deployment_id", "expected"), + [ + (None, ""), + ("", ""), + ("/home/alice/model", ""), + ("alice@example.com", ""), + ("a" * 128, ""), + ("7cd46f928e324879b82a7cabd3c0ecd", ""), + ("7cd46f928e324879b82a7cabd3c0ecdZ", ""), + ("7cd46f928e324879b82a7cabd3c0ecde", "7cd46f928e324879b82a7cabd3c0ecde"), + ("7CD46F92-8E32-4879-B82A-7CABD3C0ECDE", "7cd46f928e324879b82a7cabd3c0ecde"), + ], + ) + def test_deployment_id_validated_in_payloads( + self, monkeypatch, reporter_session, deployment_id, expected + ): + """Initial, heartbeat, and exit payloads contain only normalized UUIDs or empty IDs.""" monkeypatch.setenv("TRTLLM_DISAGG_ROLE", "context") - monkeypatch.setenv("TRTLLM_DISAGG_DEPLOYMENT_ID", "abc123") - - captured = {} - - def fake_send(payload): - captured.update(payload) - - stop_event = threading.Event() - stop_event.set() + if deployment_id is None: + monkeypatch.delenv("TRTLLM_DISAGG_DEPLOYMENT_ID", raising=False) + else: + monkeypatch.setenv("TRTLLM_DISAGG_DEPLOYMENT_ID", deployment_id) + assert usage_lib._empty_event_snapshot()["deploymentId"] == expected + captured = [] + stop_event = MagicMock() + stop_event.wait.side_effect = [False, True] with ( - patch.object(usage_lib, "_send_to_gxt", side_effect=fake_send), + patch.object(usage_lib, "_send_to_gxt", side_effect=captured.append), patch.object(usage_lib, "_HEARTBEAT_STOP", stop_event), ): usage_lib._background_reporter(None, None, "") + usage_lib.report_exit(usage_lib.TerminalOutcome("clean", exit_code_known=True)) - assert captured, "No payload was captured" - params = captured["events"][0]["parameters"] - assert params["disaggRole"] == "context" - assert params["deploymentId"] == "abc123" + events = [event for payload in captured for event in payload["events"]] + assert [event["name"] for event in events] == [ + "trtllm_initial_report", + "trtllm_heartbeat", + "trtllm_exit_report", + ] + for event in events: + assert event["parameters"]["disaggRole"] == "context" + assert event["parameters"]["deploymentId"] == expected + if deployment_id and not expected: + assert deployment_id not in json.dumps(captured) def test_disagg_payload_includes_llm_api_config_json(self, monkeypatch, reporter_session): """Disagg payloads retain sanitized LLM API config JSON fields.""" @@ -630,7 +653,8 @@ class _DisaggTelemetryArgs(BaseModel): ) monkeypatch.setenv("TRTLLM_DISAGG_ROLE", "generation") - monkeypatch.setenv("TRTLLM_DISAGG_DEPLOYMENT_ID", "deploy123") + deployment_id = usage_lib.uuid.uuid4() + monkeypatch.setenv("TRTLLM_DISAGG_DEPLOYMENT_ID", deployment_id.hex) captured = {} @@ -649,7 +673,7 @@ def fake_send(payload): assert captured, "No payload was captured" params = captured["events"][0]["parameters"] assert params["disaggRole"] == "generation" - assert params["deploymentId"] == "deploy123" + assert params["deploymentId"] == deployment_id.hex assert json.loads(params["llmApiConfigJson"]) == {"tensor_parallel_size": 2} meta = json.loads(params["llmApiConfigMetaJson"]) assert meta["capture_succeeded"] is True @@ -1286,12 +1310,17 @@ def test_terminal_event_uses_shared_session_id(self, monkeypatch, enable_telemet assert event["parameters"]["ingressPoint"] == "cli_serve" assert event["parameters"]["exitCode"] == 130 + @pytest.mark.parametrize("valid_id", [False, True]) def test_terminal_event_contains_correlation_and_counter_snapshot( - self, monkeypatch, enable_telemetry + self, monkeypatch, enable_telemetry, valid_id ): """Terminal-only sessions carry enough state for independent analysis.""" monkeypatch.setenv("TRTLLM_DISAGG_ROLE", "ctx0") - monkeypatch.setenv("TRTLLM_DISAGG_DEPLOYMENT_ID", "deployment") + deployment_id = usage_lib.uuid.uuid4() + monkeypatch.setenv( + "TRTLLM_DISAGG_DEPLOYMENT_ID", + str(deployment_id) if valid_id else "/home/alice/deployment", + ) assert usage_lib.record_llm_initialization_attempt(default_usage_context="cli_serve") usage_lib.record_llm_initialization_failure() @@ -1308,7 +1337,8 @@ def test_terminal_event_contains_correlation_and_counter_snapshot( ) params = sent[0]["events"][0]["parameters"] - assert params["deploymentId"] == "deployment" + assert params["deploymentId"] == (deployment_id.hex if valid_id else "") + assert "/home/alice/deployment" not in json.dumps(sent) assert params["disaggRole"] == "ctx0" assert params["llmInitializationAttempts"] == 1 assert params["llmInitializationFailures"] == 1