Fix stage_activation: use Triton staging kernel instead of broken simple quantize

src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py

@@ -162,49 +162,36 @@ def nvfp4_mega_moe_l2(
 def stage_activation(x_bf16):
     """Quantize BF16 activation to FP4 (E2M1) with UE4M3 block16 scales.
     
-    Simple per-token FP4 quantization using torch operations.
-    This is used for re-quantizing L1 output before the L2 GEMM.
-    
-    Returns: (x_fp4, x_sf) where
-        x_fp4: (M, K//2) int8 packed E2M1
-        x_sf: (M, K//16) float8_e4m3fn UE4M3 block scales
+    Uses the Triton staging kernel from the vLLM deepseek_v4 patch.
+    Allocates output buffers matching the kernel's expected format:
+      x_fp4: (M, K//2) int8 packed E2M1
+      x_sf: (M, K//64) uint32 packed UE4M3 (4 per uint32)
     """
-    # Per-token absmax quantization to FP4
+    from vllm.model_executor.models.staging_kernel import (
+        _stage_deepseek_v4_mega_moe_inputs,
+    )
+    
     M, K = x_bf16.shape
-    assert K % 16 == 0, f"K={K} not divisible by 16"
+    K_half = K // 2
+    K_sf = K // 64  # 4 UE4M3 per uint32, 16 values per group → K/(16*4) = K//64
     
-    x_f32 = x_bf16.float()
+    # Allocate output buffers
+    x_fp4 = torch.empty(M, K_half, dtype=torch.int8, device=x_bf16.device)
+    x_sf = torch.empty(M, K_sf, dtype=torch.int32, device=x_bf16.device)
     
-    # Reshape into blocks of 16 for block-wise scaling
-    x_blocks = x_f32.reshape(M, K // 16, 16)  # (M, K//16, 16)
+    # Create dummy topk tensors (the staging kernel writes them but we don't need them)
+    topk_weights = torch.ones(M, 1, dtype=torch.float32, device=x_bf16.device)
+    topk_ids = torch.zeros(M, 1, dtype=torch.int32, device=x_bf16.device)
+    topk_idx_out = torch.empty(M, 1, dtype=torch.int32, device=x_bf16.device)
+    topk_weights_out = torch.empty(M, 1, dtype=torch.float32, device=x_bf16.device)
     
-    # Compute per-block scale (absmax)
-    block_max = x_blocks.abs().amax(dim=-1, keepdim=True)  # (M, K//16, 1)
-    
-    # Clamp to UE4M3 representable range [0, 448]
-    block_max = block_max.clamp(min=1e-8, max=448.0)
-    
-    # Convert scale to UE4M3 (float8_e4m3fn)
-    x_sf = block_max.squeeze(-1).to(torch.float8_e4m3fn)  # (M, K//16)
-    
-    # Scale and quantize to E2M1 (4-bit, values in [-6, 6] step 1)
-    # Simple approach: scale to [-6, 6] range and round to nearest E2M1 value
-    scale_f32 = block_max.float() / 6.0  # Normalize to [-6, 6]
-    x_scaled = x_blocks / scale_f32.clamp(min=1e-8)
-    
-    # Quantize to 4-bit (E2M1 approximation: round to nearest 0.5 step)
-    x_quant = (x_scaled * 2).round() / 2  # Step of 0.5
-    x_quant = x_quant.clamp(-6, 6)
-    
-    # Pack 2 values per byte (high nibble = first, low nibble = second)
-    x_quant_4bit = (x_quant * 2).round().to(torch.int8)  # -12 to 12
-    x_4bit = x_quant_4bit.reshape(M, K // 2, 2)
-    
-    # Pack: high nibble = first value, low nibble = second
-    high = (x_4bit[:, :, 0].clamp(0, 15)).to(torch.uint8)  # E2M1 unsigned
-    low = (x_4bit[:, :, 1].clamp(0, 15)).to(torch.uint8)
-    x_fp4 = (high << 4 | low).to(torch.int8)  # (M, K//2)
+    _stage_deepseek_v4_mega_moe_inputs(
+        x_bf16, topk_weights, topk_ids,
+        x_fp4, x_sf,
+        topk_idx_out, topk_weights_out,
+    )
     
+    # x_sf is uint32 packed but stored as int32 — cast for consistency
     return x_fp4, x_sf