Prefetch block scales before TMA waits

Author: drisspg

transformer_nuggets/cute/blockscaled_tma.py

@@ -339,13 +339,29 @@ def compute_warp_consume_stage(
         enable_profiling: cutlass.Constexpr,
     ) -> tuple[pipeline.PipelineConsumer, list]:
         """Wait for one stage and accumulate the rows owned by this compute warp."""
+        scale_k = lane * self.scale_blocks_per_lane + k_tile * self.scale_blocks_per_tile
         with self.profile_scope(
             prof_buf,
             ProfileTag.TMA_WAIT,
             pid_n,
             2 + 3 * k_tile,
             enable_profiling,
         ):
+            input_scales = self.load_scale_values(
+                mSFX,
+                0,
+                scale_k,
+                input_scale_atom,
+                scale_layout,
+            )
+            prepared_weight_scales = self.prepare_weight_scale_values(
+                mSFW,
+                n0 + owned_row_start,
+                k_tile,
+                lane,
+                weight_scale_atom,
+                scale_layout,
+            )
             full = consumer.wait_and_advance()
 
         with self.profile_scope(
@@ -355,7 +371,6 @@ def compute_warp_consume_stage(
             3 + 3 * k_tile,
             enable_profiling,
         ):
-            scale_k = lane * self.scale_blocks_per_lane + k_tile * self.scale_blocks_per_tile
             x_values = self.load_lane_values(
                 sX,
                 0,
@@ -364,21 +379,6 @@ def compute_warp_consume_stage(
                 smem_atom,
                 chunk_layout,
             )
-            input_scales = self.load_scale_values(
-                mSFX,
-                0,
-                scale_k,
-                input_scale_atom,
-                scale_layout,
-            )
-            prepared_weight_scales = self.prepare_weight_scale_values(
-                mSFW,
-                n0 + owned_row_start,
-                k_tile,
-                lane,
-                weight_scale_atom,
-                scale_layout,
-            )
             for local_row in cutlass.range_constexpr(self.rows_per_warp):
                 cta_row = owned_row_start + local_row
                 global_row = n0 + cta_row