[Perf][B300] Scope fused linear-attention normalization finalizers to chunk callers

[heiheiha798]

[Perf][B300] Scope fused linear-attention normalization finalizers to chunk callers

Context and lineage

This is the single bounded salvage continuation of source issue #15 and its recorded exact-head result: #15 (comment).

The source cycle produced a fast exact-SM103 normalization finalizer for normalized linear attention, but the draft head is not correct for every shared caller. Review job 5418 exercised the actual BF16 fused_recurrent_linear_attn path at B1/T35/H2/K64/V80 and found dq error 2.124023% above the frozen 2.000000% oracle. That is a real candidate correctness failure, so the old draft PR #17 must not be reused or presented upstream.

The same evidence also showed that the finalizer is valuable on the originally selected chunk_linear_attn and fused_chunk_linear_attn training paths. This salvage starts independently from current upstream/main at 27967b970eaaf982a6960abf6cba8add9c34c7cc; it must recreate the fixture and implementation rather than cherry-pick or extend the c14 branch.

Historical evidence

• Untouched frozen fixture: 42/42 passed on B300 (job 5353).
• Final direct helper/fallback/resource gate: 42/42 passed; retained kernels used 15-21 registers with no stack or local spills (job 5384).
• Representative chunk fwd/fwdbwd geometric means were 1.10189x/1.11970x. Explicit normalized public-layer fwdbwd was 1.16148x for chunk and 1.14671x for fused chunk; dense and packed fused-chunk controls were positive (job 5385).
• All six existing BF16 registry rows improved. Fwd/fwdbwd geometric means were 1.27242x/1.23112x (job 5389).
• The representative full forward launch count fell 9 to 5 and fwdbwd 28 to 15. The normalization fragment fell from 5 to 1 forward launches and 10 to 1 backward launches. NCU measured normalization DRAM traffic down 34.81% forward and 6.89% backward, with 16/19 registers and zero spills (job 5393).
• Exact-head review then found the reachable BF16 fused-recurrent dq failure above. FP16 recurrent, public normalized short-sequence routing, dependent suites, canonical frozen-base native commands, and broader caller controls were not completed (job 5418).

The salvage hypothesis is therefore not that the failed shared helper is acceptable. It is that the proven finalizer can be made safe by narrowing dispatch to the two selected chunk callers while retaining the exact eager normalization boundary for fused recurrent and all other callers.

Bounded implementation

• Recreate the private SM103 forward/backward row finalizers in fla/ops/linear_attn/utils.py, preserving the historical eager dtype-rounding boundaries and retained two-warp-through-padded-width-128 schedule.
• Add a private, non-public caller opt-in to normalize_with_z_state / its autograd function. The default must be False.
• Only chunk_linear_attn and fused_chunk_linear_attn may pass the opt-in. fused_recurrent_linear_attn must remain source-compatible and always use the unchanged eager normalization path, regardless of sequence length.
• The fused path remains restricted to exact SM103, contiguous FP16/BF16 inputs, reverse=False, z_init=None, matched dtypes, and K/V no larger than 256. Noncontiguous backward gradients and every unsupported case keep the eager fallback.
• Prefix/reverse cumsums, z-state handling, numerator kernels, public signatures, normalization math, tolerances, precision flags, and other architectures remain unchanged.

A sequence-length-only guard is not sufficient because direct fused-recurrent calls can use longer T. Caller ownership must be explicit and private. The implementation must demonstrate that no fused normalization kernel launches from fused recurrent, including the historical failing row and long direct controls.

Tracked scope

Expected tracked files only:

• fla/ops/linear_attn/utils.py
• fla/ops/linear_attn/chunk.py
• fla/ops/linear_attn/fused_chunk.py
• fixture-first, then frozen tests/ops/test_linear_attn.py

No edit is expected in fla/ops/linear_attn/fused_recurrent.py: the private opt-in defaults to false. Benchmark registry/runner files, naive/reference code, cumsum/common/simple-GLA kernels, layers/models, public APIs, and other modules are excluded.

Fixture-first correctness contract

