Refactor attention sink examples to simplify index calculations

- Updated index handling in `example_gqa_sink_bwd_bhsd.py` and `example_mha_sink_bwd_bhsd.py` to eliminate unnecessary local allocations and streamline logic for determining start and end indices.
- Improved readability by using direct calculations instead of local variables for index bounds in pipelined loops.

Author: LeiWang1999

examples/attention_sink/example_gqa_sink_bwd_bhsd.py

@@ -81,13 +81,10 @@ def flash_fwd(
                 sinks[i] = Sinks[by]
 
             end = T.min(T.ceildiv(seq_len, block_N), T.ceildiv((bx + 1) * block_M, block_N))
-            start = T.alloc_local([1], 'int32')
-            if window_size is not None:
-                start[0] = T.max(0, (bx * block_M - window_size) // block_N)
-            else:
-                start[0] = 0
+            start = T.max(0,
+                          (bx * block_M - window_size) // block_N) if window_size is not None else 0
 
-            for k in T.Pipelined(start[0], end, num_stages=num_stages):
+            for k in T.Pipelined(start, end, num_stages=num_stages):
                 T.copy(K[bz, by // groups, k * block_N:(k + 1) * block_N, :], K_shared)
                 for i, j in T.Parallel(block_M, block_N):
                     q_idx = bx * block_M + i

examples/attention_sink/example_mha_sink_bwd_bhsd.py

@@ -267,14 +267,10 @@ def flash_bwd(
             T.clear(dk)
 
             loop_st = T.floordiv(by * block_M, block_N)
-            loop_ed = T.alloc_local([1], 'int32')
-            if window_size is not None:
-                loop_ed[0] = T.min(
-                    T.ceildiv((by + 1) * block_M + window_size, block_N),
-                    T.ceildiv(seq_len, block_N))
-            else:
-                loop_ed[0] = T.ceildiv(seq_len, block_N)
-            for k in T.Pipelined(loop_st, loop_ed[0], num_stages=num_stages):
+            loop_ed = T.min(
+                T.ceildiv((by + 1) * block_M + window_size, block_N), T.ceildiv(
+                    seq_len, block_N)) if window_size is not None else T.ceildiv(seq_len, block_N)
+            for k in T.Pipelined(loop_st, loop_ed, num_stages=num_stages):
                 T.copy(Q[bz, bx, k * block_N:(k + 1) * block_N, :], q)
                 T.clear(qkT)
                 T.gemm(K_shared, q, qkT, transpose_B=True, policy=T.GemmWarpPolicy.FullRow)

src/transform/legalize_negative_index.cc

@@ -51,8 +51,8 @@ class NegativeIndexAnalyzer : public IRVisitorWithAnalyzer {
       states.push_back(IndexSignState::kUnknown);
       needs_record = true;
       DLOG(WARNING) << "LegalizeNegativeIndex: cannot prove non-negative index "
-                   << simplified << " for buffer " << load->buffer->name
-                   << " (axis " << i << ").";
+                    << simplified << " for buffer " << load->buffer->name
+                    << " (axis " << i << ").";
     }
 
     if (needs_record) {

tilelang/intrinsics/wgmma_macro_generator.py

@@ -273,6 +273,7 @@ def wgmma(self,
         def _warp_mma(A_ptr, B_ptr, C_buf):
             tx, warp_n, warp_m = self.extract_thread_binding(thread_binding)
 
+            scale_out = T.alloc_var("int32")
             desc_a = T.alloc_wgmma_desc()
             desc_b = T.alloc_wgmma_desc()
             T.initialize_wgmma_descriptor(desc_a, A_ptr, a_swizzle_mode,
@@ -283,12 +284,16 @@ def _warp_mma(A_ptr, B_ptr, C_buf):
                                           int(b_stride_byte_offset >> 4))
             T.warpgroup_fence_operand(C_buf, num_regs=accum_regs)
             T.warpgroup_arrive()
-            for j in T.serial(num_inst_n):
-                for i in T.serial(num_inst_m):
-                    for ki in T.serial(k_dim // micro_size_k):
+
+            if clear_accum:
+                scale_out = 0
+            else:
+                scale_out = 1
+            for j in T.unroll(num_inst_n):
+                for i in T.unroll(num_inst_m):
+                    for ki in T.unroll(k_dim // micro_size_k):
                         warp_i = (warp_m // 4) * num_inst_m + i
                         warp_j = warp_n * num_inst_n + j
-                        scale_out = T.if_then_else(ki != 0, 1, T.if_then_else(clear_accum, 0, 1))
                         A_offset = (
                             ki % ak_atom_size
                         ) * micro_size_k + warp_i * 64 * a_swizzle_atom_elems + (
@@ -305,6 +310,9 @@ def _warp_mma(A_ptr, B_ptr, C_buf):
                                        (A_offset * elems_in_bytes) >> 4, desc_b.data,
                                        (B_offset * elems_in_bytes) >> 4, C_buf.data, C_offset,
                                        scale_out, scale_in_a, scale_in_b)
+                if clear_accum:
+                    scale_out = 1
+
             T.warpgroup_commit_batch()
             if wg_wait >= 0:
                 T.warpgroup_wait(wg_wait)
@@ -387,6 +395,7 @@ def wgmma_rs(self,
         def _warp_mma(A_buf, B_ptr, C_buf):
             tx, warp_n, warp_m = self.extract_thread_binding(thread_binding)
 
+            scale_out = T.alloc_var("int32")
             desc_b = T.alloc_wgmma_desc()
             T.initialize_wgmma_descriptor(desc_b, B_ptr, b_swizzle_mode,
                                           int(b_leading_byte_offset >> 4),
@@ -395,11 +404,16 @@ def _warp_mma(A_buf, B_ptr, C_buf):
             T.warpgroup_fence_operand(C_buf, num_regs=accum_regs)
             T.warpgroup_arrive()
 
-            for j in T.serial(0, num_inst_n):
-                for i in T.serial(num_inst_m):
-                    for ki in T.serial(0, (k_dim // micro_size_k)):
+            if clear_accum:
+                scale_out = 0
+            else:
+                scale_out = 1
+
+            for j in T.unroll(0, num_inst_n):
+                for i in T.unroll(num_inst_m):
+                    for ki in T.unroll(0, (k_dim // micro_size_k)):
                         warp_j = warp_n * num_inst_n + j
-                        scale_out = T.if_then_else(ki != 0, 1, T.if_then_else(clear_accum, 0, 1))
+
                         A_offset = ki * warp_rows * local_size_a + i * local_size_a
                         B_offset = (
                             ki // bk_atom_size
@@ -425,6 +439,9 @@ def _warp_mma(A_buf, B_ptr, C_buf):
                             scale_in_a,
                             scale_in_b,
                         )
+                        if clear_accum:
+                            scale_out = 1
+
             T.warpgroup_commit_batch()
             if wg_wait >= 0:
                 T.warpgroup_wait(wg_wait)