fix(cute_dsl): avoid MoE wrapper runner reference cycle (#3340)
<!-- .github/pull_request_template.md --> ## ๐ Description `CuteDslMoEWrapper` currently passes `self._forward_with_tactic` as a bound method into `CuteDslFusedMoENvfp4Runner`, creating a strong reference cycle: `wrapper -> runner -> bound method -> wrapper`. When the wrapper is used with `use_cuda_graph=True`, this can keep wrapper-owned CUDA graph resources alive after user code has dropped the wrapper, until Python cyclic GC eventually runs. This PR replaces that bound-method callback with a weakref trampoline. The runner can still call into a live wrapper, but it no longer owns the wrapper lifetime. This prevents stale wrapper CUDA resources from surviving across same-process tests or later autotune runs. ## ๐ Related Issues https://github.com/flashinfer-ai/flashinfer/pull/3286 https://github.com/flashinfer-ai/flashinfer/pull/3301 https://github.com/flashinfer-ai/flashinfer/pull/3252 ## ๐ 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 Adds a focused regression test that warms a CUDA-graph wrapper, verifies it is finalized before cyclic GC, and then runs a subsequent autotuned wrapper call to ensure the output remains NaN-free. - [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. --> <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Bug Fixes** * Improved handling and cleanup of CUDA-graph wrappers to prevent resource leaks and provide a clear error when a wrapper is no longer available. * **Tests** * Added lifetime tests covering CUDA-graph wrappers before and after autotune; verify stable, non-NaN outputs during autotune. * **Documentation** * Updated comment about cold-L2 cache behavior and noted follow-up to re-enable it once a related issue is addressed. <!-- review_stack_entry_start --> [](https://app.coderabbit.ai/change-stack/flashinfer-ai/flashinfer/pull/3340?utm_source=github_walkthrough&utm_medium=github&utm_campaign=change_stack) <!-- review_stack_entry_end --> <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
L
Lee Nau committed
18f45345ebb40529b52baad2a19a8a237d5edcd5
Parent: 3a81c3e
Committed by GitHub <noreply@github.com>
on 5/21/2026, 4:56:58 PM