diff --git a/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/cutlass_nvfp4_gemm.cu b/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/cutlass_nvfp4_gemm.cu
index 5917b1f7..66024447 100644
--- a/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/cutlass_nvfp4_gemm.cu
+++ b/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/cutlass_nvfp4_gemm.cu
@@ -168,8 +168,8 @@ int cutlass_nvfp4_gemm_run(
     using ArrayElementB = typename Gemm::GemmKernel::CollectiveMainloop::ArrayElementB;
     using ElementSF = typename Gemm::GemmKernel::CollectiveMainloop::ElementSF;
 
-    int sfa_size = cute::cosize(layout_SFA);
-    int sfb_size = cute::cosize(layout_SFB);
+    int sfa_size = cute::size(cute::filter_zeros(layout_SFA));
+    int sfb_size = cute::size(cute::filter_zeros(layout_SFB));
     int K_sf = K / InputSFVectorSize;
 
     cutlass::device_memory::allocation<ElementSF> sfa_cutlass(sfa_size);
@@ -225,7 +225,7 @@ extern "C" int cutlass_nvfp4_sfb_size(
 ) {
     using Sm1xxBlkScaledConfig = typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig;
     LayoutSFB layout_SFB = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(cute::make_shape(M, N, K, 1));
-    *out_size = cute::cosize(layout_SFB);
+    *out_size = cute::size(cute::filter_zeros(layout_SFB));
     return 0;
 }
 
@@ -239,7 +239,7 @@ extern "C" int cutlass_nvfp4_prepack_sfb_run(
     LayoutSFB layout_SFB = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(cute::make_shape(M, N, K, 1));
     using ElementSF = typename Gemm::GemmKernel::CollectiveMainloop::ElementSF;
 
-    int sfb_size = cute::cosize(layout_SFB);
+    int sfb_size = cute::size(cute::filter_zeros(layout_SFB));
     int K_sf = K / InputSFVectorSize;
 
     cudaMemsetAsync(static_cast<ElementSF*>(SFB_cutlass_ptr), 0, sfb_size * sizeof(ElementSF), stream);
@@ -281,7 +281,7 @@ extern "C" int cutlass_nvfp4_gemm_run_prepacked_sfb(
     using ArrayElementB = typename Gemm::GemmKernel::CollectiveMainloop::ArrayElementB;
     using ElementSF = typename Gemm::GemmKernel::CollectiveMainloop::ElementSF;
 
-    int sfa_size = cute::cosize(layout_SFA);
+    int sfa_size = cute::size(cute::filter_zeros(layout_SFA));
     int K_sf = K / InputSFVectorSize;
 
     // Only remap SFA (activation scales) — SFB is prepacked