[Bug] JIT: bgmv_moe rebuilds from scratch in every process because shutil.copy resets the copied sources' mtime

## Summary `gen_bgmv_moe_module()` copies its CUDA sources and headers into the JIT `generated/` directory with `shutil.copy`, which does not preserve mtime. The copies therefore receive the current time on every process start, ninja sees them as newer than the previously built objects, and the module is fully recompiled — even when `FLASHINFER_WORKSPACE_BASE` is shared across runs and nothing has changed. Measured on one machine with the default (release) JIT configuration: every fresh process pays ~217 s of rebuild. With `shutil.copy2` the second and later processes load from cache in ~4.3 s, i.e. ~213 s saved per process. This is easy to miss because the cache *looks* correct — the generated directory, the cache key and `build.ninja` are all stable across runs; only the mtimes move. ## Root cause https://github.com/flashinfer-ai/flashinfer/blob/main/flashinfer/jit/bgmv_moe.py#L100 https://github.com/flashinfer-ai/flashinfer/blob/main/flashinfer/jit/bgmv_moe.py#L108 `JitSpecNvcc.try_load()` returns `None` for JIT-path modules (its comment states that it delegates JIT-path freshness to ninja), so `build_and_load()` always invokes ninja and the freshness decision rests entirely on mtime comparison. Resetting the copied sources' mtime therefore guarantees a full rebuild every time. ## Evidence ### 1. mtime of the copied sources The source file mtime is constant. For `moe_bgmv_binding.cu` in the generated directory, across three separate processes: With `shutil.copy` (current behaviour) — a different mtime each time, equal to the moment of copying: ``` run 1: 2026-08-27 16:08:17.288354236 run 2: 2026-08-27 16:11:53.781513130 run 3: 2026-08-27 16:15:32.347126118 ``` With `shutil.copy2` — identical to the source, to the nanosecond: ``` run 1: 2026-08-27 15:51:52.591185647 run 2: 2026-08-27 15:51:52.591185647 run 3: 2026-08-27 15:51:52.591185647 ``` ### 2. ninja dry run, compiling nothing Starting from an already-built cache and only changing *how* the sources are copied, then asking `ninja -n` for its decision: ``` step 0, untouched cache generated source mtime 16:15:32 / object mtime 16:15:37 ninja: no work to do. step 1, replay shutil.copy() generated source mtime 16:20:32 / object mtime 16:15:37 <-- sources now look newer ninja: 8 build actions pending bgmv_moe_moe_bgmv_binding.cuda.o bgmv_moe_moe_bgmv_bf16_bf16_bf16.cuda.o ... all 7 objects ... bgmv_moe.so step 2, replay shutil.copy2() generated source mtime 15:51:52.591185647 / object mtime 16:15:37 ninja: no work to do. ``` The decision flips back and forth with nothing but the copy call changing, which isolates mtime as the sole cause. ### 3. Wall time, three consecutive independent processes Default release configuration, cache directory wiped before each arm. "rebuilt" was determined by comparing the mtimes of the 7 `.o` files plus `bgmv_moe.so`, not by parsing log text. | | run 1 | run 2 | run 3 | |---|---|---|---| | `shutil.copy` (current) | 216.3 s, built | 218.5 s, all 8 artifacts rebuilt | 217.2 s, all 8 artifacts rebuilt | | `shutil.copy2` | 217.3 s, built (cold cache, expected) | **4.3 s, nothing rebuilt** | **4.3 s, nothing rebuilt** | ## Reproduce ```bash export FLASHINFER_WORKSPACE_BASE=/tmp/fi-cache rm -rf "$FLASHINFER_WORKSPACE_BASE" for i in 1 2 3; do /usr/bin/time -f "run $i: %e s" \ python3 -c 'from flashinfer.jit.bgmv_moe import load_bgmv_moe_module; load_bgmv_moe_module()' done ``` Every run takes the full compile time. Applying the patch below makes runs 2 and 3 return almost immediately. ## Suggested fix ```diff --- a/flashinfer/jit/bgmv_moe.py +++ b/flashinfer/jit/bgmv_moe.py @@ -97,7 +97,7 @@ dest_path = gen_directory / fname - shutil.copy(src_path, dest_path) + shutil.copy2(src_path, dest_path) sources.append(dest_path) @@ -105,4 +105,4 @@ - shutil.copy(src_path, gen_directory / fname) + shutil.copy2(src_path, gen_directory / fname) ``` The same pattern appears in two other JIT modules, which presumably have the same behaviour: - `flashinfer/jit/rmsnorm_silu.py:420` - `flashinfer/jit/monomoe.py:110` (structurally almost identical to `bgmv_moe.py`) For what it is worth, `flashinfer/aot.py:953` already uses `shutil.copy2`, so the codebase is internally inconsistent on this point. ## Impact Any setup that invokes `bgmv_moe` repeatedly in separate processes pays the full compile cost every time. We hit this in a benchmarking harness that runs the routine once per commit per configuration, where it added tens of minutes of pure recompilation per job and caused time-limited jobs to be killed. ## Environment - flashinfer `0.6.18`, commit `9caa448`, built from source - Blackwell, `sm_120`, `TORCH_CUDA_ARCH_LIST=12.0a` - CUDA 13.2 driver, PyTorch 2.12 nightly - Behaviour verified against current `main` by inspection; the two `shutil.copy` call sites are unchanged there