[Perf][B300] Pack PaTH backward transformed-query history

[heiheiha798]

Status

This issue records an unmeasured B300 memory-optimization hypothesis. The byte formulas and thresholds below are acceptance criteria to verify, not observed performance claims.

Baseline: upstream/main@e47d5d20aeb5989b58a3738b872e7c288a9fb75f.

Problem

The public PaTH training path always calls parallel_path_attn. Its backward currently reconstructs transformed queries with transform_q_fwd_fn, which allocates a zero-filled tensor shaped [B, T, ceil(T / 512), HQ, K]. The producer writes only the sequence-valid transformed-query snapshots, and the two inter-chunk backward consumers load only those snapshots. Snapshot zero, future snapshots, and varlen cells added solely by the global maximum split count are rectangular holes.

For a sequence of length L, split size S=512, and N=ceil(L/S), the required history contains

history_cells(L) = sum(L - r*S for r in 1..N-1)

token-snapshot cells, rather than L*N. At the default bf16 MHA geometry HQ=32, K=64, this changes the transformed-query workspace from 512 MiB to 240 MiB at dense T=8192 (46.875% of the current allocation). For a flattened 8192-token pack with segment lengths [512, 1536, 2048, 4096], it changes 256 MiB to 74 MiB (28.91%). These are exact allocation formulas, not measured end-to-end peak-memory results.

PaTH is a packaged model and layer, and its default config uses 32 heads with head dimension 64. This workspace is therefore on the ordinary dense and packed-varlen training path rather than an optional backend or narrow model configuration.

Bounded method

First create one fixture-only commit containing targeted boundary tests and repository-native benchmark registrations. Production source and fla/ops/path_attn/naive.py must remain byte-for-byte upstream in that commit. Run the complete untouched B300 oracle green, record the fixture commit as BENCH_BASE, then freeze tests, references, seeds, dtypes, tolerances, numeric flags, and benchmark inputs for the optimization loop. A red untouched B300 gate is a terminal no-go for this cycle.

Replace only the internal transformed-query history layout:

1. Allocate a packed history with logical shape [sum(history_cells(L_i)), HQ, K] in the existing query dtype. For dense batches, derive each sequence base analytically. For packed varlen, build one compact int64 prefix-offset tensor over per-sequence history-cell counts.
2. Store snapshot r contiguously for local tokens [r*S, L). Every packed cell must be producer-written exactly once before either consumer runs, so allocator NaN poisoning can validate use of torch.empty rather than zero fill.
3. Pass the packed buffer and internal offset metadata to parallel_path_bwd_dkv_fn and parallel_path_bwd_dq_fn. Change only their transformed-query pointer calculation; preserve the loop bounds, transformations, dot products, accumulation order and dtype, atomics, tiles, and all other inputs/outputs.
4. Treat L <= S as a valid zero-history case: no transformed-history load or store may execute, and the public output/gradients must remain unchanged.
5. Cast program IDs, sequence IDs, split IDs, sequence bases, triangular-prefix terms, token offsets, strides, and all derived addresses to tl.int64 before arithmetic. Do not add tl.make_block_ptr or tl.advance.

The varlen allocation needs one host-visible total size. That scalar read must replace the synchronization currently incurred by get_max_num_splits; it must not add a second device-host synchronization or a per-snapshot host loop. Keep the existing rectangular layout only as an implementation-debugging baseline, not as a new runtime dispatch. If a correct packed layout requires a public API change, numerical change, additional synchronization, recomputation, or architecture-specific fallback, record no-go instead of widening scope.

Expected production files:

• fla/ops/path_attn/transform_q.py
• fla/ops/path_attn/parallel_path_bwd_inter_dkv.py
• fla/ops/path_attn/parallel_path_bwd_inter_dqh.py
• fla/ops/path_attn/parallel.py

Fixture files are limited to tests/ops/test_path_attn.py and benchmarks/ops/registry.py. Do not change layer/model/config/public API files, forward or decoding kernels, naive.py, attention math, precision, S/BT/BS, backend dispatch, or unrelated PaTH intermediates.

