From 74af6d6549dc05d9ba73850fd75f33cff01b0382 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=8C=83=E5=9D=A4?= Date: Mon, 3 Aug 2026 01:22:18 +0800 Subject: [PATCH 1/7] perf(kda): use raw bulk copies for SM90 workspace --- cula/ops/kda/sm90/k1.py | 82 ++++++++++++----------------------------- cula/ops/kda/sm90/k2.py | 73 ++++++++++++------------------------ 2 files changed, 48 insertions(+), 107 deletions(-) diff --git a/cula/ops/kda/sm90/k1.py b/cula/ops/kda/sm90/k1.py index 8b2fdd1..024fe5f 100644 --- a/cula/ops/kda/sm90/k1.py +++ b/cula/ops/kda/sm90/k1.py @@ -45,16 +45,13 @@ def k1_kernel( tma_tensor_k: cute.Tensor, tma_atom_g: cute.CopyAtom, tma_tensor_g: cute.Tensor, - tma_atom_ws_qd: cute.CopyAtom, - tma_tensor_ws_qd: cute.Tensor, - tma_atom_ws_kd: cute.CopyAtom, - tma_tensor_ws_kd: cute.Tensor, - tma_atom_ws_kr: cute.CopyAtom, - tma_tensor_ws_kr: cute.Tensor, tma_atom_ws_inv: cute.CopyAtom, tma_tensor_ws_inv: cute.Tensor, tma_atom_ws_mqk: cute.CopyAtom, tma_tensor_ws_mqk: cute.Tensor, + ws_qd: cute.Tensor, + ws_kd: cute.Tensor, + ws_kr: cute.Tensor, a_log: cute.Tensor, dt_bias: cute.Tensor, beta: cute.Tensor, @@ -148,30 +145,6 @@ def k1_kernel( cute.group_modes(gSrc_g, 0, 2), ) - gDst_qd = cute.local_tile(tma_tensor_ws_qd, (CHUNK, D), (None, None, None)) - tQDws_s, tQDws_g = cpasync.tma_partition( - tma_atom_ws_qd, - 0, - cute.make_layout(1), - cute.group_modes(s_q_decayed, 0, 2), - cute.group_modes(gDst_qd, 0, 2), - ) - gDst_kd = cute.local_tile(tma_tensor_ws_kd, (CHUNK, D), (None, None, None)) - tKDws_s, tKDws_g = cpasync.tma_partition( - tma_atom_ws_kd, - 0, - cute.make_layout(1), - cute.group_modes(s_k_decayed, 0, 2), - cute.group_modes(gDst_kd, 0, 2), - ) - gDst_kr = cute.local_tile(tma_tensor_ws_kr, (CHUNK, D), (None, None, None)) - tKRws_s, tKRws_g = cpasync.tma_partition( - tma_atom_ws_kr, - 0, - cute.make_layout(1), - cute.group_modes(s_k_restored, 0, 2), - cute.group_modes(gDst_kr, 0, 2), - ) gDst_inv = cute.local_tile(tma_tensor_ws_inv, (CHUNK, CHUNK), (None, None, None)) tINVws_s, tINVws_g = cpasync.tma_partition( tma_atom_ws_inv, @@ -189,6 +162,18 @@ def k1_kernel( cute.group_modes(gDst_mqk, 0, 2), ) ws_slot = head_idx * total_tiles + tile_idx + raw_copy_atom = cute.make_copy_atom(cpasync.CopyBulkS2GOp(), cutlass.BFloat16) + raw_smem_layout = cute.make_layout((CHUNK * D,), stride=(1,)) + raw_gmem_layout = cute.make_layout( + (CHUNK * D, total_tiles * H), + stride=(1, CHUNK * D), + ) + sQDws_raw = cute.make_tensor(s_q_decayed.iterator, raw_smem_layout) + sKDws_raw = cute.make_tensor(s_k_decayed.iterator, raw_smem_layout) + sKRws_raw = cute.make_tensor(s_k_restored.iterator, raw_smem_layout) + gQDws_raw = cute.make_tensor(ws_qd.iterator, raw_gmem_layout) + gKDws_raw = cute.make_tensor(ws_kd.iterator, raw_gmem_layout) + gKRws_raw = cute.make_tensor(ws_kr.iterator, raw_gmem_layout) if warp_idx == 0: with cute.arch.elect_one(): @@ -533,12 +518,14 @@ def k1_kernel( cute.arch.fence_view_async_shared() cute.arch.barrier() - # TMA bulk store all 5 workspace tensors (one elect_one, one thread). + # Preserve the physical K_INTER byte image for qd/kd/kr; inv/mqk remain + # layout-aware tensor TMA stores. if warp_idx == 0: + # CuTeDSL 4.6 elects the issuing lane inside direct bulk-copy atoms. + cute.copy(raw_copy_atom, sQDws_raw, gQDws_raw[(None, ws_slot)]) + cute.copy(raw_copy_atom, sKDws_raw, gKDws_raw[(None, ws_slot)]) + cute.copy(raw_copy_atom, sKRws_raw, gKRws_raw[(None, ws_slot)]) with cute.arch.elect_one(): - cute.copy(tma_atom_ws_qd, tQDws_s[(None,)], tQDws_g[(None, 0, 0, ws_slot)]) - cute.copy(tma_atom_ws_kd, tKDws_s[(None,)], tKDws_g[(None, 0, 0, ws_slot)]) - cute.copy(tma_atom_ws_kr, tKRws_s[(None,)], tKRws_g[(None, 0, 0, ws_slot)]) cute.copy(tma_atom_ws_inv, tINVws_s[(None,)], tINVws_g[(None, 0, 0, ws_slot)]) cute.copy(tma_atom_ws_mqk, tMQKws_s[(None,)], tMQKws_g[(None, 0, 0, ws_slot)]) cute.arch.cp_async_bulk_commit_group() @@ -571,9 +558,6 @@ def run_k1( stream: cuda_drv.CUstream, ): smem_layout_qk = cute.make_layout((CHUNK, D), stride=(D, 1)) - # K_INTER swizzled layout — must match kernel SMEM layout for TMA stores. - kinter_atom = make_smem_layout_atom(SmemLayoutAtomKind.K_INTER, cutlass.BFloat16) - smem_layout_qk_kinter = cute.tile_to_shape(kinter_atom, (CHUNK, D), order=(0, 1)) def make_atom(t): view = cute.make_tensor( @@ -587,18 +571,6 @@ def make_atom(t): (CHUNK, D), ) - def make_ws_store_atom(t): - view = cute.make_tensor( - t.iterator, - cute.make_layout((CHUNK, D, total_tiles * H), stride=(D, 1, CHUNK * D)), - ) - return cpasync.make_tiled_tma_atom( - cpasync.CopyBulkTensorTileS2GOp(), - view, - smem_layout_qk_kinter, - (CHUNK, D), - ) - # (CHUNK, CHUNK) bf16 plain layout for ws_inv / ws_mqk TMA bulk store. smem_layout_cc = cute.make_layout((CHUNK, CHUNK), stride=(CHUNK, 1)) @@ -620,9 +592,6 @@ def make_ws_cc_store_atom(t): tma_atom_q, tma_tensor_q = make_atom(q) tma_atom_k, tma_tensor_k = make_atom(k) tma_atom_g, tma_tensor_g = make_atom(g) - tma_atom_ws_qd, tma_tensor_ws_qd = make_ws_store_atom(ws_qd) - tma_atom_ws_kd, tma_tensor_ws_kd = make_ws_store_atom(ws_kd) - tma_atom_ws_kr, tma_tensor_ws_kr = make_ws_store_atom(ws_kr) tma_atom_ws_inv, tma_tensor_ws_inv = make_ws_cc_store_atom(ws_inv) tma_atom_ws_mqk, tma_tensor_ws_mqk = make_ws_cc_store_atom(ws_mqk) @@ -635,16 +604,13 @@ def make_ws_cc_store_atom(t): tma_tensor_k, tma_atom_g, tma_tensor_g, - tma_atom_ws_qd, - tma_tensor_ws_qd, - tma_atom_ws_kd, - tma_tensor_ws_kd, - tma_atom_ws_kr, - tma_tensor_ws_kr, tma_atom_ws_inv, tma_tensor_ws_inv, tma_atom_ws_mqk, tma_tensor_ws_mqk, + ws_qd, + ws_kd, + ws_kr, a_log, dt_bias, beta, diff --git a/cula/ops/kda/sm90/k2.py b/cula/ops/kda/sm90/k2.py index 12ccbe3..8bb89c8 100644 --- a/cula/ops/kda/sm90/k2.py +++ b/cula/ops/kda/sm90/k2.py @@ -63,16 +63,13 @@ def _make_out_kinter_one_stage(): def k2_kernel( tma_atom_v: cute.CopyAtom, tma_tensor_v: cute.Tensor, - tma_atom_kd: cute.CopyAtom, - tma_tensor_kd: cute.Tensor, - tma_atom_qd: cute.CopyAtom, - tma_tensor_qd: cute.Tensor, - tma_atom_kr: cute.CopyAtom, - tma_tensor_kr: cute.Tensor, tma_atom_inv: cute.CopyAtom, tma_tensor_inv: cute.Tensor, tma_atom_mqk: cute.CopyAtom, tma_tensor_mqk: cute.Tensor, + ws_qd: cute.Tensor, + ws_kd: cute.Tensor, + ws_kr: cute.Tensor, tma_atom_out: cute.CopyAtom, tma_tensor_out: cute.Tensor, out_gmem: cute.Tensor, @@ -156,30 +153,6 @@ def k2_kernel( cute.group_modes(sV, 0, 2), cute.group_modes(gSrc_v, 0, 2), ) - gSrc_kd = cute.local_tile(tma_tensor_kd, (CHUNK, D), (None, None, None)) - tKDs, tKDg = cpasync.tma_partition( - tma_atom_kd, - 0, - cute.make_layout(1), - cute.group_modes(sKd, 0, 2), - cute.group_modes(gSrc_kd, 0, 2), - ) - gSrc_qd = cute.local_tile(tma_tensor_qd, (CHUNK, D), (None, None, None)) - tQDs, tQDg = cpasync.tma_partition( - tma_atom_qd, - 0, - cute.make_layout(1), - cute.group_modes(sQd, 0, 2), - cute.group_modes(gSrc_qd, 0, 2), - ) - gSrc_kr = cute.local_tile(tma_tensor_kr, (CHUNK, D), (None, None, None)) - tKRs, tKRg = cpasync.tma_partition( - tma_atom_kr, - 0, - cute.make_layout(1), - cute.group_modes(sKr, 0, 2), - cute.group_modes(gSrc_kr, 0, 2), - ) gSrc_inv = cute.local_tile(tma_tensor_inv, (CHUNK, CHUNK), (None, None, None)) tIs, tIg = cpasync.tma_partition( tma_atom_inv, @@ -220,6 +193,21 @@ def k2_kernel( cute.group_modes(sBeta, 0, 2), cute.group_modes(gSrc_beta, 0, 2), ) + raw_copy_atom = cute.make_copy_atom(cpasync.CopyBulkG2SOp(), cutlass.BFloat16) + raw_stage_layout = cute.make_layout( + (CHUNK * D, STAGES), + stride=(1, CHUNK * D), + ) + raw_gmem_layout = cute.make_layout( + (CHUNK * D, total_tiles * H), + stride=(1, CHUNK * D), + ) + sKD_raw = cute.make_tensor(sKd.iterator, raw_stage_layout) + sQD_raw = cute.make_tensor(sQd.iterator, raw_stage_layout) + sKR_raw = cute.make_tensor(sKr.iterator, raw_stage_layout) + gKD_raw = cute.make_tensor(ws_kd.iterator, raw_gmem_layout) + gQD_raw = cute.make_tensor(ws_qd.iterator, raw_gmem_layout) + gKR_raw = cute.make_tensor(ws_kr.iterator, raw_gmem_layout) # Load initial_state -> sState[K_in, D_out] if has_initial_state: @@ -360,9 +348,9 @@ def k2_kernel( cute.copy(tma_atom_v, tVg_seq[(None, t, 0, head_idx)], tVs_seq[(None, s_dyn_l)], tma_bar_ptr=bar_l) else: cute.copy(tma_atom_v, tVg[(None, tg_l, 0, head_idx)], tVs[(None, s_dyn_l)], tma_bar_ptr=bar_l) - cute.copy(tma_atom_kd, tKDg[(None, 0, 0, wt_l)], tKDs[(None, s_dyn_l)], tma_bar_ptr=bar_l) - cute.copy(tma_atom_qd, tQDg[(None, 0, 0, wt_l)], tQDs[(None, s_dyn_l)], tma_bar_ptr=bar_l) - cute.copy(tma_atom_kr, tKRg[(None, 0, 0, wt_l)], tKRs[(None, s_dyn_l)], tma_bar_ptr=bar_l) + cute.copy(raw_copy_atom, gKD_raw[(None, wt_l)], sKD_raw[(None, s_dyn_l)], mbar_ptr=bar_l) + cute.copy(raw_copy_atom, gQD_raw[(None, wt_l)], sQD_raw[(None, s_dyn_l)], mbar_ptr=bar_l) + cute.copy(raw_copy_atom, gKR_raw[(None, wt_l)], sKR_raw[(None, s_dyn_l)], mbar_ptr=bar_l) cute.copy(tma_atom_inv, tIg[(None, 0, 0, wt_l)], tIs[(None, s_dyn_l)], tma_bar_ptr=bar_l) cute.copy(tma_atom_mqk, tMg[(None, 0, 0, wt_l)], tMs[(None, s_dyn_l)], tma_bar_ptr=bar_l) cute.copy(tma_atom_gt, tGTg[(None, 0, 0, wt_l)], tGTs[(None, s_dyn_l)], tma_bar_ptr=bar_l) @@ -633,13 +621,6 @@ def make_thd_atom(t, op, t_total: cutlass.Int32): ) return cpasync.make_tiled_tma_atom(op, view, kinter_smem, (CHUNK, D)) - def make_ws_qkd_atom(t): - view = cute.make_tensor( - t.iterator, - cute.make_layout((CHUNK, D, total_tiles * H), stride=(D, 1, CHUNK * D)), - ) - return cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileG2SOp(), view, kinter_smem, (CHUNK, D)) - def make_ws_cc_atom(t): view = cute.make_tensor( t.iterator, @@ -649,9 +630,6 @@ def make_ws_cc_atom(t): tma_atom_v, tma_tensor_v = make_thd_atom(v, cpasync.CopyBulkTensorTileG2SOp(), V_T_total) tma_atom_out, tma_tensor_out = make_thd_atom(out, cpasync.CopyBulkTensorTileS2GOp(), O_T_total) - tma_atom_kd, tma_tensor_kd = make_ws_qkd_atom(ws_kd) - tma_atom_qd, tma_tensor_qd = make_ws_qkd_atom(ws_qd) - tma_atom_kr, tma_tensor_kr = make_ws_qkd_atom(ws_kr) tma_atom_inv, tma_tensor_inv = make_ws_cc_atom(ws_inv) tma_atom_mqk, tma_tensor_mqk = make_ws_cc_atom(ws_mqk) @@ -701,16 +679,13 @@ def make_beta_atom(t): k2_kernel( tma_atom_v, tma_tensor_v, - tma_atom_kd, - tma_tensor_kd, - tma_atom_qd, - tma_tensor_qd, - tma_atom_kr, - tma_tensor_kr, tma_atom_inv, tma_tensor_inv, tma_atom_mqk, tma_tensor_mqk, + ws_qd, + ws_kd, + ws_kr, tma_atom_out, tma_tensor_out, out, From 33f8da18594af07b2e7ab52103b2c6f280e516e3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=8C=83=E5=9D=A4?= Date: Mon, 3 Aug 2026 11:08:08 +0800 Subject: [PATCH 2/7] fix(kda): align intracard pre-scan workspace loads --- benchmarks/bench_kda_sm90_prefill.py | 1 + cula/ops/kda/sm90/cp/pre_scan.py | 54 ++++++++++------------------ cula/ops/kda/sm90/k1.py | 3 +- cula/ops/kda/sm90/k2.py | 2 ++ 4 files changed, 24 insertions(+), 36 deletions(-) diff --git a/benchmarks/bench_kda_sm90_prefill.py b/benchmarks/bench_kda_sm90_prefill.py index d02baca..89d466c 100644 --- a/benchmarks/bench_kda_sm90_prefill.py +++ b/benchmarks/bench_kda_sm90_prefill.py @@ -107,6 +107,7 @@ def run_cula(q, k, v, g, beta, scale, A_log, dt_bias, init_state, cu_seqlens, lo use_gate_in_kernel=True, safe_gate=True, lower_bound=lower_bound, + use_intracard_cp=False, ) diff --git a/cula/ops/kda/sm90/cp/pre_scan.py b/cula/ops/kda/sm90/cp/pre_scan.py index 900b333..0e233e3 100644 --- a/cula/ops/kda/sm90/cp/pre_scan.py +++ b/cula/ops/kda/sm90/cp/pre_scan.py @@ -48,16 +48,14 @@ def pre_scan_kernel( tma_atom_v: cute.CopyAtom, tma_tensor_v: cute.Tensor, - tma_atom_kd: cute.CopyAtom, - tma_tensor_kd: cute.Tensor, - tma_atom_kr: cute.CopyAtom, - tma_tensor_kr: cute.Tensor, tma_atom_inv: cute.CopyAtom, tma_tensor_inv: cute.Tensor, tma_atom_gt: cute.CopyAtom, tma_tensor_gt: cute.Tensor, tma_atom_beta: cute.CopyAtom, tma_tensor_beta: cute.Tensor, + ws_kd: cute.Tensor, + ws_kr: cute.Tensor, H: cutlass.Constexpr[int], total_tiles: cutlass.Int32, T_total: cutlass.Int32, @@ -118,22 +116,6 @@ def pre_scan_kernel( cute.group_modes(sV, 0, 2), cute.group_modes(gSrc_v, 0, 2), ) - gSrc_kd = cute.local_tile(tma_tensor_kd, (CHUNK, D), (None, None, None)) - tKDs, tKDg = cpasync.tma_partition( - tma_atom_kd, - 0, - cute.make_layout(1), - cute.group_modes(sKd, 0, 2), - cute.group_modes(gSrc_kd, 0, 2), - ) - gSrc_kr = cute.local_tile(tma_tensor_kr, (CHUNK, D), (None, None, None)) - tKRs, tKRg = cpasync.tma_partition( - tma_atom_kr, - 0, - cute.make_layout(1), - cute.group_modes(sKr, 0, 2), - cute.group_modes(gSrc_kr, 0, 2), - ) gSrc_inv = cute.local_tile(tma_tensor_inv, (CHUNK, CHUNK), (None, None, None)) tIs, tIg = cpasync.tma_partition( tma_atom_inv, @@ -158,6 +140,19 @@ def pre_scan_kernel( cute.group_modes(sBeta, 0, 2), cute.group_modes(gSrc_beta, 0, 2), ) + raw_copy_atom = cute.make_copy_atom(cpasync.CopyBulkG2SOp(), cutlass.BFloat16) + raw_stage_layout = cute.make_layout( + (CHUNK * D, STAGES), + stride=(1, CHUNK * D), + ) + raw_gmem_layout = cute.make_layout( + (CHUNK * D, total_tiles * H), + stride=(1, CHUNK * D), + ) + sKD_raw = cute.make_tensor(sKd.iterator, raw_stage_layout) + sKR_raw = cute.make_tensor(sKr.iterator, raw_stage_layout) + 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 if tidx < D: @@ -273,8 +268,8 @@ def pre_scan_kernel( cute.copy(tma_atom_v, tVg_seq[(None, t, 0, head_idx)], tVs_seq[(None, s_dyn_l)], tma_bar_ptr=bar_l) else: cute.copy(tma_atom_v, tVg[(None, tg_l, 0, head_idx)], tVs[(None, s_dyn_l)], tma_bar_ptr=bar_l) - cute.copy(tma_atom_kd, tKDg[(None, 0, 0, wt_l)], tKDs[(None, s_dyn_l)], tma_bar_ptr=bar_l) - cute.copy(tma_atom_kr, tKRg[(None, 0, 0, wt_l)], tKRs[(None, s_dyn_l)], tma_bar_ptr=bar_l) + cute.copy(raw_copy_atom, gKD_raw[(None, wt_l)], sKD_raw[(None, s_dyn_l)], mbar_ptr=bar_l) + cute.copy(raw_copy_atom, gKR_raw[(None, wt_l)], sKR_raw[(None, s_dyn_l)], mbar_ptr=bar_l) cute.copy(tma_atom_inv, tIg[(None, 0, 0, wt_l)], tIs[(None, s_dyn_l)], tma_bar_ptr=bar_l) cute.copy(tma_atom_gt, tGTg[(None, 0, 0, wt_l)], tGTs[(None, s_dyn_l)], tma_bar_ptr=bar_l) cute.copy(tma_atom_beta, tBg[(None, 0, 0, wt_l)], tBs[(None, s_dyn_l)], tma_bar_ptr=bar_l) @@ -487,13 +482,6 @@ def make_thd_atom(t, op): ) return cpasync.make_tiled_tma_atom(op, view, kinter_smem, (CHUNK, D)) - def make_ws_qkd_atom(t): - view = cute.make_tensor( - t.iterator, - cute.make_layout((CHUNK, D, total_tiles * H), stride=(D, 1, CHUNK * D)), - ) - return cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileG2SOp(), view, kinter_smem, (CHUNK, D)) - def make_ws_cc_atom(t): view = cute.make_tensor( t.iterator, @@ -502,8 +490,6 @@ def make_ws_cc_atom(t): return cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileG2SOp(), view, cc_smem, (CHUNK, CHUNK)) tma_atom_v, tma_tensor_v = make_thd_atom(v, cpasync.CopyBulkTensorTileG2SOp()) - tma_atom_kd, tma_tensor_kd = make_ws_qkd_atom(ws_kd) - tma_atom_kr, tma_tensor_kr = make_ws_qkd_atom(ws_kr) tma_atom_inv, tma_tensor_inv = make_ws_cc_atom(ws_inv) gt_smem = cute.make_layout((D, 1), stride=(1, D)) @@ -551,16 +537,14 @@ def make_beta_atom(t): pre_scan_kernel( tma_atom_v, tma_tensor_v, - tma_atom_kd, - tma_tensor_kd, - tma_atom_kr, - tma_tensor_kr, tma_atom_inv, tma_tensor_inv, tma_atom_gt, tma_tensor_gt, tma_atom_beta, tma_tensor_beta, + ws_kd, + ws_kr, H, total_tiles, T_total, diff --git a/cula/ops/kda/sm90/k1.py b/cula/ops/kda/sm90/k1.py index 024fe5f..56f2ca9 100644 --- a/cula/ops/kda/sm90/k1.py +++ b/cula/ops/kda/sm90/k1.py @@ -519,7 +519,8 @@ def k1_kernel( cute.arch.barrier() # Preserve the physical K_INTER byte image for qd/kd/kr; inv/mqk remain - # layout-aware tensor TMA stores. + # layout-aware tensor TMA stores. This raw-workspace transport idea comes + # from Flash-Flash-KDA: https://github.com/Itssshikhar/Flash-Flash-KDA if warp_idx == 0: # CuTeDSL 4.6 elects the issuing lane inside direct bulk-copy atoms. cute.copy(raw_copy_atom, sQDws_raw, gQDws_raw[(None, ws_slot)]) diff --git a/cula/ops/kda/sm90/k2.py b/cula/ops/kda/sm90/k2.py index 8bb89c8..7395984 100644 --- a/cula/ops/kda/sm90/k2.py +++ b/cula/ops/kda/sm90/k2.py @@ -348,6 +348,8 @@ def k2_kernel( cute.copy(tma_atom_v, tVg_seq[(None, t, 0, head_idx)], tVs_seq[(None, s_dyn_l)], tma_bar_ptr=bar_l) else: cute.copy(tma_atom_v, tVg[(None, tg_l, 0, head_idx)], tVs[(None, s_dyn_l)], tma_bar_ptr=bar_l) + # Restore the byte-identical K_INTER images produced by K1. The + # raw-workspace transport idea is credited there to Flash-Flash-KDA. cute.copy(raw_copy_atom, gKD_raw[(None, wt_l)], sKD_raw[(None, s_dyn_l)], mbar_ptr=bar_l) cute.copy(raw_copy_atom, gQD_raw[(None, wt_l)], sQD_raw[(None, s_dyn_l)], mbar_ptr=bar_l) cute.copy(raw_copy_atom, gKR_raw[(None, wt_l)], sKR_raw[(None, s_dyn_l)], mbar_ptr=bar_l) From 698738643134f80d6079360b441d269d58c4a279 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=8C=83=E5=9D=A4?= Date: Mon, 3 Aug 2026 11:41:39 +0800 Subject: [PATCH 3/7] chore(bench): preserve prefill dispatch behavior --- benchmarks/bench_kda_sm90_prefill.py | 1 - 1 file changed, 1 deletion(-) diff --git a/benchmarks/bench_kda_sm90_prefill.py b/benchmarks/bench_kda_sm90_prefill.py index 89d466c..d02baca 100644 --- a/benchmarks/bench_kda_sm90_prefill.py +++ b/benchmarks/bench_kda_sm90_prefill.py @@ -107,7 +107,6 @@ def run_cula(q, k, v, g, beta, scale, A_log, dt_bias, init_state, cu_seqlens, lo use_gate_in_kernel=True, safe_gate=True, lower_bound=lower_bound, - use_intracard_cp=False, ) From 4e67e4f55d241af600235898b55cde09bd90c165 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=8C=83=E5=9D=A4?= Date: Sun, 9 Aug 2026 19:04:18 +0800 Subject: [PATCH 4/7] fix(kda): support bulk-copy election across CuTeDSL versions --- cula/ops/kda/sm90/_common.py | 41 ++++++++++++++++++++++++++++++++ cula/ops/kda/sm90/cp/pre_scan.py | 6 ++--- cula/ops/kda/sm90/k1.py | 11 +++++---- cula/ops/kda/sm90/k2.py | 8 +++---- 4 files changed, 54 insertions(+), 12 deletions(-) diff --git a/cula/ops/kda/sm90/_common.py b/cula/ops/kda/sm90/_common.py index 8044a4c..977bbba 100644 --- a/cula/ops/kda/sm90/_common.py +++ b/cula/ops/kda/sm90/_common.py @@ -3,13 +3,54 @@ """Shared low-level helpers for the SM90 FlashKDA kernels.""" +import re +from importlib.metadata import PackageNotFoundError, version as package_version + import cutlass +import cutlass.cute as cute import torch from cutlass import Int32 from cutlass._mlir.dialects import llvm as _llvm from cutlass.cutlass_dsl import T as _T +def _parse_cutedsl_version(raw_version: str) -> tuple[int, int, int]: + """Return the numeric CuTeDSL version from a release or dev version string.""" + match = re.match(r"^(\d+)\.(\d+)(?:\.(\d+))?", raw_version) + if match is None: + raise RuntimeError(f"Unable to parse the installed CuTeDSL version: {raw_version!r}") + return tuple(int(component or 0) for component in match.groups()) + + +def _installed_cutedsl_version() -> tuple[int, int, int]: + """Read the CuTeDSL version at runtime without relying on a private symbol.""" + raw_version = getattr(cutlass, "__version__", None) + if raw_version is None: + try: + raw_version = package_version("nvidia-cutlass-dsl") + except PackageNotFoundError as exc: + raise RuntimeError("nvidia-cutlass-dsl is required by the SM90 FlashKDA backend") from exc + return _parse_cutedsl_version(raw_version) + + +# CuTeDSL 4.6.0 added the missing elect_one inside cute.copy for async bulk +# atoms. Older releases need an explicit elect_one, while nesting one around +# cute.copy is incorrect in 4.6+. Keep this decision as a compile-time +# constant after detecting the installed runtime version, so the same source +# supports both API behaviours. +_CUTEDSL_VERSION = _installed_cutedsl_version() +_CUTE_COPY_AUTO_ELECTS_BULK = _CUTEDSL_VERSION >= (4, 6, 0) + + +def copy_async_bulk(atom, src, dst, **kwargs) -> None: + """Issue a CuTeDSL async bulk copy across supported elect_one APIs.""" + if cutlass.const_expr(_CUTE_COPY_AUTO_ELECTS_BULK): + cute.copy(atom, src, dst, **kwargs) + else: + with cute.arch.elect_one(): + cute.copy(atom, src, dst, **kwargs) + + def _stream_key(device: torch.device) -> tuple[str, int]: return str(device), int(torch.cuda.current_stream(device).cuda_stream) diff --git a/cula/ops/kda/sm90/cp/pre_scan.py b/cula/ops/kda/sm90/cp/pre_scan.py index 0e233e3..da40f66 100644 --- a/cula/ops/kda/sm90/cp/pre_scan.py +++ b/cula/ops/kda/sm90/cp/pre_scan.py @@ -30,7 +30,7 @@ from cutlass.cute.nvgpu.warpgroup import SmemLayoutAtomKind, make_smem_layout_atom from cutlass.cute.runtime import make_fake_compact_tensor, make_fake_stream -from cula.ops.kda.sm90._common import movm_t_b16 +from cula.ops.kda.sm90._common import copy_async_bulk, movm_t_b16 from cula.ops.kda.sm90.k2 import ( CHUNK, D, @@ -268,8 +268,8 @@ def pre_scan_kernel( cute.copy(tma_atom_v, tVg_seq[(None, t, 0, head_idx)], tVs_seq[(None, s_dyn_l)], tma_bar_ptr=bar_l) else: cute.copy(tma_atom_v, tVg[(None, tg_l, 0, head_idx)], tVs[(None, s_dyn_l)], tma_bar_ptr=bar_l) - cute.copy(raw_copy_atom, gKD_raw[(None, wt_l)], sKD_raw[(None, s_dyn_l)], mbar_ptr=bar_l) - cute.copy(raw_copy_atom, gKR_raw[(None, wt_l)], sKR_raw[(None, s_dyn_l)], mbar_ptr=bar_l) + copy_async_bulk(raw_copy_atom, gKD_raw[(None, wt_l)], sKD_raw[(None, s_dyn_l)], mbar_ptr=bar_l) + copy_async_bulk(raw_copy_atom, gKR_raw[(None, wt_l)], sKR_raw[(None, s_dyn_l)], mbar_ptr=bar_l) cute.copy(tma_atom_inv, tIg[(None, 0, 0, wt_l)], tIs[(None, s_dyn_l)], tma_bar_ptr=bar_l) cute.copy(tma_atom_gt, tGTg[(None, 0, 0, wt_l)], tGTs[(None, s_dyn_l)], tma_bar_ptr=bar_l) cute.copy(tma_atom_beta, tBg[(None, 0, 0, wt_l)], tBs[(None, s_dyn_l)], tma_bar_ptr=bar_l) diff --git a/cula/ops/kda/sm90/k1.py b/cula/ops/kda/sm90/k1.py index 56f2ca9..34aa1e4 100644 --- a/cula/ops/kda/sm90/k1.py +++ b/cula/ops/kda/sm90/k1.py @@ -30,7 +30,7 @@ from cutlass.cute.nvgpu.warpgroup import SmemLayoutAtomKind, make_smem_layout_atom from cutlass.cute.runtime import make_fake_compact_tensor, make_fake_stream -from cula.ops.kda.sm90._common import _stream_key, add_f16x2_u32, movm_t_b16 +from cula.ops.kda.sm90._common import _stream_key, add_f16x2_u32, copy_async_bulk, movm_t_b16 CHUNK: int = 16 D: int = 128 @@ -522,10 +522,11 @@ def k1_kernel( # layout-aware tensor TMA stores. This raw-workspace transport idea comes # from Flash-Flash-KDA: https://github.com/Itssshikhar/Flash-Flash-KDA if warp_idx == 0: - # CuTeDSL 4.6 elects the issuing lane inside direct bulk-copy atoms. - cute.copy(raw_copy_atom, sQDws_raw, gQDws_raw[(None, ws_slot)]) - cute.copy(raw_copy_atom, sKDws_raw, gKDws_raw[(None, ws_slot)]) - cute.copy(raw_copy_atom, sKRws_raw, gKRws_raw[(None, ws_slot)]) + # copy_async_bulk supplies elect_one only for pre-4.6 CuTeDSL. CuTeDSL + # 4.6+ emits it inside cute.copy, where an outer elect_one is invalid. + copy_async_bulk(raw_copy_atom, sQDws_raw, gQDws_raw[(None, ws_slot)]) + copy_async_bulk(raw_copy_atom, sKDws_raw, gKDws_raw[(None, ws_slot)]) + copy_async_bulk(raw_copy_atom, sKRws_raw, gKRws_raw[(None, ws_slot)]) with cute.arch.elect_one(): cute.copy(tma_atom_ws_inv, tINVws_s[(None,)], tINVws_g[(None, 0, 0, ws_slot)]) cute.copy(tma_atom_ws_mqk, tMQKws_s[(None,)], tMQKws_g[(None, 0, 0, ws_slot)]) diff --git a/cula/ops/kda/sm90/k2.py b/cula/ops/kda/sm90/k2.py index 7395984..1cb60cc 100644 --- a/cula/ops/kda/sm90/k2.py +++ b/cula/ops/kda/sm90/k2.py @@ -44,7 +44,7 @@ def _make_state_smem_layout(): return cute.tile_to_shape(atom, (D, D), (0, 1)) -from cula.ops.kda.sm90._common import _stream_key, movm_t_b16 # noqa: E402 +from cula.ops.kda.sm90._common import _stream_key, copy_async_bulk, movm_t_b16 # noqa: E402 def _make_out_kinter_one_stage(): @@ -350,9 +350,9 @@ def k2_kernel( cute.copy(tma_atom_v, tVg[(None, tg_l, 0, head_idx)], tVs[(None, s_dyn_l)], tma_bar_ptr=bar_l) # Restore the byte-identical K_INTER images produced by K1. The # raw-workspace transport idea is credited there to Flash-Flash-KDA. - cute.copy(raw_copy_atom, gKD_raw[(None, wt_l)], sKD_raw[(None, s_dyn_l)], mbar_ptr=bar_l) - cute.copy(raw_copy_atom, gQD_raw[(None, wt_l)], sQD_raw[(None, s_dyn_l)], mbar_ptr=bar_l) - cute.copy(raw_copy_atom, gKR_raw[(None, wt_l)], sKR_raw[(None, s_dyn_l)], mbar_ptr=bar_l) + copy_async_bulk(raw_copy_atom, gKD_raw[(None, wt_l)], sKD_raw[(None, s_dyn_l)], mbar_ptr=bar_l) + copy_async_bulk(raw_copy_atom, gQD_raw[(None, wt_l)], sQD_raw[(None, s_dyn_l)], mbar_ptr=bar_l) + copy_async_bulk(raw_copy_atom, gKR_raw[(None, wt_l)], sKR_raw[(None, s_dyn_l)], mbar_ptr=bar_l) cute.copy(tma_atom_inv, tIg[(None, 0, 0, wt_l)], tIs[(None, s_dyn_l)], tma_bar_ptr=bar_l) cute.copy(tma_atom_mqk, tMg[(None, 0, 0, wt_l)], tMs[(None, s_dyn_l)], tma_bar_ptr=bar_l) cute.copy(tma_atom_gt, tGTg[(None, 0, 0, wt_l)], tGTs[(None, s_dyn_l)], tma_bar_ptr=bar_l) From 3ad91f2b5d97255decd2450b70c62aff2b748e26 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=8C=83=E5=9D=A4?= Date: Sun, 9 Aug 2026 19:20:46 +0800 Subject: [PATCH 5/7] docs(bench): record SM90 raw workspace A/B --- BENCHMARK_SM90_RAW_WORKSPACE.md | 76 +++++++++++++++++++++++++++++++++ README.md | 8 ++++ 2 files changed, 84 insertions(+) create mode 100644 BENCHMARK_SM90_RAW_WORKSPACE.md diff --git a/BENCHMARK_SM90_RAW_WORKSPACE.md b/BENCHMARK_SM90_RAW_WORKSPACE.md new file mode 100644 index 0000000..5d1a9de --- /dev/null +++ b/BENCHMARK_SM90_RAW_WORKSPACE.md @@ -0,0 +1,76 @@ +# SM90 Raw Workspace Transport A/B + +This benchmark isolates the FlashKDA workspace-transport change from +intracard context parallelism. Both runs use the same process configuration, +inputs, and GPU; only the cuLA checkout changes. + +| Item | Value | +|---|---| +| Baseline | `d78f0075c69135fe191358fe8daf10e7383009d4` | +| Optimized | `4e67e4f55d241af600235898b55cde09bd90c165` | +| GPU | NVIDIA H200, SM90, 143,771 MiB device memory | +| Driver / CUDA | 570.148.08 / CUDA 13.0.48 | +| PyTorch | `2.8.0a0+34c6371d24.nv25.8` | +| CuTeDSL | `nvidia-cutlass-dsl==4.6.1` | +| FLA | `flash-linear-attention==0.5.0` | +| Inputs | BF16 Q/K/V/gate, FP32 state, H=64, D=128, `safe_gate=True` | +| CP mode | Serial K1+K2 (`use_intracard_cp=None`, planner resolves to a trivial plan) | +| Timing | 25 warmup + 100 measured iterations, IQR-mean CUDA events | +| Command | `python benchmarks/bench_kda_sm90_prefill.py --mode both` | + +The test machine is an **H200**, and these results belong to the H200 SM90 +benchmark set. + +## Fixed length + +| B | T | Baseline cuLA (ms) | Optimized cuLA (ms) | A/B | +|---:|---:|---:|---:|---:| +| 1 | 512 | 0.1241 | 0.1276 | 0.973x | +| 1 | 1024 | 0.1831 | 0.1583 | **1.157x** | +| 1 | 4096 | 0.6480 | 0.5704 | **1.136x** | +| 1 | 8192 | 1.2716 | 1.1108 | **1.145x** | +| 1 | 16384 | 2.5193 | 2.2082 | **1.141x** | +| 2 | 512 | 0.1365 | 0.1294 | **1.055x** | +| 2 | 1024 | 0.2431 | 0.1972 | **1.233x** | +| 2 | 4096 | 0.8974 | 0.7174 | **1.251x** | +| 2 | 8192 | 1.7648 | 1.4348 | **1.230x** | +| 2 | 16384 | 3.5085 | 2.8760 | **1.220x** | + +## Variable length + +| Distribution | Sequences / total T | Baseline cuLA (ms) | Optimized cuLA (ms) | A/B | +|---|---:|---:|---:|---:| +| uniform | 10 / 4096 | 0.4746 | 0.3842 | **1.235x** | +| random | 10 / 4096 | 0.5537 | 0.4550 | **1.217x** | +| skewed | 10 / 4096 | 0.6640 | 0.5754 | **1.154x** | +| uniform | 20 / 4096 | 0.4621 | 0.3752 | **1.232x** | +| random | 20 / 4096 | 0.5487 | 0.4406 | **1.245x** | +| skewed | 20 / 4096 | 0.6717 | 0.5904 | **1.138x** | +| uniform | 10 / 8192 | 0.9028 | 0.7349 | **1.229x** | +| random | 10 / 8192 | 1.0406 | 0.8582 | **1.213x** | +| skewed | 10 / 8192 | 1.2627 | 1.0924 | **1.156x** | +| uniform | 20 / 8192 | 0.8594 | 0.6886 | **1.248x** | +| random | 20 / 8192 | 1.0135 | 0.8260 | **1.227x** | +| skewed | 20 / 8192 | 1.2605 | 1.1314 | **1.114x** | +| uniform | 10 / 16384 | 1.7411 | 1.4488 | **1.202x** | +| random | 10 / 16384 | 2.0374 | 1.6995 | **1.199x** | +| skewed | 10 / 16384 | 2.4653 | 2.1679 | **1.137x** | +| uniform | 20 / 16384 | 1.6536 | 1.3662 | **1.210x** | +| random | 20 / 16384 | 1.9532 | 1.5986 | **1.222x** | +| skewed | 20 / 16384 | 2.4597 | 2.2102 | **1.113x** | + +## Summary and numerical accuracy + +- 27/28 configurations are faster; the only regression is B=1, T=512, + where the difference is within small-kernel launch noise. +- Geometric-mean A/B speedup is **1.1778x**; summing all measured cuLA + latencies gives a **15.45%** reduction. +- The optimized run reports `relative_rms_error` 0.004573–0.004924, + `rel_max` 0.007772–0.015000, and `mean_diff` 1.1e-5–1.3e-5 against + FLA. The baseline run reports the same accuracy values to the printed + precision, so the raw byte transport does not change numerical behavior. +- Flash-Flash-KDA reports isolated H100 workspace-transport reductions of + 23%, 34%, and 37% for fixed, uneven packed, and uniform packed CHUNK=16 + inputs. This cuLA result is a complete CP-off prefill A/B on a different + GPU and workload mix; it is therefore a comparable direction-of-gain + check, not an exact reproduction of those percentages. diff --git a/README.md b/README.md index c5e145f..6063c04 100644 --- a/README.md +++ b/README.md @@ -129,11 +129,19 @@ See [BENCHMARK_GB200_CUDA_130.md](BENCHMARK_GB200_CUDA_130.md) tested with CUDA See [BENCHMARK_H200.md](BENCHMARK_H200.md) tested with CUDA 12.9 for detailed results. +The raw workspace transport A/B is documented separately in +[BENCHMARK_SM90_RAW_WORKSPACE.md](BENCHMARK_SM90_RAW_WORKSPACE.md). On the +H200 validation container it makes 27/28 serial (non-intracard-CP) configurations +faster, with a 1.18x geometric-mean speedup and 15.45% latency reduction. + **Highlights:** - **KDA Modular Forward (Blackwell):** **avg 1.33x** speedup on fixed-length, **avg 1.35x** on variable-length (18 configs, uniform/skewed/random). - **Lightning Attention Prefill (Blackwell):** up to **2.08x** speedup (B=2). - **Lightning Attention Varlen (Blackwell):** **avg 1.47x** speedup across 126 configs (uniform/skewed/random). - **KDA Fused Forward (Hopper):** **avg 1.58x** speedup across fixed-length and variable-length sequences. +- **SM90 raw workspace transport:** **1.18x** geometric-mean A/B speedup + (15.45% weighted latency reduction, 27/28 configurations) with intracard CP + disabled; see [BENCHMARK_SM90_RAW_WORKSPACE.md](BENCHMARK_SM90_RAW_WORKSPACE.md). To regenerate benchmarks: From a3f2e39176802f1c531c123478502b35882a3b74 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=8C=83=E5=9D=A4?= Date: Sun, 9 Aug 2026 19:31:54 +0800 Subject: [PATCH 6/7] docs(bench): add SM90 NCU spot check --- BENCHMARK_SM90_RAW_WORKSPACE.md | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/BENCHMARK_SM90_RAW_WORKSPACE.md b/BENCHMARK_SM90_RAW_WORKSPACE.md index 5d1a9de..f304ead 100644 --- a/BENCHMARK_SM90_RAW_WORKSPACE.md +++ b/BENCHMARK_SM90_RAW_WORKSPACE.md @@ -74,3 +74,25 @@ benchmark set. inputs. This cuLA result is a complete CP-off prefill A/B on a different GPU and workload mix; it is therefore a comparable direction-of-gain check, not an exact reproduction of those percentages. + +## Nsight Compute spot check + +NCU 2025.3.0 (`--set full`, `--launch-count 2`, filtering the generated K1/K2 +names) profiled the first fixed shape, B=1, T=512, H=64. The six `ctc__*` +metrics were unavailable on this H200 setup, but the kernel duration and +L1-TEX/XBAR counters below were collected in both reports. + +| Metric | Baseline | Optimized | Change | +|---|---:|---:|---:| +| K1 GPU duration | 41.152 us | 27.680 us | **-32.7%** | +| K2 GPU duration | 76.416 us | 75.072 us | **-1.8%** | +| K1 + K2 duration | 117.568 us | 102.752 us | **-12.6%** | +| K1 global-op-TMA-store XBAR bytes | 52.43 MB | 27.26 MB | **-48.0%** | +| K2 global-op-TMA-load XBAR bytes | 70.32 MB | 45.15 MB | **-35.8%** | + +These counters show the same signature as the latency A/B: the main saving is +in K1's workspace stores, while K2's total duration changes little because +the recurrence and retained TMA loads remain unchanged. NCU does not expose +the internal TensorMap segment count directly, so the evidence is the +observable duration and transport counters rather than an inferred segment +number. From 38093d23e2ef9bc0eeb0316508fd005a8f029bf1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=8C=83=E5=9D=A4?= Date: Sun, 9 Aug 2026 19:33:38 +0800 Subject: [PATCH 7/7] style(kda): satisfy ruff import formatting --- cula/ops/kda/sm90/_common.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/cula/ops/kda/sm90/_common.py b/cula/ops/kda/sm90/_common.py index 977bbba..a24beb6 100644 --- a/cula/ops/kda/sm90/_common.py +++ b/cula/ops/kda/sm90/_common.py @@ -4,7 +4,8 @@ """Shared low-level helpers for the SM90 FlashKDA kernels.""" import re -from importlib.metadata import PackageNotFoundError, version as package_version +from importlib.metadata import PackageNotFoundError +from importlib.metadata import version as package_version import cutlass import cutlass.cute as cute