[Perf][B300] Fuse RWKV6 intra and inter chunk backward to remove FP32 dQ/dK round trips

[heiheiha798]

Problem

chunk_rwkv6_bwd currently runs chunk_rwkv6_bwd_kernel_intra, materializes full FP32 dq and dk intermediates, and then runs chunk_rwkv6_bwd_kernel_inter, which immediately reloads those tensors before writing the final gradients.

That staging adds one kernel launch and about 16 bytes of avoidable global traffic per q/k element: two FP32 tensors are each written once and read once. It also retains two full-size temporary allocations throughout this part of backward. The registered chunk_rwkv6 operator and the default chunk-mode RWKV6 training path both exercise this flow for dense and packed-variable-length inputs.

This issue evaluates a B300 optimization hypothesis. No speedup is claimed until the frozen correctness gate, repository-native benchmarks, and profiler evidence below pass.

Proposed optimization

Fold the existing off-diagonal intra-chunk dQ/dK calculation into the inter kernel for a profile-proven specialization. Accumulate those terms into the inter kernel's existing FP32 dQ/dK accumulators and write the final gradients directly, eliminating the standalone intra launch and full FP32 intermediate round trip on the optimized path.

Preserve the current mathematical and numerical sequence exactly:

1. Compute the inter/state dQ/dK contributions in FP32.
2. Form the existing state-derived b_dgk term before adding intra dK.
3. Add the intra off-diagonal contributions using the current BC=16 block-local gate centers, strict within-subblock masks, exponent placement, and FP32 accumulation.
4. Construct dG from the combined inter plus off-diagonal intra dQ/dK values.
5. Apply the current true-diagonal u bonus to final dQ/dK after dG construction, and preserve dU handling.

Do not reassociate gate exponentials around a different center, reduce accumulation precision, enable different numeric flags, or substitute the optional TileLang backend. Use tl.int64 for program/grid-derived indices and all tensor base/stride/varlen address arithmetic.

The current two-kernel implementation remains the exact fallback for unsupported or losing shapes. If B300 profiling shows that the fused register schedule is architecture-specific, guard it with an exact (10, 3) capability check using the repository's device-helper pattern. Shape dispatch is allowed only for resource/profile-proven buckets; do not create a narrow lookup table to rescue isolated rows.

Scope and lane isolation

Start from exact upstream baseline 27967b970eaaf982a6960abf6cba8add9c34c7cc in the registered lane-b worktree. The fork default is currently behind this baseline, so branch/PR work must wait for MAIN's separate safe fork-default synchronization gate.

Expected tracked paths:

• fla/ops/rwkv6/chunk.py for the production optimization and fallback dispatch.
• tests/ops/test_rwkv6.py only in an initial oracle-fixture commit that adds missing BF16 coverage, then frozen for the entire optimization loop.

Do not edit benchmarks/ops/registry.py or benchmarks/ops/run.py. chunk_rwkv6 is already registered, and those files are owned by the concurrently active lane-a GSA cycle. Do not edit the RWKV6 public API, forward kernels, recurrent/naive reference, model/layer files, tolerances, precision policy, TileLang backend, or common chunk kernels.

Use a cycle/lane-specific ignored evidence directory such as profile/c12-lane-b-rwkv6-bwd-fusion-opt/; do not read or write lane-a's worktree or artifacts.

Fixture-first frozen correctness contract

Before production edits:

1. Run the complete untouched tests/ops/test_rwkv6.py on B300 at the exact source baseline. A deterministic baseline failure is a terminal no-go, not permission to narrow the gate.
2. Add focused BF16 cases for every proposed optimized dispatch boundary. At minimum cover a dense partial chunk and an unequal packed-varlen case, the intended BT=64 path, and the D classes that may be promoted. If BT=16/32 will be optimized, add those BF16 cases before freeze as well.
3. Run the full file green against untouched production code and commit this test-only fixture as immutable BENCH_BASE.
4. Freeze tests/ops/test_rwkv6.py, fla/ops/rwkv6/fused_recurrent.py, all reference behavior, tolerances, shapes, seeds, dtypes, and numeric flags for every kernel iteration.

