RSI-Bench
Tasks/Chips & Compute Systems/GPU Kernels

Noncausal Global Extreme-GQA Ragged Attention Full VJP

Variable-length attention with the full backward pass

flashattention_varlen_feature_full_vjp_speedup Chips & Compute Systems GPU Kernels
instruction.mdthis is what the agent is given

Optimize an inherited eager implementation of packed ragged GQA cross-attention with extreme query-head sharing, bidirectionally unequal Q/K lengths, per-sequence ALiBi, softcapping, non-dense packed views, and complete Q/K/V vector-Jacobian products. Maximize raw speedup; final code is re-run on sealed hidden shapes and seeds from the same preregistered family.

Hard Constraints

  • Edit only /app/methods/main/solver.py.
  • Keep attention_varlen_feature_full_vjp(q, k, v, cu_q, cu_k, alibi_slopes, dout, causal, window_left, window_right, softcap, scale_multiplier) unchanged.
  • Return (out, dq, dk, dv) in the original packed layouts and dtypes.
  • Do not call PyTorch SDPA, private fused attention operators, or external attention packages.
  • Do not mutate inputs, inspect verifier paths, replace timers, cache by call order, or specialize on hidden identities.
  • K/V gradients reduce over all 48-96 query heads that share the single K/V head.
  • The soft cap applies to the scaled dot product before ALiBi and before masking.

What You Have

  • Every case is exact noncausal global attention: causal=False, window_left=window_right=-1.
  • Q has 48, 64, 80, or 96 heads; K/V have one head. Head dimension is 160, and activations are FP16 or BF16.
  • Every case has 18-52 independent ragged sequences. Some have Q much shorter than K; others reverse the imbalance.
  • Q/K/V and dout are views with feature padding and token steps of two or three. The final feature axis remains contiguous, but token/head strides are non-dense and must be honored.
  • ALiBi slopes are FP32 [batch, q_heads].
  • Scaled logits use (0.75 / sqrt(head_dim)) * QK, followed by 3 * tanh(logit / 3), then ALiBi and softmax.
  • ALiBi uses -slope * abs(q_position + key_length - query_length - k_position).
  • The inherited baseline and visible panel are under /app/methods and /app/problems; /app/selfcheck.py checks all outputs and raw speedup.

What You Submit

Submit /app/methods/main/ containing solver.py and local PyTorch or Triton helpers. The entry point must return the forward output and all Q/K/V gradients.

Before every evaluator run, append one contiguous canonical row (v0, v1, ...) to /app/methods/experiment_log.md and save the exact evaluated code as /app/methods/versions/vN/solver.py. Preserve failed and reverted attempts. Before finishing, make sure /app/methods/main/solver.py is byte-identical to one recorded version; the release audit rejects an unsnapshotted final edit. Record every successful aggregate as geometric mean <value>x.

How It Is Judged

Each correct sealed case scores frozen eager baseline latency / candidate median latency; reward is the geometric-mean raw speedup.

Metric

geometric-mean raw speedup over the six sealed cases · higher is better

Per case, frozen eager baseline latency / candidate median latency on an RTX A6000. All six must pass.

anchorvisible setheld-outreward
inherited eager baseline1.00x1.00x0.00
Shuman flash-attention v2.7.4—13.78x0.60
Uoptimistic upper proxy—210.39x1.00
normalisation
m <= 10
1 < m <= S0.6 * (m - 1) / (S - 1)
m > Smin(1, 0.6 + 0.4 * (m - S) / (U - S))

m = this run's held-out metric  ·  S = human flash-attention v2.7.4  ·  U = optimistic upper proxy

Linear in speedup. Each case maps on its own S and U from the sealed manifest, then the six rewards average. Any failed case: 0.

Rollouts

