[JIT Kernel] Fix NVFP4 multi-arch compilation failure (#20874)
This commit is contained in:
committed by
GitHub
parent
6ad78b008b
commit
3d749c49ca
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user