[Perf][B300] Generate Comba decay matrices with a product scan

[heiheiha798]

Status

This issue records an unmeasured B300 optimization hypothesis. The method and thresholds below are acceptance criteria to test, not observed GPU performance claims.

Baseline: upstream/main@e47d5d20aeb5989b58a3738b872e7c288a9fb75f.

Problem

The default public Comba model uses attn_mode="chunk" with head dimension 256, which routes dense and packed-varlen training through chunk_comba. In fla/ops/comba/wy_fast.py, both chunk_scaled_dot_comba_pkt_fwd_kernel and prepare_wy_repr_bwd_kernel form a full fp32 [64,64] pairwise gate factor with

exp2(g0[:, None] - g[None, :])

before masking to the strict lower triangle. At chunk size 64, each site evaluates 4,096 exponents per chunk/head even though the decay comes from a one-dimensional sequence.

Comba's cumsum kernel defines g0[i] as the exclusive prefix and g[i] as the inclusive prefix. If

d[t] = exp2(g[t] - g0[t])

then for i > j the current factor is exactly

exp2(g0[i] - g[j]) = product(d[t] for t in j+1 .. i-1).

The product is empty and equals one on the first subdiagonal. Comba gate increments are non-positive, so every d[t] is in [0,1]; this construction avoids the large inverse-decay intermediates of a similarity transform.

The hypothesis is that one length-64 vector exponent plus a reverse product scan over broadcast factors reduces special-function work enough to improve the public forward and training endpoints on B300 without changing launches, workspace, precision, or Comba's WY representation.

Baseline and profile-first gates

Before production edits, run the complete untouched B300 gate and registry benchmark. A deterministic baseline test failure, out-of-resources error, or D256 compile failure is a terminal no-go for this cycle.

Create one fixture-only commit that adds the targeted bf16 correctness cases below while production source and fla/ops/comba/naive.py remain byte-for-byte upstream. Run the complete fixture green on untouched production, record that commit as BENCH_BASE, then freeze all tests, references, seeds, dtypes, tolerances, numeric flags, registry inputs, and benchmark shapes.

Use torch profiler and a baseline full/source NCU capture before the source change. Record the share of chunk_scaled_dot_comba_pkt_fwd_kernel in full chunk_comba forward, the combined share of that kernel plus prepare_wy_repr_bwd_kernel in full forward+backward, and their share in both required Comba model training steps. Stop before implementation if the measured target work cannot arithmetically support the endpoint thresholds, if the forward kernel is below 10% of full operator forward, if the two kernels are below 5% of full operator forward+backward, if they are below 3% of either required model step, or if source evidence shows the pairwise exponent is not a material instruction/stall cost.

Bounded method

Add one small inlined @triton.jit helper in fla/ops/comba/wy_fast.py that reconstructs the strict-lower decay matrix from g0, g, and the valid-token mask:

1. Compute the fp32 vector d = exp2(g - g0), using one for invalid tail lanes.
2. Broadcast the shifted per-token factors and use tl.cumprod(..., reverse=True) along the column axis to form product(d[j+1:i]) for every row.
3. Explicitly set the strict lower triangle from the scan result and all diagonal, upper-triangle, and invalid-tail cells to zero.
4. Reuse the helper at the existing forward gate multiplication and backward dA gate multiplication. Preserve the current sign and cast point in backward.

Do not change the KKT/dA dot products, loop or launch geometry, BT/BK/BV, autotune space, fp32 accumulation, input/output dtypes, numeric flags, public signature, stored A representation, or backward contract. Do not add a global buffer, kernel launch, atomics, solve/KKT fusion, similarity transform, or unrelated refactor.

Compile and run the worst registered D256 shape immediately after the single helper change. Record Triton resource diagnostics and NCU registers, tensor/shared/local memory, spills, occupancy, and eligible warps. A scan out-of-resources failure, material spill, or occupancy collapse is terminal for this method, not permission to split/fuse adjacent kernels or widen the scope.

