Gemma4 FFPA attention backend support for accelerating Gemma4-family training

[Butterfingrz]

NeMo-Automodel recently added FFPA attention backend support for Gemma4 full-attention layers with head_dim=512. In the Gemma4-31B training path, these full-attention layers can now use ffpa-attn CuTeDSL kernels instead of the existing SDPA or FlexAttention implementations.

This targets Gemma4's expensive head_dim=512 full-attention layers. Sliding-window layers have head_dim=256 and windowed, vision-aware masks, so they are explicitly routed to FlexAttention instead of FFPA. In our end-to-end tests, FlexAttention was noticeably faster than SDPA for those sliding-window layers.

For reference, these are the operators observed when profiling Gemma4 with the SDPA and FlexAttention backends:

Backend	Layer type	head_dim	mask object	Dispatched aten/HOP op	Actual CUDA kernel
sdpa	sliding	256	4D bool Tensor	aten::_scaled_dot_product_cudnn_attention	cudnn_generated_fort_native_sdpa_sm90_flash_{fprop,bprop}_wgmma_f16
sdpa	full (no padding)	512	None (takes the is_causal fast path)	MATH (no fused op)	decomposes into aten::bmm x6 + aten::_softmax (materializes S x S)
sdpa	full (with padding)	512	4D bool Tensor	aten::_scaled_dot_product_efficient_attention	cutlass mem-efficient: fmha_cutlassF/B_bf16_aligned_*_sm80
sdpa	sliding (with padding)	256	4D bool Tensor	aten::_scaled_dot_product_cudnn_attention	same cuDNN flash as above
flex	sliding / full	256 / 512	BlockMask	torch.ops.higher_order.flex_attention (compiled)	inductor triton_tem_fused_flex_attention*

Integration

For non-CP training, attn_implementation: ffpa registers "ffpa" as a Hugging Face attention backend during model build. setup_ffpa_backend(...) validates this path before preload overrides run: direct "ffpa" requires cp_size == 1 and rejects packed sequences. This path also requires the ffpa-attn dependency, exposed through the ffpa extra.

At runtime, ffpa_attention_forward() uses FFPA only for eligible full-attention calls: head_dim=512, fp16/bf16, no attention dropout, a non-null scale, and an available FFPA CuTeDSL kernel. The dense no-mask path uses the high-level ffpa_attn_func; a 2D bool padding mask uses the FFPA varlen path. Unsupported full-attention cases fall back to SDPA or eager, with softcap forced to eager. Sliding-window and other non-causal mask functions are routed through a FlexAttention BlockMask, so those layers use FlexAttention rather than SDPA. Gemma4 MoE still uses Hugging Face Gemma4Attention, so the Gemma4-26B-A4B attention dispatcher can use the same "ffpa" backend.

For context parallelism, FFPA is wired through Gemma4's model-owned ring path instead. attn_implementation: ffpa is rejected when cp_size > 1, because the Hugging Face backend calls the FFPA op directly and would bypass the ring SDPA swap. CP runs should keep attn_implementation: sdpa and set text_config.cp_full_attn_backend: ffpa. The ring path uses a rank-uniform gate, checking causal full attention, no sliding window, _packed_seq_ids, head_dim=512, fp16/bf16, scale, and FFPA availability. Eligible full-attention ring chunks use FFPA, choosing dense or varlen per chunk; everything else stays on the compiled FlexAttention ring path.

Joint Finetuning Recipes

For non-CP Gemma4-31B fine-tuning, enable the Hugging Face attention backend directly:

model:
  _target_: nemo_automodel.NeMoAutoModelForImageTextToText.from_pretrained
  pretrained_model_name_or_path: google/gemma-4-31B-it
  torch_dtype: torch.bfloat16
  attn_implementation: ffpa
  text_config:
    use_cache: false

With this setting, head_dim=512 full-attention layers dispatch to FFPA when the layer is eligible. Layers that do not match the FFPA constraints fall back to SDPA, eager, or FlexAttention as appropriate.

Do not combine the direct Hugging Face "ffpa" backend with CP or packed_sequence; AutoModel fails fast for both cases. Packed long-sequence FFPA runs should use the CP ring configuration below.

For context-parallel Gemma4-31B fine-tuning, keep the Hugging Face attention implementation as SDPA and opt the Gemma4 ring path into FFPA:

model:
  _target_: nemo_automodel.NeMoAutoModelForImageTextToText.from_pretrained
  pretrained_model_name_or_path: google/gemma-4-31B-it
  torch_dtype: torch.bfloat16
  attn_implementation: sdpa
  text_config:
    use_cache: false
    cp_full_attn_backend: ffpa

This is the required CP setup because attn_implementation="ffpa" calls the FFPA op directly and bypasses Gemma4's model-owned CP ring hook. In CP mode, cp_full_attn_backend: ffpa routes full-attention ring chunks through FFPA, while sliding-window layers continue to use the compiled FlexAttention ring path. The CP batch path supplies _packed_seq_ids, so the ring backend can use dense chunks for single-document shards and varlen chunks when packing or cross-document masking is needed.

Results

Upstream H200 E2E result from PR #2436

The PR description also reported an end-to-end Gemma4-31B run on H200. At the same reported peak-memory level, FFPA improved TPS/GPU by about 40-50%, while the loss curve overlapped the SDPA baseline within bf16 noise.

Environment and hyperparameters:

• Hardware: 1 node x 8x H200
• Strategy: FSDP2 + activation checkpointing
• Config: DP=8, GBS=8, local_bs=1, grad-accum=4, L=8192

