D1: paired atoms epilogue (no TMEM round-trip)
Replace NO-OP round-trip + normalize + epilogue_tma_store with: - get_tmem_load_op + get_smem_store_op paired atoms - One-way TMEM→reg (normalize) →SMEM→GMEM - Eliminates ~3% error from TMEM layout mismatch - O rescale disabled (single KV tile only for now) - Pre-computed TMA partitions outside if blocks
This commit is contained in:
@@ -335,67 +335,72 @@ class FmhaKernel:
|
||||
# Wait for MMA's PV[N-1] to commit before reading O.
|
||||
final_o_bar.arrive_and_wait()
|
||||
|
||||
# === NO-OP TMEM round-trip: re-map O from MMA layout to epilog layout ===
|
||||
tTMrO_noop = cute.make_rmem_tensor(
|
||||
(tTMEM_LOADcO.shape, 128 // corr_tile_size), self.acc_dtype
|
||||
)
|
||||
for i in range(n_corr_tiles):
|
||||
tTMrO_i_ = tTMrO_noop[None, i]
|
||||
tTMrO_i_layout = cute.composition(
|
||||
tTMrO_i_.layout, cute.make_layout(tTMrO_noop.shape[0])
|
||||
)
|
||||
tTMrO_i = cute.make_tensor(tTMrO_i_.iterator, tTMrO_i_layout)
|
||||
tTMEM_LOADtO_i = cute.make_tensor(
|
||||
tTMEM_LOADtO.iterator + i * corr_tile_size,
|
||||
tTMEM_LOADtO.layout,
|
||||
)
|
||||
tTMEM_STOREtO_i = cute.make_tensor(
|
||||
tTMEM_STOREtO.iterator + i * corr_tile_size,
|
||||
tTMEM_STOREtO.layout,
|
||||
)
|
||||
cute.copy(tiled_tmem_load_o, tTMEM_LOADtO_i, tTMrO_i)
|
||||
cute.copy(tiled_tmem_store_o, tTMrO_i, tTMEM_STOREtO_i)
|
||||
cute.arch.fence_view_async_tmem_store()
|
||||
|
||||
# === Final O normalization: O *= 1/row_sum ===
|
||||
# === Correction epilog: one-way TMEM → reg (normalize) → SMEM → GMEM ===
|
||||
# Uses get_tmem_load_op + get_smem_store_op paired atoms.
|
||||
# NO TMEM round-trip — hand-constructed atoms corrupt data.
|
||||
inv_row_sum = Float32(1.0) / row_sum
|
||||
|
||||
tTMrO = cute.make_rmem_tensor(
|
||||
(tTMEM_LOADcO.shape, 128 // corr_tile_size), self.acc_dtype
|
||||
tCtO_base = cute.make_tensor(tmem_ptr + self.tmem_o0_offset, tCtO_fake.layout)
|
||||
tCtO = utils.gemm.sm100.transform_partitioned_tensor_layout(tCtO_base)
|
||||
tiled_copy_t2r, tTR_tO, tTR_rO = utils.gemm.sm100.epilogue_tmem_copy_and_partition(
|
||||
self, tidx, tCtO, tCgC, epi_tile, self.use_2cta_instrs
|
||||
)
|
||||
tTR_rC = cute.make_rmem_tensor(tTR_rO.shape, self.c_dtype)
|
||||
tiled_copy_r2s, tRS_rC, tRS_sC = utils.gemm.sm100.epilogue_smem_copy_and_partition(
|
||||
self, tiled_copy_t2r, tTR_rC, tidx, sC
|
||||
)
|
||||
tCgC_epi = cute.flat_divide(tCgC, epi_tile)
|
||||
bSG_sC, bSG_gC_partitioned = cpasync.tma_partition(
|
||||
tma_c, 0, cute.make_layout(1),
|
||||
cute.group_modes(sC, 0, 2),
|
||||
cute.group_modes(tCgC_epi, 0, 2),
|
||||
)
|
||||
epilog_sync_bar = pipeline.NamedBarrier(
|
||||
barrier_id=self.epilog_sync_bar_id,
|
||||
num_threads=32 * len(self.epilogue_warp_id),
|
||||
)
|
||||
|
||||
for i in range(n_corr_tiles):
|
||||
tTMrO_i_ = tTMrO[None, i]
|
||||
tTMrO_i_layout = cute.composition(
|
||||
tTMrO_i_.layout, cute.make_layout(tTMrO.shape[0])
|
||||
)
|
||||
tTMrO_i = cute.make_tensor(tTMrO_i_.iterator, tTMrO_i_layout)
|
||||
tTMEM_LOADtO_i = cute.make_tensor(
|
||||
tTMEM_LOADtO.iterator + i * corr_tile_size, tTMEM_LOADtO.layout
|
||||
)
|
||||
tTMEM_STOREtO_i = cute.make_tensor(
|
||||
tTMEM_STOREtO.iterator + i * corr_tile_size, tTMEM_STOREtO.layout
|
||||
)
|
||||
|
||||
cute.copy(tiled_tmem_load_o, tTMEM_LOADtO_i, tTMrO_i)
|
||||
for j in cutlass.range(cute.size(tTMrO_i), vectorize=True):
|
||||
tTMrO_i[j] = tTMrO_i[j] * inv_row_sum
|
||||
cute.copy(tiled_tmem_store_o, tTMrO_i, tTMEM_STOREtO_i)
|
||||
|
||||
cute.arch.fence_view_async_tmem_store()
|
||||
|
||||
# Epilogue: TMEM → SMEM → GMEM via TMA store.
|
||||
tCtO_base = cute.make_tensor(tmem_ptr + self.tmem_o0_offset, tCtO_fake.layout)
|
||||
acc_cons_st = pipeline.make_pipeline_state(
|
||||
pipeline.PipelineUserType.Consumer, self.num_acc_stage
|
||||
)
|
||||
c_grp = pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id))
|
||||
c_pipe = pipeline.PipelineTmaStore.create(num_stages=self.num_c_stage, producer_group=c_grp)
|
||||
acc_cons_st = utils.gemm.sm100.epilogue_tma_store(
|
||||
self, tidx, warp_idx, tma_c, tCtO_base, sC, tCgC, epi_tile,
|
||||
0, const_expr(lambda x: x), (0, 0, 0),
|
||||
acc_cons_st, acc_pipe, c_pipe,
|
||||
c_pipe = pipeline.PipelineTmaStore.create(
|
||||
num_stages=self.num_c_stage,
|
||||
producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)),
|
||||
)
|
||||
acc_pipe.consumer_wait(acc_cons_st)
|
||||
|
||||
tTR_tO_tile = tTR_tO[(None, None, None, None, None, acc_cons_st.index)]
|
||||
bSG_gC = bSG_gC_partitioned[(None, None, None, Int32(0), Int32(0), Int32(0))]
|
||||
tTR_tO_tile = cute.group_modes(tTR_tO_tile, 3, cute.rank(tTR_tO_tile))
|
||||
bSG_gC = cute.group_modes(bSG_gC, 1, cute.rank(bSG_gC))
|
||||
|
||||
subtile_cnt = cute.size(tTR_tO_tile.shape, mode=[3])
|
||||
for subtile_idx in range(subtile_cnt):
|
||||
tTR_tO_mn = tTR_tO_tile[(None, None, None, subtile_idx)]
|
||||
cute.copy(tiled_copy_t2r, tTR_tO_mn, tTR_rO)
|
||||
|
||||
# Normalize: multiply by inv_row_sum, then convert to BF16
|
||||
for j in cutlass.range(cute.size(tTR_rO), vectorize=True):
|
||||
tTR_rO[j] = tTR_rO[j] * inv_row_sum
|
||||
acc_vec = tiled_copy_r2s.retile(tTR_rO).load()
|
||||
acc_vec = acc_vec.to(self.c_dtype)
|
||||
tRS_rC.store(acc_vec)
|
||||
|
||||
c_buffer = subtile_cnt * 0 + subtile_idx
|
||||
c_buffer = c_buffer % self.num_c_stage
|
||||
cute.copy(tiled_copy_r2s, tRS_rC, tRS_sC[(None, None, None, c_buffer)])
|
||||
cute.arch.fence_proxy("async.shared", space="cta")
|
||||
epilog_sync_bar.arrive_and_wait()
|
||||
|
||||
if warp_idx == self.epilogue_warp_id[0]:
|
||||
cute.copy(tma_c, bSG_sC[(None, c_buffer)], bSG_gC[(None, subtile_idx)])
|
||||
c_pipe.producer_commit()
|
||||
c_pipe.producer_acquire()
|
||||
epilog_sync_bar.arrive_and_wait()
|
||||
|
||||
epilog_sync_bar.arrive_and_wait()
|
||||
acc_pipe.consumer_release(acc_cons_st)
|
||||
acc_cons_st.advance()
|
||||
c_pipe.producer_tail()
|
||||
|
||||
tmem.relinquish_alloc_permit()
|
||||
|
||||
@@ -1,10 +1,11 @@
|
||||
"""D1: Quick test at hd=128 to narrow down the breakage."""
|
||||
"""D1 sweep: paired atoms epilogue (no TMEM round-trip)."""
|
||||
import torch, math
|
||||
import cutlass.cute as cute
|
||||
import cutlass.torch as ct
|
||||
import cuda.bindings.driver as cuda
|
||||
from dsv4.kernels.attention.fmha import FmhaKernel
|
||||
|
||||
# Single KV tile (n=128) only — O rescale not needed
|
||||
for hd in [64, 128, 256]:
|
||||
torch.manual_seed(42)
|
||||
n = 128; m = 128
|
||||
@@ -18,7 +19,6 @@ for hd in [64, 128, 256]:
|
||||
attn = qf @ kf.T * scale; attn = torch.softmax(attn, dim=-1)
|
||||
ref = attn @ v.float()
|
||||
|
||||
# For hd>256, we'd need N-tiling, but 128 is fine as single tile
|
||||
v_kernel = v.unsqueeze(-1)
|
||||
mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q))
|
||||
mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k))
|
||||
@@ -34,4 +34,4 @@ for hd in [64, 128, 256]:
|
||||
|
||||
out = c[:,:,0].float()
|
||||
cos = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)).item()
|
||||
print(f'hd={hd}: cos {cos:.6f} {"PASS" if cos >= 0.97 else "FAIL"}')
|
||||
print(f'hd={hd}: cos {cos:.6f} {"PASS" if cos >= 0.99 else "FAIL (need >=0.99)"}')
|
||||
|
||||
Reference in New Issue
Block a user