[KDA][SM90][Accuracy] Intracard CP accumulates large errors on 64K slow-decay heads #126

Summary

The SM90 KDA intracard-CP path can accumulate very large numerical errors on 64K slow-decay heads. The error increases with the number of segments and becomes progressively larger in later sequence windows.

This is different from #89: no inf/nan is involved here. The issue appears to be finite-error accumulation while composing segment affine transforms/carries.

Environment

• GPU: NVIDIA H200 (SM90, 141 GiB)
• Python: 3.12
• PyTorch: 2.13.0a0+9186a08b2c.nv26.07
• CUDA runtime/toolkit: 13.3
• cuLA: official cuda-linear-attention==0.2.0+cu130 wheel from v0.2.0

Real-model-derived results

The metric is per-head relative L2. serial-vs-FP32 and CP-vs-FP32 use an independent FP32 naive recurrence. CP-vs-serial compares the two cuLA execution paths directly.

The inputs use real public model weights and real book text. For clarity, the nonzero-layer hidden input is embedding-derived and then passed through the target layer's real normalization/projection/conv weights; it is a proxy input, not a hidden state captured from a full model forward pass.

The error grows toward the end of the sequence:

• GLM last 8K tokens: CP-vs-FP32 output 2.39427, serial-vs-FP32 0.0398373.
• Kimi-K3 last 8K tokens: CP-vs-FP32 output 2.43918, serial-vs-FP32 0.0961614.

A manual split sweep also isolates the behavior:

• With 2 segments, CP is bit-identical to serial for both target heads.
• With 4, 8, and 16 segments, the error increases and starts only in later segments.
• This points to repeated segment-transform/carry composition rather than the first local recurrence.

Minimal public reproduction

Install the official wheel:

pip install "cuda-linear-attention==0.2.0+cu130" \
  -f https://github.com/inclusionAI/cuLA/releases/expanded_assets/v0.2.0

Run:

python repro_synthetic.py --tokens 65536 --heads 1 --split 16 --gate -10

#!/usr/bin/env python3
import argparse
import torch

D = 128
SCALE = D ** -0.5
LOWER = -5.0


def rel_l2(a, b):
    return ((a.float() - b.float()).norm() /
            b.float().norm().clamp_min(1e-30)).item()


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--tokens", type=int, default=65536)
    p.add_argument("--heads", type=int, default=1)
    p.add_argument("--split", type=int, default=16)
    p.add_argument("--gate", type=float, default=-10.0)
    args = p.parse_args()

    from cula.ops.kda.sm90.cp import intracard_prefill
    from cula.ops.kda.sm90.fwd import flash_kda_fwd

    torch.manual_seed(0)
    shape = (1, args.tokens, args.heads, D)
    q = torch.randn(shape, device="cuda", dtype=torch.bfloat16)
    k = torch.randn(shape, device="cuda", dtype=torch.bfloat16)
    v = torch.randn(shape, device="cuda", dtype=torch.bfloat16)
    g = torch.full(shape, args.gate, device="cuda", dtype=torch.bfloat16)
    beta = torch.zeros(shape[:-1], device="cuda", dtype=torch.bfloat16)
    A_log = torch.zeros(args.heads, device="cuda", dtype=torch.float32)
    dt_bias = torch.zeros(args.heads, D, device="cuda", dtype=torch.float32)

    def alloc():
        return (
            torch.empty_like(v),
            torch.empty(1, args.heads, D, D,
                        device="cuda", dtype=torch.float32),
        )

    out_s, state_s = alloc()
    flash_kda_fwd(
        q, k, v, g, beta, SCALE, out_s, A_log, dt_bias, LOWER,
        final_state=state_s,
    )

    out_c, state_c = alloc()
    intracard_prefill(
        q, k, v, g, beta, SCALE, out_c, A_log, dt_bias, LOWER,
        final_state=state_c,
        s_split=args.split,
        allow_fallback=False,
    )
    torch.cuda.synchronize()

    print(f"T={args.tokens} H={args.heads} split={args.split} gate={args.gate}")
    print(f"output relative-L2 CP-vs-serial: {rel_l2(out_c, out_s):.8e}")
    print(f"state  relative-L2 CP-vs-serial: {rel_l2(state_c, state_s):.8e}")
    window = args.tokens // 8
    for i in range(8):
        lo = i * window
        hi = args.tokens if i == 7 else (i + 1) * window
        print(
            f"tokens [{lo:5d},{hi:5d}): "
            f"{rel_l2(out_c[:, lo:hi], out_s[:, lo:hi]):.8e}"
        )


if __name__ == "__main__":
    main()

Observed output:

T=65536 H=1 split=16 gate=-10.0
output relative-L2 CP-vs-serial: 6.86958805e-03
state  relative-L2 CP-vs-serial: 7.09873484e-03
tokens [    0, 8192): 0.00000000e+00
tokens [ 8192,16384): 7.33159157e-03
...
tokens [57344,65536): 7.35112419e-03

This synthetic case is intentionally simple. It reproducibly exceeds the current TOL_RMSE = 4e-3, while the much larger errors require the real slow-decay head distribution described above.

Hypothesis

In the current SM90 CP implementation, local sState/sM computation in pre-scan uses BF16. Although b_seg, m_seg, and carries are stored as FP32 and merge uses TF32 inputs with FP32 accumulation, the segment affine transforms have already been formed through BF16 local recurrence. Repeatedly composing those transforms may amplify their approximation error when the effective transition is close to non-contracting.

This is only a hypothesis based on the source and the segment-count experiment.

Expected behavior / suggested coverage

Intracard CP should remain reasonably close to serial on 64K slow-decay heads. It would be useful to add tests covering:

• 64K sequence length;
• deliberately slow decay;
• per-head relative L2 rather than only aggregate tolerance;
• comparison against an independent FP32 recurrence;
• multiple segment counts, including more than two segments.

I can provide the two serialized single-head inputs (about 64 MiB each) if they would help reproduce the real-model-derived cases.