Repository navigation
[https://nvbugs/6665906][fix] Restore Mistral Large 3 multimodal inference - #19788
Conversation
…rence on Blackwell Share Pixtral encoder scheduling across native and HF processors. Preserve per-image geometry and use FMHA v2 for SM100 packed head-104 QKV. Cover processor behavior, kernel generation, and bounded attention workspace. Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com>
|
@coderabbitai review |
✅ Action performedReview finished.
|
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configuration
📒 Files selected for processing (1)
Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 11 remain after this review. WalkthroughFMHA kernel generation and dispatch now support selected packed-QKV head-size-104 cases. Mistral input processors share Pixtral encoder behavior, and native image batching retains individual image sizes. Tests cover these changes. The default integration test list no longer skips the specified Mistral Large 3 latency test. ChangesSM100 FMHA head-size-104 support
Mistral input processing
Test-list and warmup fixture updates
Priority: ➖ Normal Estimated code review effort: 3 (Moderate) | ~25 minutes Change: Bug fix Merge Risk: ⚪ Minimal · up to This incremental change supplies the mock timer needed by the compiled warmup path and establishes no new user-facing behavior or merge-blocking risk. 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
Full details: Docstring CoverageExplanation Docstring coverage is 47.06% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 51 functions across 11 files. (4 skipped: 4 unsupported.)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
…erage Cover representative frontend geometry and image batching cases. Assert exact scheduling budgets and successful dummy image metadata. Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com>
|
/bot run |
|
PR_Github #76095 [ run ] triggered by Bot. Commit: |
|
@coderabbitai review |
✅ Action performedReview finished.
|
|
/bot run |
|
PR_Github #76116 [ run ] triggered by Bot. Commit: |
|
PR_Github #76095 [ run ] completed with state |
|
PR_Github #76116 [ run ] completed with state
|
brnguyen2
left a comment
There was a problem hiding this comment.
Approving — the comments below are optional touch-ups, not blockers.
Both halves of the fix check out against the code paths they depend on. The vision sub-model is constructed with quant_config=None and has no KV cache, so the new dataType == dataTypeKv clause in the dispatcher is always true for Pixtral and the head-104 packed-QKV case reliably lands on fmha v2. The _has_vision_encoder gate is also sound: config_loader.py only synthesizes vision_config when vision_encoder_args is present, so audio-only and text-only native checkpoints never reach the geometry helpers. The native processor change is a real correctness fix beyond the Blackwell issue: with return_tensors="pt", transformers' MistralCommonTokenizer np.stacks the images, which raised on unequal sizes and produced a single image_sizes row for N images.
Two non-blocking items:
- Release-branch targeting. This is a fix for a bug filed against main, landing on
release/1.3. Please state the status of the main counterpart (#19507 or its successor) in the description so the branches do not diverge on the processor refactor. - SM90/SM120 scope. See the inline note on
setup.py: the padding-mask variant for head 104 is still pruned on Hopper and SM120, where the dispatcher already uses fmha v2. If that is a deliberate scoping decision, say so in the description; otherwise a follow-up would let Pixtral-Large avoid the unfused path there too.
Merge gate as stated in the description: the GB200-4_GPUs-PyTorch-1 stage must be green on this PR, since the waiver removal is only justified by that run.
…r output Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com>
|
@coderabbitai review |
|
Generate the FP16/BF16 padding variants selected by each architecture. Exercise attention parity and workspace bounds on Hopper and RTX Blackwell. Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com>
|
/bot run |
|
PR_Github #76777 [ run ] triggered by Bot. Commit: |
Limit the SM100 packed-QKV exception to FP16/BF16 FULL and CAUSAL attention. Enable SM90/SM120 FP16 FP32-accumulation padding kernels. Cover ragged numerical parity and both force flags in native CI tests. Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com>
Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com>
Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com>
|
/bot run |
|
PR_Github #76784 [ run ] triggered by Bot. Commit: |
|
PR_Github #76777 [ run ] completed with state |
|
PR_Github #76784 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #76830 [ run ] triggered by Bot. Commit: |
|
PR_Github #76830 [ run ] completed with state
|
mikeiovine
left a comment
There was a problem hiding this comment.
Stamp on behalf of runtime devs, delegating proper review to @NVIDIA/trt-llm-devs; please ping me if you think this is not accurate
|
/bot run |
|
PR_Github #76997 [ run ] triggered by Bot. Commit: |
|
PR_Github #76997 [ run ] completed with state |
….3: Restore Mistral Large 3 multimodal inference Forward-port of release/1.3 PR NVIDIA#19788 (release/1.3 commit 48c796c) to main. Adapted to main: adopts main's MultimodalItemMetadata contract from NVIDIA#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 <simengl@nvidia.com>
….3: Restore Mistral Large 3 multimodal inference Forward-port of release/1.3 PR NVIDIA#19788 (release/1.3 commit 48c796c) to main. Adapted to main: adopts main's MultimodalItemMetadata contract from _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. Also: the HF tokenizer load in MistralHFInputProcessor follows the trust_remote_code argument like the processor load next to it, and build_fmha_head104_tests requests dynamic NVRTC linking and passes plain architecture numbers (NVIDIA#20137) like the other configure-only gtest fixtures, so the target does not depend on a configure cache left by an earlier fixture. 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 <simengl@nvidia.com>
….3: Restore Mistral Large 3 multimodal inference Forward-port of release/1.3 PR NVIDIA#19788 (release/1.3 commit 48c796c) to main. Adapted to main: adopts main's MultimodalItemMetadata contract from _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. Also: the HF tokenizer load in MistralHFInputProcessor follows the trust_remote_code argument like the processor load next to it, and build_fmha_head104_tests requests dynamic NVRTC linking and passes plain architecture numbers (NVIDIA#20137) like the other configure-only gtest fixtures, so the target does not depend on a configure cache left by an earlier fixture. It also regenerates the FMHA sources, so a stale generation stamp cannot hide the head-104 generator changes, and fmhaHead104Test is registered with add_tllm_gtest (OWNER, LABELS, FULL_STACK) as main asks for new C++ tests (NVIDIA#21022). 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 <simengl@nvidia.com>
….3: Restore Mistral Large 3 multimodal inference Forward-port of release/1.3 PR NVIDIA#19788 (release/1.3 commit 48c796c) to main. Adapted to main: adopts main's MultimodalItemMetadata contract from _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. Also: the HF tokenizer load in MistralHFInputProcessor follows the trust_remote_code argument like the processor load next to it, and build_fmha_head104_tests requests dynamic NVRTC linking and passes plain architecture numbers (NVIDIA#20137) like the other configure-only gtest fixtures, so the target does not depend on a configure cache left by an earlier fixture. It also regenerates the FMHA sources, so a stale generation stamp cannot hide the head-104 generator changes, and fmhaHead104Test is registered with add_tllm_gtest (OWNER, LABELS, FULL_STACK) as main asks for new C++ tests (NVIDIA#21022). 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 <simengl@nvidia.com>
….3: Restore Mistral Large 3 multimodal inference Forward-port of release/1.3 PR NVIDIA#19788 (release/1.3 commit 48c796c) to main. Adapted to main: adopts main's MultimodalItemMetadata contract from _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. Also: the HF tokenizer load in MistralHFInputProcessor follows the trust_remote_code argument like the processor load next to it, and build_fmha_head104_tests requests dynamic NVRTC linking and passes plain architecture numbers (NVIDIA#20137) like the other configure-only gtest fixtures, so the target does not depend on a configure cache left by an earlier fixture. It also regenerates the FMHA sources, so a stale generation stamp cannot hide the head-104 generator changes, and fmhaHead104Test is registered with add_tllm_gtest (OWNER, LABELS, FULL_STACK) as main asks for new C++ tests (NVIDIA#21022). 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 <simengl@nvidia.com>
….3: Restore Mistral Large 3 multimodal inference Forward-port of release/1.3 PR NVIDIA#19788 (release/1.3 commit 48c796c) to main. Adapted to main: adopts main's MultimodalItemMetadata contract from _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. Also: the HF tokenizer load in MistralHFInputProcessor follows the trust_remote_code argument like the processor load next to it, and build_fmha_head104_tests requests dynamic NVRTC linking and passes plain architecture numbers (NVIDIA#20137) like the other configure-only gtest fixtures, so the target does not depend on a configure cache left by an earlier fixture. It also regenerates the FMHA sources, so a stale generation stamp cannot hide the head-104 generator changes, and fmhaHead104Test is registered with add_tllm_gtest (OWNER, LABELS, FULL_STACK) as main asks for new C++ tests (NVIDIA#21022). 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 <simengl@nvidia.com>
Description
Restore Mistral Large 3 multimodal inference on
release/1.3. Share the Pixtral encoder scheduling contract between native and Hugging Face input processors, preserve each native image's processed geometry when batching unequal sizes, and retain neutral behavior for text-only checkpoints.Provide head-size-104 packed-QKV padding-mask FMHA kernels on SM90, SM100, and SM120. The SM100 dispatch exception is limited to matching FP16/BF16 inputs, Q/K/V head size 104, and FULL/CAUSAL attention. SM90 and SM120 enable FP16, BF16, and forced-FP32 FP16 padding variants selected by their existing dispatcher: warp-specialized on SM90 and non-tiled on SM120. FP8 and other SM100 mask/layout combinations retain their existing routing.
This replaces closed, unmerged #19507 and removes the original four-GPU MMMU waiver. The main counterpart remains outstanding.
GB200-4_GPUs-PyTorch-1must pass before merge.Test Coverage
90-real;100-real;120-real. On H100, B200, and RTX PRO 6000 Blackwell,fmhaHead104Testpasses all six cases, covering FP16/BF16 FULL/CAUSAL attention and fixed/runtime forced-FP32 flags with short and long ragged sequences. The tests call the fused dispatcher directly and compare against an independent FP32 reference. H100 uses the existing kernel suite; B200, GB202, and RTX PRO 6000 have explicit native-test entries.pytest tests/unittest/_torch/attention/test_pixtral_head104_attention.py: two cases pass on each GPU, including the workspace bound below 64 MiB. Additional native probes pass 24 cases on H100 and 30 each on B200 and RTX PRO 6000, with ragged lengths[31, 7],[257, 65], and[2048, 32]. Baseline H100/SM120 forced-FP32 FULL requests reproduce missing kernels; the final build passes.pytest tests/unittest/_torch/modeling/test_modeling_mistral.py tests/unittest/_torch/attention/test_fmha_v2_head_size_coverage.py -m cpu_only: 16 passed with CUDA unavailable. On RTX PRO 6000, the remaining 31 Mistral cases andtest_modeling_pixtral.py -k test_pixtral_vision_model_vs_hfalso pass. All selected final tests pass without skips; local style/DCO hooks pass.accuracy/test_llm_api_pytorch_multimodal.py::TestMistralLarge3_675B::test_nvfp4_4gpus[latency_moe_trtllm]passes 900 MMMU samples at 53.111%, above the unchanged 43.123% threshold, with the original 16,384-token budget. This result does not replace the required GB200 CI gate.Validation uses Python 3.12, PyTorch 2.14.0a0 (26.08), CUDA 13.4, and
mistral-common==1.12.0with OpenCV.PR Checklist
Dev Engineer Review
cpp/kernels/fmha_v2/setup.py; verify that the intended kernel-generation changes are present in the reviewed commit.QA Engineer Review
fmhaHead104TestC++ gtest and its build and execution fixtures.l0_gb202.yml,l0_rtx_pro_6000.yml, andl0_b200.yml.Per-File QA Perspective
cpp/kernels/fmha_v2/setup.py: Changes FMHA kernel-generation eligibility and metadata formatting. Verify supported architecture, dtype, masking, warp-specialization, and tiling combinations.cpp/tensorrt_llm/kernels/fmhaDispatcher.cpp: Adds head-104 FMHA v2 dispatch conditions. Verify fallback behavior for unsupported dtypes, masks, architectures, and head sizes.tensorrt_llm/_torch/models/modeling_mistral.py: Changes image batching, encoder scheduling, dummy-data behavior, and text-only processing. Verify native/Hugging Face parity and text-only regressions.tests/unittest/_torch/modeling/test_modeling_mistral.py: Covers native geometry preservation, scheduling, and text-only behavior. It is a unit test and is not directly listed intest-dblists.tests/unittest/_torch/attention/test_fmha_v2_head_size_coverage.py: Covers supported and excluded generated head-104 kernel variants. It is a unit test and is listed inl0_gb202.ymlandl0_rtx_pro_6000.yml.tests/unittest/_torch/attention/test_pixtral_head104_attention.py: Covers ragged FP16/BF16 numerical parity and workspace limits. It is listed inl0_gb202.ymlandl0_rtx_pro_6000.yml.cpp/tests/unit_tests/kernels/fmhaHead104Test.cpp: Covers packed-QKV head-104 attention with ragged lengths, masks, dtypes, and FP32 accumulation. It is built by the CMake target and executed through the C++ integration test.tests/integration/defs/cpp/conftest.py: Builds the head-104 gtest for the requested CUDA architecture.tests/integration/defs/cpp/test_unit_tests.py: Executes the head-104 gtest for SM100 and SM120. The test is listed throughl0_b200.yml.tests/integration/test_lists/waives.txt: Removes the Mistral Large 3 four-GPU waiver. Verify that the unwaived test has suitable CI capacity.tests/integration/test_lists/test-db/l0_gb202.yml: Adds FMHA head-104 unit and Pixtral attention coverage.tests/integration/test_lists/test-db/l0_rtx_pro_6000.yml: Adds the same FMHA and Pixtral coverage for RTX PRO 6000.tests/integration/test_lists/test-db/l0_b200.yml: Adds the SM100 C++ head-104 unit test for B200 coverage.