[Perf][B300] Add a group-parallel Taylor-moment scan for Based attention

[heiheiha798]

[Perf][B300] Add a group-parallel Taylor-moment scan for Based attention

Problem

BasedLinearAttention is a public layer whose default mode="parallel" path calls parallel_based. The default feature geometry is K=16, 12 heads, and V=64 for a hidden size of 768.

The current kernels in fla/ops/based/parallel.py recompute every causal prefix in forward and dQ backward, and every causal suffix in dK/dV backward. The resulting work is quadratic in sequence length. The repository's fused_chunk_based implementation already demonstrates the exact zero-, first-, and second-order Taylor moments needed to evaluate the same polynomial attention in linear work, but its sequence-serial grid exposes only 24 long-running CTAs for the default B1/H12/K16/V64 shape. That leaves a long-context concurrency problem on a 148-SM B300.

This issue is a performance hypothesis. No speedup is claimed until the profile, correctness, resource, operator, and public-layer gates below pass.

Proposed Scope

Add an exact-SM103, dense, long-context path inside parallel_based that uses the existing Taylor algebra with a group-parallel scan schedule:

1. Split the sequence into deterministic coarse groups. The initial proof compares 512- and 1024-token groups at K=16; only measured profitable choices may be retained.
2. Compute one additive FP32 summary per (batch, head, group, value tile), containing the zero-, first-, and second-order denominator and numerator moments.
3. Run an exclusive prefix scan over the small group-summary axis, tiled across independent state dimensions.
4. Replay each group from its scanned prefix and emit its local causal outputs. The added group grid supplies long-context CTA parallelism without materializing token-by-token states.
5. Use the analogous reverse group summaries, exclusive reverse scan, and group-local replay for dQ/dK/dV.

The public signature, score 1 + s + 0.5 * s^2, scale, normalization epsilon, dtype, accumulation policy, and gradients must stay unchanged. The current kernel remains the exact fallback for non-SM103 devices, short sequences, K other than 16, unsupported dtypes/layouts, and any measured losing bucket.

Do not change fla/layers/based.py, fla/ops/based/naive.py, ReBased, generic linear-attention code, cache/state behavior, public APIs, or mathematical semantics. Do not turn this into a generic scan framework or autotuning project.

Memory Bound

The FP32 group-prefix buffer is bounded by:

B * H * NG * (1 + K + K^2) * (1 + V) * 4 bytes.

For B1/H12/T8192/K16/V64 with 512-token groups this is about 13 MiB. Backward may use an equivalently bounded transient reverse buffer. No O(T^2) allocation and no per-token K^2 * V state are allowed.

Fixture-First Correctness Contract

The current tests/ops/test_based.py does not provide a usable oracle: it constructs obsolete head-first public inputs, and its output/gradient assertions are guarded by dtype == torch.float32 even though its parametrized cases are FP16. Before any production edit:

1. Make one fixture-only commit that corrects public input layout, transposes only around the head-first naive reference, and actually checks output plus dQ/dK/dV.
2. Cover FP16, BF16, and a small FP32 control; K=8/16/32; V=60/64/100/128; T=63/111/257/1024; H=1/4/12; normalization on/off; default and unit scale; non-power-of-two tails; and supported noncontiguous inputs.
3. Add a small BasedLinearAttention(mode="parallel") oracle using the shipped scale=1 behavior and the same projected q/k/v weights. Do not use or repair the pre-existing forward_reference semantic mismatch in this performance issue.
4. Establish all dtype-specific tolerances against untouched production, record the fixture commit as BENCH_BASE, then freeze the test, reference, seeds, inputs, dtypes, numeric flags, and tolerances for the complete optimization loop.

Allocator NaN poisoning remains enabled. If untouched production fails the corrected fixture deterministically, stop and record no-go; do not repair semantics or relax the oracle in this issue.

Promotion requires the complete frozen tests/ops/test_based.py, all tests returned by scripts/find_dependent_tests.py, changed-file lint/header/syntax checks, int64 pointer/grid arithmetic review, and the repository block-pointer ban. Dense is the applicable layout; parallel_based has no varlen/state/cache API.

Profile-First and Staged Proof

On untouched B300 production, profile both the direct operator and public BasedLinearAttention(hidden_size=768, feature_dim=16, num_heads=12, mode="parallel") at B1/H12/K16/V64/T4096 and T8192.

Stop before production work unless the Based kernels are at least 30% of public-layer CUDA time at a representative long endpoint and the measured Amdahl bound can support at least a 10% layer gain.

The first production proof is forward-only at B1/H12/K16/V64/T4096:

• Compile and time the complete summary + scan + apply sequence, including workspace allocation.
• Require at least 1.30x over the current forward kernel before implementing backward.
• Stop on frozen-oracle failure, compile/resource failure, workspace above the analytic bound, local-memory spill or occupancy collapse, or a stage sum that misses the proof threshold.

Only after this proof passes may the reverse summaries/scan/local backward be implemented. Correctness-green representative performance shortfalls use at most two reviewer-approved bounded depth iterations, with all original gates unchanged.

Representative Benchmarks

Add one isolated parallel_based OpConfig in benchmarks/ops/registry.py in the fixture commit. It may add only this operator and its shapes; it must not change runner behavior, generic registry schema, input-distribution infrastructure, or absorb work from upstream PR fla-org#1144.

Primary BF16 and FP16 shapes:

