[Perf][B300] Batch query groups in NSA compression dKV backward

[heiheiha798]

Hypothesis

parallel_nsa_compression_bwd_kernel_dkv assigns one CTA to a compressed KV block and processes the causal query suffix one token at a time. For grouped-query attention, each of its four contractions therefore has an MMA free dimension of only G = HQ / H (16 in the shipped NSA configuration), while a long-sequence CTA executes thousands of serial loop bodies.

This issue will test whether batching contiguous query groups into an N = BQ * G tile improves tensor-core utilization and reduces dynamic loop work on B300 without changing the mathematical work, CTA ownership, output layout, or public API. No speedup is claimed until the profile-first and frozen-oracle gates below pass.

The default NSAConfig uses 64 query heads, 4 KV heads, head dimension 32, block size 64, and enables compression, selected-block, and sliding-window branches. The compression backward is therefore on the shipped dense and packed training path.

Bounded kernel design

Keep one CTA responsible for one [BC, K/V] compressed KV tile. Replace the scalar query loop with a compile-time BQ schedule:

• flatten BQ consecutive query tokens and their G grouped heads into BQ * G columns;
• load Q, dO, LSE, and delta for those columns;
• evaluate the existing score, probability, dV, dP, and dK contractions with the wider N dimension;
• apply the exact per-compressed-token visibility predicate using each query lane;
• retain FP32 dK/dV accumulation, disjoint stores, and dense/varlen sequence ownership.

Autotune BQ from {1, 2, 4, 8} and only consider 16 if a resource proof justifies it. Key the schedule on at least BS, BK, BV, and G, and prune oversized BQ * G tiles using the existing selected-NSA rule. The current BQ=1 warp/stage configurations must remain available unchanged as a genuine resource/performance fallback; adding wider schedules must not replace that baseline choice.

All program IDs, sequence bases, query positions, and pointer offsets must use int64 arithmetic. Do not add atomics, CSR, block pointers, recomputation, or replacement workspace. Do not modify selected attention, top-k, pooling, sliding-window attention, shared utilities, public APIs, or numerical policy. An exact-SM103 schedule guard is allowed only if profiling proves a B300-specific bucket; otherwise prefer the architecture-neutral autotuned kernel with the preserved BQ=1 fallback.

Expected tracked files:

• fla/ops/nsa/compression.py for the dKV schedule;
• tests/ops/test_nsa.py only in the initial fixture commit, then frozen;
• benchmarks/ops/registry.py only in the initial fixture commit for compression-enabled native endpoints.

fla/ops/nsa/naive.py, model/layer files, selected-attention kernels, and benchmark runner internals are frozen and out of scope.

Profile-first viability gate

Start from exact upstream baseline 27967b970eaaf982a6960abf6cba8add9c34c7cc. Before production kernel edits:

1. Run the complete untouched tests/ops/test_nsa.py on B300. A deterministic red/OOR baseline is a terminal no-go.
2. Validate the fixture-only commit against unchanged production code and record it as immutable BENCH_BASE.
3. Profile BF16 compression and full compression-enabled NSA at shipped geometry (H=4, HQ=64, G=16, K=V=32, BS=64) for T=8K/16K/32K, plus K64/V128 and unequal packed-varlen controls.
4. Attribute parallel_nsa_compression_bwd_kernel_dkv separately in compression fwdbwd and the full public path. Record its duration/share, loop count, tensor-pipe utilization, registers/spills, and peak allocation.

Stop before production edits if dKV is below 25% of compression fwdbwd on the T16K/T32K geomean, below 12% of full compression-enabled parallel_nsa fwdbwd on both long rows, or the measured Amdahl bound cannot plausibly support 1.15x compression and 1.08x full-path fwdbwd with a credible 3x dKV improvement.

After the smallest BQ proof, stop rather than widen scope if BQ * G fails to compile, introduces material local-memory spilling or occupancy collapse, or dKV is below 2.5x on representative T16K/T32K G16 rows.

Frozen correctness contract

Before kernel edits, extend the existing naive-reference matrix with focused cases for:

• fp16 and bf16 at G=16 and G=32;
• dense non-aligned T with K64/V128;
• unequal packed varlen with K128/V64;
• BS32/64, partial compressed blocks, and BQ-tail lengths;
• forward output/LSE and every reachable dq/dk/dv gradient.

Validate these cases on untouched production, commit them with the benchmark registrations, then freeze tests/ops/test_nsa.py, fla/ops/nsa/naive.py, tolerances, inputs, seeds, dtypes, skips, and numeric flags for the entire loop. Existing dense/varlen, decoding, non-power-of-two dimension, large-G, and full-NSA tests remain mandatory. NaN allocator poisoning remains enabled.

