From 1da6726a86a24f8c29abf49fc781bf450c47f409 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Fri, 15 May 2026 10:14:29 +0000 Subject: [PATCH] fix: assert float8_e4m3fn dtype in _prepack_weight_sf Non-prepacked path handles uint32 unpacking, but prepack sends weight_sf directly to _C.prepack_sfb(). If scales aren't float8, this would silently produce wrong layout/values. --- src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py b/src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py index 63289236..95cb1dee 100644 --- a/src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py +++ b/src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py @@ -109,6 +109,8 @@ def _prepack_weight_sf(weight_sf, N, K, tag): if cache_key in _prepack_weight_sf._cache: return _prepack_weight_sf._cache[cache_key] + assert weight_sf.dtype == torch.float8_e4m3fn, weight_sf.dtype + from nvfp4_megamoe_kernel.cutlass_nvfp4_gemm.kernel import prepack_sfb E = weight_sf.shape[0]