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

Fix compile test: add warmup for activation global scales

Browse Source
This commit is contained in:
biondizzle
2026-05-19 01:57:16 +00:00
parent bba3bca4d3
commit 00fe63b56f
1 changed files with 6 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
tests/test_compile_custom_op.py
View File

@@ -136,6 +136,12 @@ def main():
topk_ids = torch.tensor([[0, 1]] * 4, dtype=torch.int32, device=DEVICE)
topk_weights = torch.tensor([[0.6, 0.4]] * 4, dtype=torch.float32, device=DEVICE)
# 1. Warmup: compute activation global scales
print("\n[0] Computing activation global scales (warmup)...")
runner.compute_activation_global_scales(hidden_states, topk_weights, topk_ids)
print(f" L1 gs: {runner._l1_activation_global_scale:.6f}")
print(f" L2 gs: {runner._l2_activation_global_scale:.6f}")
# 1. Eager mode (baseline)
print("\n[1/2] Running eager mode (baseline)...")
runner._ensure_stacked()