[perf] FP8: fuse the dynamic activation quantization into Triton kernels (#1861)
Co-authored-by: SolitaryThinker <wlsaidhi@gmail.com> Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Keith
co-authored by
SolitaryThinker
Claude Opus 5.5
SolitaryThinker
Claude Opus 5.5
parent
299d0f6838
commit
cdfbd64b04
@@ -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)
|
||||
|
||||
@@ -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
|
||||
@@ -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)
|
||||
Reference in New Issue
Block a user