Promotion requires full unfiltered gates, including:

python -m benchmarks.ops.verify --op parallel_nsa_compression --base <BENCH_BASE> --modes fwd fwdbwd
python -m benchmarks.ops.verify --op parallel_nsa_full --base <BENCH_BASE> --modes fwd fwdbwd
python scripts/find_dependent_tests.py fla/ops/nsa/compression.py benchmarks/ops/registry.py

Run the complete tests/ops/test_nsa.py, every returned dependent test, and tests/models/test_modeling_nsa.py when the required FlashAttention dependency is present. Missing FlashAttention may be diagnosed, but it cannot be bypassed to promote a default-window model claim. No tolerance, precision, reference, skip, or environment-suppression change is allowed after BENCH_BASE.

Representative B300 evidence

Register two public endpoints in the immutable fixture:

• direct parallel_nsa_compression, with compressed K/V shapes constructed outside the timed call;
• full compression-enabled parallel_nsa, with compression and selected gates active and a window-free native control.

The full shipped window path is measured separately through the public layer/model with FlashAttention installed. Primary BF16 rows are T8K/16K/32K at H4/HQ64/G16/K=V32 or 64, plus K64/V128 and unequal packed varlen. T64K is optional only after the baseline resource gate.

Compare BENCH_BASE and candidate in the same isolated B300 allocation with identical inputs, warmed autotune state, clocks/idle evidence, and at least five interleaved rounds. Report every row's median, mean, standard deviation, min, p10/p20/p80/p90, sample count, equal-weight geomean, and peak allocation.

Integration must exercise the shipped NativeSparseAttention geometry and all compression/selection/window gates in dense and packed form. Run a one/two-layer public-layer harness and the repository NSA training runner where memory permits; record tokens/s and peak memory. Do not present the existing selected-only parallel_nsa registry timing as evidence for this issue.

Collect NCU full/PM-sampling and source/SourceCounters reports for baseline and candidate dKV on dense model D32/T16K, K64/V128/T32K, and one unequal varlen row. The evidence must show wider BQ * G MMAs and fewer loop bodies with unchanged mathematical/output traffic, without spills, skipped work, atomics, or hidden workspace.

Acceptance and stopping conditions

• Complete frozen operator/dependent/model correctness gates pass with unchanged API, layout, dtype, and numerical policy.
• Candidate dKV improves at least 2.5x on representative G16 rows with profiler attribution.
• Compression fwdbwd equal-weight geomean improves at least 1.15x across T8K/T16K/T32K, with T16K and T32K each at least 1.12x; forward remains within 2% noise.
• Full compression-enabled public parallel_nsa fwdbwd geomean improves at least 1.08x on T16K/T32K, every retained long row is positive, and no retained row regresses more than 3% outside measured noise.
• Dense and unequal varlen are both positive. A shipped-configuration NSA layer improves at least 5% median fwdbwd on one representative long workload and is neutral within 2% on its counterpart; the two-layer training path is positive and peak memory does not regress.

Record no-go rather than widening scope if the baseline/dependency gate is red, numerical parity needs relaxed tolerance, only an artificial large-G or small shape wins, the full public endpoint misses 1.08x, the required default-window integration cannot be exercised, or success requires edits to selected NSA, top-k, pooling, window attention, or shared utilities.

If correctness is green but representative performance misses these unchanged thresholds, submit a complete performance-depth dossier. At most two further iterations may proceed, and only after reviewer DEEPEN approval; each must remain inside this compression dKV schedule and carry explicit early resource/performance stopping evidence.

Prior art, isolation, and mergeability

Merged upstream PR fla-org#961 batches scattered selected-attention queries through CSR and reports about 5x dKV improvement at HQ64/H4/D128, but it leaves parallel_nsa_compression_bwd_kernel_dkv on its original one-query loop. This issue reuses only the validated BQ * G layout and pruning principle; compression queries are a contiguous causal suffix and require no CSR.

Current open upstream/fork inventories and remote branches contain no compression-dKV batching implementation. Upstream fla-org#1029/fla-org#1032 cover non-divisible GQA correctness, not this schedule. Lane-b issue #15 is restricted to normalized linear-attention epilogues in fla/ops/linear_attn/utils.py and tests/ops/test_linear_attn.py; it has no tracked-file, benchmark, or model overlap with this NSA work.

If the gates pass, the contribution is one substantive kernel schedule change on a shipped long-context training branch, with frozen dense/varlen correctness, direct operator evidence, and full public/model validation.

[heiheiha798]

Cycle 15 result: performance no-go

The B300 evaluation stopped at the frozen BQ=2 proof gate. The batching mechanism is correct and measurably faster, but it does not approach the required dKV-stage threshold, and the fixed reviewer returned STOP after depth review.