The existing oracle checks output/final state and dq/dk/dv/dw/du/dh0 against fused_recurrent_rwkv6, including dense tails, D=60/64/256, T=15 through 2048, chunk sizes 16/32/64, initial/final state, and irregular packed varlen D=64/100. Allocator NaN poisoning remains mandatory. The fixture closes the current gap that repository benchmarks use BF16 while the test file currently covers only FP16/FP32.

Per iteration, run the full unfiltered correctness gate before retaining or benchmarking a candidate. Promotion requires:

python -m benchmarks.ops.verify --op chunk_rwkv6 --base <BENCH_BASE> --modes fwdbwd
python scripts/find_dependent_tests.py fla/ops/rwkv6/chunk.py
ruff check fla/ops/rwkv6/chunk.py tests/ops/test_rwkv6.py

Run all dependent tests returned by the repository helper and the required changed-file pre-commit/build checks. No tolerance, test, reference, skip, environment suppression, or precision edit is allowed after BENCH_BASE.

Profile-first B300 gate

Before production edits, capture the untouched B300 backward timeline and allocation profile for representative dense BF16 D=64/128/256 rows and one unequal packed-varlen case. Attribute:

• standalone intra-kernel duration;
• inter-kernel duration;
• FP32 dQ/dK temporary allocation and write/read traffic;
• the complete public chunk_rwkv6 fwdbwd endpoint share.

Stop before broad implementation if the launch plus removable round-trip cost cannot plausibly support the 1.08x full-endpoint bar. After a minimal proof kernel, stop rather than widen scope if D=128/256 cannot compile, local-memory spills appear, occupancy collapses, or the fused schedule loses enough parallelism to erase the traffic benefit.

Repository-native benchmark plan

Use the existing six-row chunk_rwkv6 registry in BF16 and fwdbwd mode, comparing production HEAD with immutable BENCH_BASE in the same isolated B300 allocation. The matrix spans D=64/128/256, B=1..8, T=1024..16384, and H=8..96.

Add no tracked benchmark adapter. Measure varlen through an ignored, reproducible harness that calls the public chunk_rwkv6 endpoint on the frozen unequal packed cases with identical baseline/candidate inputs and timing discipline.

Final comparisons must use warmed autotune caches and at least five interleaved baseline/candidate rounds. Report each row's median, mean, standard deviation, min, p10, p20, p80, p90, sample count, and equal-weight geomean. Measure peak allocated memory and confirm the two FP32 q/k-sized temporaries disappear only on the optimized dispatch.

After operator promotion, benchmark the public RWKV6 model training path on identical dense and --varlen workloads using benchmarks/benchmark_training_throughput.py --name rwkv6. Use the same model shape/layer count, inputs, warmup, steps, clocks, and allocation for baseline and candidate. Report tokens/s distributions and peak memory; disclose any layer-count reduction required by B300 memory.

NCU evidence

Collect both required Nsight Compute passes for representative dense D=128 and D=256 optimized rows and one packed-varlen row:

• --set full --section PmSampling --section PmSampling_WarpStates
• --set source --section SourceCounters

Compare the summed baseline intra plus inter kernels with the fused kernel: duration, launch count, DRAM read/write bytes and throughput, L2 sectors/hit behavior, SM/tensor SOL, achieved occupancy, registers, local-memory spills, active/eligible warps, and dominant stalls. Confirm the intra launch and FP32 intermediate loads/stores disappear without skipped work, hidden replacement kernels, or numeric-flag changes.

Keep raw reports and logs outside git under the lane-specific evidence directory. Record exact baseline/candidate SHAs, commands, Slurm job IDs, GPU capability, clocks/idle evidence, and CUDA/PyTorch/Triton/NCU versions.

Acceptance and stop conditions

• The complete frozen RWKV6 gate and all dependent checks pass unchanged.
• Equal-weight BF16 fwdbwd geomean across all six registered rows improves at least 1.08x.
• Every D=64/128/256 class has at least one positive optimized-path result, and no retained measured row regresses more than 3% outside noise.
• Dense and packed-varlen optimized-path endpoint measurements are both non-regressing; the representative dense RWKV6 model training run improves at least 3% median tokens/s, while the varlen model run is neutral within 2% or better.
• Peak allocation decreases by the measured two FP32 intermediate sizes on optimized rows, with no asymptotically comparable replacement scratch.
• NCU attributes the result to removed launch/global traffic without spill or occupancy collapse.

