From 3d749c49ca16e5093153afebfe24f9b4d9f48d69 Mon Sep 17 00:00:00 2001
From: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com>
Date: Fri, 20 Mar 2026 10:30:04 +0800
Subject: [PATCH] [JIT Kernel] Fix NVFP4 multi-arch compilation failure
 (#20874)

---
 python/sglang/jit_kernel/nvfp4.py | 31 ++++---------------------------
 1 file changed, 4 insertions(+), 27 deletions(-)

diff --git a/python/sglang/jit_kernel/nvfp4.py b/python/sglang/jit_kernel/nvfp4.py
index 2fc663175..e13c8ac61 100644
--- a/python/sglang/jit_kernel/nvfp4.py
+++ b/python/sglang/jit_kernel/nvfp4.py
@@ -70,19 +70,6 @@ def _nvfp4_cuda_flags() -> list[str]:
     ]
 
 
-def _parse_cuda_version() -> tuple[int, int]:
-    v = torch.version.cuda
-    if not v:
-        return (0, 0)
-    parts = v.split(".")
-    if len(parts) < 2:
-        return (0, 0)
-    try:
-        return int(parts[0]), int(parts[1])
-    except ValueError:
-        return (0, 0)
-
-
 def _get_nvfp4_cuda_arch_list() -> str:
     if not torch.cuda.is_available():
         raise RuntimeError("NVFP4 JIT kernels require CUDA.")
@@ -93,20 +80,10 @@ def _get_nvfp4_cuda_arch_list() -> str:
         )
     # NVFP4 kernels use architecture-family-specific instructions and must be
     # compiled for `sm_*a` targets (e.g. sm_100a), not plain sm_100.
-    archs = [f"{major}.{minor}a"]
-    cuda_major, _cuda_minor = _parse_cuda_version()
-    if cuda_major >= 13 and "10.3a" not in archs:
-        # Match sgl-kernel AOT fatbin behavior on CUDA 13+ for Blackwell.
-        archs.append("10.3a")
-    # Preserve order while de-duplicating.
-    seen = set()
-    ordered_archs: list[str] = []
-    for arch in archs:
-        if arch in seen:
-            continue
-        seen.add(arch)
-        ordered_archs.append(arch)
-    return " ".join(ordered_archs)
+    # JIT compilation targets only the current device, unlike AOT fat-binaries;
+    # adding extra architectures here would clash with the single SGL_CUDA_ARCH
+    # value injected by load_jit().
+    return f"{major}.{minor}a"
 
 
 @contextmanager