Frozen state and correctness

• Source baseline: 27967b970eaaf982a6960abf6cba8add9c34c7cc.
• Fixture-only BENCH_BASE and branch head: 8d6b4cbeef3e153476edda08af95ded1c85f07bb.
• fla/ops/nsa/compression.py and fla/ops/nsa/naive.py remain byte-identical to the source baseline. No production candidate commit or PR was created.
• Untouched B300 gate, Slurm job 5368: 40 passed, 8 existing flash_attn skips.
• Frozen fixture gate, Slurm job 5386: 42 passed, 8 existing skips. This includes dense K64/V128 and unequal packed K128/V64 output, LSE, dq, dk, and dv oracle coverage at the existing tolerance.

Baseline profile and NCU

Jobs 5394 and 5399 establish that dKV is material:

• Direct compression fwdbwd dKV share: 91.12% at T16K and 88.91% at T32K.
• Full registered compression-enabled path dKV share: 46.96% at T16K and 43.95% at T32K.
• A 3x dKV stage improvement therefore had sufficient Amdahl headroom for the issue thresholds.

Job 5420 retained four raw baseline NCU reports for the exact selected parallel_nsa_compression_bwd_kernel_dkv launch:

• T16K/T32K grids: 16/32 CTAs over 148 SMs.
• Duration: 17.29/34.47 ms.
• 50 registers per thread, 15.67 KiB dynamic shared memory, zero local spills.
• 6.25% achieved occupancy, 0.22 eligible warps per scheduler, and about 78% no-eligible cycles.
• Source evidence shows the existing contractions already lower to UTCHMMA.

Evidence-plumbing failures are retained but are not kernel failures: job 5394 stopped on a Torch profiler API mismatch, job 5399 stopped because an outer Python-thread NVTX range did not include autograd-engine backward launches, and job 5420 exited after all four valid captures because its final CSV parser did not account for a missing trailing empty field in source-report rows.

BQ=2 proof

The ignored standalone BQ=2 proof in job 5443 preserves compressed-KV CTA ownership, the four contractions, FP32 accumulators, causal/scale semantics, output dtypes, int64 addressing, and disjoint dk/dv stores. Its dk and dv are bitwise identical to production, finite, and deterministic. It uses no global scratch.

Five interleaved rounds and 50 samples per variant produced:

• T16K production: median 10.028240 ms, mean 10.041957, std 0.024198, min 10.018880, p10 10.020643, p90 10.077536.
• T16K BQ=2: median 6.560784 ms, mean 6.570420, std 0.015504, min 6.554176, p10 6.556861, p90 6.592982. Speedup: 1.528512x.
• T32K production: median 20.185040 ms, mean 20.187256, std 0.025572, min 20.147873, p10 20.153821, p90 20.216746.
• T32K BQ=2: median 13.152992 ms, mean 13.159435, std 0.015847, min 13.131616, p10 13.141155, p90 13.179728. Speedup: 1.534635x.

Both results are far below the frozen >=2.5x dKV proof threshold. Passing would require another 1.6356x/1.6291x reduction over BQ=2.

The reviewer found that the measured BQ-scalable fraction is about 69%. Perfect scaling of that fraction projects BQ=4 at only about 2.08-2.09x and BQ=8 at only about 2.53-2.56x before additional masking, live-tile, synchronization, or resource costs. BQ=2 already raises compiled shared allocation to 23,096 bytes while leaving the 16/32-CTA global underfill unchanged. The evidence does not justify BQ=4, BQ=8, production autotuning, or scope expansion.

The absence of flash_attn is a separate blocker for the required shipped-window model promotion evidence, but it is not the reason for this stop. The schedule evidence alone is sufficient. No candidate full-op or model speedup conclusion is claimed, and there is no salvage candidate.

Durable evidence

• profile/parallel_nsa_compression-c15-lane-a-opt/trace/baseline-profile-job-5394.log
• profile/parallel_nsa_compression-c15-lane-a-opt/trace/baseline-profile-replacement-job-5399.log
• profile/parallel_nsa_compression-c15-lane-a-opt/trace/ncu-only-job-5420.log
• profile/parallel_nsa_compression-c15-lane-a-opt/trace/bq2-proof-job-5443.log
• profile/parallel_nsa_compression-c15-lane-a-opt/trace/job-5394/
• profile/parallel_nsa_compression-c15-lane-a-opt/trace/job-5399/
• profile/parallel_nsa_compression-c15-lane-a-opt/trace/job-5420/
• profile/parallel_nsa_compression-c15-lane-a-opt/trace/job-5443/

The issue remains open as a reproducible no-go record. No standalone checksums were generated.