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

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)