[Perf][B300] Evaluate KV-head-owned GQA dKV backward

[heiheiha798]

Status

This issue records an unmeasured B300 optimization hypothesis. The mechanism and thresholds below are acceptance criteria, not observed performance claims.

Baseline: upstream/main@e47d5d20aeb5989b58a3738b872e7c288a9fb75f.

Problem

The default Triton parallel_attn backward is query-head-owned. For GQA (HQ > H), parallel_attn_bwd_kernel_dkv launches over B * HQ, writes fp32 dK and dV tensors shaped with HQ, and then reduces both tensors to the real KV-head count H in separate framework operations.

At B=1, T=16384, H=4, HQ=32, K=V=128, the two fp32 intermediates occupy 512 MiB. The final bf16/fp16 KV gradients occupy 64 MiB, so the current path carries 448 MiB of avoidable allocation plus the associated per-query-head K/V loads, stores, and reduction traffic. The hypothesis is that a KV-head-owned dKV kernel can remove this amplification and reuse each K/V tile across a GQA group on B300.

This targets the dependency-light Triton backend, forced with FLA_DISABLE_BACKEND_DISPATCH=1; it does not compete with or change the optional TileLang backend.

Bounded scope

First create one fixture-only commit containing targeted correctness additions and repository-native GQA benchmark registrations. Production source and the naive reference must remain byte-for-byte upstream in this commit. Run the complete baseline gates green, record this commit as BENCH_BASE, then freeze tests, references, seeds, dtypes, tolerances, and numeric flags for the optimization loop.

Profile BENCH_BASE before production edits. Record the complete forward+backward timeline, current dKV duration, both post-dKV reductions, and peak allocated/reserved memory for one dense and one sliding-window G=8 endpoint. Stop before implementation if dKV plus reductions is below 25% of full forward+backward on both endpoints, if the reductions are negligible and profiling shows no credible K/V reuse opportunity, or if an initial G=8 proof kernel cannot compile without material spill or occupancy collapse.

If the viability gate passes:

1. Add a GQA-only Triton dKV kernel owned by (key-token tile, KV head, batch) rather than query head. One program owns every final dK/dV element for its KV head and token tile.
2. Load the K/V tile once, process the compile-time GQA group (G in {2,4,8}) in fp32 accumulation, and write directly to final input-dtype [B,T,H,K] and [B,T,H,V] buffers with one cast at the store. Do not use atomics or add an fp32 scratch tensor comparable to the current HQ-shaped intermediates.
3. Preserve the existing per-query-head dg_cumsum_k contributions when gates are present, including dense and packed-varlen addressing and sliding-window bounds.
4. Use int64 program IDs, batch/head bases, BOS/EOS offsets, token positions, strides, and derived addresses.
5. Dispatch the new path only for fp16/bf16, K == V, D in {64,128}, and G in {2,4,8} after the measured resource gate. Keep the current kernel and reductions unchanged for MHA, fp32, K != V, other dimensions, G greater than 8, unsupported platforms, and measured-losing buckets.

Do not change forward attention, dQ, sink-gradient logic, gate math, accumulation precision, public signatures, model/config code, backend dispatch infrastructure, the optional TileLang implementation, or the open @torch.compile cleanup. Do not edit fla/ops/attn/naive.py: its current output reshape cannot independently validate K != V, so unequal K/V remains a legacy-fallback boundary rather than being folded into this performance cycle.

Expected production file: fla/ops/attn/parallel.py. Fixture changes are limited to tests/ops/test_attn.py and benchmarks/ops/registry.py unless the native varlen post-init demonstrably cannot express the input.

Frozen correctness contract

The fixture should add a small, targeted set rather than a large Cartesian matrix. Against the existing fp32 naive_parallel_attn oracle and unchanged 0.005 tolerance, cover:

• optimized G=2/4/8 with fp16 and bf16 at D=64/128;
• non-tile-aligned T and B greater than 1;
• gated GQA with unequal packed lengths;
• sliding-window boundaries and packed sliding-window GQA;
• MHA, G greater than 8, fp32, non-target D, and K != V dispatch remaining on the unchanged legacy path (dispatch/source-path checks only where the naive oracle cannot represent the shape);
• complete output and dq/dk/dv/dg/dsink gradients for every reachable case, under allocator NaN poisoning.

Promotion requires the following full, unfiltered gates with the Triton backend forced:

FLA_DISABLE_BACKEND_DISPATCH=1 pytest tests/ops/test_attn.py
FLA_DISABLE_BACKEND_DISPATCH=1 pytest tests/ops/test_attn_sink.py
FLA_DISABLE_BACKEND_DISPATCH=1 pytest tests/ops/test_forgetting_attn.py
python scripts/find_dependent_tests.py fla/ops/attn/parallel.py

Run every returned dependent test plus affected-file pre-commit, header, compile, and banned-block-pointer checks. No post-freeze reference, tolerance, parametrization, skip, dtype, numeric-flag, or environment relaxation is allowed.

Repository-native B300 measurement

