[JIT Kernel] Fix NVFP4 multi-arch compilation failure (#20874)

This commit is contained in:
Mohammad Miadh Angkad
2026-03-20 10:30:04 +08:00
committed by GitHub
parent 6ad78b008b
commit 3d749c49ca

View File

@@ -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