damn clankers2

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,
-    )
+    )