[Perf][Kernel][Quantization] Triton row kernel for per-token-group FP8 quant at large token counts

tests/kernels/quantization/test_per_token_group_quant_row.py

@@ -0,0 +1,132 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Bitwise path-equivalence tests for the Triton row kernel vs
+torch.ops._C.per_token_group_fp8_quant (vllm/model_executor/layers/quantization/
+utils/fp8_utils.py).
+
+Structure follows the parametrized path-equivalence suites that reviewers asked for
+on this file's history (#21083 review: all scale layouts; #55330: bit-exact claims
+backed by torch.equal-level assertions, zero-row eps path, inf/NaN edges).
+
+Destination: tests/kernels/quantization/test_per_token_group_quant_row.py
+"""
+
+import pytest
+import torch
+
+from tests.kernels.quant_utils import FP8_DTYPE
+from vllm.model_executor.layers.quantization.utils import fp8_utils
+from vllm.platforms import current_platform
+
+GROUP_SIZE = 128
+EPS = 1e-10
+
+pytestmark = pytest.mark.skipif(
+    not current_platform.is_cuda(),
+    reason="row kernel is gated to CUDA; other platforms keep the _C/Triton paths",
+)
+
+
+def _alloc_scales(T: int, H: int, layout: str, device: str) -> torch.Tensor:
+    sf_k = H // GROUP_SIZE
+    if layout == "row":
+        return torch.empty((T, sf_k), device=device, dtype=torch.float32)
+    if layout == "col":
+        return torch.empty((sf_k, T), device=device, dtype=torch.float32).permute(
+            -1, -2
+        )
+    m = fp8_utils.get_tma_aligned_size(T, 4)
+    return torch.empty_strided((T, sf_k), (1, m), device=device, dtype=torch.float32)
+
+
+def _run(x: torch.Tensor, layout: str, use_row_kernel: bool):
+    fp8_min, fp8_max = torch.finfo(FP8_DTYPE).min, torch.finfo(FP8_DTYPE).max
+    T, H = x.shape
+    x_q = torch.empty_like(x, dtype=FP8_DTYPE)
+    x_s = _alloc_scales(T, H, layout, x.device.type and str(x.device))
+    if use_row_kernel:
+        torch.ops.vllm.per_token_group_fp8_quant_row(
+            x,
+            x_q,
+            x_s,
+            GROUP_SIZE,
+            EPS,
+            fp8_min,
+            fp8_max,
+            False,  # scale_ue8m0
+            layout != "row",  # column_major_scales
+            layout == "tma",  # tma_aligned_scales
+        )
+    else:
+        torch.ops._C.per_token_group_fp8_quant(
+            x,
+            x_q,
+            x_s,
+            GROUP_SIZE,
+            EPS,
+            fp8_min,
+            fp8_max,
+            False,  # scale_ue8m0
+            layout != "row",  # column_major
+            layout == "tma",
+        )  # tma_aligned
+    return x_q, x_s
+
+
+@pytest.mark.parametrize("T", [1, 7, 64, 512, 1452, 4096, 17083])
+@pytest.mark.parametrize("H", [2048, 4096])
+@pytest.mark.parametrize("layout", ["row", "col", "tma"])
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+@torch.inference_mode()
+def test_row_kernel_bitwise_matches_cuda(T, H, layout, dtype):
+    torch.manual_seed(0)
+    x = (torch.randn(T, H, device="cuda", dtype=dtype) * 3).contiguous()
+    if T >= 64:
+        x[T // 2] = 0.0  # all-zero row: eps floor path
+        x[0, :GROUP_SIZE] = float("inf")  # satfinite clamp path
+    q_ref, s_ref = _run(x, layout, use_row_kernel=False)
+    q_row, s_row = _run(x, layout, use_row_kernel=True)
+    assert (q_ref.view(torch.uint8) == q_row.view(torch.uint8)).all(), (
+        "quantized values differ bitwise"
+    )
+    assert (
+        s_ref.contiguous().view(torch.int32) == s_row.contiguous().view(torch.int32)
+    ).all(), "scales differ bitwise"
+
+
+@torch.inference_mode()
+def test_dispatch_fallbacks():
+    """Shapes the row kernel must refuse stay on the existing paths."""
+    x = torch.randn(2688, 2688, device="cuda", dtype=torch.bfloat16)
+    # non-power-of-two H: public API must still work (via _C), bit-identically
+    q, s = fp8_utils.per_token_group_quant_fp8(x, GROUP_SIZE, use_ue8m0=False)
+    assert q.shape == x.shape and s.shape == (2688, 2688 // GROUP_SIZE)
+    assert not fp8_utils._row_quant_eligible(x, GROUP_SIZE)
+    # UE8M0 scales fall through to the _C kernel inside the op
+    q, s = fp8_utils.per_token_group_quant_fp8(
+        torch.randn(4096, 2048, device="cuda", dtype=torch.bfloat16),
+        GROUP_SIZE,
+        use_ue8m0=True,
+    )
+    nz = s[s > 0]
+    assert torch.all(torch.log2(nz) == torch.log2(nz).round())
+
+
+def test_opcheck():
+    torch.manual_seed(0)
+    x = (torch.randn(64, 2048, device="cuda", dtype=torch.bfloat16)).contiguous()
+    x_q = torch.empty_like(x, dtype=FP8_DTYPE)
+    x_s = torch.empty((64, 2048 // GROUP_SIZE), device="cuda", dtype=torch.float32)
+    fp8_min, fp8_max = torch.finfo(FP8_DTYPE).min, torch.finfo(FP8_DTYPE).max
+    # test_schema is skipped: opcheck's cloned-output comparison calls
+    # torch.allclose on the mutated fp8 tensor, and mul_cuda has no fp8
+    # kernel — a torch harness limitation, not an op property.
+    torch.library.opcheck(
+        torch.ops.vllm.per_token_group_fp8_quant_row,
+        (x, x_q, x_s, GROUP_SIZE, EPS, fp8_min, fp8_max, False, False, False),
+        test_utils=("test_faketensor", "test_aot_dispatch_dynamic"),
+    )

vllm/compilation/passes/fusion/matcher_utils.py

@@ -42,6 +42,15 @@
 if hasattr(torch.ops._C, "per_token_group_fp8_quant"):
     QUANT_OPS[kFp8Dynamic128Sym] = torch.ops._C.per_token_group_fp8_quant.default  # noqa: E501
     QUANT_OPS[kFp8Dynamic64Sym] = torch.ops._C.per_token_group_fp8_quant.default  # noqa: E501
+if current_platform.is_cuda():
+    # On CUDA the eager path emits the size-dispatching row-quant wrapper op
+    # (same signature as the _C op, bitwise-identical outputs) instead of the
+    # _C op; patterns must match the op that actually appears in the graph.
+    # Importing fp8_utils registers torch.ops.vllm.per_token_group_fp8_quant_row.
+    from vllm.model_executor.layers.quantization.utils import fp8_utils  # noqa: F401
+
+    QUANT_OPS[kFp8Dynamic128Sym] = torch.ops.vllm.per_token_group_fp8_quant_row.default
+    QUANT_OPS[kFp8Dynamic64Sym] = torch.ops.vllm.per_token_group_fp8_quant_row.default
 
 if current_platform.is_cuda() and hasattr(torch.ops._C, "scaled_fp4_quant"):
     QUANT_OPS[kNvfp4Dynamic] = torch.ops._C.scaled_fp4_quant.out  # noqa: E501

vllm/compilation/passes/fusion/rms_quant_fusion.py

@@ -163,6 +163,14 @@ def empty_i64(*args: Any, **kwargs: Any) -> torch.Tensor:
 if hasattr(torch.ops._C, "per_token_group_fp8_quant"):
     QUANT_OPS[kFp8Dynamic128Sym] = torch.ops._C.per_token_group_fp8_quant.default  # noqa: E501
     QUANT_OPS[kFp8Dynamic64Sym] = torch.ops._C.per_token_group_fp8_quant.default  # noqa: E501
+if current_platform.is_cuda():
+    # On CUDA the eager path emits the size-dispatching row-quant wrapper op
+    # instead of the _C op (see matcher_utils.QUANT_OPS); match that op.
+    # Importing fp8_utils registers torch.ops.vllm.per_token_group_fp8_quant_row.
+    from vllm.model_executor.layers.quantization.utils import fp8_utils  # noqa: F401
+
+    QUANT_OPS[kFp8Dynamic128Sym] = torch.ops.vllm.per_token_group_fp8_quant_row.default
+    QUANT_OPS[kFp8Dynamic64Sym] = torch.ops.vllm.per_token_group_fp8_quant_row.default
 if current_platform.is_cuda() and hasattr(torch.ops._C, "scaled_fp4_quant"):
     QUANT_OPS[kNvfp4Dynamic] = torch.ops._C.scaled_fp4_quant.out
 

vllm/model_executor/layers/quantization/utils/fp8_utils.py

@@ -629,6 +629,132 @@ def _per_token_group_quant_fp8_colmajor(
     tl.store(y_s_ptr, y_s)
 
 
+@triton.jit
+def _per_token_group_quant_fp8_row(
+    x_ptr,
+    q_ptr,
+    s_ptr,
+    s_stride_row,
+    s_stride_col,
+    eps,
+    fp8_min,
+    fp8_max,
+    H: tl.constexpr,
+    GROUP: tl.constexpr,
+):
+    """One program per row: load the whole row once, compute every group's
+    scale in registers, write the row back in one coalesced store.
+
+    Reproduces torch.ops._C.per_token_group_fp8_quant bit-for-bit: fp32 absmax
+    floored at eps, IEEE round-to-nearest division (plain `/` may compile to an
+    approximate divide), round-to-nearest-satfinite fp8 conversion.
+    """
+    row = tl.program_id(0)
+    offs = tl.arange(0, H)
+    x = tl.load(x_ptr + row * H + offs).to(tl.float32)
+    g = tl.reshape(x, (H // GROUP, GROUP))
+    amax = tl.maximum(tl.max(tl.abs(g), axis=1), eps)
+    scale = tl.math.div_rn(amax, fp8_max)
+    y = tl.math.div_rn(g, scale[:, None])
+    y = tl.minimum(tl.maximum(y, fp8_min), fp8_max)
+    tl.store(q_ptr + row * H + offs, tl.reshape(y, (H,)).to(q_ptr.dtype.element_ty))
+    gi = tl.arange(0, H // GROUP)
+    tl.store(s_ptr + row * s_stride_row + gi * s_stride_col, scale)
+
+
+# Below this many tokens the CUDA kernel's lower launch cost wins; above it
+# the row kernel's full-width loads win (1.6-1.8x at prefill sizes on H200).
+_ROW_QUANT_MIN_TOKENS = 2048
+
+
+def _row_quant_eligible(x: torch.Tensor, group_size: int) -> bool:
+    H = x.shape[-1]
+    return (
+        x.dim() == 2
+        and x.is_contiguous()
+        and H % group_size == 0
+        and (H & (H - 1)) == 0  # tl.arange needs a power of two
+        and group_size <= H <= 8192
+    )
+
+
+def _per_token_group_fp8_quant_row(
+    x: torch.Tensor,
+    x_q: torch.Tensor,
+    x_s: torch.Tensor,
+    group_size: int,
+    eps: float,
+    fp8_min: float,
+    fp8_max: float,
+    scale_ue8m0: bool,
+    column_major_scales: bool,
+    tma_aligned_scales: bool,
+) -> None:
+    """Drop-in for torch.ops._C.per_token_group_fp8_quant on CUDA.
+
+    Dispatches to the row kernel at large token counts and falls back to the
+    CUDA kernel otherwise. The branch lives inside the opaque op, so
+    torch.compile never traces a data-dependent condition.
+    """
+    if (
+        scale_ue8m0
+        or x.shape[0] < _ROW_QUANT_MIN_TOKENS
+        or not _row_quant_eligible(x, group_size)
+    ):
+        torch.ops._C.per_token_group_fp8_quant(
+            x,
+            x_q,
+            x_s,
+            group_size,
+            eps,
+            fp8_min,
+            fp8_max,
+            scale_ue8m0,
+            column_major_scales,
+            tma_aligned_scales,
+        )
+        return
+    _per_token_group_quant_fp8_row[(x.shape[0],)](
+        x,
+        x_q,
+        x_s,
+        x_s.stride(0),
+        x_s.stride(1),
+        eps,
+        fp8_min,
+        fp8_max,
+        H=x.shape[-1],
+        GROUP=group_size,
+        num_warps=8,
+    )
+
+
+def _per_token_group_fp8_quant_row_fake(
+    x: torch.Tensor,
+    x_q: torch.Tensor,
+    x_s: torch.Tensor,
+    group_size: int,
+    eps: float,
+    fp8_min: float,
+    fp8_max: float,
+    scale_ue8m0: bool,
+    column_major_scales: bool,
+    tma_aligned_scales: bool,
+) -> None:
+    return None
+
+
+if current_platform.is_cuda():
+    from vllm.utils.torch_utils import direct_register_custom_op
+
+    direct_register_custom_op(
+        op_name="per_token_group_fp8_quant_row",
+        op_func=_per_token_group_fp8_quant_row,
+        mutates_args=["x_q", "x_s"],
+        fake_impl=_per_token_group_fp8_quant_row_fake,
+    )
+
+
 def per_token_group_quant_fp8(
     x: torch.Tensor,
     group_size: int,
@@ -705,6 +831,24 @@ def per_token_group_quant_fp8(
     if (
         current_platform.is_cuda_alike() or current_platform.is_xpu()
     ) and x.is_contiguous():
+        if current_platform.is_cuda():
+            # Size-dispatching wrapper around the _C op: identical signature
+            # and bitwise-identical outputs, but takes the Triton row kernel
+            # at large token counts. Fusion pattern matchers key on this op
+            # (see compilation/passes/fusion QUANT_OPS).
+            torch.ops.vllm.per_token_group_fp8_quant_row(
+                x,
+                x_q,
+                x_s,
+                group_size,
+                eps,
+                fp8_min,
+                fp8_max,
+                use_ue8m0,
+                column_major_scales,
+                tma_aligned_scales,
+            )
+            return x_q, x_s
         torch.ops._C.per_token_group_fp8_quant(
             x,
             x_q,