fix bug in tl.store mask for kernel _to_fp8_row_major_t_and_non_t

ghstack-source-id: 270cf799161c7cca6822e7c0f9c511e618321dab ghstack-comment-id: 2575984684 Pull Request resolved: #1516
pytorch · Jan 8, 2025 · 878c886 · 878c886
1 parent 367aea3
commit 878c886
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/torchao/prototype/float8nocompile/kernels/fp8_dynamic_tensorwise.py b/torchao/prototype/float8nocompile/kernels/fp8_dynamic_tensorwise.py
@@ -375,8 +375,8 @@ def _to_fp8_row_major_t_and_non_t(
         block_col_offs[:, None] * row_major_t_out_stride_row
         + block_row_offs[None, :] * row_major_t_out_stride_col
     )
-    mask = (block_row_offs[:, None] < row_major_t_num_rows) & (
-        block_col_offs[None, :] < row_major_t_num_cols
+    mask = (block_col_offs[:, None] < row_major_t_num_rows) & (
+        block_row_offs[None, :] < row_major_t_num_cols
     )
     tl.store(row_major_t_out_ptr + row_major_t_offs, fp8_vals.trans(1, 0), mask=mask)