v4.4 tag release update. (#3032)

This commit is contained in:
Junkai-Wu
2026-02-13 23:27:58 -05:00
committed by GitHub
parent 01687cfba1
commit d4bbf728ca
140 changed files with 41622 additions and 3689 deletions
@@ -200,10 +200,9 @@ def sm100_4x4x1_kernel_builder(
TMA_STORE_PIPE_DEPTH,
)
acc_shape = tiled_mma.partition_shape_C(mnk_tiler[:2])
tmem_layout = tiled_mma.make_fragment_C(
cute.append(acc_shape, EPILOGUE_STAGE_DEPTH)
).layout
tmem_layout = cute_ext.make_tmem_layout_acc(
tiled_mma, mnk_tiler, EPILOGUE_STAGE_DEPTH
)
bufferA = cute_ext.allocate(
ab_dtype,
@@ -251,11 +250,9 @@ def sm100_4x4x1_kernel_builder(
tiled_copy_t2r = cute.nvgpu.tcgen05.make_tmem_copy(copy_atom_t2r, acc_epi_div)
# Calculate the per thread destination size per iteration for output of TMEM and input of SMEM
thr_copy_t2r = tiled_copy_t2r.get_slice(tid_x)
gC_mnl_epi = cute.flat_divide(tDgD, epi_tile)
tTR_gC = thr_copy_t2r.partition_D(gC_mnl_epi)
acc_d_rmem_layout = cute.make_fragment_like(
tTR_gC[(None, None, None, 0, 0)].layout
acc_d_rmem_layout = cute_ext.make_t2r_rmem_layout(
tiled_copy_t2r, gC_mnl_epi, tid_x
)
bufferRAcc = cute_ext.allocate(