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

fix: L1 output uses unpacked E2M1 (1 byte/element) like FP8

Browse Source 
- float_e2m1_unpacksmem_t: sizeof=1, SMEM is 1 byte/element (not packed)
- TMA load unpacks 2 E2M1/global-byte → 2 SMEM bytes
- UMMA reads unpacked SMEM, packs internally for mxf4nvf4
- L1→L2 handoff: unpacked format (same byte count as FP8)
- Epilogue: 4 E2M1 bytes per uint32 STSM atom, same as FP8
- Dispatch TMA: kHidden bytes (unpacked), not kHidden/2
- Added static_assert on sizeof(a_dtype_t) and sizeof(b_dtype_t)
- Note: no bandwidth savings at L1→L2 boundary for v1
This commit is contained in:
biondizzle
2026-05-11 21:27:35 +00:00
parent 091b974736
commit 0ac73a82f9
3 changed files with 52 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
csrc/apis/mega_nvfp4.hpp
View File

@@ -28,9 +28,9 @@ get_symm_buffer_size_for_nvfp4_mega_moe(
// NVFP4 layouts: E2M1 packed (2 per byte), so token layout is K/2 bytes
// group_size=16, so SF stride is K/16 (twice as many as MXFP4)
const auto fp4_token_layout = layout::Data(hidden / 2);
const auto fp4_token_layout = layout::Data(hidden);
const auto bf16_token_layout = layout::Data(hidden * 2);
const auto fp4_intermediate_token_layout = layout::Data(intermediate_hidden / 2);
const auto fp4_intermediate_token_layout = layout::Data(intermediate_hidden);
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);
@@ -95,7 +95,7 @@ get_symm_buffer_size_for_nvfp4_mega_moe(
// NVFP4: E2M1 packed activations (2 per byte), K/2 bytes per token
auto x = torch::from_blob(
math::advance_ptr(buffer.data_ptr(), reinterpret_cast<int64_t>(input_token_buffer.base)),
{num_max_tokens_per_rank, hidden / 2},
{num_max_tokens_per_rank, hidden},
torch::TensorOptions().dtype(torch::kUInt8).device(buffer.device()));
// NVFP4 SF: K/16 bytes per token, packed as K/64 int32
auto x_sf = torch::from_blob(
@@ -113,7 +113,7 @@ get_symm_buffer_size_for_nvfp4_mega_moe(
// NVFP4: L1 output acts are E2M1 packed, intermediate_hidden/2 bytes
auto l1_acts = torch::from_blob(
math::advance_ptr(buffer.data_ptr(), reinterpret_cast<int64_t>(l1_token_buffer.base)),
{num_max_pool_tokens, hidden / 2},
{num_max_pool_tokens, hidden},
torch::TensorOptions().dtype(torch::kUInt8).device(buffer.device()));
// NVFP4 L1 SF: M-major, K/64 int32
auto l1_acts_sf = torch::from_blob(
@@ -124,7 +124,7 @@ get_symm_buffer_size_for_nvfp4_mega_moe(
// NVFP4: L2 acts are E2M1 packed, intermediate_hidden/2 bytes
auto l2_acts = torch::from_blob(
math::advance_ptr(buffer.data_ptr(), reinterpret_cast<int64_t>(l2_token_buffer.base)),
{num_max_pool_tokens, intermediate_hidden / 2},
{num_max_pool_tokens, intermediate_hidden},
torch::TensorOptions().dtype(torch::kUInt8).device(buffer.device()));
// NVFP4 L2 SF: M-major, K/64 int32
auto l2_acts_sf = torch::from_blob(