• B1/H12/K16/V64, T=128/256/2048/4096/8192/16384
• B1/H12/K16/V128, T=4096/8192
• B4/H12/K16/V64, T=2048/4096
• H4 and T=4097 controls

Use the repository-native commands with the immutable fixture commit as the comparison base:

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

Measure the public Based layer separately with an ignored deterministic harness at B1/B4 and T=2048/4096/8192. The repository has no Based pretrained model, so the public layer is the integration endpoint and no model-level claim should be made.

Compare baseline and candidate in the same isolated B300 allocation with identical inputs, environment, warmup, clocks, and repetitions. Use at least three alternating complete runs and report every retained row with median and p20/p80. NCU full/source reports must compare the current kernels against the sum of all new stages, including DRAM/L2 traffic, tensor-pipe utilization, achieved occupancy, active/eligible warps, registers, spills, waves, and source stalls.

Acceptance Gates

• Frozen output and all-gradient correctness passes without tolerance, reference, precision, or numeric-flag changes.
• Direct default B1/H12/K16/V64 reaches at least 1.50x forward and 1.30x forward+backward at both T4096 and T8192.
• Equal-weight geomean over retained long T=2048/4096/8192/16384 and V64/V128 shapes is at least 1.30x in both modes.
• Public Based layer forward+backward reaches at least 1.10x at T4096 and T8192 and at least 1.05x geomean over T2048/4096/8192.
• No retained target or fallback median regresses more than 3%; short T128/256 fallback remains within 2% noise.
• Measured workspace matches the analytic bound, and NCU attributes the result to eliminated quadratic prefix/suffix work plus useful group concurrency without material spill or skipped work.

If these gates require a one-off shape lookup, changed semantics/precision, broader architecture or module scope, or a benchmark-only result, record no-go rather than widening the issue.

Expected Files

• fla/ops/based/parallel.py
• tests/ops/test_based.py in the fixture-only commit, then frozen
• benchmarks/ops/registry.py for one isolated repository-native operator registration

Upstream PR fla-org#1144 currently edits generic registry input-profile plumbing but does not add Based coverage or touch the production/test files above. Treat this as a mechanical rebase risk: if it lands, rebase the isolated registration onto its schema and rerun the frozen gates; do not overwrite or incorporate its distribution changes.

Duplicate Audit

As of the baseline audit, no open upstream PR or issue and no fork issue implements group-parallel Based Taylor moments or touches fla/ops/based/parallel.py / tests/ops/test_based.py. Open issues fla-org#117, fla-org#229, and fla-org#411 concern generic autotuning, chunk API unification, and different blockwise masking semantics respectively. Existing fused_chunk_based supplies mathematical precedent but not this long-context group-parallel schedule.

Baseline: upstream/main@27967b970eaaf982a6960abf6cba8add9c34c7cc.

[heiheiha798]

Cycle 19 stopped at the untouched public-layer profile gate. This is not a candidate correctness or performance failure: no production kernel was edited or run.

Baseline and fixture evidence:

• Source baseline: upstream/main@27967b970eaaf982a6960abf6cba8add9c34c7cc.
• Frozen fixture/benchmark base: fcebafe711bdd5597cb904fcdea6b93204addce5.
• Slurm job 5498 ran the untouched full tests/ops/test_based.py gate on an NVIDIA B300 (SM103): 5 passed.
• Slurm job 5514 ran the frozen public-layout fixture and isolated registry smoke on the same device: 7 passed, including real output/dq/dk/dv checks and the small valid current-semantics layer oracle.
• Slurm job 5516 is invalid setup evidence only. It stopped before measurement because the ignored profiler script lacked the repository root on PYTHONPATH.

Slurm job 5519 completed the untouched 24-row direct parallel_based registry benchmark and direct T4096/T8192 timing/peak/no-Chrome profiler tables. Representative BF16 B1/H12/K16/V64 native medians were:

• T4096: 0.230160 ms forward, 0.900000 ms forward+backward.
• T8192: 0.504864 ms forward, 2.495488 ms forward+backward.

The next mandatory endpoint was the exact public BasedLinearAttention(hidden_size=768, feature_dim=16, num_heads=12, num_key_value_heads=12, mode="parallel"). Its first untouched T4096 forward deterministically failed in o_proj: the current q/k reshape produces a 192-wide parallel output while o_proj requires 768, yielding RuntimeError: mat1 and mat2 shapes cannot be multiplied (4096x192 and 768x768).

The issue contract requires this exact public integration and forbids a layer/API repair or geometry substitution. Consequently the required layer CUDA share and Amdahl gate cannot be established, and the eventual layer threshold cannot be evaluated. The cycle stopped before staged forward proof, production edits, candidate correctness/performance, model claims, or NCU. No speedup conclusion and no salvage candidate are supported.

Durable local evidence is retained at:

• profile/parallel_based_c19_lane_a-opt/logs/untouched-job-5498.log
• profile/parallel_based_c19_lane_a-opt/logs/fixture-job-5514.log
• profile/parallel_based_c19_lane_a-opt/logs/profile-job-5516.log
• profile/parallel_based_c19_lane_a-opt/logs/profile-job-5519.log
• profile/parallel_based_c19_lane_a-opt/profile/native-direct.json
• profile/parallel_based_c19_lane_a-opt/profile/direct-t{4096,8192}-{fwd,fwdbwd}-table.txt

The issue remains open as a record of the scoped hypothesis and the blocking untouched integration state.