diff --git a/src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py b/src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py index 62e96311..ec0daa1c 100644 --- a/src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py +++ b/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(