Expected production diff: fla/ops/comba/wy_fast.py only. The fixture may change tests/ops/test_comba.py only. fla/ops/comba/naive.py, registry/shapes, other Comba files, public API, model/layer/config code, and all non-Comba operators remain frozen.

The portable Triton path is preferred. If an actually observed non-NVIDIA compiler failure requires retaining the exact pairwise expression on that backend, use an existing repository device helper and keep the fallback byte-for-byte equivalent; do not invent an architecture shape table or alter other backends without evidence.

Frozen correctness contract

Untouched baseline and promotion gate:

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

The complete unfiltered tests/ops/test_comba.py remains the oracle. Existing cases cover dense and ragged varlen forward/backward, T tails, non-power-of-two D, normalizers 0.1/1/10, zero-masked gates, L2 normalization on/off, initial/final state, and every public gradient.

Before kernel edits, add only a small non-Cartesian bf16 fixture using the existing naive_chunk_comba reference and existing per-gradient tolerances:

• dense B1,T63,H1,D256, normalizer 0.1, to cover the default head dimension, harsh decay, and a partial chunk;
• one bf16 dense masked-gate D64 case;
• one bf16 uneven varlen D64 case containing a short tail and a partial 64-token chunk.

The untouched implementation must pass these cases before they are frozen. If it does not, record no-go rather than alter the reference or establish looser optimization-specific tolerances. After freeze, do not edit any test case, reference, tolerance, seed, dtype, skip, numeric flag, gate generation, or environment. NaN memory poisoning remains active and every o/ht/dq/dk/dv/dp/dg/dbeta/dh0 result must pass.

After each iteration, run the same full gate and benchmark nothing on red. Keep every attempt in ignored profile/chunk-comba-opt/OPT_LOG.md; after three consecutive non-keeps, stop and re-profile. Promotion also requires:

python scripts/find_dependent_tests.py fla/ops/comba/wy_fast.py
ruff check fla/ops/comba/wy_fast.py tests/ops/test_comba.py

Run every returned dependent test plus affected-file pre-commit, header, compile, and banned-block-pointer checks.

Repository-native B300 measurement

Use BENCH_BASE, not upstream/main, for final comparisons because it owns the frozen bf16 fixture while production is unchanged. Run all six existing registry shapes in both modes, including B8,T2048,H32,D256:

python -m benchmarks.ops.verify --op chunk_comba --base <BENCH_BASE> --modes fwd fwdbwd
python -m benchmarks.ops.run --op chunk_comba --base <BENCH_BASE> --modes fwd fwdbwd --json <ignored-path>

Use identical inputs, environment, clocks/idle state, warmup, repetitions, and fully warmed autotune state on the same isolated B300. Perform enough paired same-session comparisons to report sample count, median, mean, standard deviation, minimum, p10, p90, native p20/p80, and equal-weight geomeans for every shape; do not select best runs.

Measure the default public model path in dense and varlen training with identical layer count on base/candidate:

python benchmarks/benchmark_training_throughput.py --name comba --batch_size 1 --seq_len 8192 --warmup_steps 16 --steps 32
python benchmarks/benchmark_training_throughput.py --name comba --batch_size 4 --seq_len 2048 --context_len 2048 --varlen --warmup_steps 16 --steps 32

If memory requires reducing the layer count, apply the identical reduction to both refs and disclose it. Report tokens/s distribution and peak allocated/reserved memory.

B300 NCU evidence

Collect --set full with PM sampling and --set source --section SourceCounters for both changed kernels at representative D128 and default D256 dense shapes plus one fixed ragged varlen workload. Keep reports under ignored profile/chunk-comba-opt/trace/ with exact SHAs, commands, environment versions, and Slurm job IDs.

