v4.5 dev update. (#3153)
This commit is contained in:
@@ -624,7 +624,7 @@ class GroupedMixedInputGemmKernel:
|
||||
tile_info_empty_mbar_ptr: cute.struct.MemRange[
|
||||
cutlass.Int64, self.num_tile_info_stage
|
||||
]
|
||||
tmem_dealloc_mbar_ptr: cutlass.Int64
|
||||
tmem_dealloc_mbar: cutlass.Int64
|
||||
tmem_holding_buf: cutlass.Int32
|
||||
|
||||
self.shared_storage = SharedStorage
|
||||
@@ -824,11 +824,11 @@ class GroupedMixedInputGemmKernel:
|
||||
|
||||
# Tensor memory dealloc barrier init
|
||||
tmem = utils.TmemAllocator(
|
||||
storage.tmem_holding_buf,
|
||||
storage.tmem_holding_buf.ptr,
|
||||
barrier_for_retrieve=self.tmem_ptr_sync_barrier,
|
||||
allocator_warp_id=self.epilog_warp_id[0],
|
||||
is_two_cta=use_2cta_instrs,
|
||||
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr,
|
||||
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr,
|
||||
)
|
||||
|
||||
# Cluster arrive after barrier init
|
||||
|
||||
+3
-3
@@ -518,7 +518,7 @@ class GroupedMixedInputGemmAccScaleKernel:
|
||||
tile_info_empty_mbar_ptr: cute.struct.MemRange[
|
||||
cutlass.Int64, self.num_tile_info_stage
|
||||
]
|
||||
tmem_dealloc_mbar_ptr: cutlass.Int64
|
||||
tmem_dealloc_mbar: cutlass.Int64
|
||||
tmem_holding_buf: cutlass.Int32
|
||||
|
||||
self.shared_storage = SharedStorage
|
||||
@@ -708,11 +708,11 @@ class GroupedMixedInputGemmAccScaleKernel:
|
||||
|
||||
# Tensor memory dealloc barrier init
|
||||
tmem = utils.TmemAllocator(
|
||||
storage.tmem_holding_buf,
|
||||
storage.tmem_holding_buf.ptr,
|
||||
barrier_for_retrieve=self.tmem_ptr_sync_barrier,
|
||||
allocator_warp_id=self.epilog_warp_id[0],
|
||||
is_two_cta=use_2cta_instrs,
|
||||
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr,
|
||||
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr,
|
||||
)
|
||||
|
||||
# Cluster arrive after barrier init
|
||||
|
||||
@@ -618,7 +618,7 @@ class MixedInputGemmKernel:
|
||||
]
|
||||
acc_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage]
|
||||
acc_empty_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage]
|
||||
tmem_dealloc_mbar_ptr: cutlass.Int64
|
||||
tmem_dealloc_mbar: cutlass.Int64
|
||||
tmem_holding_buf: cutlass.Int32
|
||||
# Tensor buffers
|
||||
# (EPI_TILE_M, EPI_TILE_N, STAGE)
|
||||
@@ -820,11 +820,11 @@ class MixedInputGemmKernel:
|
||||
|
||||
# Tensor memory dealloc barrier init
|
||||
tmem = utils.TmemAllocator(
|
||||
storage.tmem_holding_buf,
|
||||
storage.tmem_holding_buf.ptr,
|
||||
barrier_for_retrieve=self.tmem_ptr_sync_barrier,
|
||||
allocator_warp_id=self.epilog_warp_id[0],
|
||||
is_two_cta=use_2cta_instrs,
|
||||
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr,
|
||||
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr,
|
||||
)
|
||||
|
||||
# Cluster arrive after barrier init
|
||||
|
||||
Reference in New Issue
Block a user