diff --git a/tokenspeed-kernel/python/tokenspeed_kernel/ops/gemm/__init__.py b/tokenspeed-kernel/python/tokenspeed_kernel/ops/gemm/__init__.py index 735aa7265..14157230c 100644 --- a/tokenspeed-kernel/python/tokenspeed_kernel/ops/gemm/__init__.py +++ b/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, diff --git a/tokenspeed-kernel/python/tokenspeed_kernel/ops/quantization/triton.py b/tokenspeed-kernel/python/tokenspeed_kernel/ops/quantization/triton.py index 43557a48d..3abf56a0f 100644 --- a/tokenspeed-kernel/python/tokenspeed_kernel/ops/quantization/triton.py +++ b/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"}), diff --git a/tokenspeed-kernel/python/tokenspeed_kernel/ops/quantization/trtllm.py b/tokenspeed-kernel/python/tokenspeed_kernel/ops/quantization/trtllm.py index 7b57cc8d4..f14fb87f6 100644 --- a/tokenspeed-kernel/python/tokenspeed_kernel/ops/quantization/trtllm.py +++ b/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"}), diff --git a/tokenspeed-kernel/test/test_kernel_api_selection.py b/tokenspeed-kernel/test/test_kernel_api_selection.py index 60965e838..b223fccb3 100644 --- a/tokenspeed-kernel/test/test_kernel_api_selection.py +++ b/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,