Spaces:
Running on Zero
Running on Zero
File size: 2,719 Bytes
28404e6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 | # Copyright 2025 Tencent Inc. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import Tuple
import torch
import triton
import triton.language as tl
# https://github.com/deepseek-ai/DeepSeek-V3/blob/main/inference/kernel.py
@triton.jit
def _fp8_per_block_quant_kernel(x_ptr, y_ptr, s_ptr, M, N, BLOCK_SIZE: tl.constexpr):
"""Quantizes FP32 tensor to FP8 format using block-wise quantization."""
pid_m = tl.program_id(axis=0)
pid_n = tl.program_id(axis=1)
n = tl.cdiv(N, BLOCK_SIZE)
offs_m = pid_m * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
offs_n = pid_n * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
offs = offs_m[:, None] * N + offs_n[None, :]
mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)
x = tl.load(x_ptr + offs, mask=mask).to(tl.float32)
max_val = tl.max(tl.abs(x))
scale = max_val / 448.0
scale = tl.where(max_val == 0.0, 1.0, scale)
y = x / scale
y = y.to(y_ptr.dtype.element_ty)
tl.store(y_ptr + offs, y, mask=mask)
tl.store(s_ptr + pid_m * n + pid_n, scale)
# triton implementation
# for weight quantization on gpu
def fp8_per_block_quant_triton(
x: torch.Tensor, block_size: int = 128
) -> Tuple[torch.Tensor, torch.Tensor]:
"""
Quantizes a FP32 2D tensor to FP8 (E4M3FN) using block-wise quantization.
For each (block_size x block_size) block:
- scale = max(abs(block)) / 448.0 (FP8 E4M3FN max magnitude)
- if block is all zeros, use scale = 1.0 to avoid div-by-zero
- scale, clamp and cast to FP8
Returns:
y: Quantized FP8 tensor, same shape as input
s: Per-block scales, shape (num_blocks_M, num_blocks_N)
"""
assert x.is_contiguous()
assert x.dim() == 2
M, N = x.size()
y = torch.empty_like(x, dtype=torch.float8_e4m3fn)
m_blocks = triton.cdiv(M, block_size)
n_blocks = triton.cdiv(N, block_size)
s = torch.empty((m_blocks, n_blocks), dtype=torch.float32, device=x.device)
def grid(meta):
return (
triton.cdiv(M, meta["BLOCK_SIZE"]),
triton.cdiv(N, meta["BLOCK_SIZE"]),
)
_fp8_per_block_quant_kernel[grid](x, y, s, M, N, BLOCK_SIZE=block_size)
return y, s
|