Fully-masked softmax rows exponentiate to 1 in the FMHA no-loop kernel

System Info

• TensorRT-LLM: 1.3.0rc23 (built from f20ea652); behaviour re-checked against main @ fbda11f5
• GPU: NVIDIA GeForce RTX 5090 (sm_120), driver 595.71.05
• PyTorch 2.11.0+cu130, CUDA 13.0, Python 3.12.13
• Ubuntu 25.10, kernel 6.17.0-41-generic
• Backend: PyTorch; single GPU, no TP/PP
• Model where observed: Qwen3.6-35B-A3B-NVFP4 with FP8 KV cache (where model-specific)

Who can help?

No response

Reproduction

In cpp/kernels/fmha_v2/src/fused_multihead_flash_attention_kernel_noloop.h, CHECK_NEG_INF is
enabled only for bidirectional sliding-window, sliding-window, and custom masks. When it is compiled
out, a row whose running maximum is still -FLT_MAX — every score in the row masked — evaluates
expf(-FLT_MAX - (-FLT_MAX)), which is expf(0) == 1. Each masked score therefore contributes
weight 1 instead of 0.

The requirement belongs to the softmax input state, not to a mask type: any instantiation that can
present an entirely masked row is affected, including ordinary causal and padding masks.

Expected behavior

A fully-masked softmax row contributes zero weight to the attention output.

actual behavior

The row's running max is still -FLT_MAX, so every masked element evaluates expf(0) == 1 and the row attends masked positions with weight 1. Silent: no error, no NaN. Presents as sequence-length-banded output corruption.

additional notes

We can demonstrate the arithmetic directly by inspection of apply_exp_with_mask. We do not have
an isolated regression test attributing our production symptom to this line alone: the fix was
deployed alongside an unrelated kernel-enumeration change, and the evidence is a multi-day production
soak rather than a controlled A/B.

A maintainer with the fmha_v2 test harness can construct the direct case far more cheaply than we
can: any causal-mask configuration whose final Q tile is partial should show masked positions
receiving weight 1 before the change and 0 after.

Questions for maintainers

1. Is enabling the guard unconditionally the fix you would want, or would you prefer it gated on the
   configurations that can actually produce a fully-masked row?
2. Is there a known-good harness in cpp/kernels/fmha_v2 you would want a regression test added to?
3. Has the performance impact of the guard on previously-unguarded instantiations been characterised
   anywhere? We have not measured it.

Per CONTRIBUTING.md this Issue Request precedes the patch, which is ready and applies cleanly to main @ fbda11f5.

Start in cpp/kernels/fmha_v2/src/fused_multihead_flash_attention_kernel_noloop.h, especially apply_exp_with_mask and the CHECK_NEG_INF guard. Use the fmha_v2 test harness to construct a causal-mask case with a partial final Q tile, then verify fully masked positions receive zero weight and add a regression test. Characterize the guard's performance impact on previously unguarded instantiations.