[Perf][B300] Parallelize ABC chunk state construction with affine summaries

[heiheiha798]

[Perf][B300] Parallelize ABC chunk state construction with affine summaries

Problem

chunk_abc is the production operator used by the shipped ABCAttention layer and ABC model. The default model geometry is H=4, K=256, M=64 slots, and V=512.

The forward chunk_abc_fwd_kernel_h and reverse chunk_abc_bwd_kernel_dh kernels launch one CTA per state tile and loop serially over every 64-token chunk. At B1, the KxM state has only 16 CTAs and the MxV state has only 32 CTAs, independent of sequence length, on a 148-SM B300. Each CTA performs a long sequence of tensor-dot updates at T=2K through 16K.

The source-level leading-dot count puts the two recurrent-state products at roughly one third of the default forward operator's principal dot work. This is only a static bound. The optimization is a hypothesis until the untouched B300 profile, correctness, full stage-sum, public-operator, and integration gates below pass.

Proposed Dataflow

Each 64-token state update is an exact diagonal affine transform:

state_out = state_in * r_chunk + c_chunk

where r_chunk is derived from the existing cumulative normalizer z and c_chunk is the existing chunk-local tensor-dot contribution.

For an exact-SM103, dense, state-free training path:

1. Launch an independent producer CTA per (chunk, state tile, batch-head) to compute each current chunk-local contribution in FP32, using the same input casts, exponentials, normalization endpoint, tensor-dot precision, scale, and masks as the existing kernel.
2. Store the fully initialized contributions in one ephemeral FP32 summary tensor.
3. Launch a lightweight ordered scan per (state tile, batch-head) that loads each summary, derives the current diagonal scale from z, writes the existing BF16/FP16 pre-chunk boundary tensor, and performs the unchanged FP32 affine update.
4. Mirror the same split for the reverse dH recurrence only after the complete forward proof passes.

This does not reduce mathematical work. It moves the tensor-dot portion from 16/32 long serial CTAs into NT times more producer CTAs, leaving only scale/add/load/store work in the low-wave scan.

Keep the current kernels as the exact fallback for non-SM103 devices, initial/final-state calls, short or partial sequences, unsupported dtype/layout/shape, B buckets not proven profitable, and every losing control. Do not introduce a common affine-scan framework, atomics, TMA dependency, block pointers, lower-precision summaries, or changes to public APIs, layers, models, references, softmax, scale, or numerical policy.

Scope

Expected tracked files:

• fla/ops/abc/chunk.py
• new tests/ops/test_abc.py in a fixture-only commit, then frozen

Explicitly excluded:

• fla/ops/abc/naive.py
• fla/layers/abc.py and all ABC model files
• common chunk/scan utilities
• benchmark registry/runner and training harnesses
• sibling RetNet work and all other operators

The initial optimized predicate is exact SM103, contiguous FP16/BF16, dense state-free calls, B in {1,2}, T divisible by 64 and at least 2048, H=4, K=256, M=64, V=512. Broaden only to frozen, measured-profitable shapes within this same operator.

Memory Bound

Only one FP32 summary is live at a time. For the larger MxV recurrence at B1, the summary is approximately 16/32/64/128 MiB at T=2048/4096/8192/16384. The KxM summary is half those sizes. B4/T16384 would require 512 MiB for the larger summary and is therefore a required memory/dispatch control, not an assumed fast-path shape.

Peak allocation and producer write + scan read traffic must be measured and charged to every result. No second comparable summary, persistent cache inflation, or uninitialized region is allowed.

Fixture-First Correctness Contract

Before any production edit, add a direct public-layout fixture that compares untouched chunk_abc with the immutable head-first naive_recurrent_abc, using explicit transposes only around the reference. Establish untouched-production dtype-specific tolerances, commit the fixture as BENCH_BASE, then freeze the fixture, naive reference, shapes, seeds, dtypes, tolerances, allocator NaN poisoning, numeric flags, and precision policy for the entire optimization loop.

Cover:

• FP16 and BF16
• B=1/2 and H=1/4
• K=64/128/256, M=32/64, V=64/256/512
• T=63/64/65/127/257 and a small aligned fast-path representative
• initial state on/off and final state on/off
• output, both final states, and independent dQ/dK/dV/dS gradients
• aligned optimized-path and short/tail/state fallback behavior

If untouched production is deterministically red, stop and record no-go. Do not modify the reference, tolerances, precision, skips, or semantics to rescue this issue.

Promotion requires the full frozen fixture, unchanged tests/models/test_modeling_abc.py, every test returned by scripts/find_dependent_tests.py, changed-file lint/header/syntax checks, int64 grid/address review, and the repository block-pointer ban.

Profile-First Gate

On untouched B300, profile public chunk_abc at BF16 B1/B2, T=4096/8192/16384, H4/K256/M64/V512. Attribute both forward h kernels and both reverse dH kernels separately and include all allocations and launches.

Record per-stage and full-op duration/share, grid/waves, achieved occupancy, registers/spills, tensor instruction count, DRAM/L2 traffic, active/eligible warps, and tensor-pipe utilization. Collect NCU full plus PM/warp sampling and source plus SourceCounters for B1/T8192 and T16384.

Stop before source edits unless a conservative producer/summary/scan model supports both at least 1.12x public forward and 1.10x public forward+backward on representative long shapes. Also stop if the low-wave stages are not material, the exact FP32 summary cannot compile, summary traffic/allocation consumes the bound, or the current kernels already expose enough useful concurrency.

Staged Proof

The first implementation proof is forward only at B1/H4/K256/M64/V512/T4096:

• Compare each old h kernel with its complete producer + allocation + scan sum.
• Require each pair to be at least 2.0x faster and the complete public forward to be at least 1.10x faster.
• Require frozen correctness, fully initialized summaries, resource-safe compilation, and no material spills or occupancy collapse.

Do not implement reverse scan/backward unless this proof passes. Correctness-green representative performance shortfalls may use at most two reviewer-approved bounded depth iterations without changing the oracle, matrix, or thresholds.

Representative Benchmarks

Use a reproducible cycle-local harness without changing benchmark infrastructure.

Primary BF16 public-op matrix:

• B={1,2}, T={2048,4096,8192,16384}, H4/K256/M64/V512
• forward and forward+backward

Breadth and fallback controls:

• FP16, B4, T4097, K128, M32, V256
• initial/final state
• short T256/1024

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

Integration measurements are the unchanged public ABCAttention(hidden_size=2048, num_heads=4, num_slots=64) at B1/T2048/4096/8192 and a deterministic local 2-4 layer ABC training step with default mixer geometry and reduced vocabulary. These are dilution/regression controls; no unmeasured model speedup claim is allowed.

NCU must show producer grids gaining useful waves and active SMs, unchanged expected tensor work, no spill/occupancy collapse, and summary/scan traffic consistent with the analytic allocation. Compare every old kernel to its complete replacement stage sum.

Acceptance Gates

• Full frozen output/final-state/all-gradient correctness and dependent/model/static gates pass unchanged.
• Public direct forward geomean is at least 1.12x and forward+backward geomean at least 1.10x over B1/B2 long primary rows.
• B1/T8192 and T16384 are each at least 1.10x in every claimed mode.
• No retained target/control regresses more than 3%; exact fallback controls remain within 2% noise.
• Default ABCAttention long training geomean is positive with no endpoint worse than 2%; the local ABC training control is non-regressing within 2%.
• Peak allocation matches one ephemeral FP32 summary plus allocator granularity, with no comparable hidden buffer.
• Profiler evidence attributes the result to useful CTA concurrency and shorter tensor-dot critical paths, not skipped work or changed numerics.

Record no-go rather than widening scope if only B1/T16384 wins, only forward wins, B2 breadth fails, summary traffic dominates, integration regresses, or correctness requires lower-precision summaries or relaxed tolerances.

Duplicate and Mergeability Audit

No open upstream or fork issue/PR was found for ABC chunk state-scan parallelization or fla/ops/abc/chunk.py. Open upstream PR fla-org#977 changes the optional short-convolution path in fla/layers/abc.py only; layer source is excluded here. No benchmark registry edit is proposed, avoiding upstream PR fla-org#1144 and sibling registry ownership.

The sibling RetNet candidate owns retention-specific superchunk replay, retention files/tests/benchmarks, and RetNet integration. This issue remains ABC-local and must not create shared scan utilities. Historical GDP, Mesa, LightNet, Parallax, NSA, normalized-linear-attention, and other fork scopes are disjoint and excluded.

Baseline: upstream/main@27967b970eaaf982a6960abf6cba8add9c34c7cc.

[heiheiha798]

PR_NO_GO_RECORDED: untouched ABC fixture is correctness-red

Cycle 22 stopped before BENCH_BASE, profiling, NCU, or production work because the frozen fixture exposed deterministic correctness failures in the untouched baseline 27967b970eaaf982a6960abf6cba8add9c34c7cc.

Untouched broad gate

Slurm job 5528 ran on one idle NVIDIA B300 SXM6 AC (SM103) at the exact clean head:

• tests/layers/test_layer_cache_layer_idx.py: 28 passed.
• tests/models/test_hybrid_attention.py: 210 passed, 2 skipped.
• tests/models/test_modeling_abc.py: all 6 cases were skipped by the existing unsupported-varlen guard; there was no failure.

Frozen fixture gate

Slurm job 5531 ran the approved eight-case public-layout chunk_abc versus immutable head-first naive_recurrent_abc oracle on the same exact head and one-file fixture scope. Result: 4 passed, 4 failed.

Failures:

• FP16 B1/T65/H4/K256/M32/V512 with initial and final state: ds absolute error 33.0, error ratio 2.162100 versus the frozen 0.005 limit.
• BF16 B2/T127/H1/K64/M64/V256 with initial state: ds absolute error 6.4375, error ratio 1.568965 versus the frozen 0.02 limit.
• FP16 B1/T257/H4/K128/M32/V64, state-free with final state: output and all input gradients passed, but hkt absolute error 277.777954 and error ratio 430.680296 versus the frozen 0.005 limit.
• BF16 B2/T257/H1/K256/M64/V512 with initial and final state: ds absolute error 9.1875, error ratio 1.362319 versus the frozen 0.02 limit.

The T63 and T64 boundary cases passed. Both exact state-free target rows also passed: BF16 B1/T2048 and FP16 B2/T2048 at H4/K256/M64/V512. Those partial passes do not satisfy the immutable full state/tail/final-state correctness contract.

No fixture tolerance, reference, seed, matrix, dtype, numeric flag, or semantics were changed. There is no BENCH_BASE commit, profile-first timing, NCU capture, source implementation, candidate, performance conclusion, commit, push, PR, or salvage candidate.

Durable evidence:

• Broad gate: profile/c22-lane-b-abc-state-scan-opt/baseline/slurm-5528.out
• Fixture gate: profile/c22-lane-b-abc-state-scan-opt/fixture/slurm-5531.out

The issue remains open. Local branch b300/c22-lane-b-abc-state-scan remains at 27967b970eaaf982a6960abf6cba8add9c34c7cc; the untracked fixture, linked worktree, and cycle evidence are preserved.