35 minwall clock
$6.51spend
8.4Mtokens
25versions, 11 kept
40 50 60 70 80 90 $0 $1 $2 $3 $4 cumulative spend on the run afv2100 microbench median latency, ms, lower is better v1 v2 v3 v4 v5 v6 v7 v8 v9 v10 v11 v12 v13 v14 v15 v16 v17 v18 v19 v20 v21 v22 v23 v24 v25
keptrevertedno scoreturning point
  1. v1Triton online-softmax forward and explicit full VJP; D128+D32 split dotsReplace eager attention with an online-softmax Triton kernel and hand-written Q/K/V gradients, splitting D160 into D128+D32 dots.selfcheck only, 2.45x5 min · $0.56
  2. v2Byte-identical v1, validated on the Q-long/K-short FP16 caseprofiler only, 60.7 ms6 min · $0.67
  3. v3Forward and dQ key tile 32 -> 64; dK/dV key tile kept at 3266.9789 min · $1.05
  4. v4dK/dV query tile 32 -> 64, forward and dQ tiles kept at 32x3262.0469 min · $1.13
  5. v5dK/dV kernel from 4 to 8 warps for its large BNxD accumulators72.92210 min · $1.26
  6. v6Split per-head dK and dV into separate kernels to halve live state69.4611 min · $1.38
  7. v7Pad the D128+D32 split to D128+D64 with feature masksTrade 20% extra arithmetic for tensor-core-friendly tile geometry instead of an odd D32 tail dot product.57.16112 min · $1.56
  8. v8Fused dK/dV key tile 32 -> 16 to cut its four-accumulator register use69.76513 min · $1.70
  9. v9Fused dK/dV key tile 32 -> 64 to reduce repeated Q/dout loads82.4514 min · $1.78
  10. v10One software-pipeline stage instead of two in all kernels61.47914 min · $1.87
  11. v11One pipeline stage for dQ only when Q tokens outnumber K tokens57.26315 min · $2.23
  12. v12int64 packed-token pointers for base spans above 2^31 elements57.417 min · $2.48
  13. v13Forward and dQ query tiles 32 -> 16 for higher occupancy86.39218 min · $2.69
  14. v14Forward and dQ query tiles 32 -> 64, dK/dV rows kept at 3243.86319 min · $2.79
  15. v15Query tiles to 128; dQ exceeds shared memory at two stagesshared memory overflow19 min · $2.89
  16. v16128-row tiles with a one-stage dQ pipeline; still 122880 B neededshared memory overflow19 min · $2.99
  17. v17Eight warps for the 64-row forward and dQ kernels52.57920 min · $3.14
  18. v18Forward and dQ key tile to 64 at the 64-row query tile45.21920 min · $3.24
  19. v19Pairwise query-head dK/dV accumulation; bad signature patchlaunch failed21 min · $3.42
  20. v20Corrected pairwise query-head dK/dV accumulationGive one program several query heads so its K/V tile is loaded once and the per-head partial-gradient storage halves.40.55822 min · $3.59
  21. v21dK/dV accumulated across groups of four query heads per program39.45922 min · $3.70
  22. v22dK/dV accumulated across groups of eight query heads per program39.2322 min · $3.80
  23. v23dK/dV accumulated across groups of sixteen query heads per program39.08223 min · $3.91
  24. v24Byte-identical v22, validated across the full six-case visible panel39.27224 min · $4.10
  25. v25Byte-identical release candidate sent to the isolated evaluatorbyte-identical to v2424 min · $4.62

25 snapshots, 24 of the 35 min, $4.62 of $6.51. v14 at $2.79 already held 43.9 ms; the eleven after it bought 4.6 ms for $1.82.

On the hidden set

held-out metricreward
inherited eager baseline1.00x0.00
human flash-attention v2.7.413.78x0.60
optimistic upper proxy210.39x1.00
this run4.13x0.1468
taskhttps://huggingface.co/datasets/autoresearch-arena/rsi-tasks-1/tree/gpu/flashattention_varlen_feature_full_vjp_speeduprollouthttps://huggingface.co/datasets/autoresearch-arena/rsi-rollout/tree/main/flashattention_varlen_feature_full_vjp_speedup/codex-gpt-5.6-sol-high/flashattention_varlen_feature_fu__D6SAqjY
PreviousCommit terminal-current numbers for a diode lot beyond the safe test windowNextMulti-stream Gated FlashFFTConv Forward