Frozen correctness contract

Before production edits, run the complete existing gate on B300:

python -m pytest -q tests/ops/test_path_attn.py

The fixture commit may add a small, non-Cartesian boundary matrix using the existing independent naive_path_attn reference and existing per-gradient tolerances. Cover dense and uneven packed-varlen behavior around S=512: T=511/512/513/1024/1025, including a pack with sub-512, exact-512, partial-tail, and multi-split segments. Include bf16 and fp16, gate on/off, default MHA plus one GQA case, D64/D128, and all reachable o/dq/dk/dv/dg/dw/dbeta gradients. Existing long dense/varlen cases remain part of the gate.

The tests must observe only public outputs and gradients, not private workspace shape. Do not change the naive reference, existing cases, tolerance values, skip conditions, seeds, numeric flags, precision, or environment to make the candidate pass. tests/conftest.py NaN poisoning remains active and must catch any unwritten packed tail or invalid boundary read.

Promotion requires:

python -m benchmarks.ops.verify --op parallel_path_attn --base <BENCH_BASE> --modes fwd fwdbwd --test-file tests/ops/test_path_attn.py
python -m pytest -q tests/ops/test_path_attn.py
python -m pytest -q tests/models/test_modeling_path_attn.py
python scripts/find_dependent_tests.py fla/ops/path_attn/transform_q.py fla/ops/path_attn/parallel_path_bwd_inter_dkv.py fla/ops/path_attn/parallel_path_bwd_inter_dqh.py fla/ops/path_attn/parallel.py

Run every returned dependent test plus affected-file pre-commit, header, compile, and banned-block-pointer checks. No benchmark may run after a failed frozen gate.

Repository-native B300 measurement

Register parallel_path_attn and parallel_path_attn_varlen, both invoking the public parallel_path_attn function and timing the first tuple output. Use bf16 q/k/v, normalized fp32 w, sigmoid(beta)*2 in fp32, and g=None for the default headline. Required shapes are:

• dense default MHA: B1,T4096/8192,H=HQ=32,D64;
• dense breadth: B1,T4096,H8,HQ32,D64;
• packed varlen: flattened B1,total-T8192,H=HQ=32,D64 with cumulative lengths [0, 512, 2048, 4096, 8192];
• small/boundary latency controls that exercise T <= 512 and a partial split.

Use BENCH_BASE, not upstream/main, for baseline comparisons because the fixture owns the registrations. Compare candidate and base in the same isolated B300 allocation with identical generated inputs, environment, clock/idle state, warmup, repetitions, and warmed compilation/autotune state. Perform at least five independent paired comparisons. Report every retained shape with sample count, median, mean, standard deviation, minimum, p10, p90, native p20/p80, and equal-weight geomean; do not select best runs.

Measure end-to-end torch.cuda.max_memory_allocated() and reserved memory in fresh processes with reset discipline for dense T4096/T8192 and the ragged 8K pack. Record absolute MiB for the public forward+backward endpoint and the exact transformed-query allocation. Baseline and candidate must use identical input lifetimes and gradient-reset behavior.

B300 profiling

Profile the untouched BENCH_BASE first at default MHA T8192. Use torch profiler with memory attribution to identify the rectangular allocation/zero fill, transform_q_fwd_kernel, parallel_path_bwd_dq_kernel, and parallel_path_bwd_dkv_kernel. The primary claim is capacity, so a small zero-fill latency share does not by itself invalidate a structurally verified peak-memory win.

Collect one full-step NCU report plus full and source-correlated captures for the three affected Triton kernels on the same dense T8192 base/candidate workload. Compare summed duration and launch count, DRAM read/write bytes and throughput, L2 sectors/hit rate, achieved occupancy, registers, local-memory spills, and source hot spots. Confirm that valid transformed-query loads/stores and math are preserved while the rectangular zero-fill traffic disappears; do not attribute a latency mechanism that the profile does not show. Keep raw reports/logs outside git with exact commands, SHAs, environment versions, and Slurm job IDs. If the user-level NCU helper is unavailable, use the repository's documented minimal full/source workflow and state that in the evidence.

