[BUG] FA2 vs NATTEN on A100: 3.5–4.8× perf gap; what is the "flash2 varlen instability" in checks.py?

[xiefan46]

Two related things on A100 (SM80):

1. NATTEN fallback is 3.5–4.8× slower than FA2 on A100

flash2/checks.py bans varlen ("banned due to instability"), so packed-sequence workloads (always varlen) silently fall back to NATTEN. On A100 that fallback lands on the legacy SM80 CUTLASS 2.X FMHA kernel, which is far slower than FA2 — and the gap grows with seq len:

seq	natten fwd	fa2 fwd	gap	natten fwd+bwd	fa2 fwd+bwd	gap
4K	56 TF/s	187 TF/s	3.4×	36 TF/s	177 TF/s	4.9×
16K	63 TF/s	216 TF/s	3.5×	42 TF/s	196 TF/s	4.7×
64K	45 TF/s	217 TF/s	4.8×	41 TF/s	198 TF/s	4.8×

(single A100-80G, bf16, varlen, 32 q / 8 kv heads, head_dim 128; both backends agree numerically to 2e-3. NATTEN plateaus at ~20% of bf16 peak and regresses at long seq; FA2 holds ~70%.)

End-to-end this roughly halves training MFU for us on A100 at 16K packed seq. H100/B200 are unaffected (flash3 / modern NATTEN kernels), but on SM80 your own default order is ["flash2", "natten"], so the varlen ban effectively removes the intended fast path.

2. What does the "flash2 varlen instability" refer to?

Could you share what the instability specifically is (which flash-attn version / API / arch), and how to reproduce it?

Perf benchmark script (goes through cosmos_framework attention frontend with backend= forced)

import os
os.environ.setdefault("SKYROBOT_ALLOW_FLASH2_VARLEN", "1")  # local env-gate lifting the ban; set before import

import torch
from cosmos_framework.model.attention.frontend import attention

DEV, DTYPE = "cuda:0", torch.bfloat16
QH, KVH, HD = 32, 8, 128
SEQS = [1024, 4096, 16384, 32768, 65536]
ITERS, WARMUP = 20, 5

cu = lambda n: torch.tensor([0, n], device=DEV, dtype=torch.int32)

def make(n, rg=False):
    g = torch.Generator(device=DEV).manual_seed(0)
    q = torch.randn(1, n, QH, HD, device=DEV, dtype=DTYPE, generator=g, requires_grad=rg)
  g = torch.Generator(device=DEV).manual_seed(0)
  q = torch.randn(1, n, QH, HD, device=DEV, dtype=DTYPE, generator=g, requires_grad=rg)
  k = torch.randn(1, n, KVH, HD, device=DEV, dtype=DTYPE, generator=g, requires_grad=rg)
  v = torch.randn(1, n, KVH, HD, device=DEV, dtype=DTYPE, generator=g, requires_grad=rg)
  return q, k, v

def call(backend, q, k, v, n):
  return attention(q, k, v, cumulative_seqlen_Q=cu(n), cumulative_seqlen_KV=cu(n),
                   max_seqlen_Q=n, max_seqlen_KV=n, backend=backend)

def bench(fn):
  for _ in range(WARMUP): fn()
  torch.cuda.synchronize(); ts = []
  for _ in range(ITERS):
      s, e = torch.cuda.Event(True), torch.cuda.Event(True)
      s.record(); fn(); e.record(); torch.cuda.synchronize()
      ts.append(s.elapsed_time(e))
  return sorted(ts)[len(ts)//2]

for n in SEQS:
  row = [f"{n:>6}"]
  for backend in ("natten", "flash2"):
      q, k, v = make(n)
      with torch.no_grad():
          ms = bench(lambda: call(backend, q, k, v, n))
      tf = 4*n*n*QH*HD / (ms/1e3) / 1e12
      qg, kg, vg = make(n, rg=True)
      dout = torch.randn(1, n, QH, HD, device=DEV, dtype=DTYPE)
      def fb():
          out = call(backend, qg, kg, vg, n)
          (out[0] if isinstance(out, tuple) else out).backward(dout)
          qg.grad = kg.grad = vg.grad = None
      msfb = bench(fb)
      tffb = 4*n*n*QH*HD*3.5 / (msfb/1e3) / 1e12
      row.append(f"{backend}: fwd {ms:.2f}ms/{tf:.0f}TF  f+b {msfb:.2f}ms/{tffb:.0f}TF")
  print(" | ".join(row), flush=True)

Env: A100-SXM4-80GB, Ubuntu 22.04, driver 575.57.08, CUDA 12.8, torch 2.10.0+cu128, flash-attn 2.7.4.post1, NATTEN 0.21.6.dev6, python 3.13.

[lfengad]

@gekurian might be able to help on answering this? Thx!