diff --git a/python/sglang/kernels/fused_op.py b/python/sglang/kernels/fused_op.py index 34ffe0aab..f896bec08 100644 --- a/python/sglang/kernels/fused_op.py +++ b/python/sglang/kernels/fused_op.py @@ -70,6 +70,7 @@ BACKEND_METHODS: Dict[KernelBackend, str] = { KernelBackend.FLASHINFER: "forward_flashinfer", KernelBackend.DEEPGEMM: "forward_deepgemm", KernelBackend.AITER: "forward_aiter", + KernelBackend.TORCH_NPU: "forward_npu", } # best -> fallback. ``torch_compile`` is deliberately absent: auto-selection @@ -82,6 +83,7 @@ DEFAULT_PRIORITY: Tuple[KernelBackend, ...] = ( KernelBackend.DEEPGEMM, KernelBackend.CUTE_DSL, KernelBackend.AITER, + KernelBackend.TORCH_NPU, KernelBackend.TRITON, KernelBackend.TORCH, ) @@ -269,6 +271,9 @@ class BaseFusedOp(ABC): def forward_aiter(self, *args, **kwargs): raise NotImplementedError(f"{self.op}: no aiter backend") + def forward_npu(self, *args, **kwargs): + raise NotImplementedError(f"{self.op}: no npu backend") + # --- selection --- def available_backends(self) -> List[KernelBackend]: diff --git a/python/sglang/kernels/ops/activation/__init__.py b/python/sglang/kernels/ops/activation/__init__.py index 0a493f251..eb0e2de96 100644 --- a/python/sglang/kernels/ops/activation/__init__.py +++ b/python/sglang/kernels/ops/activation/__init__.py @@ -180,11 +180,98 @@ class GeluTanhAndMulOp(_GatedActivationOp): return F.gelu(gate, approximate="tanh") +class ReLU2Op(BaseFusedOp): + """``out = relu(input) ** 2`` (single-input, not gated). + + The real kernel is the CUDA JIT path (``sglang.jit_kernel.activation.relu2``, + used in production on CUDA); elsewhere the torch reference runs. + """ + + op = "activation.relu2" + priority = (KernelBackend.JIT, KernelBackend.TORCH) + capabilities = {KernelBackend.JIT: _CUDA} + format_signature = FormatSignature( + supported_dtypes=_ACT_DTYPES, + description="relu(x) ** 2; returns tensor", + ) + descriptions = { + KernelBackend.JIT: "relu(x)**2 (sglang.jit_kernel).", + KernelBackend.TORCH: "relu(x)**2 (pure-torch reference).", + } + + def forward_native( + self, input: torch.Tensor, out: Optional[torch.Tensor] = None + ) -> torch.Tensor: + import torch.nn.functional as F + + x = F.relu(input) + result = x * x + if out is None: + return result + out.copy_(result) + return out + + def forward_jit( + self, input: torch.Tensor, out: Optional[torch.Tensor] = None + ) -> torch.Tensor: + from sglang.jit_kernel.activation import relu2 + + result = relu2(input) + if out is None: + return result + out.copy_(result) + return out + + +class QuickGELUOp(BaseFusedOp): + """``out = input * sigmoid(1.702 * input)`` (single-input, not gated). + + Only ROCm has a native kernel (``sgl_kernel.gelu_quick``, AOT on HIP); CUDA + uses the torch reference in production, so no CUDA backend is registered. + """ + + op = "activation.gelu_quick" + priority = (KernelBackend.AOT, KernelBackend.TORCH) + capabilities = {KernelBackend.AOT: _HIP} + format_signature = FormatSignature( + supported_dtypes=_ACT_DTYPES, + description="x * sigmoid(1.702 * x); returns tensor", + ) + descriptions = { + KernelBackend.AOT: "quick GELU (sgl_kernel wheel, ROCm).", + KernelBackend.TORCH: "quick GELU (pure-torch reference).", + } + + def forward_native( + self, input: torch.Tensor, out: Optional[torch.Tensor] = None + ) -> torch.Tensor: + import torch + + result = input * torch.sigmoid(1.702 * input) + if out is None: + return result + out.copy_(result) + return out + + def forward_aot( + self, input: torch.Tensor, out: Optional[torch.Tensor] = None + ) -> torch.Tensor: + import torch + from sgl_kernel import gelu_quick + + if out is None: + out = torch.empty(input.shape, dtype=input.dtype, device=input.device) + gelu_quick(input, out) + return out + + _SILU_AND_MUL = register_fused_op(SiluAndMulOp(), __name__, "_SILU_AND_MUL") _GELU_AND_MUL = register_fused_op(GeluAndMulOp(), __name__, "_GELU_AND_MUL") _GELU_TANH_AND_MUL = register_fused_op( GeluTanhAndMulOp(), __name__, "_GELU_TANH_AND_MUL" ) +_RELU2 = register_fused_op(ReLU2Op(), __name__, "_RELU2") +_GELU_QUICK = register_fused_op(QuickGELUOp(), __name__, "_GELU_QUICK") def silu_and_mul( @@ -208,13 +295,27 @@ def gelu_tanh_and_mul( return _GELU_TANH_AND_MUL(input, out) +def relu2(input: torch.Tensor, out: Optional[torch.Tensor] = None) -> torch.Tensor: + """``out = relu(input) ** 2``.""" + return _RELU2(input, out) + + +def gelu_quick(input: torch.Tensor, out: Optional[torch.Tensor] = None) -> torch.Tensor: + """``out = input * sigmoid(1.702 * input)``.""" + return _GELU_QUICK(input, out) + + __all__ = [ "SiluAndMulOp", "GeluAndMulOp", "GeluTanhAndMulOp", + "ReLU2Op", + "QuickGELUOp", "silu_and_mul", "gelu_and_mul", "gelu_tanh_and_mul", + "relu2", + "gelu_quick", ] diff --git a/python/sglang/kernels/ops/layernorm/__init__.py b/python/sglang/kernels/ops/layernorm/__init__.py index 8de0679d6..9f9953114 100644 --- a/python/sglang/kernels/ops/layernorm/__init__.py +++ b/python/sglang/kernels/ops/layernorm/__init__.py @@ -1,10 +1,11 @@ """Layer-normalization kernels. Each operator is a :class:`~sglang.kernels.fused_op.BaseFusedOp` with a -pure-``torch`` reference (``forward_native``) plus optimized CUDA backends, +pure-``torch`` reference (``forward_native``) plus optimized per-device backends, all behind one signature. The public module-level functions are thin wrappers -over module-level instances; auto-selection prefers the AOT ``sgl_kernel`` -implementation on CUDA and falls back to the native reference elsewhere. +over module-level instances; auto-selection follows the production default for +the live device: AOT ``sgl_kernel`` on CUDA, ``aiter`` (or rocm-triton for +gemma) on ROCm, ``torch_npu`` on Ascend, native reference otherwise. Pick a specific backend with e.g. ``_RMSNORM.forward(x, w, backend=KernelBackend.JIT)`` or globally via ``SGLANG_FORCE_FUSED_OP_BACKEND``. @@ -26,9 +27,21 @@ if TYPE_CHECKING: _NORM_DTYPES = ("float16", "bfloat16") _CUDA = frozenset({CapabilityRequirement.CUDA}) +_HIP = frozenset({CapabilityRequirement.HIP}) +_NPU = frozenset({CapabilityRequirement.NPU}) +# Unlike the gated-activation ops, sgl_kernel does *not* build the rmsnorm ops +# for ROCm (production: ``if _is_cuda or _is_xpu or _is_musa: from sgl_kernel +# import rmsnorm`` — HIP is absent), so AOT here is CUDA-only. ROCm instead has +# an ``aiter`` path, and Ascend a ``torch_npu`` path — a clean illustration that +# the same ``AOT`` provenance covers different devices per op. +# Priority (best -> fallback) is device-agnostic; per-op CapabilityRequirement +# decides eligibility, so on CUDA this resolves to AOT, on HIP to AITER, on NPU +# to TORCH_NPU, each matching the production default for that device. _NORM_PRIORITY = ( KernelBackend.AOT, KernelBackend.JIT, + KernelBackend.AITER, + KernelBackend.TORCH_NPU, KernelBackend.TORCH, ) @@ -44,6 +57,8 @@ class RMSNormOp(BaseFusedOp): capabilities = { KernelBackend.AOT: _CUDA, KernelBackend.JIT: _CUDA, + KernelBackend.AITER: _HIP, + KernelBackend.TORCH_NPU: _NPU, } format_signature = FormatSignature( supported_dtypes=_NORM_DTYPES, @@ -52,6 +67,8 @@ class RMSNormOp(BaseFusedOp): descriptions = { KernelBackend.AOT: "RMS normalization (sgl_kernel wheel).", KernelBackend.JIT: "RMS normalization (sglang.jit_kernel).", + KernelBackend.AITER: "RMS normalization (aiter rmsnorm2d_fwd, ROCm).", + KernelBackend.TORCH_NPU: "RMS normalization (torch_npu, Ascend).", KernelBackend.TORCH: "RMS normalization (pure-torch reference).", } @@ -103,6 +120,40 @@ class RMSNormOp(BaseFusedOp): jit_rmsnorm(input, weight, out, eps) return out + def forward_aiter( + self, + input: torch.Tensor, + weight: torch.Tensor, + eps: float = 1e-6, + out: Optional[torch.Tensor] = None, + enable_pdl: Optional[bool] = None, + ) -> torch.Tensor: + import torch + from aiter import rmsnorm2d_fwd + + # Mirrors production srt/layers/layernorm.py: rmsnorm2d_fwd(out, x, w, eps) + # writes the normalized result in-place into ``out`` (ROCm path). + if out is None: + out = torch.empty_like(input) + rmsnorm2d_fwd(out, input, weight, eps) + return out + + def forward_npu( + self, + input: torch.Tensor, + weight: torch.Tensor, + eps: float = 1e-6, + out: Optional[torch.Tensor] = None, + enable_pdl: Optional[bool] = None, + ) -> torch.Tensor: + import torch_npu + + result = torch_npu.npu_rms_norm(input, weight, eps)[0] + if out is None: + return result + out.copy_(result) + return out + class FusedAddRMSNormOp(BaseFusedOp): """In-place ``residual += input; input = RMSNorm(residual) * weight``. @@ -116,6 +167,8 @@ class FusedAddRMSNormOp(BaseFusedOp): capabilities = { KernelBackend.AOT: _CUDA, KernelBackend.JIT: _CUDA, + KernelBackend.AITER: _HIP, + KernelBackend.TORCH_NPU: _NPU, } format_signature = FormatSignature( supported_dtypes=_NORM_DTYPES, @@ -129,6 +182,10 @@ class FusedAddRMSNormOp(BaseFusedOp): KernelBackend.JIT: ( "Fused residual-add + RMS normalization (sglang.jit_kernel)." ), + KernelBackend.AITER: ("Fused residual-add + RMS normalization (aiter, ROCm)."), + KernelBackend.TORCH_NPU: ( + "Fused residual-add + RMS normalization (torch_npu, Ascend)." + ), KernelBackend.TORCH: ( "Fused residual-add + RMS normalization (pure-torch reference)." ), @@ -174,19 +231,65 @@ class FusedAddRMSNormOp(BaseFusedOp): return jit_fused_add_rmsnorm(input, residual, weight, eps) + def forward_aiter( + self, + input: torch.Tensor, + residual: torch.Tensor, + weight: torch.Tensor, + eps: float = 1e-6, + enable_pdl: Optional[bool] = None, + ) -> None: + import torch + from aiter import rmsnorm2d_fwd_with_add + + # aiter writes the normalized value and the new residual into separate + # out buffers (production call order: out, x, residual_out, residual, w, + # eps); copy them back to honor this op's in-place contract. + out = torch.empty_like(input) + residual_out = torch.empty_like(residual) + rmsnorm2d_fwd_with_add(out, input, residual_out, residual, weight, eps) + input.copy_(out) + residual.copy_(residual_out) + + def forward_npu( + self, + input: torch.Tensor, + residual: torch.Tensor, + weight: torch.Tensor, + eps: float = 1e-6, + enable_pdl: Optional[bool] = None, + ) -> None: + import torch_npu + + # torch_npu.npu_add_rms_norm(residual, x, w, eps) -> (normed, _, new_sum) + out, _, residual_out = torch_npu.npu_add_rms_norm(residual, input, weight, eps) + input.copy_(out) + residual.copy_(residual_out) + class GemmaRMSNormOp(BaseFusedOp): """``out = (input / RMS(input)) * (weight + 1)``; returns a tensor.""" op = "layernorm.gemma_rmsnorm" priority = _NORM_PRIORITY - capabilities = {KernelBackend.AOT: _CUDA} + # AOT (sgl_kernel) on CUDA; JIT is the ROCm rocm-triton path + # (sglang.jit_kernel.minimax_m3) — a JIT provenance pinned to HIP, distinct + # from the CUDA-only JIT on the plain rmsnorm ops; torch_npu on Ascend. + capabilities = { + KernelBackend.AOT: _CUDA, + KernelBackend.JIT: _HIP, + KernelBackend.TORCH_NPU: _NPU, + } format_signature = FormatSignature( supported_dtypes=_NORM_DTYPES, description="out = (x / RMS(x)) * (weight + 1); returns tensor", ) descriptions = { KernelBackend.AOT: "Gemma-style RMS normalization (sgl_kernel wheel).", + KernelBackend.JIT: ( + "Gemma-style RMS normalization (rocm-triton, sglang.jit_kernel)." + ), + KernelBackend.TORCH_NPU: ("Gemma-style RMS normalization (torch_npu, Ascend)."), KernelBackend.TORCH: "Gemma-style RMS normalization (pure-torch reference).", } @@ -221,13 +324,53 @@ class GemmaRMSNormOp(BaseFusedOp): return sgl_kernel.gemma_rmsnorm(input, weight, eps, out, enable_pdl) + def forward_jit( + self, + input: torch.Tensor, + weight: torch.Tensor, + eps: float = 1e-6, + out: Optional[torch.Tensor] = None, + enable_pdl: Optional[bool] = None, + ) -> torch.Tensor: + from sglang.jit_kernel.minimax_m3.rmsnorm import ( + gemma_rmsnorm as rocm_triton_gemma_rmsnorm, + ) + + result = rocm_triton_gemma_rmsnorm(input, weight, eps) + if out is None: + return result + out.copy_(result) + return out + + def forward_npu( + self, + input: torch.Tensor, + weight: torch.Tensor, + eps: float = 1e-6, + out: Optional[torch.Tensor] = None, + enable_pdl: Optional[bool] = None, + ) -> torch.Tensor: + import torch_npu + + result = torch_npu.npu_gemma_rms_norm(input, weight, eps)[0] + if out is None: + return result + out.copy_(result) + return out + class GemmaFusedAddRMSNormOp(BaseFusedOp): """In-place ``residual += input; input = GemmaRMSNorm(residual) * (weight + 1)``.""" op = "layernorm.gemma_fused_add_rmsnorm" priority = _NORM_PRIORITY - capabilities = {KernelBackend.AOT: _CUDA} + # AOT (sgl_kernel) on CUDA; JIT is the ROCm rocm-triton path on HIP. + # NPU here would use ``sgl_kernel_npu.add_gemma_rms_norm`` (a distinct AOT-npu + # wheel provenance, not torch_npu) — deferred until that provenance lands. + capabilities = { + KernelBackend.AOT: _CUDA, + KernelBackend.JIT: _HIP, + } format_signature = FormatSignature( supported_dtypes=_NORM_DTYPES, in_place=True, @@ -235,6 +378,10 @@ class GemmaFusedAddRMSNormOp(BaseFusedOp): ) descriptions = { KernelBackend.AOT: ("Gemma-style fused residual-add + RMS normalization."), + KernelBackend.JIT: ( + "Gemma-style fused residual-add + RMS normalization " + "(rocm-triton, sglang.jit_kernel)." + ), KernelBackend.TORCH: ( "Gemma-style fused residual-add + RMS normalization " "(pure-torch reference)." @@ -271,6 +418,25 @@ class GemmaFusedAddRMSNormOp(BaseFusedOp): input, residual, weight, eps, enable_pdl ) + def forward_jit( + self, + input: torch.Tensor, + residual: torch.Tensor, + weight: torch.Tensor, + eps: float = 1e-6, + enable_pdl: Optional[bool] = None, + ) -> None: + from sglang.jit_kernel.minimax_m3.rmsnorm import ( + gemma_fused_add_rmsnorm as rocm_triton_gemma_fused_add_rmsnorm, + ) + + # rocm-triton returns (normed, new_residual); honor the in-place contract. + norm_out, residual_out = rocm_triton_gemma_fused_add_rmsnorm( + input, residual, weight, eps + ) + input.copy_(norm_out) + residual.copy_(residual_out) + _RMSNORM = register_fused_op(RMSNormOp(), __name__, "_RMSNORM") _FUSED_ADD_RMSNORM = register_fused_op( diff --git a/python/sglang/kernels/spec.py b/python/sglang/kernels/spec.py index 0ca15ab15..1db35d2ec 100644 --- a/python/sglang/kernels/spec.py +++ b/python/sglang/kernels/spec.py @@ -33,7 +33,8 @@ class KernelBackend(str, Enum): (the ``sgl_kernel`` wheel, built for CUDA *and* ROCm) are both cross-device; which devices a given op supports is expressed by its :class:`CapabilityRequirement` list. Platform-specific libraries (e.g. - ``aiter`` on AMD) are just additional provenance values. + ``aiter`` on AMD, ``torch_npu`` on Ascend) are just additional provenance + values, each pinned to its device by its ``CapabilityRequirement``. """ TORCH = "torch" # pure-torch reference (forward_native) @@ -45,7 +46,8 @@ class KernelBackend(str, Enum): FLASHINFER = "flashinfer" DEEPGEMM = "deepgemm" AITER = "aiter" # AMD aiter library (device=HIP) - # TODO(RFC #29630): more provenance as needed (npu / cpu-avx, ...) + TORCH_NPU = "torch_npu" # Ascend NPU vendor runtime (device=NPU) + # TODO(RFC #29630): more provenance as needed (cpu-avx, sgl_kernel_npu, ...) class DeviceType(str, Enum): @@ -53,8 +55,9 @@ class DeviceType(str, Enum): CUDA = "cuda" HIP = "hip" + NPU = "npu" # Ascend NPU (torch_npu / sgl_kernel_npu) CPU = "cpu" - # TODO(RFC #29630): NPU / XPU / ... as backends land. + # TODO(RFC #29630): XPU / MUSA / ... as backends land. class PlatformInfo(msgspec.Struct, frozen=True): @@ -98,6 +101,9 @@ class PlatformInfo(msgspec.Struct, frozen=True): try: if torch.version.hip is not None and torch.cuda.is_available(): return cls(device_type="hip") + npu = getattr(torch, "npu", None) + if npu is not None and npu.is_available(): + return cls(device_type="npu") if torch.cuda.is_available(): major, minor = torch.cuda.get_device_capability() return cls( @@ -123,7 +129,7 @@ class CapabilityRequirement(msgspec.Struct, frozen=True): (``min_cuda_arch`` / ``max_cuda_arch`` apply only when ``device == CUDA``). The device-only cases are so common that they are exposed as class constants - (``CapabilityRequirement.CUDA`` / ``.HIP``); use :meth:`cuda` for an + (``CapabilityRequirement.CUDA`` / ``.HIP`` / ``.NPU``); use :meth:`cuda` for an arch-bounded CUDA requirement (e.g. ``CapabilityRequirement.cuda( min_sm=(10, 0))`` for SM100+). """ @@ -136,6 +142,7 @@ class CapabilityRequirement(msgspec.Struct, frozen=True): # instances of the class itself). ClassVar keeps them out of msgspec fields. CUDA: ClassVar[CapabilityRequirement] HIP: ClassVar[CapabilityRequirement] + NPU: ClassVar[CapabilityRequirement] @classmethod def cuda( @@ -160,6 +167,7 @@ class CapabilityRequirement(msgspec.Struct, frozen=True): CapabilityRequirement.CUDA = CapabilityRequirement(device=DeviceType.CUDA) CapabilityRequirement.HIP = CapabilityRequirement(device=DeviceType.HIP) +CapabilityRequirement.NPU = CapabilityRequirement(device=DeviceType.NPU) def capabilities_satisfied( diff --git a/test/registered/kernels/test_fused_op.py b/test/registered/kernels/test_fused_op.py index 3f47735d7..b0729d1a2 100644 --- a/test/registered/kernels/test_fused_op.py +++ b/test/registered/kernels/test_fused_op.py @@ -152,6 +152,8 @@ class TestBaseFusedOp(unittest.TestCase): KernelBackend.TORCH_COMPILE, KernelBackend.JIT, KernelBackend.AOT, + KernelBackend.AITER, + KernelBackend.TORCH_NPU, }, ) # Dotted targets resolve to the bound backend methods. diff --git a/test/registered/kernels/test_kernels_namespace.py b/test/registered/kernels/test_kernels_namespace.py index d81b727e4..acc516a31 100644 --- a/test/registered/kernels/test_kernels_namespace.py +++ b/test/registered/kernels/test_kernels_namespace.py @@ -28,15 +28,37 @@ EXPECTED_OPS = { "torch", "torch_compile", }, - "layernorm.rmsnorm": {"aot", "jit", "torch", "torch_compile"}, + "activation.relu2": {"jit", "torch", "torch_compile"}, + "activation.gelu_quick": {"aot", "torch", "torch_compile"}, + "layernorm.rmsnorm": { + "aot", + "jit", + "aiter", + "torch_npu", + "torch", + "torch_compile", + }, "layernorm.fused_add_rmsnorm": { + "aot", + "jit", + "aiter", + "torch_npu", + "torch", + "torch_compile", + }, + "layernorm.gemma_rmsnorm": { + "aot", + "jit", + "torch_npu", + "torch", + "torch_compile", + }, + "layernorm.gemma_fused_add_rmsnorm": { "aot", "jit", "torch", "torch_compile", }, - "layernorm.gemma_rmsnorm": {"aot", "torch", "torch_compile"}, - "layernorm.gemma_fused_add_rmsnorm": {"aot", "torch", "torch_compile"}, # curated dual/single-backend wrapper ops "gemm.fp8_scaled_mm": {"aot"}, "gemm.dsv3_fused_a_gemm": {"aot", "jit"}, @@ -249,6 +271,41 @@ class TestKernelsNamespace(unittest.TestCase): finally: fo._platform = saved + def test_layernorm_cross_device_coverage(self): + # The rmsnorm ops illustrate that the *same* provenance covers different + # devices per op: AOT (sgl_kernel) is CUDA-only here (sgl_kernel does not + # build rmsnorm for ROCm), so HIP falls to AITER and NPU to torch_npu, + # each matching the production default for that device. gemma uses a + # rocm-triton JIT path on HIP -- a JIT provenance pinned to HIP, unlike + # the CUDA-only JIT on plain rmsnorm. + import sglang.kernels.fused_op as fo + from sglang.kernels.ops.layernorm import ( + _FUSED_ADD_RMSNORM, + _GEMMA_RMSNORM, + _RMSNORM, + ) + + B = self.K.KernelBackend + cuda = self.K.PlatformInfo(device_type="cuda", cuda_arch_major=9) + hip = self.K.PlatformInfo(device_type="hip") + npu = self.K.PlatformInfo(device_type="npu") + saved = fo._platform + try: + for plat, expect in ((cuda, B.AOT), (hip, B.AITER), (npu, B.TORCH_NPU)): + fo._platform = lambda p=plat: p + self.assertEqual(_RMSNORM._resolve_backend(), expect) + self.assertEqual(_FUSED_ADD_RMSNORM._resolve_backend(), expect) + # gemma: AOT on CUDA, rocm-triton JIT on HIP, torch_npu on NPU. + for plat, expect in ((cuda, B.AOT), (hip, B.JIT), (npu, B.TORCH_NPU)): + fo._platform = lambda p=plat: p + self.assertEqual(_GEMMA_RMSNORM._resolve_backend(), expect) + # AOT rmsnorm is CUDA-only (not HIP) -- distinct from activation's AOT. + fo._platform = lambda: hip + self.assertFalse(_RMSNORM.backend_eligible(B.AOT)) + self.assertTrue(_RMSNORM.backend_eligible(B.AITER)) + finally: + fo._platform = saved + def test_selector_explicit_backend(self): spec = self.K.select_kernel( "layernorm.rmsnorm", backend=self.K.KernelBackend.JIT @@ -301,6 +358,7 @@ class TestKernelsNamespace(unittest.TestCase): # and dedup, so {CUDA, HIP} == {HIP, CUDA}. self.assertEqual(cap.CUDA, cap(device=dev.CUDA)) self.assertEqual(cap.HIP, cap(device=dev.HIP)) + self.assertEqual(cap.NPU, cap(device=dev.NPU)) self.assertEqual({cap.CUDA, cap.HIP}, {cap.HIP, cap.CUDA}) self.assertEqual(len({cap.CUDA, cap(device=dev.CUDA)}), 1) # cuda(min_sm=...) factory: an SM100+ CUDA requirement.