[feat] pack the GenRec training forward through FlashAttention-2

[WhiteSwan1]

GenRec padded every variable-length prompt before the teacher-forced forward, so every row paid for the longest row in its batch. This runs that path on packed embeddings through the FlashAttention-2 varlen kernel instead: one concatenated stream, explicit cu_seq_lens boundaries, per-row position_ids, and response-only logits gathered straight out of the packed layout.

Measured

On 1× H20, GenRecCausalLMModel over Qwen2.5-0.5B:

batch 20, real length	SDPA (the path this replaces)	FLASH_ATTENTION_2
step time	448.8 ms	233.4 ms — 1.92×
peak memory	31,515 MiB	20,147 MiB — −36%

The margin grows with sequence length: 2.55× at 2048 tokens, where SDPA peaks at 76.8 GB on a 96 GB card against FA2's 41.1 GB.

Worth reading the breakdown before optimising further: after packing, attention is only 7.0% of the step at production length, and the linear layers are 60% running at ~90% of H20 BF16 peak. The win here is mostly not paying for padding — 43% of tokens at batch 20 — rather than a faster attention kernel.

Choosing the kernel

GenRecModelConfig.common.attn_kernel picks the backend, carrying the transformers name straight through to attn_implementation:

optional string attn_kernel = 6 [default = "sdpa"];

sdpa is the default because it is the backend whose dependency always ships — building a GenRec model then needs neither a GPU nor the flash_attn wheel, which keeps the CPU lane and export working. Training runs opt into flash_attention_2 explicitly. There is no AUTO and no silent fallback: a misconfiguration surfaces as the library's own error rather than a quietly slower path.

Each kernel gets the layout it can actually serve. Only the varlen kernel reads cu_seq_lens; every other kernel rebuilds the boundaries as a dense (1, 1, T, T) mask, which costs (ΣL)² where left-padding costs Σ(B·L²). So _forward branches — packed for flash_attention_2, left-padded for everything else. The supervised window is the same absolute positions either way, so the labels are built once and only the logits differ. Measured at 8 rows of 128: packed is 1,048,576 score-pairs behind a 1 MiB mask, padded is 131,072 behind no mask at all.

flash_attn is declared in requirements/cu126|cu129|cu130.txt beside the other self-hosted wheels — three lines each, one per supported Python (3.10/3.11/3.12), matching the coverage dynamicemb, faiss and fbgemm_gpu_hstu already have.

Test plan

• Unit coverage for the packed varlen metadata, the suffix window and its labels, response-only loss and gradients, the short-row guard, and backbone construction under both kernels.
• _RowIsolationCase asserts the invariant that matters — rewrite row 0 and row 1's logits must stay bit-identical — and runs under both kernels: the sdpa arm on the CPU lane on every PR, the flash arm on the GPU lane.
• Verified the two branches agree to 5.96e-08 with identical labels.
• GPU coverage for qwen2/qwen3 packed-vs-solo forwards, projected-slot gradients under BF16/FP16 and under fp32 masters with BF16 autocast, and ragged beam decoding.

Known gaps

• The wheel is locked to the torch minor version through libtorch's C++ ABI, not just to CUDA: it links 70 c10::/at:: symbols against unversioned sonames, so a torch bump without a rebuild fails at import, not at install. Requires-Dist: torch carries no bound, so pip cannot catch the mismatch.
• cu126 has no sm100/sm120 build, unlike cu129 and cu130. Blackwell on a CUDA 12.6 install falls back to SDPA.

FA3 was measured too and is not proposed here: +2.5% at production length, no memory saving, and no wheel — it needs a from-source Hopper build against the same torch minor ABI.

[github-actions]

Review summary — overall this is a well-built change. I verified the new label-window math against the old left-padded construction: the keep positions, the masked_fill condition, and the shift-based loss land on exactly the same (row, response) pairs in both layouts, and the short-row guard is the right call for a packed stream. The bit-exact rewrite-invariance check in _RowIsolationCase is a particularly strong guard, and _check_backbone_interfaces catching unbounded responses up front covers the logits_suffix_len cast.

Findings (details inline):

1. Per-step host sync — the short-row guard's bool(...any()) blocks on GPU training every step; the collator could validate the same invariant host-side for free.
2. Cross-kernel equivalence is not pinned by a test — the PR's headline claim (packed FA2 ≡ padded SDPA) is only backed by one-off manual verification; each CI arm self-compares within a single kernel.
3. Two docstrings contradict the code — _left_pad_packed_inputs is not decode-only, and _padded_logits is not servable by "every kernel" given the CUDNN_ATTENTION exclusion.
4. Import order — the TYPE_CHECKING block placement in test_util.py likely trips ruff I001.

One behavior change worth spelling out for release notes: the short-row guard now applies to the SDPA path too. The old padded forward trained through rows shorter than logits_suffix_len (with partially ignored labels); now a single such row fails the whole batch with a ValueError. Loud and well-messaged, but existing SDPA users with empty-prompt rows will hit it on upgrade.

The wheel constraints (cp311-only, torch-ABI-locked, no AUTO/fallback) are already honestly documented in the PR body — nothing to add there.

[tiankongdeguiji]

Could we use a string for attn_kernel instead of introducing a new protobuf enum here?

This value is essentially passed through to Hugging Face's attn_implementation ("sdpa", "flash_attention_2", etc.), so using the upstream string directly would avoid adding another enum to the proto namespace and also remove the extra _ATTN_KERNEL mapping.

For example:

optional string attn_kernel = 6 [default = "sdpa"];