| |
| |
| @@ -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 |
|
|
