diff --git a/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/kernel.py b/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/kernel.py index 26fbffa6..f43ccae9 100644 --- a/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/kernel.py +++ b/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 diff --git a/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/pytorch_binding.cpp b/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/pytorch_binding.cpp index c7549088..a19ec77e 100644 --- a/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/pytorch_binding.cpp +++ b/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(M), static_cast(N), static_cast(K), - 1.0f, 0.0f, + static_cast(alpha), 0.0f, cuda_stream ); diff --git a/src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py b/src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py index 9eaa5e9b..bf543442 100644 --- a/src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py +++ b/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 diff --git a/src/nvfp4_megamoe_kernel/symm_buffer.py b/src/nvfp4_megamoe_kernel/symm_buffer.py index 4984ecd7..d974f1ec 100644 --- a/src/nvfp4_megamoe_kernel/symm_buffer.py +++ b/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} " diff --git a/vllm/patches/deepseek_v4.py b/vllm/patches/deepseek_v4.py index e5a046db..1b66350c 100644 --- a/vllm/patches/deepseek_v4.py +++ b/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)