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

revert: restore mxf4nvf4/block16 code (correct path for sm_100a)

Browse Source 
Reverted to commit 36b439e's NVFP4 kernel code:
- kGranK=16, mxf4nvf4.block_scale.scale_vec::4X
- float_ue4m3_t instruction descriptor
- Block16 SF layout (4X TMEM)
- UE4M3 L1 epilogue
- No UE4M3→UE8M0 conversion, no block16→block32 merge

The mxf4nvf4.scale_vec::4X PTX instruction compiles successfully
on both sm_100 and sm_100f with CUDA 13.0. The previous build 17
error was likely from a different cause, not the arch flag.

Python: reverted transform_nvfp4_weights_for_mega_moe to use
pack_ue4m3_to_int32 with gran_k=16, no UE8M0 conversion.
This commit is contained in:
biondizzle
2026-05-11 15:02:47 +00:00
parent e80fe9af60
commit fbdddaccf4
5 changed files with 56 additions and 160 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
csrc/apis/layout.hpp
View File

@@ -53,7 +53,6 @@ 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
// Supports gran_k=32 (MXFP4 and NVFP4-block32), 128 (FP8)
if (sf.scalar_type() == torch::kInt and gran_mn == 1 and (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);

18
csrc/apis/mega_nvfp4.hpp
View File

@@ -30,8 +30,8 @@ get_symm_buffer_size_for_nvfp4_mega_moe(
const auto fp8_token_layout = layout::Data(hidden);
const auto bf16_token_layout = layout::Data(hidden * 2);
const auto fp8_intermediate_token_layout = layout::Data(intermediate_hidden);
const auto nvfp4_sf_layout = layout::Data(hidden / 32);
const auto nvfp4_intermediate_sf_layout = layout::Data(intermediate_hidden / 32);
const auto nvfp4_sf_layout = layout::Data(hidden / 16);
const auto nvfp4_intermediate_sf_layout = layout::Data(intermediate_hidden / 16);
const auto input_topk_idx_layout = layout::Data(num_topk * sizeof(int64_t), false);
const auto input_topk_weights_layout = layout::Data(num_topk * sizeof(float), false);
const auto l1_topk_weights_layout = layout::Data(sizeof(float), false);
@@ -86,7 +86,7 @@ get_symm_buffer_size_for_nvfp4_mega_moe(
// Check SF buffer requirements
// NVFP4: hidden must be divisible by 64 (4 UE4M3 scales per int32, group_size=16)
DG_HOST_ASSERT(hidden % 128 == 0 and intermediate_hidden % 128 == 0);
DG_HOST_ASSERT(hidden % 64 == 0 and intermediate_hidden % 64 == 0);
DG_HOST_ASSERT(num_max_padded_sf_pool_tokens % 4 == 0);
// Slice function
@@ -98,7 +98,7 @@ get_symm_buffer_size_for_nvfp4_mega_moe(
// NVFP4 SF: K/16 bytes per token, packed as K/64 int32
auto x_sf = torch::from_blob(
math::advance_ptr(buffer.data_ptr(), reinterpret_cast<int64_t>(input_sf_buffer.base)),
{num_max_tokens_per_rank, hidden / 128},
{num_max_tokens_per_rank, hidden / 64},
torch::TensorOptions().dtype(torch::kInt).device(buffer.device()));
auto topk_idx = torch::from_blob(
math::advance_ptr(buffer.data_ptr(), reinterpret_cast<int64_t>(input_topk_idx_buffer.base)),
@@ -115,7 +115,7 @@ get_symm_buffer_size_for_nvfp4_mega_moe(
// NVFP4 L1 SF: M-major, K/64 int32
auto l1_acts_sf = torch::from_blob(
math::advance_ptr(buffer.data_ptr(), reinterpret_cast<int64_t>(l1_sf_buffer.base)),
{num_max_padded_sf_pool_tokens, hidden / 128},
{num_max_padded_sf_pool_tokens, hidden / 64},
{1, num_max_padded_sf_pool_tokens},
torch::TensorOptions().dtype(torch::kInt).device(buffer.device()));
auto l2_acts = torch::from_blob(
@@ -125,7 +125,7 @@ get_symm_buffer_size_for_nvfp4_mega_moe(
// NVFP4 L2 SF: M-major, K/64 int32
auto l2_acts_sf = torch::from_blob(
math::advance_ptr(buffer.data_ptr(), reinterpret_cast<int64_t>(l2_sf_buffer.base)),
{num_max_padded_sf_pool_tokens, intermediate_hidden / 128},
{num_max_padded_sf_pool_tokens, intermediate_hidden / 64},
{1, num_max_padded_sf_pool_tokens},
torch::TensorOptions().dtype(torch::kInt).device(buffer.device()));
return std::make_tuple(x, x_sf, topk_idx, topk_weights, l1_acts, l1_acts_sf, l2_acts, l2_acts_sf);
@@ -153,7 +153,7 @@ static void fp8_nvfp4_mega_moe(
// Config checks
const auto num_tokens = static_cast<int>(y.size(0));
const auto [rm, rn, rk] = recipe;
DG_HOST_ASSERT(rm == 1 and rn == 1 and rk == 32); // NVFP4 block32: group_size=32
DG_HOST_ASSERT(rm == 1 and rn == 1 and rk == 16); // NVFP4: group_size=16
DG_HOST_ASSERT(activation == "swiglu");
// Activation checks
@@ -175,8 +175,8 @@ static void fp8_nvfp4_mega_moe(
DG_HOST_ASSERT(l1_weights.is_contiguous() and l2_weights.is_contiguous());
// Check weight SF layout for UE4M3 packing, MN-major, and TMA alignment
// NVFP4 block32: kGranK=32, SF packed as int32 (4 UE4M3 bytes per int32)
constexpr int kGranMN = 1, kGranK = 32;
// NVFP4: kGranK=16, SF packed as int32 (4 UE4M3 bytes per int32)
constexpr int kGranMN = 1, kGranK = 16;
check_sf_layout(l1_weights_sf, intermediate_hidden * 2, hidden, kGranMN, kGranK,
num_experts_per_rank, true, false, torch::kInt);
check_sf_layout(l2_weights_sf, hidden, intermediate_hidden, kGranMN, kGranK,