Fused SwiGLU epilogue with granularity-8 weight interleave
- Fix interleave_l1_weights: remove //2 bug (g=granularity_bf16 for N-axis) - Apply L1 weight+SF interleave in runner._ensure_stacked() and moe_pipeline - De-interleave L1 GEMM output before gate/up split - Fused SwiGLU kernel: epi_tile=(128,8) for subtile-level pairing - Even subtiles = gate: SiLU in FP32 registers, save to register buffer - Odd subtiles = up: silu(gate)*up from buffer - Both branches produce same BF16 tensor type (CuTeDSL constraint) - run_nvfp4_moe_fused() pipeline: fused L1 + PyTorch L2 - Runner: fused_swiglu=True option for CuTeDSLMoERunner - Layertest: both fused and non-fused paths PASS (cosine 0.988) - README.md updated with current status and lessons learned
This commit is contained in:
@@ -17,6 +17,7 @@ sys.path.insert(0, REPO_ROOT)
|
||||
|
||||
from cutedsl.moe_pipeline import (
|
||||
run_nvfp4_moe,
|
||||
run_nvfp4_moe_fused,
|
||||
)
|
||||
|
||||
NVFP4_MODEL_DIR = "/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4"
|
||||
@@ -251,6 +252,32 @@ def main():
|
||||
else:
|
||||
print(f" PASS: cosine {cosine:.6f} >= {COSINE_THRESHOLD}")
|
||||
|
||||
# ══════════════════════════════════════════════════════════════
|
||||
# Fused SwiGLU pipeline test
|
||||
# ══════════════════════════════════════════════════════════════
|
||||
print(f"\n Running CuTeDSL NVFP4 MoE FUSED pipeline (first run compiles)...")
|
||||
fused_output = run_nvfp4_moe_fused(
|
||||
hidden_states, expert_ids, expert_weights,
|
||||
weights, expert_indices,
|
||||
)
|
||||
print(f" Fused: amax={fused_output.abs().max():.4f} mean={fused_output.float().mean():.6f}")
|
||||
|
||||
fused_cosine = torch.nn.functional.cosine_similarity(
|
||||
fused_output.flatten().unsqueeze(0).float(),
|
||||
ref_output.flatten().unsqueeze(0).float(),
|
||||
).item()
|
||||
fused_mse = (fused_output.float() - ref_output.float()).pow(2).mean().item()
|
||||
|
||||
print(f"\n{'=' * 70}")
|
||||
print(f" FUSED RESULT: cosine={fused_cosine:.6f} MSE={fused_mse:.6e}")
|
||||
print(f"{'=' * 70}")
|
||||
|
||||
if fused_cosine < COSINE_THRESHOLD:
|
||||
print(f" FAIL: fused cosine {fused_cosine:.6f} < {COSINE_THRESHOLD}")
|
||||
sys.exit(1)
|
||||
else:
|
||||
print(f" PASS: fused cosine {fused_cosine:.6f} >= {COSINE_THRESHOLD}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
Reference in New Issue
Block a user