Added flash_compatible mode to decoder; fixed CUDA Graph graph-break in spectrum encoder

[GauravSRC]

Motivation

Profiling Casanovo's NAR (non-autoregressive) decoder revealed two blockers preventing hardware-accelerated attention and CUDA Graph capture:

1. AnalyteTransformerDecoder.embed() unconditionally builds tgt_key_padding_mask and generates a causal tgt_mask, both of which cause PyTorch's MultiheadAttention to fall back from the FlashAttention SDPA backend to the slower memory-efficient backend.
2. SpectrumTransformerEncoder.forward() constructs the global-token mask via torch.tensor([[False]] * batch_size) --a data-dependent Python list comprehension that TorchDynamo cannot trace symbolically, fragmenting the CUDA Graph into 4 disconnected sub-graphs.

Changes

• analytes.py: add flash_compatible: bool = False to embed() and forward(). When True, skips tgt_key_padding_mask and causal tgt_mask generation. Fully backward-compatible --default is False, no existing caller is affected.
• spectra.py: replace list-comprehension mask with new_zeros() tensor factory. Numerically identical output, no behaviour change.
• Not changed: utils.py, __init__.py, all other files.

Backward Compatibility --Verified Bit-Exact

A dedicated numerical equivalence test confirmed the flash_compatible=False default path (used by every existing casanovo installation) produces zero change in output. Both code paths were run on the same decoder instance (identical weights), so any difference could only come from logic, not weights. The new default-path code was compared against a standalone reimplementation of the original embed() body across 4 test conditions (auto-generated causal mask, explicit caller-provided mask, full forward() token scores, plus a 10-trial robustness sweep across varying batch sizes/sequence lengths/padding patterns) -- 13 tests total, all bit-exact with max_diff = 0.00e+00. This is not merely within torch.allclose tolerance; outputs are identical to the last bit. Existing casanovo users are completely unaffected by this change.

Performance Validation

Profiled on real MS spectra (NVIDIA L4, PyTorch 2.7.1) with the flash_compatible=True opt-in path:

• At bs=1: 20.5ms →12.9ms (1.6×) combining the fix with torch.compile
• At bs=512: up to 2.3× speedup
• CUDA Graph fragmentation reduced from 4 compiled regions to 2, directly validating the spectra.py graph-break fix
• FlashAttention SDPA backend now activates for decoder self-attention (confirmed via profiler kernel trace)

As a side investigation, TF32 (torch.backends.cuda.matmul.allow_tf32) was also tested as an alternative to BF16 to avoid autocast's casting overhead. TF32 is fully compatible with torch.compile (comparable ~1.5–1.7× gains) but --as expected, since TF32 doesn't change tensor dtype --cannot activate FlashAttention, confirming BF16 remains the better precision choice when flash attention matters.

Testing

Existing test suite passes unchanged (8/8). New test_analyte_decoder_flash_compatible verifies the opt-in path produces correct output shape and that the AR path remains unaffected.