Skip to content

Commit 9c3b746

Browse files
nvpohanhyzh119
andauthored
bugfix: Fix FusedMoeRunner does not exist error (flashinfer-ai#1424)
Fix `AttributeError: FusedMoeRunner does not exist` error when running cutlass MoE. Broken by flashinfer-ai#1201 <!-- .github/pull_request_template.md --> ## 📌 Description <!-- What does this PR do? Briefly describe the changes and why they’re needed. --> ## 🔍 Related Issues <!-- Link any related issues here --> ## 🚀 Pull Request Checklist Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete. ### ✅ Pre-commit Checks - [x] I have installed `pre-commit` by running `pip install pre-commit` (or used your preferred method). - [x] I have installed the hooks with `pre-commit install`. - [x] I have run the hooks manually with `pre-commit run --all-files` and fixed any reported issues. > If you are unsure about how to set up `pre-commit`, see [the pre-commit documentation](https://pre-commit.com/). ## 🧪 Tests - [x] Tests have been added or updated as needed. - [x] All tests are passing (`unittest`, etc.). ## Reviewer Notes <!-- Optional: anything you'd like reviewers to focus on, concerns, etc. --> --------- Signed-off-by: Po-Han Huang <[email protected]> Co-authored-by: Zihao Ye <[email protected]>
1 parent 96337ff commit 9c3b746

File tree

1 file changed

+11
-8
lines changed

1 file changed

+11
-8
lines changed

flashinfer/fused_moe/core.py

Lines changed: 11 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -266,7 +266,7 @@ def gen_cutlass_fused_moe_sm100_module(use_fast_build: bool = False) -> JitSpec:
266266

267267
@functools.cache
268268
def get_cutlass_fused_moe_sm100_module(use_fast_build: bool = False):
269-
module = gen_cutlass_fused_moe_sm100_module(use_fast_build).build_and_load(
269+
gen_cutlass_fused_moe_sm100_module(use_fast_build).build_and_load(
270270
class_name="FusedMoeRunner"
271271
)
272272

@@ -329,14 +329,17 @@ def __init__(
329329
)
330330

331331
if instance_key not in MoERunner.runner_dict:
332-
MoERunner.runner_dict[instance_key] = module.FusedMoeRunner(
333-
x_dtype,
334-
weight_dtype,
335-
output_dtype,
336-
use_deepseek_fp8_block_scale,
337-
use_w4a8_group_scaling,
338-
use_mxfp8_act_scaling,
332+
MoERunner.runner_dict[instance_key] = (
333+
torch.classes.fused_moe_sm100.FusedMoeRunner(
334+
x_dtype,
335+
weight_dtype,
336+
output_dtype,
337+
use_deepseek_fp8_block_scale,
338+
use_w4a8_group_scaling,
339+
use_mxfp8_act_scaling,
340+
)
339341
)
342+
340343
self.fused_moe_runner = MoERunner.runner_dict[instance_key]
341344

342345
def get_valid_tactics(

0 commit comments

Comments
 (0)