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 1855efcdede3..b46cf042ad1c 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 dc26b6f98ee0..b7b2299f4dde 100644 --- a/cpp/tests/unit_tests/kernels/CMakeLists.txt +++ b/cpp/tests/unit_tests/kernels/CMakeLists.txt @@ -105,6 +105,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 1d9e2cf33ed7..bae0b654703a 100644 --- a/tensorrt_llm/_torch/models/modeling_mistral.py +++ b/tensorrt_llm/_torch/models/modeling_mistral.py @@ -325,7 +325,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) @@ -340,152 +344,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 +class MistralPixtralEncoderContract: + """MM encoder item-scheduling contract for Pixtral-based Mistral VLMs. - @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 get_mm_encoder_item_metadata( - self, - _prompt_token_ids: List[int], - multimodal_data: Dict[str, Any], - ) -> Optional[MultimodalEncoderItemMetadata]: - """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 MultimodalEncoderItemMetadata( - item_refs=item_refs, - encoder_token_lengths=encoder_token_lengths, - output_embedding_lengths=output_embedding_lengths) + `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. @@ -559,11 +466,49 @@ def _dummy_mm_data_for_size( } } + def get_mm_encoder_item_metadata( + self, + _prompt_token_ids: List[int], + multimodal_data: Dict[str, Any], + ) -> Optional[MultimodalEncoderItemMetadata]: + """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 MultimodalEncoderItemMetadata( + item_refs=item_refs, + encoder_token_lengths=encoder_token_lengths, + output_embedding_lengths=output_embedding_lengths) + 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) @@ -581,14 +526,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)} @@ -603,6 +552,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 " @@ -629,6 +582,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 @@ -653,9 +702,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 85e5da514b20..276ac07b38e3 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 730d26d9920b..3194559ecf6a 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/thop/parallel/test_w4a16_linear.py diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index cf5d0e06e5f8..4bae90905444 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -24,7 +24,6 @@ accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_nvfp4[tep4_latency_moe 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.py::TestQwen3_8B::test_w4a8_mxfp4[fp8-latency] SKIP (https://nvbugs/6566768) -accuracy/test_llm_api_pytorch_multimodal.py::TestMistralLarge3_675B::test_nvfp4_4gpus[latency_moe_trtllm] SKIP (https://nvbugs/6665906) accuracy/test_llm_api_pytorch_multimodal.py::TestNanoV3Omni::test_auto_dtype[fp8_mmmu_encoder_cuda_graph] SKIP (https://nvbugs/6842250) accuracy/test_llm_api_pytorch_multimodal.py::TestNemotron_Nano_12B_V2_VL::test_auto_dtype[forced_chunked_prefill] SKIP (https://nvbugs/6836304) cpp/test_multi_gpu.py::test_cache_transceiver[8proc-mooncake_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 b6ffceb7f14d..34b0727d4f04 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.bindings import executor as executor_lib @@ -683,25 +686,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() @@ -718,8 +755,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_encoder_item_metadata( [], {"image": {"image_sizes": [[28, 56], [56, 56]]}} @@ -730,6 +766,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_encoder_item_metadata(input_ids, extra["multimodal_data"]) + metadata.validate() + assert metadata.item_refs == [("image", i) for i in range(len(sizes))] + assert metadata.encoder_token_lengths == encoder_lengths + assert metadata.output_embedding_lengths == 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() @@ -815,3 +924,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_encoder_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_encoder_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_encoder_item_metadata([], {}) is None + assert proc.get_mm_encoder_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})