Register separate public-endpoint entries for dense GQA, sliding-window GQA, and gated unequal-varlen GQA, all calling parallel_attn, using shape_q_hq, output_is_tuple=False, and test_file='tests/ops/test_attn.py'. Use BENCH_BASE rather than upstream/main for comparisons because the baseline commit owns the new registry fixture.

Required bf16 forward+backward endpoints:

• dense: B1,T4096/8192,H4,HQ32,D128 and B2,T2048,H8,HQ32,D64;
• sliding window: B1,T8192/16384,H4,HQ32,D128,W256/512;
• gated unequal varlen: flattened totals 8192/16384 with H4,HQ32,D128;
• breadth controls for G=2/4/8 and D=64/128, plus unchanged forward, MHA, G>8, and non-target fallback controls.

Compare HEAD and BENCH_BASE in the same isolated B300 allocation with identical inputs, clocks/idle state, environment, warmup, repetitions, and warmed autotune state. Use at least five independent same-session comparisons. Report per-shape sample count, median, mean, std, min, p10, p90, native p20/p80, equal-weight geomean, kernel count, and torch.cuda.max_memory_allocated/reserved with reset discipline. Keep raw JSON/logs, commands, environment versions, SHAs, and Slurm job IDs outside git.

Collect full and source-correlated NCU profiles for one dense G=8 and one W=512 G=8 workload. Compare summed dKV plus reduction duration, launches, DRAM read/write bytes and throughput, L2 behavior, tensor/SM SOL, achieved occupancy, registers, local-memory spills, eligible/active warps, and dominant source stalls. Confirm the two post-reduction kernels disappear only on the optimized dispatch and are not replaced by hidden casts or reductions. The user-level NCU helper is unavailable in this environment, so use the repository's documented minimal full/source workflow and state that in the evidence.

Acceptance criteria

• Every frozen op and dependent gate passes with exact API/layout/dtype behavior and finite in-tolerance gradients.
• Full forward+backward equal-weight geomean is at least 1.08x across all optimized public shapes; every required G=4/8 dense, sliding-window, and varlen endpoint is at least 1.04x; no endpoint regresses by more than 3% outside measured noise.
• Unchanged forward, MHA, G>8, fp32, unequal-K/V, and other fallback controls remain within 2% noise.
• The representative G=8, T=16384 workload reduces measured peak allocation by at least 30% and removes the analytical 448 MiB HQ-shaped dK/dV excess without a comparable new scratch allocation.
• NCU explains the endpoint win through lower K/V traffic, direct final stores, and removed reductions without material spill, occupancy collapse, skipped work, or clock artifacts.

A memory-only improvement, an isolated sliding-window win, a result requiring TileLang, or a miss on the frozen correctness/performance/resource gates is a no-go. Do not widen into forward attention, dQ, model changes, optional backends, or unrelated dispatch cleanup to rescue it.

Upstream overlap

No open upstream issue, PR, or remote branch targets KV-head-owned Triton GQA dKV. Merged fla-org#260 introduced the current MHA/GQA kernel, fla-org#504 tuned A100 blocks, fla-org#824 added sliding windows, and fla-org#941 fixed the unrelated NV==1 correctness bug. Merged fla-org#846 is the closest precedent, but its optional TileLang backward remains query-head-owned and uses atomic dK/dV accumulation. Open fla-org#980 only removes the erroneous class-level @torch.compile; open fla-org#1032 adds GQA divisibility validation. Both are small rebase risks in the same file, not semantic duplicates, and neither cleanup belongs in this issue.

[heiheiha798]

B300 evaluation result: no-go at the frozen baseline gate

Cycle 7 stopped before profiling or production implementation because the new frozen oracle exposed a deterministic numerical failure in the untouched Triton implementation.

• Source baseline: e47d5d20aeb5989b58a3738b872e7c288a9fb75f
• Fixture-only commit: b259d2d5bc5bf5ea9ff314a9ea7df580298e63e4
• Environment: B300/SM103 with FLA_DISABLE_BACKEND_DISPATCH=1
• Slurm job: 5178
• Durable raw log: profile/parallel_attn_gqa-opt/fixture-5178.log

The full unfiltered tests/ops/test_attn.py gate completed with 37 passed and 1 failed. The failing case was the targeted G8 bf16 D128 gated unequal-varlen case with cumulative sequence lengths [0, 33, 161, 390]. Its dg error ratio was 0.008288 against the unchanged 0.005 tolerance, with maximum absolute error 0.257812.

The same case passed the other frozen comparisons:

• output ratio: 0.001632
• dq ratio: 0.003064
• dk ratio: 0.003087
• dv ratio: 0.002566

The allocation used set -e, so it stopped immediately after this failure and did not run tests/ops/test_attn_sink.py or tests/ops/test_forgetting_attn.py.

fla/ops/attn/parallel.py and fla/ops/attn/naive.py remained byte-for-byte identical to the source baseline. No viability profile, NCU collection, candidate kernel, production-source edit, benchmark comparison, push, or PR was performed. This is therefore a no-go caused by a red frozen untouched-production fixture, not a candidate implementation failure, and it supports no technical speedup conclusion.