8b51380466
Co-authored-by: Aryan Kumar <aryan5v@users.noreply.github.com> Co-authored-by: Aryan Kumar <aryank@Aryans-Mac-Studio.local> Co-authored-by: SolitaryThinker <wlsaidhi@gmail.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
85 lines
3.8 KiB
Python
85 lines
3.8 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Triton helpers for per-token x per-channel FP8 linears on GPUs whose rowwise ``_scaled_mm`` is slow.
|
|
|
|
On sm89 (RTX 4090 / L40S / RTX 6000 Ada) ``torch._scaled_mm`` with rowwise scales runs at
|
|
~70 TFLOPS, below bf16 (~160), while the per-tensor kernel reaches 220-305 TFLOPS. The
|
|
per-token x per-channel result is recovered exactly by running the per-tensor kernel with unit
|
|
scales and applying ``out[i, j] *= sx[i] * sw[j]`` in one pass over the output, which costs
|
|
5-10% of the GEMM instead of 2-4x.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import functools
|
|
|
|
import torch
|
|
import triton
|
|
import triton.language as tl
|
|
|
|
FP8_MAX = 448.0
|
|
FP8_MIN_SCALE = 1.0 / (FP8_MAX * 512.0)
|
|
|
|
|
|
@functools.cache
|
|
def rowwise_scaled_mm_is_slow() -> bool:
|
|
"""Ada (sm89) has no fast rowwise-scaled FP8 GEMM in torch; Hopper and Blackwell do."""
|
|
return torch.cuda.is_available() and torch.cuda.get_device_capability() == (8, 9)
|
|
|
|
|
|
@triton.jit
|
|
def _quantize_rowwise_kernel(x_ptr, q_ptr, s_ptr, K, stride_x, stride_q, BLOCK_K: tl.constexpr):
|
|
row = tl.program_id(0).to(tl.int64)
|
|
x_row = x_ptr + row * stride_x
|
|
amax = tl.zeros((BLOCK_K, ), dtype=tl.float32)
|
|
for k in range(0, K, BLOCK_K):
|
|
cols = k + tl.arange(0, BLOCK_K)
|
|
amax = tl.maximum(amax, tl.abs(tl.load(x_row + cols, mask=cols < K, other=0.0).to(tl.float32)))
|
|
scale = tl.maximum(tl.max(amax, axis=0) / 448.0, 1.0 / (448.0 * 512.0))
|
|
tl.store(s_ptr + row, scale)
|
|
inv = 1.0 / scale
|
|
for k in range(0, K, BLOCK_K):
|
|
cols = k + tl.arange(0, BLOCK_K)
|
|
v = tl.load(x_row + cols, mask=cols < K, other=0.0).to(tl.float32) * inv
|
|
v = tl.minimum(tl.maximum(v, -448.0), 448.0)
|
|
tl.store(q_ptr + row * stride_q + cols, v.to(tl.float8e4nv), mask=cols < K)
|
|
|
|
|
|
def quantize_rowwise_fp8(x_2d: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
|
"""Per-token FP8 quantization in one launch. Returns ``(x_fp8 [M, K], x_scale [M, 1] float32)``."""
|
|
x_2d = x_2d.contiguous()
|
|
M, K = x_2d.shape
|
|
q = torch.empty((M, K), device=x_2d.device, dtype=torch.float8_e4m3fn)
|
|
s = torch.empty((M, 1), device=x_2d.device, dtype=torch.float32)
|
|
if M:
|
|
_quantize_rowwise_kernel[(M, )](x_2d, q, s, K, x_2d.stride(0), q.stride(0), BLOCK_K=1024, num_warps=4)
|
|
return q, s
|
|
|
|
|
|
@triton.jit
|
|
def _scale_rows_cols_kernel(o_ptr, sx_ptr, sw_ptr, M, N, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr):
|
|
rows = tl.program_id(0).to(tl.int64) * BLOCK_M + tl.arange(0, BLOCK_M)
|
|
cols = tl.program_id(1) * BLOCK_N + tl.arange(0, BLOCK_N)
|
|
mask = (rows[:, None] < M) & (cols[None, :] < N)
|
|
# int64 offsets: a 78k-token fc_in output has 2.2e9 elements, past int32.
|
|
ptrs = o_ptr + rows[:, None] * N + cols[None, :]
|
|
v = tl.load(ptrs, mask=mask, other=0.0).to(tl.float32)
|
|
sx = tl.load(sx_ptr + rows, mask=rows < M, other=0.0)
|
|
sw = tl.load(sw_ptr + cols, mask=cols < N, other=0.0)
|
|
tl.store(ptrs, (v * sx[:, None] * sw[None, :]).to(o_ptr.dtype.element_ty), mask=mask)
|
|
|
|
|
|
def scaled_mm_token_channel(x_fp8: torch.Tensor, x_scale: torch.Tensor, w_fp8_t: torch.Tensor,
|
|
w_scale: torch.Tensor) -> torch.Tensor:
|
|
"""``(x_fp8 * x_scale) @ (w_fp8_t * w_scale)`` in bf16 via the fast per-tensor GEMM plus a scale epilogue.
|
|
|
|
The unit-scale GEMM output is at most 448^2 * K, far inside bf16 range, and its relative
|
|
precision is that of any bf16 output, so the epilogue loses nothing against rowwise scaling.
|
|
"""
|
|
one = torch.ones((), device=x_fp8.device, dtype=torch.float32)
|
|
out = torch._scaled_mm(x_fp8, w_fp8_t, scale_a=one, scale_b=one, out_dtype=torch.bfloat16)
|
|
if isinstance(out, tuple):
|
|
out = out[0]
|
|
M, N = out.shape
|
|
grid = (triton.cdiv(M, 64), triton.cdiv(N, 128))
|
|
_scale_rows_cols_kernel[grid](out, x_scale.reshape(-1), w_scale.reshape(-1), M, N, BLOCK_M=64, BLOCK_N=128)
|
|
return out
|