diff --git a/USAGE.md b/USAGE.md index e37a61b9..aa3b04c2 100644 --- a/USAGE.md +++ b/USAGE.md @@ -213,6 +213,11 @@ Pass `use_intracard_cp` to the Hopper prefill: Works with **any sequence length** (non-CHUNK-aligned is handled internally) and **dense or varlen** input. The auto decision uses the device SM count plus two tunables in `cula/ops/kda/sm90/cp/plan.py` (`CULA_KDA_CP_RERUN_RATIO`, `CULA_KDA_CP_AUTO_MIN_SEG_TILES`). +The SM90 CP pre-scan keeps its segment transition matrix canonically in FP32 +registers and maintains a BF16 mirror only for Tensor Core products. This +prevents transition error from compounding when many slow-decay segments are +composed. + ```python o, final_state = flashkda_prefill( q=q, k=k, v=v, g=g, beta=beta, A_log=A_log, dt_bias=dt_bias, diff --git a/cula/ops/kda/sm90/cp/pre_scan.py b/cula/ops/kda/sm90/cp/pre_scan.py index da40f663..4cc31b9a 100644 --- a/cula/ops/kda/sm90/cp/pre_scan.py +++ b/cula/ops/kda/sm90/cp/pre_scan.py @@ -91,6 +91,9 @@ def pre_scan_kernel( sINV = smem.allocate_tensor(cutlass.BFloat16, cc_stage_layout, 128) sState = smem.allocate_tensor(cutlass.BFloat16, state_layout, 128) sM = smem.allocate_tensor(cutlass.BFloat16, state_layout, 128) + # sState intentionally matches the serial FlashKDA BF16 recurrence. sM + # is only its Tensor Core operand mirror; the CP-only affine transform is + # kept canonically in a distributed FP32 register fragment below. sGt = smem.allocate_tensor(cutlass.Float32, cute.make_layout((D, 1, STAGES), stride=(1, D, D)), 128) sBeta = smem.allocate_tensor(cutlass.BFloat16, cute.make_layout((CHUNK, 1, STAGES), stride=(1, 64, 64)), 128) sMbar = smem.allocate_tensor(cutlass.Int64, cute.make_layout((STAGES,)), 16) @@ -154,12 +157,10 @@ def pre_scan_kernel( gKD_raw = cute.make_tensor(ws_kd.iterator, raw_gmem_layout) gKR_raw = cute.make_tensor(ws_kr.iterator, raw_gmem_layout) - # sState=0, sM=I + # sState=0. The distributed FP32 M accumulator initializes sM below. if tidx < D: for e in cutlass.range_constexpr(D): sState[tidx, e] = cutlass.BFloat16(0.0) - sM[tidx, e] = cutlass.BFloat16(0.0) - sM[tidx, tidx] = cutlass.BFloat16(1.0) cute.arch.barrier() mma_atom = warp.MmaF16BF16Op(cutlass.BFloat16, cutlass.Float32, (16, 8, 16)) @@ -223,10 +224,26 @@ def pre_scan_kernel( tCrKrA_state_blk_cv = smem_thr_copy_A_state.retile(tCrKrA_state_blk) tCrUpd_blk = thr_mma_state.make_fragment_C(tiled_mma_state.partition_shape_C((CHUNK, D))) sState_blk_tile = cute.flat_divide(sState, (CHUNK, D)) - sM_blk_tile = cute.flat_divide(sM, (CHUNK, D)) coord_state_blk = cute.make_identity_tensor((CHUNK, D)) tCcState_blk = thr_mma_state.partition_C(coord_state_blk) + # Four compute warps jointly own the complete transition matrix. Keeping + # this fragment in registers preserves the original shared-memory + # footprint and avoids round-tripping the canonical value through SMEM. + tCrKrA_m = thr_mma_state.make_fragment_A(thr_mma_state.partition_A(sKr_T_view_s0)) + tCrKrA_m_cv = smem_thr_copy_A_state.retile(tCrKrA_m) + tCrMAcc = thr_mma_state.make_fragment_C(tiled_mma_state.partition_shape_C((D, D))) + tCcM = thr_mma_state.partition_C(cute.make_identity_tensor((D, D))) + tCsM = thr_mma_state.partition_C(sM) + if warp_idx < LOAD_WARP_IDX: + for i in cutlass.range_constexpr(cute.size(tCrMAcc)): + ii: cutlass.Constexpr[int] = i + row, col = tCcM[ii] + value = cutlass.Float32(1.0) if row == col else cutlass.Float32(0.0) + tCrMAcc[ii] = value + tCsM[ii] = cutlass.BFloat16(value) + cute.arch.barrier() + tCrU_T = thr_mma.make_fragment_B(thr_mma.partition_B(sKr_T_ref)) sINV_s0 = sINV[(None, None, 0)] @@ -406,30 +423,21 @@ def pre_scan_kernel( ii: cutlass.Constexpr[int] = i tCrU_T_post_u32[ii] = movm_t_b16(cutlass.Int32(tCrU_post_u32[ii])) - # State update': M = M*gt + kr^T @ U - for mi in cutlass.range_constexpr(M_BLOCKS): - sKr_T_blk_s = sKr_T_blk_tile_s[None, None, mi, 0] - cute.copy( - smem_tiled_copy_A_state, - smem_thr_copy_A_state.partition_S(sKr_T_blk_s), - tCrKrA_state_blk_cv, - ) - tCrUpd_blk.fill(0.0) - cute.gemm(tiled_mma_state, tCrUpd_blk, tCrKrA_state_blk, tCrU_T_post, tCrUpd_blk) - - sM_blk = sM_blk_tile[None, None, mi, 0] - tCsM_blk = thr_mma_state.partition_C(sM_blk) - m_frag_blk = cute.make_fragment_like(tCsM_blk, cutlass.BFloat16) - gt_frag_blk_m = cute.make_fragment_like(tCsM_blk, cutlass.Float32) - m_off2: cutlass.Constexpr[int] = mi * CHUNK - for i in cutlass.range_constexpr(cute.size(m_frag_blk)): - ii: cutlass.Constexpr[int] = i - m_frag_blk[ii] = tCsM_blk[ii] - gt_frag_blk_m[ii] = sGt_s[m_off2 + tCcState_blk[ii][0]] - for i in cutlass.range_constexpr(cute.size(tCrUpd_blk)): - ii: cutlass.Constexpr[int] = i - old = cutlass.Float32(m_frag_blk[ii]) * gt_frag_blk_m[ii] - tCsM_blk[ii] = cutlass.BFloat16(old + tCrUpd_blk[ii]) + # State update': M = M*gt + kr^T @ U. Accumulate into the + # persistent FP32 register fragment and publish only a BF16 image + # for the next chunk's kd@M Tensor Core operand. + cute.copy( + smem_tiled_copy_A_state, + smem_thr_copy_A_state.partition_S(sKr_T_s), + tCrKrA_m_cv, + ) + for i in cutlass.range_constexpr(cute.size(tCrMAcc)): + ii: cutlass.Constexpr[int] = i + tCrMAcc[ii] = tCrMAcc[ii] * sGt_s[tCcM[ii][0]] + cute.gemm(tiled_mma_state, tCrMAcc, tCrKrA_m, tCrU_T_post, tCrMAcc) + for i in cutlass.range_constexpr(cute.size(tCrMAcc)): + ii: cutlass.Constexpr[int] = i + tCsM[ii] = cutlass.BFloat16(tCrMAcc[ii]) cute.arch.barrier(barrier_id=1, number_of_threads=128) cute.arch.fence_view_async_shared() @@ -448,7 +456,11 @@ def pre_scan_kernel( if tidx < D: for d_out in cutlass.range_constexpr(D): b_state_g[state_base_f + cutlass.Int32(d_out * D) + cutlass.Int32(tidx)] = cutlass.Float32(sState[tidx, d_out]) - m_state_g[state_base_f + cutlass.Int32(d_out * D) + cutlass.Int32(tidx)] = cutlass.Float32(sM[tidx, d_out]) + if warp_idx < LOAD_WARP_IDX: + for i in cutlass.range_constexpr(cute.size(tCrMAcc)): + ii: cutlass.Constexpr[int] = i + row, col = tCcM[ii] + m_state_g[state_base_f + cutlass.Int32(col * D + row)] = tCrMAcc[ii] @cute.jit diff --git a/tests/test_kda_sm90_intracard_cp.py b/tests/test_kda_sm90_intracard_cp.py index 5302cb9e..438018d7 100644 --- a/tests/test_kda_sm90_intracard_cp.py +++ b/tests/test_kda_sm90_intracard_cp.py @@ -67,6 +67,11 @@ def _rel_rmse(a, b): return (a - b).pow(2).mean().sqrt().item() / max(b.pow(2).mean().sqrt().item(), 1e-6) +def _rel_l2(a, b): + a, b = a.float(), b.float() + return ((a - b).norm() / b.norm().clamp_min(1e-30)).item() + + def _assert_cp_matches(actual, ref, name): rrmse = _rel_rmse(actual, ref) assert rrmse < TOL_RMSE, f"{name}: rel_rmse {rrmse:.2e} >= {TOL_RMSE}" @@ -168,6 +173,121 @@ def test_trivial_manual_plan_falls_back_to_serial(): assert torch.equal(fin_fb, fin_ref) +@needs_cuda +@pytest.mark.kda_slow +def test_slow_decay_64k_cp_relative_l2_matches_fla_scale(): + """Regression for #126 using the issue's public synthetic reproducer.""" + torch.manual_seed(0) + h = 1 + shape = (1, 65536, h, D) + q = torch.randn(shape, dtype=torch.bfloat16, device="cuda") + k = torch.randn(shape, dtype=torch.bfloat16, device="cuda") + v = torch.randn(shape, dtype=torch.bfloat16, device="cuda") + g = torch.full(shape, -10.0, dtype=torch.bfloat16, device="cuda") + beta = torch.zeros(shape[:-1], dtype=torch.bfloat16, device="cuda") + A_log = torch.zeros(h, dtype=torch.float32, device="cuda") + dt_bias = torch.zeros(h, D, dtype=torch.float32, device="cuda") + + def run(*, split): + out = torch.empty_like(v) + state = torch.empty(1, h, D, D, dtype=torch.float32, device="cuda") + if split == 1: + flash_kda_fwd(q, k, v, g, beta, SCALE, out, A_log, dt_bias, LB, final_state=state) + else: + intracard_prefill( + q, + k, + v, + g, + beta, + SCALE, + out, + A_log, + dt_bias, + LB, + final_state=state, + s_split=split, + allow_fallback=False, + ) + return out, state + + out_serial, state_serial = run(split=1) + for split in (2, 4, 8, 16): + out_cp, state_cp = run(split=split) + + if split == 2: + # There is only one carry handoff, so the CP path is exactly the + # same BF16 recurrence as serial FlashKDA. + assert torch.equal(out_cp, out_serial) + assert torch.equal(state_cp, state_serial) + else: + # This public reproducer has an almost-zero segment transition; + # it mostly measures the different BF16 state-rounding points. + # Keep the divergence bounded while the transition-sensitive + # regression below directly covers the #126 amplification. + assert _rel_l2(out_cp, out_serial) < 8e-3, f"output split={split}" + assert _rel_l2(state_cp, state_serial) < 8e-3, f"state split={split}" + + +@needs_cuda +@pytest.mark.kda_slow +def test_slow_decay_16k_cp_transition_matches_fp64(): + """The CP-only affine transition must not accumulate in BF16 (#126).""" + from fla.ops.kda import fused_recurrent_kda + + torch.manual_seed(123) + h = 1 + shape = (1, 16384, h, D) + q = torch.randn(shape, dtype=torch.bfloat16, device="cuda") + k = torch.randn(shape, dtype=torch.bfloat16, device="cuda") + v = torch.randn(shape, dtype=torch.bfloat16, device="cuda") + g = torch.full(shape, -10.0, dtype=torch.bfloat16, device="cuda") + beta = torch.full(shape[:-1], -10.0, dtype=torch.bfloat16, device="cuda") + A_log = torch.zeros(h, dtype=torch.float32, device="cuda") + dt_bias = torch.zeros(h, D, dtype=torch.float32, device="cuda") + h0 = torch.randn(1, h, D, D, dtype=torch.float32, device="cuda") * 0.1 + + out = torch.empty_like(v) + state = torch.empty_like(h0) + intracard_prefill( + q, + k, + v, + g, + beta, + SCALE, + out, + A_log, + dt_bias, + LB, + initial_state=h0, + final_state=state, + s_split=16, + allow_fallback=False, + ) + + g64 = LB * torch.sigmoid(torch.exp(A_log.double()).view(1, 1, h, 1) * (g.double() + dt_bias.double().view(1, 1, h, D))) + beta64 = torch.sigmoid(beta.double()) + ref_out, ref_state = fused_recurrent_kda( + q.double(), + k.double(), + v.double(), + g64, + beta64, + scale=SCALE, + initial_state=h0.double(), + output_final_state=True, + use_qk_l2norm_in_kernel=True, + use_gate_in_kernel=False, + transpose_state_layout=True, + ) + + # BF16 sM gives about 1.45e-1 / 5.60e-1 on this case. Keeping the + # canonical transition in FP32 reduces that to about 5.5e-2 / 1.4e-2. + assert _rel_l2(out, ref_out) < 8e-2 + assert _rel_l2(state, ref_state) < 3e-2 + + @needs_cuda @pytest.mark.kda_fast @pytest.mark.parametrize("s_split", [2, 4, 7])