[megatron] Fix FA4 backward recompilation: pass max_seqlen as int in PackedSeqParams

[xuwang0117]

Fixes #9797

PR type

• Bug Fix

PR information

get_packed_seq_params() passed max_seqlen_q/kv as 0-dim CUDA tensors (from cu_seqlens.diff().max(), no .item()).

flash-attn 4 (CuTe DSL) embeds booleans derived from max_seqlen into its backward-kernel compile-cache key. A 0-dim tensor there hashes by object identity, so the cache never hits and every micro-batch backward triggers a full JIT recompilation (tens of seconds each). With expert parallelism the other ranks stall on MoE all-to-all while any rank compiles, so the whole step is serialized on compilation. The forward key does not contain these fields, so only backward is affected.

This PR converts max_seqlen_q/kv to Python int before constructing PackedSeqParams (one microsecond-level device sync per micro-batch, negligible). This also matches the flash-attn varlen API contract (max_seqlen is documented/typed as int; transformers/modeling_flash_attention_utils.py follows the same convention).

Root-cause evidence (from hooking flash_attn.cute.interface._flash_attn_bwd.compile_cache during real training): 16 compilations of the same key within the first minute, with the key ending in

(..., tensor(False, device='cuda:0'), tensor(False, device='cuda:0'))

Experiment results

8×B200, Qwen3.5-35B-A3B full-parameter SFT (megatron sft, packing + padding_free, EP=8, max_length=8192, --attention_backend auto):

setup	per-step	GPU util	memory
no flash-attn (unfused fallback)	39.2 s/it	—	145.6 GiB
FA4 4.0.0b11, before fix	482–535 s/it	~0% (CPU-bound compiling)	81 GiB
FA4 4.0.0b11, after fix	13–14 s/it	50–70%	87.5 GiB

Loss curve identical before/after (first-step loss 1.3514 in both).

🤖 Generated with Claude Code

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.