File size: 5,077 Bytes
2a247be | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 | --- a/python/sglang/jit_kernel/flash_attn/cute/flash_fwd_sm100.py
+++ b/python/sglang/jit_kernel/flash_attn/cute/flash_fwd_sm100.py
@@ -448,6 +448,12 @@
self.load_warp_ids = (14, 15)
self.empty_warp_ids = ()
+ self.tmem_alloc_warp_ids = (
+ self.mma_warp_id,
+ *self.softmax0_warp_ids,
+ *self.softmax1_warp_ids,
+ *self.correction_warp_ids,
+ )
if self.use_correction_warps_for_epi:
self.empty_warp_ids = self.empty_warp_ids + self.epilogue_warp_ids
self.epilogue_warp_ids = self.correction_warp_ids
@@ -1732,15 +1738,7 @@
tmem_alloc_barrier = pipeline.NamedBarrier(
barrier_id=int(NamedBarrierFwdSm100.TmemPtr),
- num_threads=cute.arch.WARP_SIZE
- * len(
- (
- self.mma_warp_id,
- *self.softmax0_warp_ids,
- *self.softmax1_warp_ids,
- *self.correction_warp_ids,
- )
- ),
+ num_threads=cute.arch.WARP_SIZE * len(self.tmem_alloc_warp_ids),
)
# Tensor memory dealloc barrier init
tmem = cutlass.utils.TmemAllocator(
@@ -2207,11 +2205,11 @@
# ///////////////////////////////////////////////////////////////////////////////
# MMA
# ///////////////////////////////////////////////////////////////////////////////
+ tmem.allocate(cute.arch.get_max_tmem_alloc_cols("sm_100"))
+ if warp_idx in self.tmem_alloc_warp_ids:
+ tmem.wait_for_alloc()
if warp_idx == self.mma_warp_id:
cute.arch.setmaxregister_decrease(self.num_regs_other)
- # Alloc tensor memory buffer
- tmem.allocate(cute.arch.get_max_tmem_alloc_cols("sm_100"))
- tmem.wait_for_alloc()
tmem_ptr = tmem.retrieve_ptr(self.qk_acc_dtype)
self.mma(
tiled_mma_qk,
@@ -2278,8 +2276,6 @@
) or (const_expr(self.q_stage == 1) and warp_idx <= self.softmax0_warp_ids[-1]):
# increase register after decreasing
cute.arch.setmaxregister_increase(self.num_regs_softmax)
- # sync with mma warp before retrieving tmem ptr
- tmem.wait_for_alloc()
tmem_ptr = tmem.retrieve_ptr(self.qk_acc_dtype)
softmax_loop = partial(
self.softmax_loop,
@@ -2335,8 +2331,6 @@
# ///////////////////////////////////////////////////////////////////////////////
if warp_idx >= self.correction_warp_ids[0] and warp_idx < self.mma_warp_id:
cute.arch.setmaxregister_decrease(self.num_regs_correction)
- # sync with mma warp before retrieving tmem ptr
- tmem.wait_for_alloc()
tmem_ptr = tmem.retrieve_ptr(self.qk_acc_dtype)
self.correction_loop(
thr_mma_qk,
@@ -4740,6 +4734,7 @@
else:
row_max = None
pipeline_sm_stats.consumer_release_w_index(stage)
+ sink_rescale = Float32(1.0)
if const_expr(learnable_sink is not None):
LOG2_E = math.log2(math.e)
sink_val = learnable_sink_val[stage]
@@ -4749,18 +4744,22 @@
row_max = sink_val * (LOG2_E / softmax_scale_log2_eff)
row_sum = max_offset_scale
else:
- row_sum += cute.math.exp2(
- sink_val * LOG2_E
- - row_max * softmax_scale_log2_eff
- + max_offset,
+ old_max_log2 = row_max * softmax_scale_log2_eff
+ new_max_log2 = max(old_max_log2, sink_val * LOG2_E)
+ sink_rescale = cute.math.exp2(
+ old_max_log2 - new_max_log2, fastmath=True
+ )
+ row_sum = row_sum * sink_rescale + cute.math.exp2(
+ sink_val * LOG2_E - new_max_log2 + max_offset,
fastmath=True,
)
+ row_max = new_max_log2 / softmax_scale_log2_eff
acc_O_mn_row_is_zero_or_nan = row_sum == 0.0 or row_sum != row_sum
stats[stage] = (row_sum, row_max, acc_O_mn_row_is_zero_or_nan)
scale = cute.arch.rcp_approx(
row_sum if not acc_O_mn_row_is_zero_or_nan else 1.0
)
- scale = scale * v_descale
+ scale = scale * sink_rescale * v_descale
# Wait for the last O to be ready from the MMA warp
pipeline_o_acc.consumer_wait_w_index_phase(
stage, o_corr_consumer_phase
|