Fix [Spark unit test CI]: defer torch._dynamo.disable to avoid import-time crash in CI (#3290)
<!-- .github/pull_request_template.md --> ## ๐ Description - #3271 (feat(moe): add SM120 W4A16 b12x kernels) added a @torch._dynamo.disable decorator to current_cuda_stream() in cute_dsl/utils.py. This eagerly imports torch._dynamo at module load time, which triggers getpass.getuser() during cache-dir initialization. This crashes in CI containers running as unmapped UIDs (e.g. Spark runners with -u $(id -u):$(id -g) mapping to UID 996, which has no /etc/passwd entry). - Replaces the eager decorator with a self-replacing lazy wrapper that defers torch._dynamo.disable to the first call of current_cuda_stream(), with zero overhead on subsequent calls. ## ๐ 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. --> <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Performance** * Improved module import performance by deferring CUDA stream initialization until first use, reducing startup overhead. [](https://app.coderabbit.ai/change-stack/flashinfer-ai/flashinfer/pull/3290) <!-- end of auto-generated comment: release notes by coderabbit.ai --> Co-authored-by: Cursor <cursoragent@cursor.com>
K
Ka-Hyun Nam committed
9d282787d412218aaf6193259598c0487b4926a8
Parent: 00b4dc5
Committed by GitHub <noreply@github.com>
on 5/12/2026, 4:09:04 PM