Plumb global scale as GEMM alpha instead of folding into UE4M3

stage_activation now returns (x_fp4, x_sf, input_global_scale).
The global scale is applied as the CUTLASS GEMM alpha parameter
in the epilogue: D = alpha * A @ B, avoiding the fp32→UE4M3
round-trip that folding would introduce.

Changes:
- stage_activation: returns global scale as 3rd value
- cutlass_nvfp4_gemm C++ binding: alpha param (was hardcoded 1.0)
- cutlass_grouped_nvfp4_gemm: passes alpha to per-expert GEMM
- nvfp4_mega_moe_l1/l2: accept alpha, pass to grouped GEMM
- nvfp4_moe_full: reads symm_buffer.input_global_scale for L1,
  uses stage_activation's returned global scale for L2
- SymmBuffer: added input_global_scale field
- vllm patch: stores global scale from stage_activation

src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/kernel.py

@@ -23,11 +23,12 @@ def cutlass_nvfp4_blockscaled_gemm(
     B_packed,   # (N, K_half) int8 packed E2M1
     SFB,        # scale factors for B (float8_e4m3fn)
     M, N, K,    # Problem dimensions (K in FP4 elements)
+    alpha=1.0,  # fp32 scalar applied in epilogue: D = alpha * A @ B + beta * C
 ):
     """Single NVFP4 block-scaled GEMM using CUTLASS."""
     if not _CUTLASS_AVAILABLE:
         raise RuntimeError("CUTLASS NVFP4 GEMM extension not available")