Fallback rows, forward controls, and unsupported architectures must remain within 2% noise. If the gain survives only in an artificial microshape, requires relaxed numerics, overlaps another backend/common-kernel rewrite, or misses the endpoint/model thresholds, record no-go rather than changing the issue into a TileLang or broader RWKV6 rewrite.

Duplicate and mergeability audit

At upstream 27967b97, no open upstream or fork issue/PR and no remote branch targets RWKV6 intra/inter backward fusion or the FP32 dQ/dK staging tensors. Merged upstream PR fla-org#1045 adds an opt-in TileLang RWKV6 intra-forward kernel and explicitly leaves backward-specific kernels in Triton. Open PR fla-org#1046 changes common/GLA/Delta/GDN TileLang backward paths and does not touch RWKV6.

Lane-a issue #12 and branch b300/gsa-tile-parallel-cached-decode touch only tests/ops/test_gsa.py, benchmarks/ops/registry.py, and benchmarks/ops/run.py; this normalized lane-b scope changes RWKV6 source/test paths and uses the existing registry without editing shared benchmark files. The tracked-file ownership and timed operator/model surfaces are therefore disjoint.

If the thresholds pass, this is a focused, independently mergeable Triton training-backward optimization: one public operator, one full-tensor round trip removed, an exact fallback, no API/model-format change, and a profiler-verifiable B300 mechanism.

[heiheiha798]

B300 validation result: no-go

Validation used upstream baseline 27967b970eaaf982a6960abf6cba8add9c34c7cc and the frozen fixture commit da428a7b262aeced290ebc73f594e6797fd04d1f (BENCH_BASE).

Correctness remained green throughout the required gates:

• Untouched baseline: 13 passed (Slurm job 5295).
• Frozen BF16 dense/packed-varlen fixture gate: 17 passed (job 5303).
• Final fused candidate against the frozen gate: 17 passed (job 5333).
• Ruff, Python syntax, frozen-file, and scoped-diff checks passed.

The exact SM103/BT64 candidate fused the existing intra dQ/dK work into the inter kernel. Profiling confirmed that it removed the standalone intra launch and eliminated both full FP32 temporary round trips, with peak-memory reductions matching the expected temporary sizes exactly.

Shape	Baseline intra + inter (us)	Fused (us)	Pair speedup	Peak-memory reduction (bytes)
Dense D64	285.356	267.717	1.065887x	33,554,432
Dense D128	1,264.609	1,327.847	0.952376x	134,217,728
Dense D256	10,481.178	11,219.712	0.934175x	1,073,741,824
Packed varlen D128	160.956	139.972	1.149916x	16,769,024

The baseline profile was job 5308 and the candidate profile was job 5335, using the same harness, inputs, and environment. D64 did not reach the required 1.08x bar, while D128 and D256 regressed. This triggers the issue's explicit profile-first stop condition, so the six-row registry benchmark, model benchmarks, and NCU collection were intentionally not run.

The local prototype is retained on branch b300/c12-lane-b-rwkv6-bwd-fusion: HEAD is the frozen BENCH_BASE commit da428a7b262aeced290ebc73f594e6797fd04d1f, with the scoped production prototype remaining as an uncommitted tracked change. The branch was not pushed and no draft PR was created.

Retained evidence in the preserved cycle-12 lane-b linked worktree:

• profile/c12-lane-b-rwkv6-bwd-fusion-opt/OPT_LOG.md
• profile/c12-lane-b-rwkv6-bwd-fusion-opt/logs/c12-lane-b-fused-gate2-5333.log
• profile/c12-lane-b-rwkv6-bwd-fusion-opt/logs/c12-lane-b-profile-fused-5335.log
• profile/c12-lane-b-rwkv6-bwd-fusion-opt/trace/baseline/summary.json
• profile/c12-lane-b-rwkv6-bwd-fusion-opt/trace/fused-v1/summary.json

No salvage candidate is proposed: the bounded fused schedule failed the original performance contract, including regressions in both larger D classes. This issue remains open as the validation record; no PR was created.