Update blackwell tutorial to be compatible with 4.5-dev version (#3130)
* Update blackwell tutorial to be compatible with 4.5-dev version * update example for reverted changes * add more example fix
This commit is contained in:
@@ -219,11 +219,11 @@ def kernel(
|
||||
* len((mma_warp_id, *epilogue_warp_ids)), # 5 warps = 160 threads
|
||||
)
|
||||
tmem = utils.TmemAllocator(
|
||||
storage.tmem_holding_buffer,
|
||||
storage.tmem_holding_buffer.ptr,
|
||||
barrier_for_retrieve=tmem_alloc_barrier,
|
||||
allocator_warp_id=epilogue_warp_ids[0],
|
||||
is_two_cta=True if use_2cta_instrs else False,
|
||||
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar,
|
||||
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr,
|
||||
)
|
||||
|
||||
# Partition tensors for TMA; This requires the tensors partitioned for MMA
|
||||
|
||||
@@ -152,11 +152,11 @@ def kernel(
|
||||
* len((mma_warp_id, *epilogue_warp_ids)), # 5 warps = 160 threads
|
||||
)
|
||||
tmem = utils.TmemAllocator(
|
||||
storage.tmem_holding_buffer,
|
||||
storage.tmem_holding_buffer.ptr,
|
||||
barrier_for_retrieve=tmem_alloc_barrier,
|
||||
allocator_warp_id=epilogue_warp_ids[0],
|
||||
is_two_cta=True,
|
||||
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar,
|
||||
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr,
|
||||
)
|
||||
|
||||
num_tma_copy_bytes = (
|
||||
|
||||
@@ -159,11 +159,11 @@ def kernel(
|
||||
* len((mma_warp_id, *epilogue_warp_ids)), # 5 warps = 160 threads
|
||||
)
|
||||
tmem = utils.TmemAllocator(
|
||||
storage.tmem_holding_buffer,
|
||||
storage.tmem_holding_buffer.ptr,
|
||||
barrier_for_retrieve=tmem_alloc_barrier,
|
||||
allocator_warp_id=epilogue_warp_ids[0],
|
||||
is_two_cta=True,
|
||||
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar,
|
||||
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr,
|
||||
)
|
||||
|
||||
num_tma_copy_bytes = (
|
||||
|
||||
@@ -184,11 +184,11 @@ def cluster_specific_kernel(
|
||||
* len((mma_warp_id, *epilogue_warp_ids)), # 5 warps = 160 threads
|
||||
)
|
||||
tmem = utils.TmemAllocator(
|
||||
storage.tmem_holding_buffer,
|
||||
storage.tmem_holding_buffer.ptr,
|
||||
barrier_for_retrieve=tmem_alloc_barrier,
|
||||
allocator_warp_id=epilogue_warp_ids[0],
|
||||
is_two_cta=True,
|
||||
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar,
|
||||
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr,
|
||||
)
|
||||
|
||||
num_tma_copy_bytes = (
|
||||
|
||||
@@ -171,11 +171,11 @@ def kernel(
|
||||
* len((mma_warp_id, *epilogue_warp_ids)), # 5 warps = 160 threads
|
||||
)
|
||||
tmem = utils.TmemAllocator(
|
||||
storage.tmem_holding_buffer,
|
||||
storage.tmem_holding_buffer.ptr,
|
||||
barrier_for_retrieve=tmem_alloc_barrier,
|
||||
allocator_warp_id=epilogue_warp_ids[0],
|
||||
is_two_cta=True,
|
||||
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar,
|
||||
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr,
|
||||
)
|
||||
|
||||
num_tma_copy_bytes = (
|
||||
|
||||
@@ -214,11 +214,11 @@ def gemm(
|
||||
* len((mma_warp_id, *epilogue_warp_ids)), # 5 warps = 160 threads
|
||||
)
|
||||
tmem = utils.TmemAllocator(
|
||||
storage.tmem_holding_buffer,
|
||||
storage.tmem_holding_buffer.ptr,
|
||||
barrier_for_retrieve=tmem_alloc_barrier,
|
||||
allocator_warp_id=epilogue_warp_ids[0],
|
||||
is_two_cta=True,
|
||||
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar,
|
||||
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr,
|
||||
)
|
||||
|
||||
num_tma_copy_bytes = (
|
||||
|
||||
Reference in New Issue
Block a user