-    return _C.forward(A_packed, SFA, B_packed, SFB, M, N, K)
+    return _C.forward(A_packed, SFA, B_packed, SFB, M, N, K, alpha)
 
 
 def cutlass_grouped_nvfp4_gemm(
@@ -37,6 +38,7 @@ def cutlass_grouped_nvfp4_gemm(
     weight_sf,    # (E_per_rank, N, sf_k) float8_e4m3fn block scales
     topk_ids,     # (num_tokens, NUM_TOPK) int32
     topk_weights, # (num_tokens, NUM_TOPK) float32
+    alpha=1.0,    # fp32 scalar: D = alpha * A @ B (from stage_activation global scale)
 ):
     """Per-expert grouped GEMM for MoE dispatch using CUTLASS NVFP4.
     
@@ -77,6 +79,7 @@ def cutlass_grouped_nvfp4_gemm(
             expert_x, expert_x_sf,
             expert_w, expert_w_sf,  # Pass directly — already (N, K_half) and (N, sf_k)
             M_expert, N, K,
+            alpha=alpha,
         )  # (M_expert, N) bfloat16
         
         # Check for CUDA errors after each expert GEMM

src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/pytorch_binding.cpp

@@ -18,7 +18,8 @@ torch::Tensor cutlass_nvfp4_gemm_forward(
     torch::Tensor SFA,
     torch::Tensor B_packed,
     torch::Tensor SFB,
-    int64_t M, int64_t N, int64_t K
+    int64_t M, int64_t N, int64_t K,
+    double alpha = 1.0
 ) {
     auto D = torch::empty({M, N}, torch::dtype(torch::kBFloat16).device(A_packed.device()));
 
@@ -30,7 +31,7 @@ torch::Tensor cutlass_nvfp4_gemm_forward(
         B_packed.data_ptr(), SFB.data_ptr(),
         D.data_ptr(),
         static_cast<int>(M), static_cast<int>(N), static_cast<int>(K),
-        1.0f, 0.0f,
+        static_cast<float>(alpha), 0.0f,
         cuda_stream
     );
 

src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py

@@ -96,6 +96,7 @@ def nvfp4_mega_moe_l1(
     topk_ids,     # (num_tokens, NUM_TOPK) int32
     topk_weights, # (num_tokens, NUM_TOPK) float32
     num_experts_per_rank,
+    alpha=1.0,    # fp32 scalar from stage_activation global scale
 ):
     """L1 GEMM: gate_up_proj — Native NVFP4 block-scaled MMA.
 
@@ -121,6 +122,7 @@ def nvfp4_mega_moe_l1(
         x_fp4, x_sf_fp8,
         l1_weights, w_sf_fp8,
         topk_ids, topk_weights,
+        alpha=alpha,
     )
     return output  # (num_tokens, 6144) bfloat16
 
@@ -133,6 +135,7 @@ def nvfp4_mega_moe_l2(
     topk_ids,     # (num_tokens, NUM_TOPK) int32
     topk_weights, # (num_tokens, NUM_TOPK) float32
     num_experts_per_rank,
+    alpha=1.0,    # fp32 scalar from stage_activation global scale
 ):
     """L2 GEMM: down_proj — Native NVFP4 block-scaled MMA.
 
@@ -155,6 +158,7 @@ def nvfp4_mega_moe_l2(
         x_fp4, x_sf_fp8,
         l2_weights, w_sf_fp8,
         topk_ids, topk_weights,
+        alpha=alpha,
     )
     return output  # (num_tokens, 7168) bfloat16
 
@@ -228,10 +232,13 @@ def stage_activation(x_bf16):
        Pack as 4-bit sign-magnitude nibbles (bit3=sign, bits2:0=mag index)
        Block scale = block_max / 6.0 stored as UE4M3 (float8_e4m3fn)
     
-    The global scale is folded into the block scales (same as weight_transform
-    folds weight_scale_2 into weight_scale), so the GEMM API is unchanged:
-        dequant = e2m1_magnitude * block_scale
-        where block_scale = (block_max / 6.0) * global_scale
+    Returns (x_fp4, x_sf, input_global_scale) where:
+      x_fp4: packed E2M1 nibbles
+      x_sf: UE4M3 block scales (NOT folded with global scale)
+      input_global_scale: fp32 per-tensor scale, applied as GEMM alpha
+    
+    The GEMM applies global scale via alpha: D = alpha * (A_sf * A_fp4) @ (B_sf * B_fp4)
+    This avoids fp32→UE4M3 round-trip from folding, preserving precision.
     """
     x_f32 = x_bf16.float()
     
@@ -249,12 +256,7 @@ def stage_activation(x_bf16):
     
     x_fp4, x_sf = _quantize_to_e2m1(x_normalized)
     
-    # Fold global scale into block scales: sf_folded = sf * global_scale
-    # (same as _fold_global_scale in weight_transform.py)
-    sf_folded = x_sf.to(torch.float32) * input_global_scale
-    x_sf_folded = sf_folded.clamp(0.0, 448.0).to(torch.float8_e4m3fn)
-    
-    return x_fp4, x_sf_folded
+    return x_fp4, x_sf, input_global_scale
 
 
 def nvfp4_mega_moe_full(
@@ -302,6 +304,7 @@ def nvfp4_mega_moe_full(
     # Step 1: Read staged activation from symm_buffer
     x_fp4 = symm_buffer.x[:num_tokens]
     x_sf = symm_buffer.x_sf[:num_tokens]
+    l1_global_scale = symm_buffer.input_global_scale  # fp32, from stage_activation
     topk_ids = symm_buffer.topk_idx[:num_tokens]
     topk_weights = symm_buffer.topk_weights[:num_tokens]
 
@@ -330,6 +333,7 @@ def nvfp4_mega_moe_full(
     l1_output = nvfp4_mega_moe_l1(
         x_fp4, x_sf, l1_w, l1_sf,
         topk_ids_local, topk_weights, num_experts_per_rank,
+        alpha=l1_global_scale,
     )
 
     # NaN-trace: check L1 output
@@ -350,12 +354,13 @@ def nvfp4_mega_moe_full(
               f"abs_max={activated.abs().max().item():.4e}")
 
     # Step 4: Quantize L1 output → FP4
-    l1_fp4, l1_sf_out = stage_activation(activated)
+    l1_fp4, l1_sf_out, l2_global_scale = stage_activation(activated)
 
     # Step 5: L2 GEMM (native NVFP4 block-scaled MMA)
     l2_output = nvfp4_mega_moe_l2(
         l1_fp4, l1_sf_out, l2_w, l2_sf,
         topk_ids_local, topk_weights, num_experts_per_rank,
+        alpha=l2_global_scale,
     )
 
     # NaN-trace: check L2 output

src/nvfp4_megamoe_kernel/symm_buffer.py

@@ -68,6 +68,10 @@ class SymmBuffer:
             dtype=torch.bfloat16, device=device,
         )
         
+        # Per-tensor global scale from stage_activation (fp32 scalar)
+        # Applied as GEMM alpha: D = global_scale * (A_sf * A_fp4) @ (B_sf * B_fp4)
+        self.input_global_scale = 1.0
+        
         if MEGA_MOE_DEBUG:
             print(f"[SymmBuffer] x={self.x.shape} x_sf={self.x_sf.shape} "
                   f"topk_idx={self.topk_idx.shape} topk_weights={self.topk_weights.shape} "

vllm/patches/deepseek_v4.py

@@ -533,9 +533,10 @@ class DeepseekV4MegaMoEExperts(nn.Module):
 
         # Quantize activation using the kernel's PyTorch stage_activation
         # (same code path the kernel uses for L1→L2 requantization).
-        x_fp4, x_sf = stage_activation(hidden_states)
+        x_fp4, x_sf, input_global_scale = stage_activation(hidden_states)
         symm_buffer.x[:num_tokens].copy_(x_fp4)
         symm_buffer.x_sf[:num_tokens].copy_(x_sf)
+        symm_buffer.input_global_scale = input_global_scale
         symm_buffer.topk_idx[:num_tokens].copy_(topk_ids)
         symm_buffer.topk_weights[:num_tokens].copy_(topk_weights)