[None][feat] Support custom masks in TRTLLM attention

[yuxianq]

Description

Add custom-mask support to the PyTorch TRTLLM attention backend while keeping generation on TRTLLM-gen kernels.

Main changes:

• add a Triton custom-mask context FMHA to the phased TRTLLM backend
• compose context and generation implementations across phased FMHA providers, including mixed batches
• support Gemma4 H512 context preprocessing and direct TRTLLM-gen generation
• reuse existing attention metadata instead of introducing custom-mask-only metadata fields
• preserve explicit Gemma4 attention backend overrides
• use FP8 queries for FP8-KV TRTLLM-gen decode, selecting the faster QkvE4m3 specialization

Performance

B200, TP/PP=1, NVFP4 weights, FP8 KV cache, ISL/OSL=1000/1000:

Model	BS	TRTLLM throughput gain vs FlashInfer	TRTLLM latency reduction
Gemma-4-26B-A4B-NVFP4	1 / 32 / 64 / 128 / 256	7.1% / 6.3% / 6.9% / 7.6% / 9.3%	6.7% / 5.9% / 6.5% / 7.1% / 8.6%
Gemma-4-31B-IT-NVFP4	1 / 32 / 64 / 128 / 256	6.0% / 4.2% / 3.6% / 3.0% / 37.7%	5.6% / 4.0% / 3.5% / 2.9% / 29.9%

The 31B BS256 point is KV-capacity constrained and was reproduced with a second paired run.

Test Coverage

• B200:
  • tests/unittest/_torch/attention/test_triton_custom_mask_fmha.py
  • tests/unittest/_torch/modeling/test_gemma4_multimodal.py
  • tests/unittest/_torch/modeling/test_modeling_gemma4.py
  • Result: 117 passed, 14 skipped
• Focused new regression cases: 6 passed
• Pre-commit hooks passed for all modified files
• Gemma4 B200 performance matrix completed for both FlashInfer and TRTLLM at BS 1/32/64/128/256

PR Checklist

• PR description clearly explains what and why
• Follows TensorRT-LLM coding guidelines
• Test cases are provided for new code paths
• No public API changes
• No new dependencies
• Documentation updated for the phased attention design