SIGN IN SIGN UP

FlashInfer: Kernel Library for LLM Serving

0 0 188 Python

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 -->

[![Review Change
Stack](https://storage.googleapis.com/coderabbit_public_assets/review-stack-in-coderabbit-ui.svg)](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