Implement simple FP4 quantization for L1→L2 re-quant step (no vLLM fp4_utils dependency)

src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py

@@ -162,12 +162,50 @@ def nvfp4_mega_moe_l2(
 def stage_activation(x_bf16):
     """Quantize BF16 activation to FP4 (E2M1) with UE4M3 block16 scales.
     
-    Uses the Triton staging kernel from the vLLM deepseek_v4 patch.
+    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
     """
-    from vllm.model_executor.models.staging_kernel import (
-        _stage_deepseek_v4_mega_moe_inputs,
-    )
-    return _stage_deepseek_v4_mega_moe_inputs(x_bf16)
+    # Per-token absmax quantization to FP4
+    M, K = x_bf16.shape
+    assert K % 16 == 0, f"K={K} not divisible by 16"
+    
+    x_f32 = x_bf16.float()
+    
+    # Reshape into blocks of 16 for block-wise scaling
+    x_blocks = x_f32.reshape(M, K // 16, 16)  # (M, K//16, 16)
+    
+    # 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)
+    
+    return x_fp4, x_sf
 
 
 def nvfp4_mega_moe_full(