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.
This commit is contained in:
2026-05-15 10:14:29 +00:00
parent 489c620159
commit 1da6726a86

View File

@@ -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]