test(jit): assert BMM export symlink under GEN_SRC_DIR

[kahyunnam]

📌 Description

Fixes #4166.

tests/utils/test_gen_module_symlink_race_condition.py fails on release-v0.6.16 rc1 (unit_test_b300, cu129/cu130):

AssertionError: Expected /tmp/flashinfer_test_fused_moe_symlink_ng4ln5vs/flashinfer/trtllm/
batched_gemm/trtllmGen_bmm_export to be a symlink

Root cause

#3468 ("move JIT symlinks to writable gen dir", b9890c85) moved the trtllmGen_bmm_export symlink from FLASHINFER_CUBIN_DIR to FLASHINFER_GEN_SRC_DIR in flashinfer/jit/fused_moe.py, flashinfer/jit/moe_utils.py, and flashinfer/jit/gemm/core.py. That PR touched three production files and no tests.

The test redirects only FLASHINFER_CUBIN_DIR, then asserts the symlink under it, so it now checks a location the code no longer writes to.

There is a second, quieter consequence. Since FLASHINFER_GEN_SRC_DIR is never redirected, the symlink is created in the real workspace directory and points into the test's temp directory, which shutil.rmtree deletes on teardown. The between-iteration symlink_path.unlink() also targeted the temp path, so it never removed the symlink actually in use — meaning the race stopped being re-triggered after the first iteration, and the test would have silently under-tested even with the assertion path corrected.

Fix

• Redirect FLASHINFER_GEN_SRC_DIR alongside FLASHINFER_CUBIN_DIR in the worker, and assert the symlink there. tests/jit/test_jit_cpp_ext.py already redirects FLASHINFER_GEN_SRC_DIR this way, and flashinfer/aot.py reassigns the same attribute, so this is the established pattern.
• Keep the artifact cache and generated-source dir as separate temp subdirectories, so the assertion still fails if the symlink ever moves back under FLASHINFER_CUBIN_DIR. Pointing both at one directory would make the assertion pass either way.
• Restore the race semantics and confine all symlinks to the temp tree.
• Correct the ensure_symlink() docstring, which still cited the pre-fix: move JIT symlinks to writable gen dir #3468 CUBIN_DIR location.

No production behavior changes; the only non-test edit is a docstring.

🔍 Related Issues

• Fixes [bug][v0.6.16] test_gen_module_symlink_race_condition assert is_symlink #4166
• Root cause: fix: move JIT symlinks to writable gen dir #3468

🚀 Pull Request Checklist

✅ Pre-commit Checks

• I have installed pre-commit by running pip install pre-commit (or used my preferred method).
• I have installed the hooks with pre-commit install.
• I have run the hooks manually with pre-commit run --all-files and fixed any reported issues.

All hooks pass on the changed files (ruff check, ruff format, mypy, whitespace/EOL).

🧪 Tests

• Tests have been updated as needed.

The test requires SM100/SM12x and returns early elsewhere, so it cannot be executed end-to-end on the SM89/90 host available to me. It still needs a B300/GB200 run to confirm the CI failure is cleared.

To verify the path logic without that hardware, I exercised the real gen_trtllm_gen_fused_moe_sm100_module() with only artifact download/checksum verification stubbed (the symlink is created before the nvcc-arch check raises), and recorded the ensure_symlink() link path. A stand-in directory represents the shared workspace dir so nothing touches the real cache.

Redirecting FLASHINFER_CUBIN_DIR only, as the test did before:

[ensure_symlink] link=/tmp/verify_before_.../shared_workspace_generated/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export
  symlink present in cubin_dir (old assertion target) : False   <-- assertion fails, matches #4166
  symlink present in shared gen dir (real workspace)  : True    <-- leaks outside the temp tree
  symlink present in test gen dir (temp, isolated)    : False

Redirecting both, as this PR does:

[ensure_symlink] link=/tmp/verify_after_.../generated/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export
  symlink present in cubin_dir (old assertion target) : False
  symlink present in shared gen dir (real workspace)  : False   <-- no leak
  symlink present in test gen dir (temp, isolated)    : True    <-- assertion passes

Collection and the non-SM100 early-return path were also confirmed green: pytest -q tests/utils/test_gen_module_symlink_race_condition.py → 1 passed.

Reviewer Notes

gen_fused_moe_worker_process now takes a (cubin_dir, gen_src_dir) tuple, since Pool.map passes a single argument.

flashinfer/jit/moe_utils.py and flashinfer/jit/gemm/core.py have the same symlink layout but no test asserts their paths, so nothing else needed updating.

Summary by CodeRabbit

• Documentation

  • Corrected an example path in the symlink helper documentation.

• Tests

  • Improved race-condition coverage for generated module symlinks.
  • Reduced test flakiness by separating generated files from cached build artifacts.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: 92a63b8c-b440-49c0-9c31-e84e9fed4b35

📥 Commits

Reviewing files that changed from the base of the PR and between d536efd and 7736b6f.

📒 Files selected for processing (2)

• flashinfer/jit/cubin_loader.py
• tests/utils/test_gen_module_symlink_race_condition.py

---

📝 Walkthrough

Walkthrough

The change corrects the ensure_symlink documentation and updates the symlink race test to use separate cubin-cache and generated-source directories while validating concurrent symlink creation and header access.

Changes

Symlink race test

Layer / File(s)	Summary
Generated-source symlink path and worker validation flashinfer/jit/cubin_loader.py, tests/utils/test_gen_module_symlink_race_condition.py	The documented symlink path now uses GEN_SRC_DIR; workers patch both directory settings and validate the symlink and generated headers.
Separated-directory race orchestration tests/utils/test_gen_module_symlink_race_condition.py	The race test uses distinct cubins and generated directories and removes only the generated-source symlink between worker runs.

Estimated code review effort: 2 (Simple) | ~10 minutes

Possibly related PRs

• flashinfer-ai/flashinfer#2903: Introduces the symlink-based header workflow exercised by the updated test.
• flashinfer-ai/flashinfer#2979: Addresses related concurrency behavior in ensure_symlink().
• flashinfer-ai/flashinfer#3468: Moves the fused BMM export symlink under FLASHINFER_GEN_SRC_DIR, matching this test setup.

Suggested labels: run-ci

Suggested reviewers: saltyminty, bkryu, jimmyzho

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly summarizes the main change: asserting the BMM export symlink under GEN_SRC_DIR.
Description check	✅ Passed	The description matches the template well, with description, related issues, tests, checklist, and reviewer notes filled in.
Linked Issues check	✅ Passed	The changes address #4166 by redirecting FLASHINFER_GEN_SRC_DIR and asserting the symlink in the correct location.
Out of Scope Changes check	✅ Passed	The extra docstring tweak and test refactor are directly related to the same symlink-path fix and are not out of scope.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[kahyunnam]

/bot run tests/utils

[flashinfer-bot]

GitLab MR !1054 has been created, and the CI pipeline #59840427 is currently running. I'll report back once the pipeline job completes.