fix: 384 experts (not 256), clarify cross-rank reduce is in caller
DeepSeek-V4-Pro has 384 routed experts, 48 per rank (384/8). The cross-rank all-reduce happens in the parent DeepseekV4MoE.forward, not in our kernel. Our kernel writes local output; caller does reduce. Fixed README, nvfp4_mega_moe.py comments.
This commit is contained in:
@@ -7,7 +7,7 @@ Architecture:
|
||||
- L1 GEMM: gate_up_proj (FP4 x FP4 → BF16 with UE4M3 scales)
|
||||
- SiLU+Mul activation
|
||||
- L2 GEMM: down_proj (FP4 x FP4 → BF16 with UE4M3 scales)
|
||||
- NVLink cross-rank sync via symm buffer
|
||||
- NVLink cross-rank sync handled by caller (not this kernel)
|
||||
- Expert parallel: each rank handles NUM_EXPERTS/8 experts
|
||||
|
||||
The kernel uses native NVFP4 block-scaled MMA via tcgen05.mma
|
||||
@@ -253,7 +253,7 @@ def nvfp4_mega_moe_full(
|
||||
3. SiLU + Mul (activation)
|
||||
4. Quantize L1 output → FP4 + UE4M3 scales
|
||||
5. L2 GEMM: down_proj (native NVFP4 block-scaled MMA)
|
||||
6. NVLink sync + reduce across ranks → write to y
|
||||
6. Write to y (caller handles cross-rank all-reduce)
|
||||
|
||||
Uses tcgen05.mma.kind::mxf8f6f4.block_scale for native E2M1×E2M1
|
||||
with UE4M3 block-16 scaling in Blackwell tensor cores.
|
||||
@@ -305,5 +305,5 @@ def nvfp4_mega_moe_full(
|
||||
topk_ids, topk_weights, num_experts_per_rank,
|
||||
)
|
||||
|
||||
# Step 6: Write to output
|
||||
# Step 6: Write to output (caller handles cross-rank all-reduce)
|
||||
y.copy_(l2_output)
|
||||
|
||||
Reference in New Issue
Block a user