review comments

Signed-off-by: Lucas Wilkinson <lwilkinson@neuralmagic.com>
NVIDIA · Feb 26, 2025 · 2a9256f · 2a9256f
1 parent 60ddaa3
commit 2a9256f
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/...warp_specialized_gemm_with_blockwise_scaling/reference/host/gemm_with_groupwise_scaling.h b/...warp_specialized_gemm_with_blockwise_scaling/reference/host/gemm_with_groupwise_scaling.h
@@ -221,8 +221,10 @@ void gett_mainloop(
   const int N = cute::size<0>(mainloop_params.B.layout());
   const int ScaleGranularityM = M / cute::size<0>(mainloop_params.ScaleA);
   const int ScaleGranularityN = N / cute::size<0>(mainloop_params.ScaleB);
-  assert(M % ScaleGranularityM == 0 && "ScaleGranularityM must divide M");
-  assert(N % ScaleGranularityN == 0 && "ScaleGranularityN must divide N"); 
+  assert(ScaleGranularityM && M % ScaleGranularityM == 0 
+    && "ScaleGranularityM must divide M");
+  assert(ScaleGranularityN && N % ScaleGranularityN == 0 
+    && "ScaleGranularityN must divide N"); 
 
   cute::Tensor blockscale_A = domain_offset(
     make_coord(m / ScaleGranularityM, _0{}), mainloop_params.ScaleA(_, _, l));

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
@@ -441,7 +441,7 @@ struct CollectiveMma<
     }
 
     #pragma unroll
-    for (int i = 0; i < size(tApA_ScaleA); ++i) {
+    for (int i = 0; i < size(tBpB_ScaleB); ++i) {
       tBpB_ScaleB(i) = get<0>(tBcB_ScaleB(i)) < 
         std::min(scales_n, (n_coord + 1) * ScaleNsPerTile);
     }