From bcfa05230f271a724ef7a1a6c6a0e64de188c31b Mon Sep 17 00:00:00 2001
From: Kaixi Hou <kaixih@users.noreply.github.com>
Date: Tue, 19 May 2026 17:14:53 -0700
Subject: [PATCH] Make cubin symlink setup idempotent

---
 flashinfer/jit/cubin_loader.py   |  7 +++-
 tests/utils/test_cubin_loader.py | 57 ++++++++++++++++++++++++++++++++
 2 files changed, 63 insertions(+), 1 deletion(-)
 create mode 100644 tests/utils/test_cubin_loader.py

diff --git a/flashinfer/jit/cubin_loader.py b/flashinfer/jit/cubin_loader.py
index fd7724ff1d2..049149b5c4c 100644
--- a/flashinfer/jit/cubin_loader.py
+++ b/flashinfer/jit/cubin_loader.py
@@ -261,7 +261,12 @@ def ensure_symlink(
                 link.unlink()
             else:
                 shutil.rmtree(link)
-        link.symlink_to(target)
+        try:
+            link.symlink_to(target)
+        except FileExistsError:
+            if link.is_symlink() and link.resolve() == target.resolve():
+                return
+            raise
 
 
 def verify_symlinked_headers(
diff --git a/tests/utils/test_cubin_loader.py b/tests/utils/test_cubin_loader.py
new file mode 100644
index 00000000000..00142b6e517
--- /dev/null
+++ b/tests/utils/test_cubin_loader.py
@@ -0,0 +1,57 @@
+"""
+Copyright (c) 2026 by FlashInfer team.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+  http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+import pathlib
+
+import pytest
+
+from flashinfer.jit.cubin_loader import ensure_symlink
+
+
+def test_ensure_symlink_tolerates_concurrent_same_target(monkeypatch, tmp_path):
+    target = tmp_path / "target"
+    target.mkdir()
+    link = tmp_path / "include" / "target"
+    original_symlink_to = pathlib.Path.symlink_to
+
+    def create_link_then_raise(self, symlink_target, *args, **kwargs):
+        original_symlink_to(self, symlink_target, *args, **kwargs)
+        raise FileExistsError
+
+    monkeypatch.setattr(pathlib.Path, "symlink_to", create_link_then_raise)
+
+    ensure_symlink(link, target)
+
+    assert link.is_symlink()
+    assert link.resolve() == target.resolve()
+
+
+def test_ensure_symlink_reraises_concurrent_different_target(monkeypatch, tmp_path):
+    target = tmp_path / "target"
+    other_target = tmp_path / "other_target"
+    target.mkdir()
+    other_target.mkdir()
+    link = tmp_path / "include" / "target"
+    original_symlink_to = pathlib.Path.symlink_to
+
+    def create_wrong_link_then_raise(self, symlink_target, *args, **kwargs):
+        original_symlink_to(self, other_target, *args, **kwargs)
+        raise FileExistsError
+
+    monkeypatch.setattr(pathlib.Path, "symlink_to", create_wrong_link_then_raise)
+
+    with pytest.raises(FileExistsError):
+        ensure_symlink(link, target)