biondizzle/DeepGEMM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: gran_k=16 in transform_sf + sm_100a arch for NVFP4 mega_moe

Browse Source 
- transform_sf_into_required_layout: add gran_k=16 branch for NVFP4 UE4M3
  scales (4 per int32, group_size=16). Previously only handled 32/128.
- get_arch: always return '100a' for SM100, never '100f'. The family
  variant lacks mxf4nvf4 (NVFP4 block-scaled MMA) support, causing
  'scale_vec::4X not supported on sm_100f' errors.
- transform_nvfp4_weights_for_mega_moe: fold weight_scale_2 into block
  scales, pack UE4M3→int32, transpose MN-major, call
  transform_sf_into_required_layout with gran_k=16.
This commit is contained in:
biondizzle
2026-05-11 16:11:11 +00:00
parent fbdddaccf4
commit 86a1263f44
3 changed files with 74 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
csrc/apis/layout.hpp
View File

@@ -53,7 +53,10 @@ static torch::Tensor transform_sf_into_required_layout(const torch::Tensor& sf,
}
// (INT, 1, gran_k) on SM100: transform to TMA-aligned and MN-major
if (sf.scalar_type() == torch::kInt and gran_mn == 1 and (gran_k == 32 or gran_k == 128) and arch_major == 10)
// gran_k=16: NVFP4 UE4M3 (4 per int32, group_size=16)
// gran_k=32: MXFP4 UE8M0 (4 per int32, group_size=32)
// gran_k=128: FP32 block scale
if (sf.scalar_type() == torch::kInt and gran_mn == 1 and (gran_k == 16 or gran_k == 32 or gran_k == 128) and arch_major == 10)
return check_sf_layout(sf, mn, k, gran_mn, gran_k, num_groups, true, false, torch::kInt);
DG_HOST_UNREACHABLE("Unknown SF transformation");

9
csrc/jit/device_runtime.hpp
View File

@@ -88,10 +88,15 @@ public:
std::string get_arch(const bool& number_only = false,
const bool& support_arch_family = false) {
const auto [major, minor] = get_arch_pair();
if (major == 10 and minor != 1) {
if (major == 10) {
if (number_only)
return "100";
return support_arch_family ? "100f" : "100a";
// Always target 100a for SM100 — the 'f' family variant
// lacks mxf4nvf4 (NVFP4 block-scaled MMA) support, which
// causes 'scale_vec::4X not supported on sm_100f' errors.
// Since DeepGEMM is JIT-compiled for the exact GPU, there's
// no benefit to targeting the restricted family subset.
return "100a";
}
return std::to_string(major * 10 + minor) + (number_only ? "" : "a");
}