For each arm we hold everything fixed and flip only attn_implementation between
eager and the AOTriton sdpa fast path. Hypothesis: 3–5×.
GptOssForCausalLM has no SDPA path in transformers (recorded as error runs).The speedup shrinks with model size (1.22× → 1.13×) and reverses for full-SFT (0.92×). The LoRA GEMMs over a frozen bf16 base dominate wall-time, so swapping the attention kernel moves a smaller slice as models grow; for Gemma full-SFT the AOTriton flash backward is actually slower than eager, with no memory saving either.
The kernel path was verified to actually change, not silently fall back: the validator confirms the
flash/efficient SDPA backends run (and match the math backend numerically) only with
TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1. Corroborating a genuinely different kernel, the
sdpa LoRA runs also draw ~7 GB less peak memory — consistent with flash attention not materializing
the N×N score matrix.