v4.3 tag release update. (#2789)
This commit is contained in:
@@ -40,6 +40,7 @@ from cutlass.cute.nvgpu import cpasync, tcgen05
|
||||
import cutlass.torch as cutlass_torch
|
||||
import cutlass.utils as utils
|
||||
import cutlass.pipeline as pipeline
|
||||
from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
|
||||
import cutlass.utils.blackwell_helpers as sm100_utils
|
||||
import cutlass.utils.blockscaled_layout as blockscaled_utils
|
||||
from cutlass.cute.runtime import from_dlpack
|
||||
@@ -177,22 +178,18 @@ class Sm100GroupedBlockScaledGemmKernel:
|
||||
self.threads_per_cta = 32 * len(
|
||||
(self.mma_warp_id, self.tma_warp_id, *self.epilog_warp_id)
|
||||
)
|
||||
# Set barrier for cta sync, epilogue sync and tmem ptr sync
|
||||
self.cta_sync_barrier = pipeline.NamedBarrier(
|
||||
barrier_id=1,
|
||||
num_threads=self.threads_per_cta,
|
||||
)
|
||||
# Set barrier for epilogue sync and tmem ptr sync
|
||||
self.epilog_sync_barrier = pipeline.NamedBarrier(
|
||||
barrier_id=2,
|
||||
barrier_id=1,
|
||||
num_threads=32 * len(self.epilog_warp_id),
|
||||
)
|
||||
self.tmem_alloc_barrier = pipeline.NamedBarrier(
|
||||
barrier_id=3,
|
||||
barrier_id=2,
|
||||
num_threads=32 * len((self.mma_warp_id, *self.epilog_warp_id)),
|
||||
)
|
||||
# Barrier used by MMA/TMA warps to signal A/B tensormap initialization completion
|
||||
self.tensormap_ab_init_barrier = pipeline.NamedBarrier(
|
||||
barrier_id=4,
|
||||
barrier_id=3,
|
||||
num_threads=64,
|
||||
)
|
||||
self.smem_capacity = utils.get_smem_capacity_in_bytes("sm_100")
|
||||
@@ -646,6 +643,7 @@ class Sm100GroupedBlockScaledGemmKernel:
|
||||
cluster=(*self.cluster_shape_mn, 1),
|
||||
smem=self.shared_storage.size_in_bytes(),
|
||||
stream=stream,
|
||||
min_blocks_per_mp=1,
|
||||
)
|
||||
return
|
||||
|
||||
@@ -781,11 +779,9 @@ class Sm100GroupedBlockScaledGemmKernel:
|
||||
cute.arch.mbarrier_init(
|
||||
tmem_dealloc_mbar_ptr, num_tmem_dealloc_threads
|
||||
)
|
||||
cute.arch.mbarrier_init_fence()
|
||||
|
||||
# Cluster arrive after barrier init
|
||||
if cute.size(self.cluster_shape_mn) > 1:
|
||||
cute.arch.cluster_arrive_relaxed()
|
||||
pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True)
|
||||
|
||||
#
|
||||
# Setup smem tensor A/B/SFA/SFB/C
|
||||
@@ -944,10 +940,7 @@ class Sm100GroupedBlockScaledGemmKernel:
|
||||
#
|
||||
# Cluster wait before tensor memory alloc
|
||||
#
|
||||
if cute.size(self.cluster_shape_mn) > 1:
|
||||
cute.arch.cluster_wait()
|
||||
else:
|
||||
self.cta_sync_barrier.arrive_and_wait()
|
||||
pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn)
|
||||
|
||||
#
|
||||
# Get tensormap buffer address
|
||||
@@ -2894,6 +2887,7 @@ def run(
|
||||
tensor_of_tensormap,
|
||||
max_active_clusters,
|
||||
current_stream,
|
||||
options=f"--opt-level 2",
|
||||
)
|
||||
|
||||
# reference check
|
||||
|
||||
Reference in New Issue
Block a user