[Perf][B300] Recompute PKDA dAqk on chip in backward

[heiheiha798]

Problem

chunk_precond_kda is a registered public operator and the default chunk-mode training/prefill path for the shipped PrecondKDA model. Its backward currently computes dAqk = do @ v_new^T in chunk_precond_kda_bwd_kernel_dAv, writes a full causal FP32 [B, T, H, 64] tensor, and immediately reloads that tensor in chunk_precond_kda_bwd_kernel_intra.

At representative registered shapes, the dAqk allocation is 0.188 GiB for B=1,T=8192,H=96 and 0.250 GiB for B=4,T=4096,H=64; producer write plus consumer read doubles those traffic figures. The public model-default B=4,T=2048,H=16,D=128 path still moves 64 MiB per layer invocation. This issue evaluates whether an exact-SM103 on-chip recomputation schedule is faster than that full-tensor round trip. No speedup is claimed until the frozen correctness, native endpoint, and NCU gates below pass.

Revised optimization contract

The current intra kernel consumes two different slices of the asymmetric dAqk matrix:

• dQ needs a [BC, BT] row fragment for the owned 16-token subchunk.
• dK-preconditioned needs a [BT, BC] column fragment for the same token subchunk.

Those fragments are not transposes and one [16,64] tile cannot serve both paths. The optimized design must therefore make the tradeoff explicit:

1. Split the existing dAv producer into a lean dv-only kernel. Keep dAkk materialized and unchanged.
2. On exact SM103 only, launch one intra CTA per (chunk, head, 16-token subchunk) rather than per K tile.
3. Recompute both orthogonal FP32 fragments on chip from do and v_new, apply the existing causal mask and scale exactly, and loop over K tiles inside the CTA so each fragment is reused for all dQ/dK-preconditioned work.
4. Accumulate beta over the internal K loop and write its final per-token value directly, eliminating the NK-sized beta partial workspace and reduction when the optimized path is used.
5. Preserve the current global dAkk contract, gate centering/exp placement, safe-gate behavior, FP32 accumulation, varlen semantics, and final gradient layout.

This recomputes each required dAqk element in a row-owned and a column-owned fragment instead of pretending the asymmetric matrix is reusable by transpose. Promotion requires profiling to show that the duplicated tensor-core work is cheaper than the eliminated FP32 allocation/store/load and beta reduction.

Use an exact get_device_capability(device.index) == (10, 3) host dispatch. All other devices, unsupported/resource-heavy shapes, and measured losing buckets retain the current implementation unchanged. Do not add atomics, a full-chunk [64,64] live fragment, another asymptotically comparable workspace, a new backend dependency, or changes to the public API, ATK, context-parallel behavior, reference math, tolerances, or numeric flags.

Scope and lane isolation

Start from exact upstream baseline 27967b970eaaf982a6960abf6cba8add9c34c7cc in the registered lane-a worktree.

Expected tracked paths:

• fla/ops/precond_kda/chunk.py
• fla/ops/precond_kda/chunk_bwd.py
• fla/ops/precond_kda/chunk_intra.py
• tests/ops/test_precond_kda.py only in an initial oracle-fixture commit, then frozen

Do not edit fla/ops/precond_kda/naive.py, benchmark registry/runner files, model/layer files, ATK, common KDA, utilities, or sibling-lane files. Lane-b issue #13 is confined to RWKV6 source/test paths and the existing chunk_rwkv6 benchmark surface. Prior lane-a issue #12 is terminal and concerns GSA cached decode; it does not overlap these files or this operator.

Keep all raw logs, profiler reports, scripts, and the optimization log outside git under a cycle/lane-specific ignored directory such as profile/chunk_precond_kda-c13-lane-a-opt/.

Fixture-first frozen correctness contract

Before production edits:

1. Run the complete untouched B300 gate at the exact baseline:

   python -m benchmarks.ops.verify --op chunk_precond_kda --base 27967b970eaaf982a6960abf6cba8add9c34c7cc --modes fwdbwd

   A deterministic red/OOR baseline or an unreproducible registered shape is a terminal no-go, not permission to narrow or relax the oracle.

2. Add focused, deterministic dense and ragged-varlen K != V cases for K=64,V=128 and K=128,V=64, exercising forward and every reachable gradient against the existing naive reference. Validate them against untouched production code and commit this test-only fixture as immutable BENCH_BASE.

3. Freeze tests/ops/test_precond_kda.py, fla/ops/precond_kda/naive.py, all references, tolerances, shapes, seeds, dtypes, skips, and numeric flags for every kernel iteration.

The existing full file remains mandatory. It covers dense/varlen forward and backward, safe and ordinary gates, in-kernel and pre-gated modes, L2 normalization, recompute on/off, initial/final A state, transposed state layout, partial/non-power-of-two dimensions, and allocator NaN poisoning.