Backend	Peak memory	Tokens/s/GPU	Loss
sdpa baseline	~65 GiB	1.00x	Reference
ffpa candidate	~65 GiB	1.4-1.5x	Overlaps sdpa within bf16 noise

W&B report: Gemma4 FFPA H200 E2E comparison

Note: the small memory gap visible in this early H200 plot comes from the first FFPA SAC integration. FFPA forward ops were pinned to MUST_SAVE, so activation_checkpointing: true did not recompute the FFPA attention ops. Read the memory line here as an artifact of that checkpointing policy, not as a pure kernel-memory comparison.

The benchmark recipes used for these runs are examples/vlm_finetune/gemma4/gemma4_31b_ffpa_mock_8k.yaml for the non-CP path and examples/vlm_finetune/gemma4/gemma4_31b_ffpa_mock_packing_cp8_16k.yaml for the Ring CP path.

Perf summary for PR #2436, rerun on a local 8x H100 80GB setup at local rerun commit dac1683f7. The CP1 mock 4k and 8k FFPA rows were updated from the latest reruns. Each run used 30 train steps; TPS/GPU is averaged over steps 20-29 only.

Setting	Backend	TPS/GPU	Max mem (GiB)
CP1 mock 4k	Flex	943.60	46.81
CP1 mock 4k	SDPA	883.35	52.32
CP1 mock 4k	FFPA	1171.65	46.81
CP1 mock 8k	Flex	559.66	65.48
CP1 mock 8k	SDPA	OOM	-
CP1 mock 8k	FFPA	654.39	65.48
CP8 packed 16k	Flex	286.39	39.85
CP8 packed 16k	FFPA	317.10	40.51

FFPA speedup:

Setting	FFPA TPS/GPU	Baseline	Baseline TPS/GPU	Speedup
CP1 mock 4k	1171.65	Flex	943.60	+24.2%
CP1 mock 4k	1171.65	SDPA	883.35	+32.6%
CP1 mock 8k	654.39	Flex	559.66	+16.9%
CP1 mock 8k	654.39	SDPA	OOM	N/A
CP8 packed 16k	317.10	Flex	286.39	+10.7%

Notes:

• CP1 mock 8k SDPA OOMs because it falls back to MATH.
• Direct attn_implementation: ffpa is non-CP and non-packed only; packed FFPA uses the CP ring path.
• FFPA forward ops are included in selective activation checkpointing's must-save set, so the CuTeDSL forward kernels are not replayed during backward.
• Latest FFPA reruns used PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True.
• Latest CP1 mock 4k FFPA allocator diagnostics: 0 CUDA alloc retries, 0 CUDA OOMs.
• Latest CP1 mock 8k FFPA allocator diagnostics: 359 total CUDA alloc retries by step 29, 125 retries during steps 20-29, 0 CUDA OOMs.

Follow-up

Future work will extend this FFPA full-attention backend path across Gemma4-family models that use head_dim=512 full-attention layers.

• Wire Gemma4-MoE Ring CP into the FFPA ring backend. CP does not use the Hugging Face "ffpa" attention backend directly; it runs through Gemma4's model-owned ring hook, so each attention module must be attached with use_ffpa=True. Today only the dense Gemma4 branch reads text_config.cp_full_attn_backend and passes that flag into attach_gemma4_cp_ring_attention(...). The Gemma4-MoE decoder layer still calls attach_gemma4_cp_ring_attention(self.self_attn) with the default use_ffpa=False, so google/gemma-4-26B-A4B-it Ring CP currently falls back to the FlexAttention ring for full-attention layers instead of using FFPA.

Debugging notes

• Gemma4 full-attention layers exposed a separate SDPA debugging issue before the FFPA path: head_dim=512 plus the no-padding causal fast path can make attention_mask=None, which previously caused Transformers' sdpa_attention_forward to pass enable_gqa=True. PyTorch's fused SDPA backends cannot handle that GQA broadcast case for head_dim > 256, so dispatch silently falls back to the MATH backend and materializes the S x S score matrix, turning streaming fused attention memory into O(S^2) memory and causing long-sequence OOMs. Hugging Face Transformers PR #46960 fixes this by only enabling GQA-in-SDPA for the safe case where key/value head dimensions are equal and <= 256; Gemma4 head_dim=512 full-attention therefore takes the repeat-KV path instead of the unsupported enable_gqa=True path.

• The local 8x H100 80GB 8K run did not reach the expected ~30% TPS gain, but the bottleneck was not the FFPA kernel suddenly getting slower. At 8K, the run pushed the real reserved-memory watermark to the allocator/FSDP2 boundary. It did not OOM, but the allocator frequently flushed cache and retried allocations. Those retries usually bring synchronization and the cudaMalloc slow path, which can consume the TPS gain FFPA should provide at this sequence length. PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True helps reduce slivers and fragmentation, but it does not reduce the real live-tensor footprint, FSDP2 all-gather demand, or FFPA layout-copy demand. This is why the H200 result looks closer to the ideal case, with roughly 40-50% TPS improvement, while the local H100 CP1 mock 8K run only shows +16.9% over Flex.

---

Related PRs:

• perf: FFPA D=512 attention backend for Gemma4 (3x fwd / 6x bwd vs SDPA) #2436
• feat: add mock VLM dataset and Gemma4 pretokenize support #1682
• feat: add gemma4 moe ring CP #1914
• feat(gemma4): context parallelism for dense 31B #2592

I'm truly grateful to @HuiyingLi , @athitten and @akoumpa for their help, and to the Automodel team for their trust.