152 lines
5.0 KiB
Python
152 lines
5.0 KiB
Python
"""Fused codebook weight dequantization kernel for OrbitQuant.
|
|
|
|
Replaces _orbitquant.codebook_dequant (unpack -> offset -> codebook gather
|
|
-> per-row norms) with a single elementwise+gather kernel.
|
|
|
|
Layout matches _quant_utils.pack_int4 / kernels.triton_int4_unpack:
|
|
low nibble = even index, high nibble = odd index, two's complement.
|
|
|
|
Output dtype follows the provided `out` tensor (default float16, the
|
|
compute dtype used at the matmul). For an exact float32 parity check,
|
|
pass an fp32 `out`.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import torch
|
|
|
|
try:
|
|
import triton
|
|
import triton.language as tl
|
|
_HAS_TRITON = True
|
|
except ImportError:
|
|
_HAS_TRITON = False
|
|
|
|
_DEFAULT_BLOCK_K: int = 2048
|
|
|
|
|
|
if _HAS_TRITON:
|
|
|
|
@triton.jit
|
|
def _codebook_dequant_int4_kernel(
|
|
packed_ptr, # [N, K_packed] uint8
|
|
norms_ptr, # [N] float (per-row L2 norms)
|
|
cb_ptr, # [2**bits] float32 (Lloyd-Max centroids)
|
|
out_ptr, # [N, K]
|
|
stride_pn, stride_on,
|
|
K,
|
|
BLOCK_K: tl.constexpr,
|
|
):
|
|
pid_n = tl.program_id(0)
|
|
pid_k = tl.program_id(1)
|
|
offs = pid_k * BLOCK_K + tl.arange(0, BLOCK_K)
|
|
mask = offs < K
|
|
|
|
# 1. Unpack two's-complement nibble -> int8 -> index in [0, 15]
|
|
packed_col = offs // 2
|
|
is_high = (offs % 2).to(tl.uint8)
|
|
byte = tl.load(packed_ptr + pid_n * stride_pn + packed_col,
|
|
mask=mask, other=0).to(tl.uint8)
|
|
nibble = tl.where(is_high == 0, byte & 0x0F, (byte >> 4) & 0x0F).to(tl.int32)
|
|
idx = tl.where(nibble >= 8, nibble - 16, nibble) + 8 # offset +2^(b-1)
|
|
|
|
# 2. Codebook gather (tiny table -> L1/registers)
|
|
centroid = tl.load(cb_ptr + idx, mask=mask, other=0.0).to(tl.float32)
|
|
|
|
# 3. Per-row L2 norm (norms is contiguous 1-D after reshape in wrapper)
|
|
row_norm = tl.load(norms_ptr + pid_n).to(tl.float32)
|
|
val = centroid * row_norm
|
|
|
|
# 4. Store (cast to output dtype)
|
|
tl.store(out_ptr + pid_n * stride_on + offs,
|
|
val.to(out_ptr.dtype.element_ty), mask=mask)
|
|
|
|
|
|
@triton.jit
|
|
def _codebook_dequant_int8_kernel(
|
|
idx_ptr, # [N, K] int8 (already-unpacked indices)
|
|
norms_ptr, # [N]
|
|
cb_ptr, # [2**bits]
|
|
out_ptr, # [N, K]
|
|
stride_in, stride_on,
|
|
K,
|
|
BLOCK_K: tl.constexpr,
|
|
):
|
|
pid_n = tl.program_id(0)
|
|
pid_k = tl.program_id(1)
|
|
offs = pid_k * BLOCK_K + tl.arange(0, BLOCK_K)
|
|
mask = offs < K
|
|
|
|
raw = tl.load(idx_ptr + pid_n * stride_in + offs, mask=mask, other=0).to(tl.int32)
|
|
idx = raw + 128 # [-128,127] -> [0,255]
|
|
|
|
centroid = tl.load(cb_ptr + idx, mask=mask, other=0.0).to(tl.float32)
|
|
row_norm = tl.load(norms_ptr + pid_n).to(tl.float32)
|
|
val = centroid * row_norm
|
|
|
|
tl.store(out_ptr + pid_n * stride_on + offs,
|
|
val.to(out_ptr.dtype.element_ty), mask=mask)
|
|
|
|
|
|
def codebook_dequant_triton(
|
|
qdata: torch.Tensor,
|
|
norms: torch.Tensor,
|
|
codebook: torch.Tensor,
|
|
d_padded: int,
|
|
bits: int,
|
|
out: torch.Tensor | None = None,
|
|
) -> torch.Tensor:
|
|
"""Fused unpack + codebook gather + per-row norms (Triton, CUDA only).
|
|
|
|
Args:
|
|
qdata: packed indices. uint8 (INT4, 2/byte) or int8 (INT8).
|
|
norms: [N] or [N,1] per-row L2 norms (bf16/fp16/fp32).
|
|
codebook: [2**bits] float32 Lloyd-Max centroids.
|
|
d_padded: output last dim (never trimmed to d at inference).
|
|
bits: 4 or 8.
|
|
out: optional [N, d_padded] output (dtype controls storage;
|
|
default float16 for inference, pass fp32 for parity tests).
|
|
|
|
Returns:
|
|
[N, d_padded] tensor.
|
|
"""
|
|
if not _HAS_TRITON:
|
|
raise RuntimeError("Triton is required for codebook_dequant_triton")
|
|
|
|
assert qdata.is_cuda, "codebook_dequant_triton requires CUDA input"
|
|
N = qdata.shape[0]
|
|
K = d_padded
|
|
|
|
if out is None:
|
|
out = torch.empty((N, K), dtype=torch.float16, device=qdata.device)
|
|
else:
|
|
assert out.shape == (N, K), f"out shape {out.shape} != ({N}, {K})"
|
|
|
|
norms1d = norms.reshape(-1).contiguous()
|
|
cb = codebook if codebook.is_contiguous() else codebook.contiguous()
|
|
|
|
BLOCK_K = _DEFAULT_BLOCK_K
|
|
grid = (N, triton.cdiv(K, BLOCK_K))
|
|
|
|
if bits == 4:
|
|
assert qdata.dtype == torch.uint8, f"INT4 path expects uint8, got {qdata.dtype}"
|
|
_codebook_dequant_int4_kernel[grid](
|
|
qdata, norms1d, cb, out,
|
|
qdata.stride(0), out.stride(0),
|
|
K, BLOCK_K=BLOCK_K,
|
|
)
|
|
elif bits == 8:
|
|
assert qdata.dtype == torch.int8, f"INT8 path expects int8, got {qdata.dtype}"
|
|
_codebook_dequant_int8_kernel[grid](
|
|
qdata, norms1d, cb, out,
|
|
qdata.stride(0), out.stride(0),
|
|
K, BLOCK_K=BLOCK_K,
|
|
)
|
|
else:
|
|
raise ValueError(f"Unsupported bits={bits}; only 4 or 8 supported")
|
|
|
|
return out
|
|
|
|
|
|
__all__ = ["codebook_dequant_triton"]
|