Rewrite prefix key-padding masks as k/v truncation to keep SDPA on the flash path (PyTorch)

[arnaudlvq]

Problem

Any non-null attn_mask forces torch.nn.functional.scaled_dot_product_attention
off the fused flash/memory-efficient backends onto the math backend
(check_for_attn_mask, sdp_utils_cpp.h#L258),
which materializes the [T_q, T_src] score matrix per head — O(N²) memory.

In practice the PyTorch model OOMs around ~16k context rows on a 24 GB GPU.
At a 121k-row context, the score matrix would be ~235 GB per ICL block in bf16.
The fallback is silent: no warning unless a backend is forced via
torch.nn.attention.sdpa_kernel.

Fix

The ICL attention mask is always a [B,1,1,S] boolean prefix key-padding
mask: valid keys are the first n rows, the same n for every query and batch
element. Masking those keys is algebraically identical to truncating k/v to the
first n rows and passing attn_mask=None, which lets SDPA dispatch to the
fused backends (memory linear in N).

MultiheadAttention.forward now detects that exact pattern right before the
SDPA call and rewrites it; any other mask (ragged n per batch element, float
bias, non-4D) falls through unchanged. The masked form is presumably needed at
pretraining time, where packed batches have a different boundary per element —
that case is ragged and is untouched by this rewrite.

~15 lines in tabfm/src/pytorch/model.py.

Measured

100k-row x 20-feature regression context, bf16, n_estimators=8, Colab A100:
fit+predict completes with peak allocated VRAM 13.9 GB (weights included),
R² 0.999 on the synthetic task. Without the rewrite the mask keeps SDPA on the
math backend, which needs ~160 GB of scores per block at that size
(100,000² x 8 heads x 2 bytes) and OOMs around ~16k rows on a 24 GB GPU.
Also validated at 121k rows on a private real dataset with metric-identical
predictions vs. the small-context baseline.

Tests

tabfm/src/pytorch/prefix_mask_test.py (torch-only, CPU, <1 s):

• masked forward == truncated-k/v forward (up to float32 GEMM-shape round-off,
  observed max |diff| ~2e-7);
• a ragged per-element prefix mask is NOT rewritten and still matches
  per-element truncated forwards;
• an all-False mask is not misdetected;
• end-to-end TabFM forward with a batched train_size stays finite with the
  expected shape.

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.