v4.4 tag release update. (#3032)
This commit is contained in:
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user