[None][perf] Reduce Gemma multimodal attention mask construction peak memory

[zhaoyangwang-nvidia]

@coderabbitai summary

Description

Long Gemma3/Gemma4 multimodal prefills can OOM while constructing a dense attention mask because the existing implementation materializes several quadratic intermediate tensors and copies per-request masks into a packed batch buffer. Construct the packed boolean mask directly with a shared Triton kernel, using only linear-sized vision-blob metadata in addition to the final output.

Preserve causal/sliding-window attention, same-blob vision bidirectionality, Gemma4 cached-prefix handling, audio causality, and the existing model-specific gating. The final dense mask still requires Q*(P+Q) bytes per request.

On one B200, warmed mask-construction peak allocated memory (including the output, window 1024, one request with no cached prefix) was:

Model	Query length	Before	After
Gemma3	32K	5.00 GiB	1.00 GiB
Gemma3	64K	20.00 GiB	4.00 GiB
Gemma4	32K	3.00 GiB	1.00 GiB
Gemma4	64K	12.00 GiB	4.00 GiB

With an artificial 1.5 GiB PyTorch allocator budget, 32K Gemma4 mask construction OOMs before this change and succeeds afterward. These are isolated mask-construction measurements, not end-to-end serving memory claims.

Test Coverage

• tests/unittest/_torch/modules/test_gemma_attention_mask.py: registered in l0_b200.yml; 19 GPU cases passed using the source-isolated helper on B200. Covers exact mask values, image/video/audio, strided input and input immutability, packed requests, prefix/tile boundaries, empty inputs, CUDA graph replay, and peak allocation. The initial 10 cases also passed through the full TensorRT-LLM package before the review added boundary cases.
• tests/unittest/_torch/modeling/test_modeling_gemma3.py::TestGemma3::test_gemma3_allclose_to_hf[backend:flashinfer_config:1b] and ...::test_gemma3_allclose_to_hf[backend:flashinfer_config:27b]: passed, including image-mask prefill/decode with reduced-layer random-weight models.
• Additional local validation: 185 exact-reference cases and 6 full large-matrix comparisons passed, including a 64K x 64K mask to check 64-bit output addressing; existing model mask tests: 12 passed.
• Pre-commit and git diff --check: passed.

Validation limitations: the reused native build lacks SM100 FMHA kernels, so the default TRTLLM-backend Gemma3 sanity test failed in that environment. Full pretrained multimodal serving and Gemma4 full-model inference have not been validated; this PR is draft pending broader integration coverage.

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

• PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

• Test cases are provided for new code paths (see test instructions)

• If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.

• Any new dependencies have been scanned for license and vulnerabilities

• CODEOWNERS updated if ownership changes

• Documentation updated as needed

• Update tava architecture diagram if there is a significant design change in PR.

• The reviewers assigned automatically/manually are appropriate for the PR.

• Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

[zhaoyangwang-nvidia]

/bot run

[tensorrt-cicd]

PR_Github #78078 [ run ] triggered by Bot. Commit: c14295f Link to invocation

[tensorrt-cicd]

PR_Github #78078 [ run ] completed with state SUCCESS. Commit: c14295f
/LLM/main/L0_MergeRequest_PR pipeline #64358 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation