From 16f91ff0e175ad7205827a16b8e382b3886f7998 Mon Sep 17 00:00:00 2001
From: biondizzle <biondizzle@gmail.com>
Date: Thu, 14 May 2026 13:02:10 +0000
Subject: [PATCH] fix: rewrite stage_activation with proper E2M1 quantization
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Three bugs fixed:
1. clamp(0,15) was destroying sign bits — E2M1 is sign-magnitude 4-bit
   nibbles, not unsigned. Half the activation was zeroed.
2. Scale stored block_max but divided by block_max/6, so stored scale was
   6× too large. Now correctly stores block_max/6 (the actual dequant factor).
3. Uniform 0.5 step doesn't match E2M1 values {0,0.5,1,1.5,2,3,4,6}.
   Now snaps to nearest E2M1 representable magnitude.

New _quantize_to_e2m1 helper handles all three correctly:
- Sign-magnitude 4-bit nibble packing (bit3=sign, bits2:0=mag index)
- Correct block scale (block_max / 6.0)
- Nearest-neighbor to actual E2M1 values
---
 src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py | 95 +++++++++++++++-------
 1 file changed, 64 insertions(+), 31 deletions(-)

diff --git a/src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py b/src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py
index 87a50976..30a37a65 100644
--- a/src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py
+++ b/src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py
@@ -154,42 +154,75 @@ def nvfp4_mega_moe_l2(
     return output  # (num_tokens, 7168) bfloat16
 
 
+# E2M1 (FP4) representable magnitudes: {0, 0.5, 1, 1.5, 2, 3, 4, 6}
+# Bit patterns (3-bit, no sign): 000=0, 001=0.5, 010=1, 011=1.5, 100=2, 101=3, 110=4, 111=6
+# Full 4-bit nibble: bit 3 = sign, bits 2:0 = magnitude index
+_E2M1_MAGNITUDES = torch.tensor([0, 0.5, 1, 1.5, 2, 3, 4, 6], dtype=torch.float32)
+
+
+def _quantize_to_e2m1(x_f32):
+    """Quantize float32 values to E2M1 (FP4) nibble indices.
+    
+    Maps each value to the nearest E2M1 representable magnitude,
+    then packs as 4-bit sign-magnitude nibbles.
+    
+    Returns (nibbles, scales) where:
+      nibbles: (..., N) uint8 with 4-bit sign-magnitude per value
+      scales: (..., N//16) float8_e4m3fn block scales
+    """
+    *batch, N = x_f32.shape
+    assert N % 16 == 0, f"Last dim {N} not divisible by 16 (block size)"
+
+    # Reshape into blocks of 16 for block-wise scaling
+    x_blocks = x_f32.reshape(*batch, N // 16, 16)
+    
+    # Per-block absmax determines the scale
+    block_max = x_blocks.abs().amax(dim=-1, keepdim=True).clamp(min=1e-8, max=448.0)
+    
+    # Scale so that the max maps to 6.0 (largest E2M1 magnitude)
+    # Dequant: x_reconstructed = x_e2m1 * scale, where scale = block_max / 6.0
+    scale_f32 = block_max / 6.0
+    x_scaled = x_blocks / scale_f32.clamp(min=1e-8)
+    
+    # Find nearest E2M1 magnitude for each value
+    signs = torch.sign(x_scaled)  # +1, -1, or 0
+    abs_scaled = x_scaled.abs()   # 0..6 range
+    
+    # Nearest E2M1 magnitude: find closest in {0, 0.5, 1, 1.5, 2, 3, 4, 6}
+    mags = _E2M1_MAGNITUDES.to(device=abs_scaled.device)
+    # Distance from each value to each magnitude
+    dists = (abs_scaled.unsqueeze(-1) - mags).abs()  # (..., 16, 8)
+    idx = dists.argmin(dim=-1)  # (..., 16) — index into E2M1 magnitudes
+    
+    # Clamp to valid range (safety)
+    idx = idx.clamp(0, 7).to(torch.uint8)
+    
+    # Build 4-bit sign-magnitude nibble: bit3=sign, bits2:0=magnitude index
+    sign_bit = (signs < 0).to(torch.uint8)  # 1 if negative
+    nibbles = (sign_bit << 3) | idx  # (..., 16) uint8, values 0..15
+    
+    # Pack 2 nibbles per byte: low nibble = even index, high nibble = odd index
+    nibbles = nibbles.reshape(*batch, N // 2, 2)
+    packed = (nibbles[..., 1] << 4) | nibbles[..., 0]  # (..., N//2) uint8
+    
+    # Scale factors: what the GEMM needs to reconstruct the original values
+    # dequant = e2m1_magnitude * scale, so scale = block_max / 6.0
+    sf = scale_f32.squeeze(-1).to(torch.float8_e4m3fn)  # (..., N//16)
+    
+    return packed.to(torch.int8), sf
+
+
 def stage_activation(x_bf16):
     """Quantize BF16 activation to FP4 (E2M1) with UE4M3 block16 scales.
     
-    Uses vLLM's per_tensor_cast_to_fp4 utility for L1→L2 re-quantization.
-    This is a simplified quantization — proper E2M1 with UE4M3 block scales
-    would require the staging kernel, but the staging kernel's API is complex.
-    
-    For now, we dequant the CUTLASS L1 output and use a simple absmax quantize.
-    TODO: Use the Triton staging kernel with proper buffer allocation.
+    Proper E2M1 quantization:
+    - Per-block (16 values) absmax scaling
+    - Snap to nearest E2M1 representable value: {0, ±0.5, ±1, ±1.5, ±2, ±3, ±4, ±6}
+    - Pack as 4-bit sign-magnitude nibbles (bit3=sign, bits2:0=mag index)
+    - Block scale = block_max / 6.0 stored as UE4M3 (float8_e4m3fn)
     """
-    M, K = x_bf16.shape
-    K_half = K // 2
-    K_sf = K // 16  # 1 scale per 16 values
-    
-    # Simple per-token absmax quantization
     x_f32 = x_bf16.float()
-    
-    # Reshape into blocks of 16 for block-wise scaling
-    x_blocks = x_f32.reshape(M, K_sf, 16)
-    block_max = x_blocks.abs().amax(dim=-1, keepdim=True).clamp(min=1e-8, max=448.0)
-    
-    # Scale to E2M1 range and quantize
-    scale_f32 = block_max / 6.0
-    x_scaled = x_blocks / scale_f32.clamp(min=1e-8)
-    x_quant = (x_scaled * 2).round() / 2  # step of 0.5
-    x_quant = x_quant.clamp(-6, 6)
-    
-    # Pack 2 values per byte (E2M1 packing)
-    x_q4 = (x_quant * 2).round().to(torch.int8).reshape(M, K_half, 2)
-    high = (x_q4[:, :, 0].clamp(0, 15)).to(torch.uint8)
-    low = (x_q4[:, :, 1].clamp(0, 15)).to(torch.uint8)
-    x_fp4 = (high << 4 | low).to(torch.int8)
-    
-    # Scale factors as float8_e4m3fn
-    x_sf = block_max.squeeze(-1).to(torch.float8_e4m3fn)
-    
+    x_fp4, x_sf = _quantize_to_e2m1(x_f32)
     return x_fp4, x_sf