damn clankers

src/nvfp4_megamoe_kernel/symm_buffer.py

@@ -36,17 +36,17 @@ class SymmBuffer:
         device = torch.cuda.current_device()
         
         # NVFP4: packed E2M1 (2 values per byte), so K//2
-        sf_k_groups_hidden = hidden_size // (16 * 4)  # UE4M3 block16, 4 packed per uint32
-        sf_k_groups_inter = intermediate_size // (16 * 4)
+        sf_k_groups_hidden = hidden_size // 16  # UE4M3 block16, 1 scale per group
+        sf_k_groups_inter = intermediate_size // 16
         
-        # Staging buffers (matching DeepGEMM layout)
+        # Staging buffers (matching kernel's expected input format)
         self.x = torch.empty(
             max_num_tokens, hidden_size // 2,
             dtype=torch.int8, device=device,
         )
         self.x_sf = torch.empty(
             max_num_tokens, sf_k_groups_hidden,
-            dtype=torch.uint32, device=device,
+            dtype=torch.float8_e4m3fn, device=device,
         )
         self.topk_idx = torch.empty(
             max_num_tokens, top_k,