Acceptance criteria

• The untouched baseline, frozen full operator gate, PaTH model integration test, dependent tests, and repository checks are green with unchanged references/tolerances/numerics.
• Packed transformed-query bytes equal the analytic formula. They are at most 46.875% of the rectangular allocation at dense T8192 and at most 28.91% for the specified ragged 8K pack, with no asymptotically comparable replacement scratch.
• Public forward+backward peak allocated memory improves by at least 20% at default MHA T8192 and at least 25% for the ragged 8K pack. Report reserved-memory behavior separately.
• No required endpoint's median forward+backward latency regresses by more than 3% outside measured noise, and the equal-weight geomean is at least 0.99x baseline. Forward remains within 3% noise because its production path is unchanged.
• Profiling confirms the allocation/traffic mechanism without hidden recomputation, skipped snapshots, new synchronization, uncoalesced access collapse, or spill regression.

A miss on the frozen correctness, structural-byte, peak-memory, or latency guards is a no-go. Do not rescue it by changing precision, recomputing history, fusing unrelated kernels, adding shape tables, modifying forward/decoding, or changing public/model behavior.

Upstream overlap

No open upstream issue, PR, or remote branch targets PaTH transformed-query history packing. Historical PaTH work covers the original operator/model (fla-org#384), head-dimension refactoring (fla-org#503), correctness fixes (fla-org#581/fla-org#633), and long-addressing fixes (fla-org#769/fla-org#994/fla-org#1082), but the rectangular zero-filled history remains in current main. Open fla-org#1144 changes benchmark infrastructure and may cause a mechanical registry rebase; it does not touch PaTH kernels or implement this optimization.

This change is independent of prior local ATK, HGRN, RWKV7, cross-entropy, GRPO, MoBA, and parallel-attention GQA scopes. It must start from the stated upstream baseline and must not modify or depend on preserved historical branches, issues, PRs, or evidence.

[heiheiha798]

B300 evaluation result: no-go at the untouched baseline gate

Cycle 8 stopped before the fixture or implementation stages because the existing PaTH backward does not pass its full untouched test file on B300/SM103.

• Source baseline and branch HEAD: e47d5d20aeb5989b58a3738b872e7c288a9fb75f
• Slurm job: 5216
• Hardware: NVIDIA B300 SXM6 AC, compute capability 10.3
• Software: PyTorch 2.13.0+cu130, CUDA 13.0, Triton 3.7.1
• Durable raw log: profile/parallel_path_attn_history-opt/untouched-gate-5216.log

The exact full python -m pytest tests/ops/test_path_attn.py -q command completed with 4 passed and 5 failed in 321.50 seconds. Every failure occurred when the existing parallel_path_bwd_dq_kernel was launched from parallel_path_bwd_inter_dqh.py: Triton reported OutOfResources: tensor memory, Required: 544, Hardware limit: 512.

The five failing cases were:

• dense B5/T512/H2/HQ8/D128, bf16, forget gate enabled;
• dense B1/T4000/H1/HQ2/D128, bf16, forget gate disabled;
• packed H2/HQ4/D128, fp16, cu_seqlens=[0,15,333,2048], forget gate disabled;
• packed H2/HQ4/D128, fp16, cu_seqlens=[0,15,333,2048], forget gate enabled;
• packed H2/HQ16/D128, fp16, cu_seqlens=[0,500,1023,2000,3000,4096], forget gate enabled.

tests/ops/test_path_attn.py, all fla/ops/path_attn production files, and benchmarks/ops/registry.py remained byte-for-byte identical to the baseline. No fixture-only commit or BENCH_BASE was created, and no profile, NCU collection, packed-history implementation, benchmark comparison, push, or PR was performed.

This is an untouched B300 baseline resource failure, not a packed transformed-query-history candidate failure. Under the frozen gate contract, the cycle ended without retry, kernel workaround, test relaxation, or scope expansion, and it supports no technical performance conclusion.