fix(kernel): use Triton FP8 quantization on SM120

tokenspeed-kernel/python/tokenspeed_kernel/ops/gemm/__init__.py

@@ -30,7 +30,7 @@
 import tokenspeed_kernel.ops.gemm.triton  # noqa: F401
 import tokenspeed_kernel.ops.gemm.trtllm  # noqa: F401
 import torch
-from tokenspeed_kernel.platform import Platform
+from tokenspeed_kernel.platform import ArchVersion, Platform
 from tokenspeed_kernel.profiling import ShapeCapture, kernel_scope
 from tokenspeed_kernel.selection import select_kernel
 from tokenspeed_kernel.signature import (
@@ -163,6 +163,21 @@ def _online_quantize_mxfp8(
     """
     block_k = block_size[1]
 
+    if (
+        kernel_name in {"flashinfer_mm_fp8_blockscale", "triton_mm_fp8_blockscale"}
+        and _platform.is_nvidia
+        and _platform.arch_version == ArchVersion(12, 0)
+    ):
+        from tokenspeed_kernel.ops.quantization import quantize_fp8_with_scale
+
+        return quantize_fp8_with_scale(
+            A,
+            granularity="token_group",
+            group_size=block_k,
+            scale_encoding="float32",
+            solution="triton",
+        )
+
     def ensure_row_major_scales(
         qA: torch.Tensor,
         A_scales: torch.Tensor,

tokenspeed-kernel/python/tokenspeed_kernel/ops/quantization/triton.py

@@ -294,7 +294,7 @@ def triton_quantize_fp8(
     "fp8_with_scale",
     name="triton_quantize_fp8_with_scale",
     solution="triton",
-    capability=CapabilityRequirement(vendors=frozenset({"amd"})),
+    capability=CapabilityRequirement(vendors=frozenset({"amd", "nvidia"})),
     signatures=format_signatures("x", "dense", {torch.bfloat16, torch.float16}),
     traits={
         "granularity": frozenset({"token_group_128"}),

tokenspeed-kernel/python/tokenspeed_kernel/ops/quantization/trtllm.py

@@ -20,7 +20,11 @@
 from __future__ import annotations
 
 import torch
-from tokenspeed_kernel.platform import current_platform
+from tokenspeed_kernel.platform import (
+    ArchVersion,
+    CapabilityRequirement,
+    current_platform,
+)
 from tokenspeed_kernel.registry import Priority, error_fn, register_kernel
 from tokenspeed_kernel.signature import format_signatures
 
@@ -64,6 +68,10 @@ def trtllm_fp8_tensor(x: torch.Tensor) -> torch.Tensor:
         "fp8_with_scale",
         name="trtllm_quantize_fp8_with_scale",
         solution="trtllm",
+        capability=CapabilityRequirement(
+            max_arch_version=ArchVersion(10, 9),
+            vendors=frozenset({"nvidia"}),
+        ),
         signatures=format_signatures("x", "dense", {torch.bfloat16, torch.float16}),
         traits={
             "granularity": frozenset({"tensor", "token", "token_group_128"}),

tokenspeed-kernel/test/test_kernel_api_selection.py

@@ -47,6 +47,10 @@
 import tokenspeed_kernel.ops.moe.flashinfer as _moe_flashinfer
 import tokenspeed_kernel.ops.moe.gluon as _moe_gluon
 import tokenspeed_kernel.ops.moe.triton as _moe_triton
+import tokenspeed_kernel.ops.quantization as _quantization_pkg
+import tokenspeed_kernel.ops.quantization.flashinfer as _quantization_flashinfer
+import tokenspeed_kernel.ops.quantization.triton as _quantization_triton
+import tokenspeed_kernel.ops.quantization.trtllm as _quantization_trtllm
 import tokenspeed_kernel.ops.sampling as _sampling_pkg
 import tokenspeed_kernel.ops.sampling.cute_dsl as _sampling_cute_dsl
 import tokenspeed_kernel.ops.sampling.gluon as _sampling_gluon
@@ -135,6 +139,11 @@
     _moe_triton_mxfp4,
     _moe_triton,
     _moe_pkg,
+    # Quantization registration modules.
+    _quantization_flashinfer,
+    _quantization_triton,
+    _quantization_trtllm,
+    _quantization_pkg,
     # Sampling registration modules.
     _sampling_cute_dsl,
     _sampling_gluon,