leavelet/cutlass
1
0
Fork 0
Code Issues Pull Requests Actions 27 Packages Projects Releases Wiki Activity

v4.3.4 update. (#2892)

Browse Source
This commit is contained in:
Junkai-Wu
2025-12-21 11:49:12 -05:00
committed by GitHub
parent 331e2f451c
commit 7f5fe3edf1
31 changed files with 839 additions and 240 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/python/CuTeDSL/blackwell/mixed_input_fmha/mixed_input_fmha_decode.py
+19 -16
View File
@@ -401,11 +401,7 @@ class MixedInputFusedMultiHeadAttentionDecode:
d_per_blk = 128
d_blks = cute.ceil_div(d, d_per_blk)
reduction(
o, m, l,
o_partial, m_partial, l_partial,
scale_o
).launch(
self.reduction(o, m, l, o_partial, m_partial, l_partial, scale_o).launch(
grid=[d_blks, h_q, b],
block=[d_per_blk, 1, 1],
cluster=[1, 1, 1],
@@ -1173,7 +1169,7 @@ class MixedInputFusedMultiHeadAttentionDecode:
# Reduce colmax in smem
if lane_store_max:
smem_fmax(tSsM.iterator + tSsM.layout(lane_idx), tSrM_lane)
self.smem_fmax(tSsM.iterator + tSsM.layout(lane_idx), tSrM_lane)
# Wait for colmax then load
cute.arch.barrier(barrier_id=softmax_nbar_id, number_of_threads=warpgroup_threads)
@@ -1259,7 +1255,7 @@ class MixedInputFusedMultiHeadAttentionDecode:
# Reduce cluster colmax
if warpgroup_widx == 0:
if lane_store_max:
dsmem_fmax(
self.dsmem_fmax(
sM_cluster.iterator + sM_layout((0, lane_idx)),
sM[(0, lane_idx)],
m_cluster_full_ptr
@@ -1280,7 +1276,7 @@ class MixedInputFusedMultiHeadAttentionDecode:
else:
# other splits copy cluster colmax into local smem
if lane_store_max:
sM_cluster[0, lane_idx] = dsmem_load(
sM_cluster[0, lane_idx] = self.dsmem_load(
sM_cluster.iterator + sM_layout((0, lane_idx))
)
@@ -1300,8 +1296,10 @@ class MixedInputFusedMultiHeadAttentionDecode:
sM_lane = sM_lane * scale_qs
# Store colsum and colmax
gmem_fadd(gL_partial.iterator + gL_partial.layout(lane_idx), sL_lane)
gmem_fmax(gM.iterator + gM.layout(lane_idx), sM_lane)
self.gmem_fadd(
gL_partial.iterator + gL_partial.layout(lane_idx), sL_lane
)
self.gmem_fmax(gM.iterator + gM.layout(lane_idx), sM_lane)
if kv_split_in_cluster == 0:
gM_partial[lane_idx] = sM_lane
@@ -1350,7 +1348,7 @@ class MixedInputFusedMultiHeadAttentionDecode:
# Store colsum and colmax
gL_partial[lane_idx] = sL_lane
gM_partial[lane_idx] = sM_lane
gmem_fmax(gM.iterator + gM.layout(lane_idx), sM_lane)
self.gmem_fmax(gM.iterator + gM.layout(lane_idx), sM_lane)
o_handle = o_consumer.wait_and_advance()
cute.copy(thr_load_s, tStO, tSrO)
@@ -1378,6 +1376,7 @@ class MixedInputFusedMultiHeadAttentionDecode:
return
@staticmethod
@cute.kernel
def reduction(
o : cute.Tensor,
@@ -1414,6 +1413,7 @@ class MixedInputFusedMultiHeadAttentionDecode:
return
@staticmethod
@cute.jit
def _mapa(ptr : Pointer, cta_rank_in_cluster : Int32 = 0):
llvm_ptr = ptr.llvm_ptr
@@ -1424,8 +1424,8 @@ class MixedInputFusedMultiHeadAttentionDecode:
)
@cute.jit
def dsmem_load(val_ptr : Pointer):
val_llvm_ptr = _mapa(val_ptr, 0)
def dsmem_load(self, val_ptr: Pointer):
val_llvm_ptr = self._mapa(val_ptr, 0)
ret = llvm.inline_asm(
Float32.mlir_type,
@@ -1439,6 +1439,7 @@ class MixedInputFusedMultiHeadAttentionDecode:
return Float32(ret)
@staticmethod
@cute.jit
def warp_fmax(val : Float32):
ret = llvm.inline_asm(
@@ -1472,10 +1473,10 @@ class MixedInputFusedMultiHeadAttentionDecode:
)
@cute.jit
def dsmem_fmax(val_ptr : Pointer, val : Float32, mbar_ptr : Pointer):
def dsmem_fmax(self, val_ptr: Pointer, val: Float32, mbar_ptr: Pointer):
expect_tx_bytes = Int32(Float32.width // 8)
val_llvm_ptr = _mapa(val_ptr, 0)
mbar_llvm_ptr = _mapa(mbar_ptr, 0)
val_llvm_ptr = self._mapa(val_ptr, 0)
mbar_llvm_ptr = self._mapa(mbar_ptr, 0)
nvvm.mbarrier_txn(
mbar_llvm_ptr,
@@ -1499,6 +1500,7 @@ class MixedInputFusedMultiHeadAttentionDecode:
asm_dialect=llvm.AsmDialect.AD_ATT,
)
@staticmethod
@cute.jit
def gmem_fmax(ptr : Pointer, val : Float32):
llvm.inline_asm(
@@ -1529,6 +1531,7 @@ class MixedInputFusedMultiHeadAttentionDecode:
asm_dialect=llvm.AsmDialect.AD_ATT,
)
@staticmethod
@cute.jit
def gmem_fadd(ptr : Pointer, val : Float32):
llvm.inline_asm(