biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix hidden_size: shared expert uses 7168, not HC_DIM 28672

Browse Source
This commit is contained in:
biondizzle
2026-05-18 20:10:32 +00:00
parent 70f50a1ec6
commit f07643791e
1 changed files with 7 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
tests/test_shared_expert.py
View File

@@ -14,9 +14,7 @@ from safetensors import safe_open
MODEL_PATH = "/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4"
DEVICE = "cuda:0"
LAYER_IDX = 0
HIDDEN_SIZE = 7168
HC_MULT = 4
HC_DIM = HC_MULT * HIDDEN_SIZE # 28672
HIDDEN_SIZE = 7168 # shared expert input dim (from checkpoint weight shapes)
INTERMEDIATE_SIZE = 3072
SWIGLU_LIMIT = 10.0
NUM_TOKENS = 4
@@ -78,9 +76,9 @@ def main():
down_sf = P(f"{prefix}.down_proj.weight_scale")
down_gs = P(f"{prefix}.down_proj.weight_scale_2").item()
print(f"gate_proj: shape={gate_w.shape} gs={gate_gs:.8f}")
print(f"up_proj: shape={up_w.shape} gs={up_gs:.8f}")
print(f"down_proj: shape={down_w.shape} gs={down_gs:.8f}")
print(f"gate_proj: shape={gate_w.shape} gs={gate_gs:.8f} sf_shape={gate_sf.shape}")
print(f"up_proj: shape={up_w.shape} gs={up_gs:.8f} sf_shape={up_sf.shape}")
print(f"down_proj: shape={down_w.shape} gs={down_gs:.8f} sf_shape={down_sf.shape}")
# Stack gate + up into gate_up_proj (same format as MoE L1)
# gate/up weights are (intermediate, hidden) uint8 packed
@@ -104,7 +102,7 @@ def main():
# Create runner
runner = CuTeDSLSharedExpertRunner(
hidden_size=HC_DIM,
hidden_size=HIDDEN_SIZE,
intermediate_size=INTERMEDIATE_SIZE,
max_num_tokens=8192,
device=DEVICE,
@@ -119,7 +117,7 @@ def main():
runner.finalize_weights()
# Warmup to compute activation global scales
dummy = torch.randn(NUM_TOKENS, HC_DIM, dtype=torch.bfloat16, device=DEVICE) * 2.0
dummy = torch.randn(NUM_TOKENS, HIDDEN_SIZE, dtype=torch.bfloat16, device=DEVICE) * 2.0
runner._ensure_initialized()
runner.compute_activation_global_scales(dummy)
print(f"Warmup gs: L1={runner._l1_activation_global_scale:.6f} "
@@ -127,7 +125,7 @@ def main():
# Run CuTeDSL
print("\n--- CuTeDSL Forward ---")
hidden = torch.randn(NUM_TOKENS, HC_DIM, dtype=torch.bfloat16, device=DEVICE) * 2.0
hidden = torch.randn(NUM_TOKENS, HIDDEN_SIZE, dtype=torch.bfloat16, device=DEVICE) * 2.0
with torch.no_grad():
output = runner.run(hidden)