From cdfbd64b042bd0586e3bf7ec63e847573f2bc31c Mon Sep 17 00:00:00 2001 From: Keith Date: Wed, 7 Oct 2026 18:20:04 +0800 Subject: [PATCH] [perf] FP8: fuse the dynamic activation quantization into Triton kernels (#1861) Co-authored-by: SolitaryThinker Co-authored-by: Claude Opus 5.5 (1M context) --- fastvideo/layers/quantization/fp8_config.py | 11 ++ .../layers/quantization/fp8_quant_kernels.py | 119 ++++++++++++++++++ .../quantization/test_fp8_quant_kernels.py | 113 +++++++++++++++++ 3 files changed, 243 insertions(+) create mode 100644 fastvideo/layers/quantization/fp8_quant_kernels.py create mode 100644 fastvideo/tests/ops/quantization/test_fp8_quant_kernels.py diff --git a/fastvideo/layers/quantization/fp8_config.py b/fastvideo/layers/quantization/fp8_config.py index 5b7e3f3cb..5198b92b5 100644 --- a/fastvideo/layers/quantization/fp8_config.py +++ b/fastvideo/layers/quantization/fp8_config.py @@ -20,6 +20,11 @@ from fastvideo.layers.quantization.base_config import ( QuantizationConfig, QuantizeMethodBase, ) +from fastvideo.layers.quantization.fp8_quant_kernels import ( + quantize_rowwise_fused, + quantize_tensorwise_fused, + triton_quant_available, +) from fastvideo.models.utils import set_weight_attrs logger = logging.getLogger(__name__) @@ -73,6 +78,8 @@ def _supports_fp8_compute() -> bool: def _quantize_tensorwise(x_2d: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: """Returns ``(x_fp8 [M, K], x_scale [1] float32)``.""" + if triton_quant_available(x_2d): + return quantize_tensorwise_fused(x_2d) x_absmax = x_2d.abs().amax().float() x_scale = (x_absmax / FP8_MAX).clamp(min=FP8_MIN_SCALE) x_fp8 = (x_2d / x_scale.to(x_2d.dtype)).clamp(-FP8_MAX, FP8_MAX).to(FP8_DTYPE) @@ -88,6 +95,10 @@ def _rowwise_scaled_mm_is_slow() -> bool: def _quantize_rowwise(x_2d: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: """Returns ``(x_fp8 [M, K], x_scale [M, 1] float32)``.""" + # The NaN-preserving fused kernel runs on every GPU with a recent Triton; the + # sm89 kernel from fp8_kernels stays as the fast path where it cannot run. + if triton_quant_available(x_2d): + return quantize_rowwise_fused(x_2d) if x_2d.is_cuda and _rowwise_scaled_mm_is_slow(): from fastvideo.layers.quantization.fp8_kernels import quantize_rowwise_fp8 return quantize_rowwise_fp8(x_2d) diff --git a/fastvideo/layers/quantization/fp8_quant_kernels.py b/fastvideo/layers/quantization/fp8_quant_kernels.py new file mode 100644 index 000000000..53b0041c1 --- /dev/null +++ b/fastvideo/layers/quantization/fp8_quant_kernels.py @@ -0,0 +1,119 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Fused FP8 (e4m3fn) activation quantization for the generic FP8 linear path. + +The eager formulas in ``fp8_config`` (``abs`` -> ``amax`` -> ``div`` -> ``clamp`` -> +``to(float8)``) are four to five memory passes over the activation. At MiniMax-H3's +shapes that costs about as much as the FP8 ``_scaled_mm`` it feeds (1.06 ms next +to a 1.25 ms GEMM for a 65,536 x 5,376 activation on MI355X). These kernels do the +same math in one pass (rowwise) or an ``aminmax`` reduction plus one pass +(tensorwise): 3x (tensorwise) and 4-5x (rowwise) faster at those shapes. + +NaN handling follows the eager formulas: a NaN anywhere in a row (rowwise) or in +the tensor (tensorwise) makes its scale and its FP8 values NaN. Triton's default +``maximum``/``minimum`` and ``tl.max`` return the non-NaN operand, so every max, +min and clamp below asks for NaN propagation. +""" +from __future__ import annotations + +import torch + +FP8_DTYPE = torch.float8_e4m3fn +FP8_MAX = float(torch.finfo(FP8_DTYPE).max) # 448.0 +FP8_MIN_SCALE = 1.0 / (FP8_MAX * 512.0) + +try: + import triton + import triton.language as tl + # The kernels need propagate_nan= on tl.maximum/minimum (tl.PropagateNan), + # which older Tritons lack; without it fall back to the eager formulas. + _HAS_TRITON = hasattr(tl, "PropagateNan") +except Exception: # noqa: BLE001 + _HAS_TRITON = False + +if _HAS_TRITON: + + @triton.jit + def _max_keep_nan(a, b): + return tl.maximum(a, b, propagate_nan=tl.PropagateNan.ALL) + + @triton.jit + def _clamp_keep_nan(x, FP8_MAX: tl.constexpr): + y = tl.maximum(x, -FP8_MAX, propagate_nan=tl.PropagateNan.ALL) + return tl.minimum(y, FP8_MAX, propagate_nan=tl.PropagateNan.ALL) + + @triton.jit + def _rowwise_quant_kernel(x_ptr, out_ptr, scale_ptr, K, stride_x, stride_out, FP8_MAX: tl.constexpr, + MIN_SCALE: tl.constexpr, BLOCK: tl.constexpr): + """One program per row: amax over the row, then scale/clamp/cast in a second sweep. + + The row is re-read from L2 in the second sweep, so K is not limited by + the register tile and the kernel stays one launch.""" + # int64 offsets: M * K can pass 2^31 elements on long videos. + row = tl.program_id(0).to(tl.int64) + x_row = x_ptr + row * stride_x + amax = tl.zeros([BLOCK], dtype=tl.float32) + for start in range(0, K, BLOCK): + offs = start + tl.arange(0, BLOCK) + x = tl.load(x_row + offs, mask=offs < K, other=0.0).to(tl.float32) + amax = _max_keep_nan(amax, tl.abs(x)) + row_amax = tl.reduce(amax, 0, _max_keep_nan) + scale = tl.maximum(row_amax / FP8_MAX, MIN_SCALE, propagate_nan=tl.PropagateNan.ALL) + inv = 1.0 / scale + out_row = out_ptr + row * stride_out + for start in range(0, K, BLOCK): + offs = start + tl.arange(0, BLOCK) + x = tl.load(x_row + offs, mask=offs < K, other=0.0).to(tl.float32) + y = _clamp_keep_nan(x * inv, FP8_MAX) + tl.store(out_row + offs, y.to(out_ptr.dtype.element_ty), mask=offs < K) + tl.store(scale_ptr + row, scale) + + @triton.jit + def _scale_cast_kernel(x_ptr, out_ptr, inv_scale_ptr, n_elements, FP8_MAX: tl.constexpr, BLOCK: tl.constexpr): + """Tensorwise second pass: one read, one FP8 write.""" + pid = tl.program_id(0).to(tl.int64) # int64 offsets: n can pass 2^31 elements + offs = pid * BLOCK + tl.arange(0, BLOCK) + mask = offs < n_elements + inv = tl.load(inv_scale_ptr) + x = tl.load(x_ptr + offs, mask=mask, other=0.0).to(tl.float32) + y = _clamp_keep_nan(x * inv, FP8_MAX) + tl.store(out_ptr + offs, y.to(out_ptr.dtype.element_ty), mask=mask) + + +def triton_quant_available(x: torch.Tensor) -> bool: + return _HAS_TRITON and x.is_cuda and x.dtype in (torch.bfloat16, torch.float16) + + +def quantize_rowwise_fused(x_2d: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Returns ``(x_fp8 [M, K], x_scale [M, 1] float32)``; one kernel launch.""" + x_2d = x_2d.contiguous() + m, k = x_2d.shape + out = torch.empty((m, k), dtype=FP8_DTYPE, device=x_2d.device) + scale = torch.empty((m, 1), dtype=torch.float32, device=x_2d.device) + block = min(4096, triton.next_power_of_2(k)) + if m: + _rowwise_quant_kernel[(m, )](x_2d, + out, + scale, + k, + x_2d.stride(0), + out.stride(0), + FP8_MAX=FP8_MAX, + MIN_SCALE=FP8_MIN_SCALE, + BLOCK=block, + num_warps=8) + return out, scale + + +def quantize_tensorwise_fused(x_2d: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Returns ``(x_fp8 [M, K], x_scale [1] float32)``; an ``aminmax`` reduction + one pass.""" + x_2d = x_2d.contiguous() + # aminmax, maximum and clamp all propagate NaN, as the eager amax does. + lo, hi = torch.aminmax(x_2d) + amax = torch.maximum(lo.abs(), hi.abs()).float() + scale = (amax / FP8_MAX).clamp(min=FP8_MIN_SCALE).view(1) + inv = 1.0 / scale + out = torch.empty_like(x_2d, dtype=FP8_DTYPE) + n = x_2d.numel() + block = 4096 + _scale_cast_kernel[(triton.cdiv(n, block), )](x_2d, out, inv, n, FP8_MAX=FP8_MAX, BLOCK=block, num_warps=8) + return out, scale diff --git a/fastvideo/tests/ops/quantization/test_fp8_quant_kernels.py b/fastvideo/tests/ops/quantization/test_fp8_quant_kernels.py new file mode 100644 index 000000000..fa8147c9e --- /dev/null +++ b/fastvideo/tests/ops/quantization/test_fp8_quant_kernels.py @@ -0,0 +1,113 @@ +# SPDX-License-Identifier: Apache-2.0 +"""The fused FP8 activation quantizers must reproduce the eager formulas: identical +scales, and dequantized values within one FP8 rounding step (the fused kernels scale +in FP32 where the eager path scales in the activation dtype). NaN inputs must give +NaN scales and NaN FP8 values, as they do in the eager path.""" +import pytest +import torch + +from fastvideo.layers.quantization import fp8_config, fp8_quant_kernels +from fastvideo.layers.quantization.fp8_quant_kernels import ( + FP8_DTYPE, + FP8_MAX, + quantize_rowwise_fused, + quantize_tensorwise_fused, + triton_quant_available, +) + +# Triton is optional: without it the eager formulas run on the GPU too. +needs_gpu_and_triton = pytest.mark.skipif(not (torch.cuda.is_available() and fp8_quant_kernels._HAS_TRITON), + reason="needs a GPU with Triton") + + +def _eager_tensorwise(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + scale = (x.abs().amax().float() / FP8_MAX).clamp(min=fp8_config.FP8_MIN_SCALE) + return (x / scale.to(x.dtype)).clamp(-FP8_MAX, FP8_MAX).to(FP8_DTYPE), scale.view(1) + + +def _eager_rowwise(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + scale = (x.abs().amax(dim=-1, keepdim=True).float() / FP8_MAX).clamp(min=fp8_config.FP8_MIN_SCALE) + return (x / scale.to(x.dtype)).clamp(-FP8_MAX, FP8_MAX).to(FP8_DTYPE), scale + + +def _assert_eager_path(x: torch.Tensor) -> None: + assert not triton_quant_available(x) + for quantize, eager in ((fp8_config._quantize_tensorwise, _eager_tensorwise), + (fp8_config._quantize_rowwise, _eager_rowwise)): + q, s = quantize(x) + q_ref, s_ref = eager(x) + assert torch.equal(q.view(torch.uint8), q_ref.view(torch.uint8)) and torch.equal(s, s_ref) + + +def test_cpu_inputs_keep_the_eager_path() -> None: + _assert_eager_path(torch.randn(8, 64, dtype=torch.bfloat16)) + + +def test_cpu_empty_rows_keep_the_eager_path() -> None: + x = torch.empty((0, 96), dtype=torch.bfloat16) + q, s = fp8_config._quantize_rowwise(x) + q_ref, s_ref = _eager_rowwise(x) + assert q.shape == (0, 96) and s.shape == (0, 1) + assert torch.equal(q.view(torch.uint8), q_ref.view(torch.uint8)) and torch.equal(s, s_ref) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") +def test_gpu_without_triton_keeps_the_eager_path(monkeypatch) -> None: + monkeypatch.setattr(fp8_quant_kernels, "_HAS_TRITON", False) + # On sm89 rowwise quantization would otherwise take the fp8_kernels path, + # which scales in FP32 and so is not bitwise equal to the eager formula. + monkeypatch.setattr(fp8_config, "_rowwise_scaled_mm_is_slow", lambda: False) + _assert_eager_path(torch.randn(8, 64, device="cuda", dtype=torch.bfloat16)) + + +@needs_gpu_and_triton +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.parametrize( + ("m", "k", "magnitude"), + [ + (1000, 5376, 3.0), + (4096, 14336, 3.0), + (7, 96, 3.0), + # Scales near FP8_MIN_SCALE are subnormal in fp16, where the eager path + # rounds the scale to the activation dtype before dividing. + (64, 5376, 1e-4), + ]) +def test_fused_matches_eager_on_gpu(m: int, k: int, magnitude: float, dtype: torch.dtype) -> None: + torch.manual_seed(0) + x = (torch.randn(m, k, device="cuda") * magnitude).to(dtype) + x[0, 0] = 200.0 * magnitude / 3.0 # an outlier separates the tensorwise and rowwise scales + assert triton_quant_available(x) + for fused, eager in ((quantize_tensorwise_fused, _eager_tensorwise), (quantize_rowwise_fused, _eager_rowwise)): + q, s = fused(x) + q_ref, s_ref = eager(x) + assert q.dtype == FP8_DTYPE and q.shape == x.shape and s.shape == s_ref.shape + torch.testing.assert_close(s, s_ref, rtol=1e-6, atol=0) + dequant, dequant_ref = q.float() * s, q_ref.float() * s_ref + # one FP8 ulp at the row/tensor maximum, scaled back to the activation's range + ulp = (s * FP8_MAX / 2**3).max() + assert (dequant - dequant_ref).abs().max() <= ulp + + +@needs_gpu_and_triton +def test_fused_rowwise_empty_rows_on_gpu() -> None: + x = torch.empty((0, 96), device="cuda", dtype=torch.bfloat16) + q, s = quantize_rowwise_fused(x) + q_ref, s_ref = _eager_rowwise(x) + assert q.shape == q_ref.shape and s.shape == s_ref.shape + assert torch.equal(q.view(torch.uint8), q_ref.view(torch.uint8)) and torch.equal(s, s_ref) + + +@needs_gpu_and_triton +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +def test_fused_keeps_nan_like_eager_on_gpu(dtype: torch.dtype) -> None: + torch.manual_seed(0) + x = torch.randn(4, 96, device="cuda").to(dtype) + x[1, 5] = float("nan") + for fused, eager in ((quantize_tensorwise_fused, _eager_tensorwise), (quantize_rowwise_fused, _eager_rowwise)): + q, s = fused(x) + q_ref, s_ref = eager(x) + assert s_ref.isnan().any() # the eager behaviour that the fused kernels must keep + assert torch.equal(s.isnan(), s_ref.isnan()) + assert torch.equal(q.float().isnan(), q_ref.float().isnan()) + finite = ~s_ref.isnan() + torch.testing.assert_close(s[finite], s_ref[finite], rtol=1e-6, atol=0)