chore: import upstream snapshot with attribution

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

@@ -0,0 +1,376 @@
+# Copyright (c) 2026 LightSeek Foundation
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+from __future__ import annotations
+
+from typing import Literal
+
+import torch
+from tokenspeed_kernel.profiling import ShapeCapture, kernel_scope
+from tokenspeed_kernel.selection import select_kernel
+from tokenspeed_kernel.signature import dense_tensor_format, format_signature
+
+__all__ = [
+    "quantize_fp8",
+    "quantize_fp8_with_scale",
+    "quantize_mxfp8",
+    "quantize_nvfp4",
+    "quantize_mxfp4",
+]
+
+
+def quantize_fp8(
+    x: torch.Tensor,
+    scale: float | torch.Tensor | None = None,
+    # kernel options
+    enable_pdl: bool = False,
+    # dispatch options
+    override: str | None = None,
+    solution: str | None = None,
+) -> torch.Tensor:
+    """Quantize x to same-shape FP8 with an optional scalar scale.
+
+    This API covers static activation-scale cases such as GPT-OSS MXFP4 MoE
+    W4A8 input quantization and plain FP8 casts. If scale is provided, it is
+    the actual quantization scale, so the backend computes x / scale before
+    casting. If scale is omitted, the backend performs a pure FP8 cast.
+
+    Args:
+        x: Input tensor.
+        scale: Optional scalar scale, as a Python value or scalar tensor. If
+            provided, the backend computes x / scale before casting. If omitted,
+            the backend performs a pure FP8 cast.
+        enable_pdl: Whether to request Programmatic Dependent Launch support.
+        override: Optional exact kernel name or solution override.
+        solution: Optional registered solution to select.
+
+    Returns:
+        Quantized FP8 tensor with the same shape as x.
+
+    Non-scalar scales belong in dynamic quantization APIs and should be
+    on-device tensors.
+
+    """
+
+    traits = {
+        "has_scale": scale is not None,
+    }
+    signature = format_signature(x=dense_tensor_format(x.dtype))
+    kernel = select_kernel(
+        "quantization",
+        "fp8",
+        signature,
+        traits=traits,
+        solution=solution,
+        override=override,
+    )
+    shape_params = {
+        "shape": tuple(x.shape),
+        "has_scale": scale is not None,
+    }
+    ShapeCapture.get().record("quantization", "fp8", kernel.name, x.dtype, shape_params)
+    with kernel_scope(
+        "quantization", "fp8", x.dtype, kernel_name=kernel.name, **shape_params
+    ):
+        return kernel(
+            x,
+            scale=scale,
+            enable_pdl=enable_pdl,
+        )
+
+
+def quantize_fp8_with_scale(
+    x: torch.Tensor,
+    # quantization options
+    granularity: Literal["tensor", "token", "token_group"] = "tensor",
+    group_size: int | None = None,
+    scale_encoding: Literal["float32", "ue8m0", "packed_ue8m0"] = "float32",
+    # kernel options
+    enable_pdl: bool = False,
+    # dispatch options
+    override: str | None = None,
+    solution: str | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Quantize x to FP8 while dynamically computing scales.
+
+    Use granularity="tensor" for one scale over the whole tensor,
+    granularity="token" for one scale per row/token, and
+    granularity="token_group" for one scale per row/token and contiguous group
+    along the last dimension.
+
+    Args:
+        x: Input tensor.
+        granularity: Scale granularity. Supported values are tensor, token, and
+            token_group.
+        group_size: Number of contiguous values per scale group along the last
+            dimension. Required for token_group granularity.
+        scale_encoding: Scale encoding for token_group granularity, such as
+            float32, ue8m0, or packed_ue8m0.
+        enable_pdl: Whether to request Programmatic Dependent Launch support.
+        override: Optional exact kernel name or solution override.
+        solution: Optional registered solution to select.
+
+    Returns:
+        Tuple of quantized FP8 tensor and scale tensor.
+
+    The expected scale shapes are [1] for tensor granularity, [M, 1] for token
+    granularity, and [M, ceil(K / group_size)] or a backend-specific layout for
+    token_group granularity, where M = x.reshape(-1, x.shape[-1]).shape[0].
+    Returned scales use float32 dtype for scale_encoding="float32" and a
+    backend-specific encoded integer dtype for non-float encodings such as
+    "ue8m0".
+    """
+
+    if granularity not in {"tensor", "token", "token_group"}:
+        raise ValueError(f"unsupported FP8 dynamic granularity: {granularity!r}")
+    if granularity == "token_group":
+        if group_size is None or group_size <= 0:
+            raise ValueError(
+                f"token_group granularity requires positive group_size, got {group_size}"
+            )
+        granularity_trait = f"token_group_{group_size}"
+    else:
+        granularity_trait = granularity
+    traits = {
+        "granularity": granularity_trait,
+        "scale_encoding": scale_encoding,
+    }
+    signature = format_signature(x=dense_tensor_format(x.dtype))
+    kernel = select_kernel(
+        "quantization",
+        "fp8_with_scale",
+        signature,
+        traits=traits,
+        solution=solution,
+        override=override,
+    )
+    shape_params = {
+        "shape": tuple(x.shape),
+        "granularity": granularity_trait,
+        "group_size": group_size,
+        "scale_encoding": scale_encoding,
+    }
+    ShapeCapture.get().record(
+        "quantization", "fp8_with_scale", kernel.name, x.dtype, shape_params
+    )
+    with kernel_scope(
+        "quantization",
+        "fp8_with_scale",
+        x.dtype,
+        kernel_name=kernel.name,
+        **shape_params,
+    ):
+        return kernel(
+            x,
+            granularity=granularity,
+            group_size=group_size,
+            scale_encoding=scale_encoding,
+            enable_pdl=enable_pdl,
+        )
+
+
+def quantize_mxfp8(
+    x: torch.Tensor,
+    # kernel options
+    enable_pdl: bool = False,
+    # dispatch options
+    override: str | None = None,
+    solution: str | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Quantize x to MXFP8 format.
+
+    MXFP8 uses FP8 data plus encoded vector scales, commonly one scale per 32
+    values along the last dimension.
+
+    Args:
+        x: Input tensor.
+        enable_pdl: Whether to request Programmatic Dependent Launch support.
+        override: Optional exact kernel name or solution override.
+        solution: Optional registered solution to select.
+
+    Returns:
+        Tuple of quantized MXFP8 tensor and encoded scale tensor.
+
+    """
+
+    traits = {}
+    signature = format_signature(x=dense_tensor_format(x.dtype))
+    kernel = select_kernel(
+        "quantization",
+        "mxfp8",
+        signature,
+        traits=traits,
+        solution=solution,
+        override=override,
+    )
+    shape_params = {
+        "shape": tuple(x.shape),
+    }
+    ShapeCapture.get().record(
+        "quantization", "mxfp8", kernel.name, x.dtype, shape_params
+    )
+    with kernel_scope(
+        "quantization", "mxfp8", x.dtype, kernel_name=kernel.name, **shape_params
+    ):
+        return kernel(
+            x,
+            enable_pdl=enable_pdl,
+        )
+
+
+def quantize_nvfp4(
+    x: torch.Tensor,
+    scale: float | torch.Tensor | None = None,
+    # quantization options
+    scale_layout: Literal["linear", "swizzled"] = "swizzled",
+    # kernel options
+    enable_pdl: bool = False,
+    # dispatch options
+    override: str | None = None,
+    solution: str | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Quantize x to packed NVFP4.
+
+    NVFP4 uses packed E2M1x2 data with one E4M3 scale factor per 16 values.
+    The quantized output is usually shaped [M, K/2].
+
+    scale is the actual input scale. Backend adapters should handle any
+    backend-specific inverse-scale convention internally.
+
+    Args:
+        x: Input tensor.
+        scale: Optional scalar input scale.
+        scale_layout: Scale-factor layout. "linear" returns unswizzled scales;
+            "swizzled" requests the backend-specific layout used by FP4 GEMM.
+        enable_pdl: Whether to request Programmatic Dependent Launch support.
+        override: Optional exact kernel name or solution override.
+        solution: Optional registered solution to select.
+
+    Returns:
+        Tuple of packed NVFP4 tensor and scale-factor tensor.
+
+    """
+
+    traits = {
+        "scale_layout": scale_layout,
+        "has_scale": scale is not None,
+    }
+    signature = format_signature(x=dense_tensor_format(x.dtype))
+    kernel = select_kernel(
+        "quantization",
+        "nvfp4",
+        signature,
+        traits=traits,
+        solution=solution,
+        override=override,
+    )
+    shape_params = {
+        "shape": tuple(x.shape),
+        "scale_layout": scale_layout,
+        "has_scale": scale is not None,
+    }
+    ShapeCapture.get().record(
+        "quantization", "nvfp4", kernel.name, x.dtype, shape_params
+    )
+    with kernel_scope(
+        "quantization", "nvfp4", x.dtype, kernel_name=kernel.name, **shape_params
+    ):
+        return kernel(
+            x,
+            scale=scale,
+            scale_layout=scale_layout,
+            enable_pdl=enable_pdl,
+        )
+
+
+def quantize_mxfp4(
+    x: torch.Tensor,
+    # quantization options
+    global_scale: float | None = None,
+    scale_size: int = 32,
+    scale_layout: Literal["linear", "128x4", "8x4"] = "128x4",
+    # kernel options
+    enable_pdl: bool = False,
+    # dispatch options
+    override: str | None = None,
+    solution: str | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Quantize x to packed MXFP4.
+
+    MXFP4 uses packed E2M1x2 data with one UE8M0 scale factor per scale_size
+    values, usually scale_size=32. The quantized output is usually shaped
+    [M, K/2].
+
+    global_scale is optional because some backends compute the global scale
+    internally from the input. If provided, it is the actual global scale.
+
+    Args:
+        x: Input tensor.
+        global_scale: Optional scalar global scale.
+        scale_size: Number of values per scale-factor vector.
+        scale_layout: Scale-factor layout, such as linear, 128x4, or 8x4.
+        enable_pdl: Whether to request Programmatic Dependent Launch support.
+        override: Optional exact kernel name or solution override.
+        solution: Optional registered solution to select.
+
+    Returns:
+        Tuple of packed MXFP4 tensor and scale-factor tensor.
+
+    """
+
+    if scale_size <= 0:
+        raise ValueError(f"scale_size must be positive, got {scale_size}")
+
+    traits = {
+        "scale_size": scale_size,
+        "scale_layout": scale_layout,
+        "has_global_scale": global_scale is not None,
+        "scale_encoding": "ue8m0",
+    }
+    signature = format_signature(x=dense_tensor_format(x.dtype))
+    kernel = select_kernel(
+        "quantization",
+        "mxfp4",
+        signature,
+        traits=traits,
+        solution=solution,
+        override=override,
+    )
+    shape_params = {
+        "shape": tuple(x.shape),
+        "scale_size": scale_size,
+        "scale_layout": scale_layout,
+        "has_global_scale": global_scale is not None,
+    }
+    ShapeCapture.get().record(
+        "quantization", "mxfp4", kernel.name, x.dtype, shape_params
+    )
+    with kernel_scope(
+        "quantization", "mxfp4", x.dtype, kernel_name=kernel.name, **shape_params
+    ):
+        return kernel(
+            x,
+            global_scale=global_scale,
+            scale_size=scale_size,
+            scale_layout=scale_layout,
+            enable_pdl=enable_pdl,
+        )
+
+
+# Backend registration (side-effect imports).
+import tokenspeed_kernel.ops.quantization.flashinfer  # noqa: E402,F401
+import tokenspeed_kernel.ops.quantization.triton  # noqa: E402,F401
+import tokenspeed_kernel.ops.quantization.trtllm  # noqa: E402,F401

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

@@ -0,0 +1,27 @@
+# Copyright (c) 2026 LightSeek Foundation
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+"""CUDA quantization kernels."""
+
+from tokenspeed_kernel.registry import error_fn
+
+try:
+    from tokenspeed_kernel.thirdparty.cuda.marlin import gptq_marlin_repack
+except ImportError:
+    gptq_marlin_repack = error_fn
+
+__all__ = ["gptq_marlin_repack"]

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

@@ -0,0 +1,105 @@
+# Copyright (c) 2026 LightSeek Foundation
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+from __future__ import annotations
+
+import torch
+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
+
+platform = current_platform()
+
+fp4_quantize = error_fn
+flashinfer_quantize_mxfp8 = error_fn
+flashinfer_quantize_nvfp4 = error_fn
+mxfp8_quantize = error_fn
+nvfp4_block_scale_interleave = error_fn
+fp8_blockscale_quantize_runner_sm90 = error_fn
+
+if platform.is_nvidia:
+    from flashinfer import mxfp8_quantize
+
+    if platform.is_hopper:
+        from flashinfer.gemm.gemm_base import (
+            get_fp8_blockscale_gemm_runner_sm90 as fp8_blockscale_quantize_runner_sm90,
+        )
+
+    @register_kernel(
+        "quantization",
+        "mxfp8",
+        name="flashinfer_quantize_mxfp8",
+        solution="flashinfer",
+        signatures=format_signatures("x", "dense", {torch.bfloat16, torch.float16}),
+        traits={},
+        priority=Priority.PERFORMANT,
+    )
+    def flashinfer_quantize_mxfp8(
+        x: torch.Tensor,
+        enable_pdl: bool = False,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        return mxfp8_quantize(x, False)
+
+
+if platform.is_nvidia and platform.is_blackwell:
+    from flashinfer import (
+        fp4_quantize,
+        nvfp4_block_scale_interleave,
+    )
+
+    @register_kernel(
+        "quantization",
+        "nvfp4",
+        name="flashinfer_quantize_nvfp4",
+        solution="flashinfer",
+        capability=CapabilityRequirement(
+            min_arch_version=ArchVersion(10, 0),
+            vendors=frozenset({"nvidia"}),
+        ),
+        signatures=format_signatures("x", "dense", {torch.bfloat16, torch.float16}),
+        traits={
+            "has_scale": frozenset({True}),
+        },
+        priority=Priority.PERFORMANT,
+    )
+    def flashinfer_quantize_nvfp4(
+        x: torch.Tensor,
+        scale: float | torch.Tensor | None = None,
+        scale_layout: str = "swizzled",
+        enable_pdl: bool = False,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        # The public quantization API uses the actual scale; FlashInfer's FP4
+        # helper expects the inverse scale used before packing.
+        scale_inv = 1.0 / scale
+        return fp4_quantize(
+            x,
+            global_scale=scale_inv,
+            sf_vec_size=16,
+            is_sf_swizzled_layout=scale_layout == "swizzled",
+            enable_pdl=enable_pdl,
+        )
+
+
+__all__ = [
+    "fp4_quantize",
+    "mxfp8_quantize",
+    "nvfp4_block_scale_interleave",
+    "fp8_blockscale_quantize_runner_sm90",
+]

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

@@ -0,0 +1,484 @@
+# Copyright (c) 2026 LightSeek Foundation
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+from typing import Optional
+
+import torch
+from tokenspeed_kernel._triton import tl, triton
+from tokenspeed_kernel.platform import CapabilityRequirement, current_platform
+from tokenspeed_kernel.registry import Priority, register_kernel
+from tokenspeed_kernel.signature import format_signatures
+
+platform = current_platform()
+
+
+@triton.jit
+def _fp8_quantize_kernel(
+    x_ptr,
+    out_ptr,
+    scale,
+    M,
+    N,
+    x_row_stride,
+    out_row_stride,
+    BLOCK_N: tl.constexpr,
+    EVEN_N: tl.constexpr,
+    FP8_DTYPE: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    HAS_SCALE: tl.constexpr,
+    HAS_SCALE_TENSOR: tl.constexpr,
+    ENABLE_PDL: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    m_idx = pid * BLOCK_M + tl.arange(0, BLOCK_M)
+    m_mask = m_idx < M
+    n_idx = tl.arange(0, BLOCK_N)
+
+    # PDL: wait for the producer kernel (e.g., kv_b_proj GEMM) to drain before
+    # we read its output. No-op when disabled.
+    if ENABLE_PDL:
+        tl.extra.cuda.gdc_wait()
+
+    if EVEN_N:
+        load_mask = m_mask[:, None]
+    else:
+        load_mask = m_mask[:, None] & (n_idx[None, :] < N)
+
+    x_off = m_idx[:, None] * x_row_stride + n_idx[None, :]
+    x = tl.load(x_ptr + x_off, mask=load_mask)
+
+    x = x.to(tl.float32)
+    if HAS_SCALE:
+        if HAS_SCALE_TENSOR:
+            scale = tl.load(scale)
+        x = x * (1.0 / scale)
+    x_fp8 = x.to(FP8_DTYPE)
+
+    out_off = m_idx[:, None] * out_row_stride + n_idx[None, :]
+    tl.store(out_ptr + out_off, x_fp8, mask=load_mask)
+
+    # PDL: signal that dependents (e.g., FMHA) can begin their preamble.
+    if ENABLE_PDL:
+        tl.extra.cuda.gdc_launch_dependents()
+
+
+def _flatten_to_2d(x: torch.Tensor):
+    """Flatten leading dims into a single M, returning (M, N, row_stride).
+
+    Requires stride(-1) == 1 and that all leading dims pack contiguously
+    onto the row stride — i.e. ``stride(d) == shape(d+1) * stride(d+1)`` for
+    every ``d < ndim - 2``. This holds for fully-contiguous tensors and for
+    last-dim slices like ``kv[..., qk_nope:]`` where the leading dims still
+    pack onto a uniform row stride.
+    """
+    assert x.stride(-1) == 1, f"expected stride-1 inner dim, got stride={x.stride(-1)}"
+    N = x.shape[-1]
+    if x.ndim == 1:
+        return 1, N, N
+    M = x.numel() // N
+    row_stride = x.stride(-2)
+    # Validate that every leading dim packs onto the next.
+    for d in range(x.ndim - 2):
+        expected = x.shape[d + 1] * x.stride(d + 1)
+        if x.stride(d) != expected:
+            raise ValueError(
+                f"cannot flatten dim {d}: stride={x.stride(d)} but expected "
+                f"shape[{d+1}]*stride[{d+1}]={expected}. Tensor shape={tuple(x.shape)}, "
+                f"stride={tuple(x.stride())}."
+            )
+    return M, N, row_stride
+
+
+def fp8_quantize(
+    x: torch.Tensor,
+    scale: float | torch.Tensor | None = None,
+    out: Optional[torch.Tensor] = None,
+    fp8_dtype: torch.dtype = torch.float8_e4m3fn,
+    enable_pdl: bool = False,
+) -> torch.Tensor:
+    """Cast a BF16/FP16 tensor to FP8 with an optional per-tensor scale.
+
+    Computes ``out = saturate((x / scale) -> fp8)`` element-wise when scale is
+    provided. When scale is omitted, this is a pure FP8 cast.
+
+    Args:
+        x: BF16 or FP16 tensor. Must have stride(-1) == 1; leading dims must
+           pack uniformly onto the row stride (true for contiguous tensors and
+           for last-dim slice views like ``kv[..., qk_nope:]``).
+        scale: optional scalar divisor applied before the cast. Python values
+           are passed as plain kernel args; scalar tensors are loaded on device.
+        out: optional pre-allocated FP8 output. Same shape as ``x``. If not
+           provided, allocated as contiguous.
+        fp8_dtype: ``torch.float8_e4m3fn`` (default), ``torch.float8_e5m2`` or
+           ``torch.float8_e4m3fnuz`` (the bias-8 e4m3 used on AMD CDNA3).
+        enable_pdl: opt into Programmatic Dependent Launch (Hopper+). Caller
+           must also pass ``launch_pdl=True`` upstream / downstream as needed.
+
+    Returns:
+        FP8 tensor with the same shape as ``x``.
+    """
+    assert x.dtype in (
+        torch.bfloat16,
+        torch.float16,
+    ), f"fp8_quantize input must be bf16/fp16, got {x.dtype}"
+    assert fp8_dtype in (
+        torch.float8_e4m3fn,
+        torch.float8_e5m2,
+        torch.float8_e4m3fnuz,
+    ), f"fp8_quantize unsupported fp8 dtype: {fp8_dtype}"
+    has_scale = scale is not None
+    has_scale_tensor = isinstance(scale, torch.Tensor)
+    if has_scale_tensor:
+        assert scale.numel() == 1, "scale tensor must be scalar"
+        scale = scale.contiguous()
+
+    M, N, x_row_stride = _flatten_to_2d(x)
+
+    if out is None:
+        out = torch.empty(x.shape, dtype=fp8_dtype, device=x.device)
+    else:
+        assert out.shape == x.shape and out.dtype == fp8_dtype
+    out_M, _, out_row_stride = _flatten_to_2d(out)
+    assert out_M == M
+
+    if fp8_dtype is torch.float8_e4m3fn:
+        fp8_dtype_const = tl.float8e4nv
+    elif fp8_dtype is torch.float8_e5m2:
+        fp8_dtype_const = tl.float8e5
+    else:
+        fp8_dtype_const = tl.float8e4b8
+
+    # Block-size heuristic — picked from per-shape best configs in an
+    # nsys-driven sweep on B200 (kv_a [s,512] and v [s,h,128] for K2.5).
+    # Pattern: num_warps=4, num_stages=2 win universally; BLOCK_M ramps with
+    # M to amortize launch as the grid grows.
+    # See tasks/k2.5_optimization/{tune_fp8_quantize_nsys,parse_tune_fp8_quantize_nsys}.py
+    if M <= 2048:
+        block_m = 4
+    elif M <= 16384:
+        block_m = 16
+    else:
+        block_m = 32
+    num_warps = 4
+    num_stages = 2
+
+    grid = (triton.cdiv(M, block_m),)
+
+    block_n = max(1, triton.next_power_of_2(N))
+    even_n = block_n == N
+
+    # ``launch_pdl`` is a NVIDIA-only Triton runtime kwarg (Hopper+ Programmatic
+    # Dependent Launch). The HIP backend rejects unknown kwargs, so only forward
+    # it when PDL is actually requested.
+    extra_kwargs = {"launch_pdl": True} if enable_pdl else {}
+
+    _fp8_quantize_kernel[grid](
+        x,
+        out,
+        1.0 if scale is None else scale,
+        M,
+        N,
+        x_row_stride,
+        out_row_stride,
+        BLOCK_N=block_n,
+        EVEN_N=even_n,
+        FP8_DTYPE=fp8_dtype_const,
+        BLOCK_M=block_m,
+        HAS_SCALE=has_scale,
+        HAS_SCALE_TENSOR=has_scale_tensor,
+        ENABLE_PDL=enable_pdl,
+        num_warps=num_warps,
+        num_stages=num_stages,
+        **extra_kwargs,
+    )
+    return out
+
+
+@triton.jit
+def _fp8_token_group_quantize_kernel(
+    x_ptr,
+    out_ptr,
+    scale_ptr,
+    group_size,
+    eps,
+    bit8_min,
+    bit8_max,
+    BLOCK: tl.constexpr,
+):
+    group_id = tl.program_id(0)
+    cols = tl.arange(0, BLOCK)
+    mask = cols < group_size
+    offsets = group_id.to(tl.int64) * group_size + cols
+
+    x = tl.load(x_ptr + offsets, mask=mask, other=0.0).to(tl.float32)
+    scale = tl.maximum(tl.max(tl.abs(x), axis=0), eps) / bit8_max
+    out = tl.clamp(x / scale, bit8_min, bit8_max).to(out_ptr.dtype.element_ty)
+
+    tl.store(out_ptr + offsets, out, mask=mask)
+    tl.store(scale_ptr + group_id, scale)
+
+
+def _fp8_token_group_quantize(
+    x: torch.Tensor,
+    group_size: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if x.shape[-1] % group_size != 0:
+        raise ValueError(
+            f"the last dimension of x must be divisible by group_size, got "
+            f"shape={tuple(x.shape)}, group_size={group_size}"
+        )
+    if not x.is_contiguous():
+        raise ValueError("x must be contiguous")
+
+    out_dtype = platform.fp8e4m3fn.dtype
+    out = torch.empty_like(x, device=x.device, dtype=out_dtype)
+    scales = torch.empty(
+        x.shape[:-1] + (x.shape[-1] // group_size,),
+        device=x.device,
+        dtype=torch.float32,
+    )
+
+    groups = x.numel() // group_size
+    block = triton.next_power_of_2(group_size)
+    num_warps = min(max(block // 256, 1), 8)
+    bit8_max = platform.fp8e4m3fn.max
+    bit8_min = -bit8_max
+
+    _fp8_token_group_quantize_kernel[(groups,)](
+        x,
+        out,
+        scales,
+        group_size,
+        1e-10,
+        bit8_min=bit8_min,
+        bit8_max=bit8_max,
+        BLOCK=block,
+        num_warps=num_warps,
+        num_stages=1,
+    )
+    return out, scales
+
+
+@register_kernel(
+    "quantization",
+    "fp8",
+    name="triton_quantize_fp8",
+    solution="triton",
+    signatures=format_signatures("x", "dense", {torch.bfloat16, torch.float16}),
+    traits={"has_scale": frozenset({True, False})},
+    priority=Priority.PORTABLE,
+)
+def triton_quantize_fp8(
+    x: torch.Tensor,
+    scale: float | torch.Tensor | None = None,
+    enable_pdl: bool = False,
+) -> torch.Tensor:
+    return fp8_quantize(x, scale=scale, enable_pdl=enable_pdl)
+
+
+@register_kernel(
+    "quantization",
+    "fp8_with_scale",
+    name="triton_quantize_fp8_with_scale",
+    solution="triton",
+    capability=CapabilityRequirement(vendors=frozenset({"amd", "nvidia"})),
+    signatures=format_signatures("x", "dense", {torch.bfloat16, torch.float16}),
+    traits={
+        "granularity": frozenset({"token_group_128"}),
+        "scale_encoding": frozenset({"float32"}),
+    },
+    priority=Priority.PORTABLE,
+)
+def triton_quantize_fp8_with_scale(
+    x: torch.Tensor,
+    granularity: str = "tensor",
+    group_size: int | None = None,
+    scale_encoding: str = "float32",
+    enable_pdl: bool = False,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if granularity != "token_group" or group_size != 128:
+        raise ValueError(
+            "triton FP8 dynamic quantization currently supports only "
+            f"granularity='token_group' with group_size=128, got "
+            f"granularity={granularity!r}, group_size={group_size}."
+        )
+    if scale_encoding != "float32":
+        raise ValueError(
+            "triton FP8 dynamic quantization currently requires "
+            f"scale_encoding='float32', got {scale_encoding!r}."
+        )
+    return _fp8_token_group_quantize(x.contiguous(), 128)
+
+
+@triton.jit
+def _mxfp4_quantize_block(x):
+    max_normal: tl.constexpr = 6
+    min_normal: tl.constexpr = 1
+    amax = tl.max(tl.abs(x), axis=0)
+    amax = amax.to(tl.int32, bitcast=True)
+    amax = (amax + 0x200000).to(tl.uint32, bitcast=True) & 0xFF800000
+    amax = amax.to(tl.float32, bitcast=True)
+    scale_e8m0_unbiased = tl.log2(amax).floor() - 2
+    scale_e8m0_unbiased = tl.clamp(scale_e8m0_unbiased, min=-127, max=127)
+    scale_byte = scale_e8m0_unbiased.to(tl.uint8) + 127
+    qx = x * tl.exp2(-scale_e8m0_unbiased)
+    qx = qx.to(tl.uint32, bitcast=True)
+
+    sign = qx & 0x80000000
+    qx = qx ^ sign
+    qx_fp32 = qx.to(tl.float32, bitcast=True)
+    saturate_mask = qx_fp32 >= max_normal
+    denormal_mask = (not saturate_mask) & (qx_fp32 < min_normal)
+    normal_mask = not (saturate_mask | denormal_mask)
+
+    denorm_exp: tl.constexpr = (127 - 1) + (23 - 1) + 1
+    denorm_mask_int: tl.constexpr = denorm_exp << 23
+    denorm_mask_float: tl.constexpr = tl.cast(denorm_mask_int, tl.float32, bitcast=True)
+    denormal_x = qx_fp32 + denorm_mask_float
+    denormal_x = denormal_x.to(tl.uint32, bitcast=True)
+    denormal_x -= denorm_mask_int
+    denormal_x = denormal_x.to(tl.uint8)
+
+    normal_x = qx
+    mant_odd = (normal_x >> (23 - 1)) & 1
+    normal_x += 0xC11FFFFF
+    normal_x += mant_odd
+    normal_x = normal_x >> (23 - 1)
+    normal_x = normal_x.to(tl.uint8)
+
+    e2m1 = tl.full(x.shape, 0x7, dtype=tl.uint8)
+    e2m1 = tl.where(normal_mask, normal_x, e2m1)
+    e2m1 = tl.where(denormal_mask, denormal_x, e2m1)
+    sign_lp = sign >> (23 + 8 - 1 - 2)
+    sign_lp = sign_lp.to(tl.uint8)
+    e2m1 = e2m1 | sign_lp
+    e2m1 = tl.reshape(e2m1, [16, 2])
+    evens, odds = tl.split(e2m1)
+    return evens | (odds << 4), scale_byte
+
+
+@triton.jit
+def _mxfp4_quantize_kernel(
+    x_ptr,
+    out_ptr,
+    scale_ptr,
+    M,
+    x_row_stride,
+    out_row_stride,
+    scale_row_stride,
+    N: tl.constexpr,
+):
+    pid_m = tl.program_id(0)
+    pid_g = tl.program_id(1)
+    offs_k = pid_g * 32 + tl.arange(0, 32)
+    x = tl.load(x_ptr + pid_m * x_row_stride + offs_k).to(tl.float32)
+
+    packed, scale_byte = _mxfp4_quantize_block(x)
+    pack_idx = tl.arange(0, 16)
+    tl.store(out_ptr + pid_m * out_row_stride + pid_g * 16 + pack_idx, packed)
+    tl.store(scale_ptr + pid_m * scale_row_stride + pid_g, scale_byte)
+
+
+def mxfp4_quantize(
+    x: torch.Tensor,
+    out: Optional[torch.Tensor] = None,
+    scales: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Quantize a BF16/FP16 tensor to packed MXFP4.
+
+    The last dimension is quantized in groups of 32 values. The returned data
+    packs two E2M1 values per byte and returns one uint8 E8M0 scale per group.
+    """
+    assert x.dtype in (
+        torch.bfloat16,
+        torch.float16,
+    ), f"mxfp4_quantize input must be bf16/fp16, got {x.dtype}"
+    M, N, x_row_stride = _flatten_to_2d(x)
+    if N % 32 != 0:
+        raise ValueError("mxfp4_quantize requires the last dimension divisible by 32")
+
+    out_shape = (*x.shape[:-1], N // 2)
+    scale_shape = (*x.shape[:-1], N // 32)
+    if out is None:
+        out = torch.empty(out_shape, dtype=torch.uint8, device=x.device)
+    else:
+        assert out.shape == out_shape and out.dtype == torch.uint8
+    if scales is None:
+        scales = torch.empty(scale_shape, dtype=torch.uint8, device=x.device)
+    else:
+        assert scales.shape == scale_shape and scales.dtype == torch.uint8
+
+    out_M, _, out_row_stride = _flatten_to_2d(out)
+    scales_M, _, scale_row_stride = _flatten_to_2d(scales)
+    assert out_M == M and scales_M == M
+
+    grid = (M, N // 32)
+    _mxfp4_quantize_kernel[grid](
+        x,
+        out,
+        scales,
+        M,
+        x_row_stride,
+        out_row_stride,
+        scale_row_stride,
+        N=N,
+        num_warps=1,
+    )
+    return out, scales
+
+
+@register_kernel(
+    "quantization",
+    "mxfp4",
+    name="triton_quantize_mxfp4",
+    solution="triton",
+    capability=CapabilityRequirement(vendors=frozenset({"amd"})),
+    signatures=format_signatures("x", "dense", {torch.bfloat16, torch.float16}),
+    traits={
+        "scale_size": frozenset({32}),
+        "scale_layout": frozenset({"linear"}),
+        "has_global_scale": frozenset({False}),
+        "scale_encoding": frozenset({"ue8m0"}),
+    },
+    priority=Priority.PORTABLE,
+)
+def triton_quantize_mxfp4(
+    x: torch.Tensor,
+    global_scale: float | None = None,
+    scale_size: int = 32,
+    scale_layout: str = "linear",
+    enable_pdl: bool = False,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if global_scale is not None:
+        raise ValueError("triton MXFP4 quantization does not support global_scale")
+    if scale_size != 32:
+        raise ValueError(
+            f"triton MXFP4 quantization requires scale_size=32, got {scale_size}"
+        )
+    if scale_layout != "linear":
+        raise ValueError(
+            "triton MXFP4 quantization requires scale_layout='linear', "
+            f"got {scale_layout!r}"
+        )
+    return mxfp4_quantize(x)
+
+
+__all__ = [
+    "fp8_quantize",
+    "mxfp4_quantize",
+    "triton_quantize_mxfp4",
+    "triton_quantize_fp8_with_scale",
+]

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

@@ -0,0 +1,117 @@
+# Copyright (c) 2026 LightSeek Foundation
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+from __future__ import annotations
+
+import torch
+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
+
+platform = current_platform()
+
+trtllm_fp8_token_group_128 = error_fn
+trtllm_fp8_token = error_fn
+trtllm_fp8_tensor = error_fn
+
+if platform.is_nvidia:
+    from tokenspeed_kernel.thirdparty.trtllm import (
+        per_tensor_quant_fp8 as _trtllm_per_tensor_quant_fp8,
+    )
+    from tokenspeed_kernel.thirdparty.trtllm import (
+        per_token_group_quant_8bit as _trtllm_per_token_group_quant_8bit,
+    )
+    from tokenspeed_kernel.thirdparty.trtllm import (
+        per_token_quant_fp8 as _trtllm_per_token_quant_fp8,
+    )
+
+    _FP8_DTYPE = platform.fp8e4m3fn.dtype
+
+    def trtllm_fp8_token_group_128(x: torch.Tensor) -> torch.Tensor:
+        qweight, _scale = _trtllm_per_token_group_quant_8bit(x, group_size=128)
+        return qweight.float()
+
+    def trtllm_fp8_token(x: torch.Tensor) -> torch.Tensor:
+        output = torch.empty_like(x, dtype=_FP8_DTYPE)
+        scale = torch.empty(x.size(0), dtype=torch.float32, device=x.device)
+        _trtllm_per_token_quant_fp8(x, output, scale)
+        return output.float()
+
+    def trtllm_fp8_tensor(x: torch.Tensor) -> torch.Tensor:
+        output = torch.empty_like(x, dtype=_FP8_DTYPE)
+        scale = torch.zeros(1, dtype=torch.float32, device=x.device)
+        _trtllm_per_tensor_quant_fp8(x, output, scale)
+        return output.float()
+
+    @register_kernel(
+        "quantization",
+        "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"}),
+            "scale_encoding": frozenset({"float32", "ue8m0"}),
+        },
+        priority=Priority.PERFORMANT,
+    )
+    def trtllm_quantize_fp8_with_scale(
+        x: torch.Tensor,
+        granularity: str = "tensor",
+        group_size: int | None = None,
+        scale_encoding: str = "float32",
+        enable_pdl: bool = False,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if granularity in {"tensor", "token"}:
+            if scale_encoding != "float32":
+                raise ValueError(f"TRT-LLM {granularity} FP8 requires float32 scales")
+
+            q = torch.empty_like(x, dtype=_FP8_DTYPE)
+            if granularity == "tensor":
+                scale = torch.empty(1, dtype=torch.float32, device=x.device)
+                _trtllm_per_tensor_quant_fp8(x, q, scale)
+            else:
+                scale = torch.empty(x.shape[:-1], dtype=torch.float32, device=x.device)
+                _trtllm_per_token_quant_fp8(x, q, scale)
+                scale = scale.unsqueeze(-1)
+            return q, scale
+
+        if granularity == "token_group":
+            return _trtllm_per_token_group_quant_8bit(
+                x,
+                group_size=group_size,
+                use_ue8m0=scale_encoding == "ue8m0",
+            )
+
+        raise ValueError(f"unsupported TRT-LLM FP8 granularity: {granularity!r}")
+
+
+__all__ = [
+    "trtllm_fp8_token_group_128",
+    "trtllm_fp8_token",
+    "trtllm_fp8_tensor",
+]