Before production edits, recreate the source fixture on untouched current main, add the missing caller-isolation coverage, run it on one B300, and commit the test-only result as the new immutable BENCH_BASE. Then freeze references, tolerances, inputs, seeds, dtypes, numeric flags, and skips.

Coverage must include:

• the historical positive-feature FP16/BF16 dense K64/K128/K256 and unequal packed-varlen output/dq/dk/dv cases;
• the exact failing BF16 fused-recurrent B1/T35/H2/K64/V80 row against naive_recurrent_linear_attn, plus its FP16 counterpart;
• fused-recurrent direct controls at T=1/35/64/65/257, dense and applicable packed-varlen/state/reverse cases, proving the eager path remains selected;
• public normalized LinearAttention short routes at T<=64, including output, input gradient, and parameter gradients;
• chunk and fused-chunk dense/packed rows that explicitly prove the finalizers launch only under the private opt-in;
• contiguous and noncontiguous upstream gradients, z_init, reverse, FP32, non-SM103, K/V>256, and allocator NaN-poison fallback controls;
• complete output/gradient writes, finite values, and deterministic repetition.

The historical recurrent failure must pass without changing its reference, tolerance, dtype, scale, seed, or feature distribution. Any fused-recurrent launch of _normalize_fwd_kernel or _normalize_bwd_kernel is a correctness/scope failure.

Every iteration must run the full frozen tests/ops/test_linear_attn.py, all tests returned by scripts/find_dependent_tests.py for the four tracked files, the relevant linear-attention model/layer tests, and unchanged syntax/lint/header/int64/block-pointer checks. No new skip or numerical relaxation is allowed.

Measurement and profiler plan

Use the existing repository chunk_linear_attn registration without tracked benchmark edits. Compare immutable BENCH_BASE and candidate in separate pinned roots on one exclusive B300 with identical inputs, warmed state, at least five interleaved rounds, and complete sample distributions.

Measure:

• all six existing BF16 chunk_linear_attn registry rows in fwd and fwdbwd;
• the historical dense D64/D128/D256 and unequal packed D128/V96 representative rows;
• normalized fused_chunk_linear_attn dense/packed controls;
• explicit normalized public LinearAttention chunk and fused-chunk endpoints;
• fused-recurrent and public T<=64 fallback controls, which must remain within noise of the eager baseline.

Report median, mean, standard deviation, min, p10/p20/p80/p90, sample count, equal-weight geometric means, peak allocation, exact commits/environment, and job IDs. Run the canonical repository verify/run commands against the new immutable base; do not substitute an ignored microbenchmark for the native gate.

Capture full and source NCU plus a launch timeline at BF16 B4/T2048/H16/K128/V128. Compare the complete operator and normalization fragment, including launch count, DRAM/L2 traffic, duration, registers, shared/local memory, spills, occupancy, active/eligible warps, and source-correlated loads/reductions/stores. Timeline evidence must also show zero finalizer launches from the fused-recurrent controls.

Unchanged acceptance thresholds

• Full frozen direct/caller/dependent/model correctness green with no oracle, tolerance, precision, seed, flag, or skip change.
• Existing six-row BF16 registry fwd geometric mean at least 1.05x and fwdbwd geometric mean at least 1.08x.
• Every retained D class has a positive optimized row; no representative row regresses more than 3%.
• Dense and packed chunk/fused-chunk endpoints are non-regressing.
• Explicit normalized public integration improves at least 3% on a dense endpoint and remains within 2% on the packed endpoint.
• Fused-recurrent, public short, and all other fallback controls remain within 2% noise, with no finalizer launch.
• Full/source NCU corroborates reduced launches/traffic without spills or hidden work, and peak memory does not regress.

These are the original source-issue thresholds, not relaxed salvage gates. If correctness is green but representative performance misses, use the mandatory performance depth-review process. At most two DEEPEN decisions may adjust only the same row-finalizer schedule or caller-local launch/dataflow; the module, architecture, correctness contract, and thresholds remain fixed. This salvage attempt cannot create another salvage candidate.

Duplicate, overlap, and mergeability boundaries

