From 09d1307d78a0b000cf48d50e0820d09f5937e25b Mon Sep 17 00:00:00 2001 From: biondizzle Date: Thu, 14 May 2026 20:34:51 +0000 Subject: [PATCH] damn clankers2 --- src/nvfp4_megamoe_kernel/symm_buffer.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/src/nvfp4_megamoe_kernel/symm_buffer.py b/src/nvfp4_megamoe_kernel/symm_buffer.py index 7c46658f..4984ecd7 100644 --- a/src/nvfp4_megamoe_kernel/symm_buffer.py +++ b/src/nvfp4_megamoe_kernel/symm_buffer.py @@ -35,11 +35,16 @@ class SymmBuffer: device = torch.cuda.current_device() - # NVFP4: packed E2M1 (2 values per byte), so K//2 - sf_k_groups_hidden = hidden_size // 16 # UE4M3 block16, 1 scale per group + # NVFP4 packed E2M1: 2 FP4 values per byte → K//2 bytes per token. + # Scales are UE4M3 (float8_e4m3fn), one per 16-element group → K//16 + # bytes per token, UNPACKED. This is what `stage_activation` produces + # and what the CUTLASS NVFP4 block-scaled GEMM consumes directly. + # (The DeepGEMM API packed 4 UE4M3 into one uint32 — we don't, because + # our CUTLASS kernel reads scales as float8_e4m3fn.) + sf_k_groups_hidden = hidden_size // 16 sf_k_groups_inter = intermediate_size // 16 - - # Staging buffers (matching kernel's expected input format) + + # Staging buffers self.x = torch.empty( max_num_tokens, hidden_size // 2, dtype=torch.int8, device=device, @@ -84,4 +89,4 @@ def get_symm_buffer_for_nvfp4_mega_moe( return SymmBuffer( group, num_experts, max_num_tokens, top_k, hidden_size, intermediate_size, - ) + ) \ No newline at end of file