Per-iteration subsets may provide fast signal, but promotion requires the full unfiltered fixture gate, all model tests, dependent tests, and repository compile/lint/pre-commit checks. No tolerance, reference, skip, precision, or environment-suppression edit is allowed after BENCH_BASE.

Profile-first B300 gate

Before production edits, profile the untouched exact baseline at the public model-default geometry B=4,T=2048,H=16,K=V=128 and representative registered D64/D128 shapes. Attribute the complete fwdbwd endpoint, dAv producer, intra consumer, dAqk allocation/store/load traffic, and beta partial reduction.

Proceed only if the affected stages are at least 20% of fwdbwd and the specifically removable traffic/reduction has an Amdahl bound capable of a 10% full-op improvement. The full dAv duration is not all removable because dv remains.

After the smallest proof kernel, stop rather than widen scope if either orthogonal fragment or the internal K loop fails to compile, spills materially, collapses occupancy, loses too much NK parallelism, or duplicated dAqk dot work erases the traffic benefit. Test D256 and ragged tails as resource/fallback boundaries before broad tuning.

Repository-native benchmark plan

Compare production HEAD with immutable BENCH_BASE in the same isolated B300 allocation:

python -m benchmarks.ops.verify --op chunk_precond_kda --base <BENCH_BASE> --modes fwd fwdbwd
python -m benchmarks.ops.run --op chunk_precond_kda --base <BENCH_BASE> --modes fwd fwdbwd

Use all valid registered rows, plus native custom D64 and model-default B=4,T=2048,H=16,D=128 shapes. Measure a frozen unequal-varlen public-op case with the same inputs and timing discipline. Forward is an unchanged-path control.

For public-path evidence, compare identical two-layer precond_kda training runs at B=1,T=2048 in dense and --varlen --context_len 2048 modes. Baseline and candidate must use the same seed, environment, warmup, steps, clocks, and allocation.

Final evidence requires at least three interleaved baseline/candidate blocks and per-shape median, mean, standard deviation, min, p10, p90, sample count, and equal-weight geomean, plus measured peak allocation.

NCU evidence

Collect both required Nsight Compute passes on B=4,T=2048,H=16,D=128 and one ragged-varlen optimized row:

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

Compare the summed baseline dAv/intra/beta-reduction sequence with the candidate dv-only/on-chip-intra sequence: duration, launches, DRAM read/write bytes and throughput, L2 traffic, tensor instruction count, SM/tensor SOL, achieved occupancy, registers, local-memory spills, active/eligible warps, and dominant stalls. The report must show the global dAqk store/load and NK beta workspace/reduction disappear while making the duplicated dot work visible; timing alone is insufficient.

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

Acceptance and stop conditions

• Full frozen op fixture, K != V coverage, model tests, and dependent checks pass unchanged.
• Combined replaced backward stages improve at least 1.25x median at the public D128 shape.
• Public-default fwdbwd improves at least 1.10x; optimized D64/D128 rows achieve at least 1.08x equal-weight fwdbwd geomean.
• Every registered row, D256/resource fallback, varlen endpoint, and forward control is non-regressing within 2% outside demonstrated noise.
• Dense and varlen two-layer training are each neutral within 2%, with at least one improving at least 1.05x.
• Peak allocation confirms removal of the full FP32 dAqk tensor and NK beta partials without asymptotically comparable replacement scratch.
• NCU attributes the result to lower global traffic/reduction work without material spill or occupancy collapse.

Stop and record no-go after three logged, materially distinct schedule attempts if these endpoints fail. Do not convert this issue into ATK work, standard KDA work, a benchmark-only tune, a TileLang rewrite, or a broader backward refactor.

Duplicate and mergeability audit

At upstream 27967b97, no open upstream/fork issue or PR and no current remote branch targets PKDA backward or on-chip dAqk consumption. The old feature/kda_bwd_improvement branch is fully behind main and predates PKDA. Standard-KDA work such as fla-org#1053/fla-org#1054, fla-org#1112, and fla-org#1128 targets fla/ops/kda/** or alternate backends, while fla-org#1046 covers other linear-attention families. PKDA was introduced by merged fla-org#950 and its source has not had a follow-up performance rewrite.

If the evidence gates pass, this is a focused, independently mergeable B300 optimization: one registered/default-model training path, one full FP32 recurrence workspace eliminated, exact architecture and shape fallback, no API/checkpoint change, and a profiler-verifiable recompute-versus-HBM systems result.

[heiheiha798]

B300 evaluation result: performance no-go

The exact-SM103 on-chip dAqk prototype was evaluated against source baseline 27967b970eaaf982a6960abf6cba8add9c34c7cc and immutable fixture commit cf63e06f9875002d565bc1e1f6ad51cb03ce5d30 on an NVIDIA B300 SXM6 AC (capability 10.3, CUDA 13.0, PyTorch 2.13.0+cu130, Triton 3.7.1).

