bench fix

2025-09-19 12:34:08 -04:00
parent 1e7302d30a
commit 42aadc5069
3 changed files with 27 additions and 5 deletions
--- a/scripts/debug_qwen2_experts.py
+++ b/scripts/debug_qwen2_experts.py
@@ -33,7 +33,7 @@ def main() -> None:

    block = Qwen2MoeSparseMoeBlock(cfg).to("cuda", dtype=torch.bfloat16)
    experts = block.experts
-    setattr(experts, "_ax_parent_block", block)
+    experts._ax_parent_block = block

    impls = _iter_expert_impls(experts)
    print(f"impl count: {len(impls)}")