fix < 128

Signed-off-by: Lucas Wilkinson <lwilkinson@neuralmagic.com>
NVIDIA · Feb 25, 2025 · 7f541db · 7f541db
1 parent b42a7b4
commit 7f541db
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/...de/cutlass/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized_fp8_blockwise_scaling.hpp b/...de/cutlass/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized_fp8_blockwise_scaling.hpp
@@ -457,12 +457,14 @@ struct CollectiveMma<
 
     #pragma unroll
     for (int i = 0; i < size(tApA_ScaleA); ++i) {
-      tApA_ScaleA(i) = get<0>(tAcA_ScaleA(i)) < scales_m;
+      tApA_ScaleA(i) = get<0>(tAcA_ScaleA(i)) < 
+        std::min(scales_m, (m_coord + 1) * ScaleMsPerTile);
     }
 
     #pragma unroll
     for (int i = 0; i < size(tApA_ScaleA); ++i) {
-      tBpB_ScaleB(i) = get<0>(tBcB_ScaleB(i)) < scales_n;
+      tBpB_ScaleB(i) = get<0>(tBcB_ScaleB(i)) < 
+        std::min(scales_n, (n_coord + 1) * ScaleNsPerTile);
     }
 
     // Mainloop