The prototype implemented the issue's revised ownership exactly: a dv-only producer, one intra CTA per chunk/head/16-token subchunk, both FP32 row and column dAqk fragments recomputed from do and v_new, an internal K-tile loop, direct beta accumulation, and exact legacy fallback outside the guarded SM103 fp16/bf16 K/V<=128 path. It retained materialized dAkk, unchanged APIs/numerics/references, int64 addressing, and introduced no atomics, block pointers, or comparable scratch.

Correctness and viability

• Untouched gate job 5311: complete registered verifier green.
• Frozen-fixture job 5321: 49/49 tests green against untouched production, including the added dense K64/V128 and uneven packed K128/V64 cases.
• Baseline profile/NCU job 5337: affected stages were 30.20% of public D128 fwdbwd and 33.73% of long D128 fwdbwd. Public dAv/intra/beta totaled 1.357376 ms of 4.494112 ms, so the original Amdahl gate supported a bounded proof.
• Dual-fragment proof job 5340: public, D256, and ragged cases compiled and were deterministic/finite; the synthetic proof kernel used 78 registers, had no spills, and reached 34.70% public achieved occupancy.
• Production validation job 5342: completed 0:0; the frozen repository-native gate passed 49/49 (49 passed, 15 warnings) before performance measurement.

No oracle, tolerance, dtype, skip, seed, numeric flag, or reference was changed after the fixture commit.

Same-session performance

Five interleaved baseline/candidate rounds were collected with 30 samples per round (n=150 per source and shape). Speedup is baseline median divided by candidate median.

Shape	Baseline median (ms)	Candidate median (ms)	Speedup	Baseline p90 (ms)	Candidate p90 (ms)
Public B4/T2048/H16/K128/V128	4.479392	6.248864	0.716833x	4.497792	6.272333
Registered B8/T1024/H8/K64/V64	1.870960	2.321696	0.805859x	1.916208	2.361619
Long B1/T8192/H96/K128/V128	23.608640	34.771423	0.678967x	23.656915	34.826191
Dense asymmetric B2/T513/H4/K64/V128	1.519408	1.466448	1.036114x	1.584122	1.527798
Ragged asymmetric total-T513/H4/K128/V64	1.774816	1.722960	1.030097x	1.841050	1.782557
D256 legacy fallback B1/T2048/H16	3.268768	3.269536	0.999765x	3.289846	3.290659

The public/D64/long representative geomean was 0.731996x, versus the required >=1.08x optimized-row geomean and >=1.10x public endpoint. Repository-native rows independently reproduced 0.68x to 0.83x for every dispatched D64/D128 class. The two approximately 1.03x small asymmetric wins are not representative and do not satisfy the public/default contract.

The intended storage removal was real:

• Public D128 peak allocated delta fell 1,147,150,336 -> 1,112,547,840 bytes (-3.02%).
• Registered D64 fell 295,703,040 -> 278,925,824 bytes (-5.67%).
• Long D128 fell 6,882,902,016 -> 6,675,284,480 bytes (-3.02%).
• D256 fallback allocation was unchanged.

The active bound is structural. At K128 the candidate removes the legacy NK-owned program dimension, serializes four K-tile recurrence bodies in each CTA, and evaluates row plus column coverage of the full do @ v_new domain twice. The public grid still contains about 8,192 intra CTAs, so simple global underfill does not explain the regression. Holding the profiled unaffected public time fixed estimates the candidate replacement at about 3.112 ms versus the baseline affected sequence's 1.357 ms. Meeting the unchanged public threshold would require roughly a 3.33x speedup of the current replacement and then a 1.495x improvement over the already faster materialized baseline sequence.

The performance depth review returned STOP at depth iteration 0. Halving the fragment/subchunk size would retain duplicated dot work and internal K serialization; splitting row and column consumers would retain both full-domain recomputations and K loops while adding a launch/global gradient handoff. Neither bounded follow-up has evidence capable of recovering the measured 28-47% representative regressions and then clearing the original promotion margin.

Candidate model integration and candidate NCU were not run after the hard operator threshold failed. The prototype remains uncommitted and unpushed on the local cycle branch; no PR was created and no salvage candidate is proposed. This is a reasoned performance no-go, not a correctness failure and not a claim that the general problem is invalid. The issue remains open.

Durable ignored evidence is retained under:

• profile/chunk_precond_kda-c13-lane-a-opt/profile/job-5337/
• profile/chunk_precond_kda-c13-lane-a-opt/ncu/job-5337/
• profile/chunk_precond_kda-c13-lane-a-opt/proof/job-5340/
• profile/chunk_precond_kda-c13-lane-a-opt/candidate/job-5342/aggregate.json
• profile/chunk_precond_kda-c13-lane-a-opt/candidate-validation-5342.log