• Current upstream main remains unchanged on this mechanism. No open upstream issue/PR or remote branch implements the SM103 normalized linear-attention finalizers. The old l2norm remote branch is historical and has no changes to fla/ops/linear_attn or its tests from its merge base.
• Draft PR Fuse SM103 linear-attention normalization epilogues #17 is the rejected source head and evidence only. Do not update, reuse, or stack it.
• Active salvage issue [Perf][B300] Productionize BQ=2 tiling for NSA compression dKV backward #26 owns NSA compression dKV, fla/ops/nsa/compression.py, tests/ops/test_nsa.py, optional NSA registry metadata, and NSA integration/dependencies. This issue owns no NSA or benchmark infrastructure and must not use that cycle's worktree, environment, or artifacts.
• Open benchmark PR [Test] Add realistic benchmark input profiles fla-org/flash-linear-attention#1144 is irrelevant because this issue makes no benchmark-file change.

The principal risks are caller-dispatch leakage, fused-recurrent numerical interaction, private autograd signature plumbing, and historical performance not surviving the current independent implementation. The explicit default-false caller gate, frozen recurrent oracle, native endpoint measurements, and unchanged eager fallback bound those risks.

[heiheiha798]

PR_NO_GO_RECORDED — c27 caller-isolation fixture gate

Cycle 27 was a fresh-current-main salvage of #15/#17, scoped to opt-in SM103 normalization finalizers for chunk and fused-chunk callers while keeping fused-recurrent/public-short routes on the eager boundary. It stopped before BENCH_BASE and before any production candidate.

Untouched baseline

B300/SM103 job 5566 ran at exact clean head 27967b970eaaf982a6960abf6cba8add9c34c7cc and passed the complete pre-fixture gate:

• tests/ops/test_linear_attn.py: 30 passed
• tests/layers/test_layer_cache_layer_idx.py: 28 passed
• tests/models/test_hybrid_attention.py: 210 passed, 2 skipped
• tests/models/test_modeling_linear_attn.py: 1 passed, 5 existing unsupported-varlen skips

The durable log ends BASELINE_GATE_PASS at profile/c27-lane-b-linear-attn-optin-finalizers-opt/baseline/c27-lane-b-baseline-5566.log.

Frozen fixture result

B300/SM103 job 5568 ran the new uncommitted caller/oracle fixture against untouched production at the same exact head. It ended 14 failed, 45 passed; set -e correctly stopped before the three dependent suites.

Ten direct eager-boundary rows failed inside the test-local reference construction before comparing production outputs: torch.autograd.grad reported that reference k was unused because the head-first chunk_global_cumsum expression has no autograd edge to k. Those are fixture-reference construction reds, not candidate correctness results.

Four independently valid untouched fused_recurrent_linear_attn versus naive_recurrent_linear_attn rows exceeded the frozen dQ ratios:

• BF16 T1: diff 0.000732, ratio 15048.624349 > 0.02
• FP16 T64: diff 0.031433, ratio 0.323746 > 0.005
• FP16 T257: diff 0.054443, ratio 0.863499 > 0.005
• FP16 reverse T65: diff 0.054077, ratio 0.390213 > 0.005

The exact historical regression rows were green on untouched production: BF16 B1/T35/H2/K64/V80 dQ ratio 0.012006 < 0.02, and its FP16 counterpart dQ ratio 0.000000 < 0.005. BF16 T65, packed, and initial-state controls also passed. All six dense/packed caller opt-in/isolation controls passed, and all four FP16/BF16 public T<=64 chunk/fused-chunk versus recurrent output/input/parameter-gradient controls passed exactly with zero finalizer launches.

The durable fixture log is profile/c27-lane-b-linear-attn-optin-finalizers-opt/fixture/c27-lane-b-fixture-5568.log.

Disposition

The frozen baseline/oracle contract is not green, and the authorized no-correction/no-retry boundary is exhausted. No BENCH_BASE, production finalizer, source correctness, performance/profile/NCU result, commit, push, or PR was produced in c27, so this records no candidate correctness or speedup claim. No further salvage is proposed. The fork issue remains OPEN; the local branch stays at exact baseline with only the uncommitted fixture, and the c27 worktree/environment/logs are preserved.