Compare kernel duration, launch count, SM/SOL throughput, achieved occupancy, registers, tensor/shared/local memory, spills, eligible/active warps, dominant stalls, and source-correlated special-function versus scan instructions. Source/disassembly evidence must confirm that each pairwise matrix exponent was replaced by one vector exponent; when an executed special-function counter is available, require at least 90% fewer exponent-lane instructions at the two changed sites. Do not invent unavailable SM103 metric aliases. The changed kernels must improve enough to explain the full endpoint result rather than merely moving work into an opaque scan lowering.

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 bf16 fixture, full dense/varlen forward/backward gate, dependent tests, and repository checks are green with unchanged references/tolerances/numerics.
• Full operator forward geomean is at least 1.08x and forward+backward geomean is at least 1.03x across all registry shapes. No individual median or p90 regresses by more than 3% outside measured noise.
• Both dense and varlen Comba model runs improve median throughput by at least 3%, with no peak-memory regression.
• The two changed kernels each improve by at least 1.15x on representative target shapes, and NCU/source evidence attributes the endpoint result to removal of pairwise exponent work without resource, spill, occupancy, or hidden-work regressions.

Failure of the frozen correctness gate, D256 compile/resource gate, operator/model thresholds, or NCU mechanism is a no-go. Do not rescue the result by changing precision, tolerances, tests, autotune buckets, chunk size, other kernels, model settings asymmetrically, or scope.

Upstream overlap

No open upstream issue, PR, or remote branch implements a Comba product-scan decay matrix. Open fla-org#797 changes GDN only and uses an ungated similarity transform/fused inference path; it does not touch Comba, and its inverse-decay formulation is deliberately excluded here. Open fla-org#1144 changes benchmark infrastructure and may create a mechanical registry rebase, but it does not change Comba kernels or this decay computation. Open fla-org#1145 is an Ascend solve_tril change and is unrelated.

Historical Comba work includes the general exp2 migration but no product scan. This cycle must remain independent of all prior local ATK, HGRN, RWKV7, cross-entropy, GRPO, MoBA, parallel-attention GQA, and PaTH branches/issues/evidence.

[heiheiha798]

B300 evaluation result: no-go at the frozen fixture gate

Cycle 9 stopped before profiling or production implementation because the fixture-first correctness gate was red on untouched production.

Reproducible evidence

• Source baseline: e47d5d20aeb5989b58a3738b872e7c288a9fb75f
• Fixture-only commit: 2ae9d3a12d704d11eefeeade07055a80d4cc7845
• Slurm job: 5221
• Command: python -m pytest tests/ops/test_comba.py -q
• Result: 27 passed, 1 failed in 758.16 seconds on an NVIDIA B300 SXM6 (SM103).
• Durable log: profile/chunk_comba_product_scan-opt/fixture-5221.log

The sole failure was the required bf16 partial-chunk fixture:

test_chunk[B1-T63-H1-D256-scale1-gate_logit_normalizer0.1-mask_p0-use_qk_l2norm_in_kernelFalse-torch.bfloat16]

The unchanged db oracle tolerance is 0.005; untouched production reported an absolute difference of 0.031250 and an error ratio of 0.005499. The case completed and reported its other output, final-state, and gradient comparisons. The new masked dense D64 and uneven packed D64 bf16 fixtures passed.

The fixture commit changes only tests/ops/test_comba.py. fla/ops/comba/, including naive.py, and the benchmark registry remain byte-identical to the source baseline. The preceding untouched gate in job 5220 passed 25/25 tests, and all six registered chunk_comba fwd/fwdbwd shapes, including D256, compiled and ran.

Per the frozen-oracle stop rule, I did not change the fixture, reference, tolerances, or numeric settings. No production kernel was implemented; no profile, model run, NCU collection, post-fixture benchmark, push, or PR was performed. Therefore this cycle supports no candidate correctness or performance conclusion. The issue remains open for future evaluation from a green frozen baseline.