nodes/document97-Comfy-INT4-HIPtest
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
Files
main
document97-Comfy-INT4-HIPtest/native_int4.py
T

731 lines
27 KiB
Python
Raw Permalink Blame History

import logging
import os
import time
import torch
import triton
import triton.language as tl
from .convrot import build_hadamard as _build_hadamard
_LOGGED_NATIVE_PATH = False
_TARGET_CHECKED = False
_IS_GFX1103_TARGET = False
_TARGET_DESCRIPTION = "uninitialized"
_PROFILE_ENABLED = os.environ.get("COMFY_INT4_PROFILE", "0").lower() in ("1", "true", "yes", "on")
_PROFILE_FAILED = False
_PROFILE_COUNTS: dict[tuple, int] = {}
@triton.jit
def _round_to_nearest_even(x):
lower = tl.floor(x)
fraction = x - lower
lower_int = lower.to(tl.int32)
tie = tl.where((lower_int & 1) != 0, lower + 1.0, lower)
return tl.where(fraction > 0.5, lower + 1.0, tl.where(fraction < 0.5, lower, tie))
@triton.jit
def _convrot_matmul_kernel(
x_ptr,
h_ptr,
out_ptr,
group_rows,
GROUP_SIZE: tl.constexpr,
BLOCK_M: tl.constexpr,
BLOCK_N: tl.constexpr,
BLOCK_K: tl.constexpr,
):
pid_m = tl.program_id(0)
pid_n = tl.program_id(1)
offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
offs_k = tl.arange(0, BLOCK_K)
x_ptrs = x_ptr + offs_m[:, None] * GROUP_SIZE + offs_k[None, :]
h_ptrs = h_ptr + offs_k[:, None] * GROUP_SIZE + offs_n[None, :]
accumulator = tl.zeros((BLOCK_M, BLOCK_N), tl.float32)
for k_base in range(0, GROUP_SIZE, BLOCK_K):
x = tl.load(x_ptrs, mask=(offs_m[:, None] < group_rows) & (offs_k[None, :] + k_base < GROUP_SIZE), other=0.0)
h = tl.load(h_ptrs, mask=(offs_k[:, None] + k_base < GROUP_SIZE) & (offs_n[None, :] < GROUP_SIZE), other=0.0)
accumulator += tl.dot(x, h)
x_ptrs += BLOCK_K
h_ptrs += BLOCK_K * GROUP_SIZE
tl.store(
out_ptr + offs_m[:, None] * GROUP_SIZE + offs_n[None, :],
accumulator,
mask=(offs_m[:, None] < group_rows) & (offs_n[None, :] < GROUP_SIZE),
)
@triton.jit
def _quantize_i4_rowwise_kernel(x_ptr, q_ptr, scale_ptr, k_size, BLOCK_SIZE: tl.constexpr):
row = tl.program_id(0)
offsets = tl.arange(0, BLOCK_SIZE)
values = tl.load(x_ptr + row * k_size + offsets, mask=offsets < k_size, other=0.0).to(tl.float32)
absmax = tl.max(tl.abs(values), axis=0)
scale = tl.maximum(absmax / 7.0, 1e-10)
tl.store(scale_ptr + row, scale)
pair_offsets = tl.arange(0, BLOCK_SIZE // 2)
even_offsets = pair_offsets * 2
odd_offsets = even_offsets + 1
even = tl.load(x_ptr + row * k_size + even_offsets, mask=even_offsets < k_size, other=0.0).to(tl.float32) / scale
odd = tl.load(x_ptr + row * k_size + odd_offsets, mask=odd_offsets < k_size, other=0.0).to(tl.float32) / scale
even = _round_to_nearest_even(even)
odd = _round_to_nearest_even(odd)
even = tl.minimum(tl.maximum(even, -7.0), 7.0).to(tl.int32)
odd = tl.minimum(tl.maximum(odd, -7.0), 7.0).to(tl.int32)
packed = (even & 0xF) | ((odd & 0xF) << 4)
tl.store(q_ptr + row * (k_size // 2) + pair_offsets, packed.to(tl.int8), mask=pair_offsets < k_size // 2)
@triton.jit
def _wmma_i4_accumulate(
a0,
a1,
b0,
b1,
c0,
c1,
c2,
c3,
c4,
c5,
c6,
c7,
):
return tl.inline_asm_elementwise(
asm=(
"v_mov_b32 v8, $8\n\t"
"v_mov_b32 v9, $9\n\t"
"v_mov_b32 v10, $10\n\t"
"v_mov_b32 v11, $11\n\t"
"v_mov_b32 v16, $12\n\t"
"v_mov_b32 v17, $13\n\t"
"v_mov_b32 v18, $14\n\t"
"v_mov_b32 v19, $15\n\t"
"v_mov_b32 v20, $16\n\t"
"v_mov_b32 v21, $17\n\t"
"v_mov_b32 v22, $18\n\t"
"v_mov_b32 v23, $19\n\t"
"v_wmma_i32_16x16x16_iu4 v[24:31], v[8:9], v[10:11], v[16:23] neg_lo:[1,1,0]\n\t"
"v_mov_b32 $0, v24\n\t"
"v_mov_b32 $1, v25\n\t"
"v_mov_b32 $2, v26\n\t"
"v_mov_b32 $3, v27\n\t"
"v_mov_b32 $4, v28\n\t"
"v_mov_b32 $5, v29\n\t"
"v_mov_b32 $6, v30\n\t"
"v_mov_b32 $7, v31"
),
constraints=(
"=v,=v,=v,=v,=v,=v,=v,=v,"
"v,v,v,v,v,v,v,v,v,v,v,v,"
"~{v8},~{v9},~{v10},~{v11},"
"~{v16},~{v17},~{v18},~{v19},~{v20},~{v21},~{v22},~{v23},"
"~{v24},~{v25},~{v26},~{v27},~{v28},~{v29},~{v30},~{v31}"
),
args=[a0, a1, b0, b1, c0, c1, c2, c3, c4, c5, c6, c7],
dtype=(tl.int32,) * 8,
is_pure=True,
pack=1,
)
@triton.jit
def _load_i4_word(ptr, byte_offset, row_mask):
x0 = tl.load(ptr + byte_offset, mask=row_mask, other=0).to(tl.uint8).to(tl.uint32)
x1 = tl.load(ptr + byte_offset + 1, mask=row_mask, other=0).to(tl.uint8).to(tl.uint32)
x2 = tl.load(ptr + byte_offset + 2, mask=row_mask, other=0).to(tl.uint8).to(tl.uint32)
x3 = tl.load(ptr + byte_offset + 3, mask=row_mask, other=0).to(tl.uint8).to(tl.uint32)
return (x0 | (x1 << 8) | (x2 << 16) | (x3 << 24)).to(tl.int32, bitcast=True)
@triton.jit
def _w4a4_gemm_kernel(
a_ptr,
w_ptr,
a_scale_ptr,
w_scale_ptr,
bias_ptr,
out_ptr,
m_size,
n_size,
k_size,
stride_am,
stride_wn,
stride_om,
n_start,
HAS_BIAS: tl.constexpr,
):
tile_m = tl.program_id(0) * 16
tile_n = n_start + tl.program_id(1) * 16
wave_lane = tl.arange(0, 32)
lane = wave_lane % 16
lane_group = wave_lane // 16
row_a = tile_m + lane
row_w = tile_n + lane
mask_a = row_a < m_size
mask_w = row_w < n_size
c0 = tl.zeros((32,), tl.int32)
c1 = tl.zeros((32,), tl.int32)
c2 = tl.zeros((32,), tl.int32)
c3 = tl.zeros((32,), tl.int32)
c4 = tl.zeros((32,), tl.int32)
c5 = tl.zeros((32,), tl.int32)
c6 = tl.zeros((32,), tl.int32)
c7 = tl.zeros((32,), tl.int32)
for k_base in range(0, k_size, 16):
byte_k = k_base // 2
a_byte = row_a * stride_am + byte_k
w_byte = row_w * stride_wn + byte_k
a0 = _load_i4_word(a_ptr, a_byte, mask_a)
a1 = _load_i4_word(a_ptr, a_byte + 4, mask_a)
b0 = _load_i4_word(w_ptr, w_byte, mask_w)
b1 = _load_i4_word(w_ptr, w_byte + 4, mask_w)
c0, c1, c2, c3, c4, c5, c6, c7 = _wmma_i4_accumulate(
a0,
a1,
b0,
b1,
c0,
c1,
c2,
c3,
c4,
c5,
c6,
c7,
)
# Convert metadata in registers. Keeping the source tensor in its stored
# dtype avoids a separate eager dtype-conversion kernel on every forward.
w_scales = tl.load(w_scale_ptr + tile_n + lane, mask=mask_w, other=0.0).to(tl.float32)
if HAS_BIAS:
bias = tl.load(bias_ptr + tile_n + lane, mask=mask_w, other=0.0).to(tl.float32)
else:
bias = tl.zeros((32,), tl.float32)
accumulators = (c0, c1, c2, c3, c4, c5, c6, c7)
for element in tl.static_range(8):
row_in_tile = element * 2 + lane_group
row = tile_m + row_in_tile
col = tile_n + lane
a_scale = tl.load(a_scale_ptr + row, mask=row < m_size, other=0.0).to(tl.float32)
value = accumulators[element].to(tl.float32)
value *= a_scale * w_scales
value += bias
tl.store(out_ptr + row * stride_om + col, value, mask=(row < m_size) & (col < n_size))
@triton.jit
def _w4a4_gemm_32x32_kernel(
a_ptr,
w_ptr,
a_scale_ptr,
w_scale_ptr,
bias_ptr,
out_ptr,
m_size,
n_size,
k_size,
stride_am,
stride_wn,
stride_om,
n_start,
HAS_BIAS: tl.constexpr,
):
tile_m = tl.program_id(0) * 32
tile_n = n_start + tl.program_id(1) * 32
wave_lane = tl.arange(0, 32)
lane = wave_lane % 16
lane_group = wave_lane // 16
row_a0 = tile_m + lane
row_a1 = tile_m + 16 + lane
row_w0 = tile_n + lane
row_w1 = tile_n + 16 + lane
mask_a0 = row_a0 < m_size
mask_a1 = row_a1 < m_size
mask_w0 = row_w0 < n_size
mask_w1 = row_w1 < n_size
tl0 = tl.zeros((32,), tl.int32)
tl1 = tl.zeros((32,), tl.int32)
tl2 = tl.zeros((32,), tl.int32)
tl3 = tl.zeros((32,), tl.int32)
tl4 = tl.zeros((32,), tl.int32)
tl5 = tl.zeros((32,), tl.int32)
tl6 = tl.zeros((32,), tl.int32)
tl7 = tl.zeros((32,), tl.int32)
tr0 = tl.zeros((32,), tl.int32)
tr1 = tl.zeros((32,), tl.int32)
tr2 = tl.zeros((32,), tl.int32)
tr3 = tl.zeros((32,), tl.int32)
tr4 = tl.zeros((32,), tl.int32)
tr5 = tl.zeros((32,), tl.int32)
tr6 = tl.zeros((32,), tl.int32)
tr7 = tl.zeros((32,), tl.int32)
bl0 = tl.zeros((32,), tl.int32)
bl1 = tl.zeros((32,), tl.int32)
bl2 = tl.zeros((32,), tl.int32)
bl3 = tl.zeros((32,), tl.int32)
bl4 = tl.zeros((32,), tl.int32)
bl5 = tl.zeros((32,), tl.int32)
bl6 = tl.zeros((32,), tl.int32)
bl7 = tl.zeros((32,), tl.int32)
br0 = tl.zeros((32,), tl.int32)
br1 = tl.zeros((32,), tl.int32)
br2 = tl.zeros((32,), tl.int32)
br3 = tl.zeros((32,), tl.int32)
br4 = tl.zeros((32,), tl.int32)
br5 = tl.zeros((32,), tl.int32)
br6 = tl.zeros((32,), tl.int32)
br7 = tl.zeros((32,), tl.int32)
for k_base in range(0, k_size, 16):
byte_k = k_base // 2
a0_byte = row_a0 * stride_am + byte_k
a1_byte = row_a1 * stride_am + byte_k
w0_byte = row_w0 * stride_wn + byte_k
w1_byte = row_w1 * stride_wn + byte_k
at0 = _load_i4_word(a_ptr, a0_byte, mask_a0)
at1 = _load_i4_word(a_ptr, a0_byte + 4, mask_a0)
ab0 = _load_i4_word(a_ptr, a1_byte, mask_a1)
ab1 = _load_i4_word(a_ptr, a1_byte + 4, mask_a1)
wl0 = _load_i4_word(w_ptr, w0_byte, mask_w0)
wl1 = _load_i4_word(w_ptr, w0_byte + 4, mask_w0)
wr0 = _load_i4_word(w_ptr, w1_byte, mask_w1)
wr1 = _load_i4_word(w_ptr, w1_byte + 4, mask_w1)
tl0, tl1, tl2, tl3, tl4, tl5, tl6, tl7 = _wmma_i4_accumulate(
at0, at1, wl0, wl1, tl0, tl1, tl2, tl3, tl4, tl5, tl6, tl7
)
tr0, tr1, tr2, tr3, tr4, tr5, tr6, tr7 = _wmma_i4_accumulate(
at0, at1, wr0, wr1, tr0, tr1, tr2, tr3, tr4, tr5, tr6, tr7
)
bl0, bl1, bl2, bl3, bl4, bl5, bl6, bl7 = _wmma_i4_accumulate(
ab0, ab1, wl0, wl1, bl0, bl1, bl2, bl3, bl4, bl5, bl6, bl7
)
br0, br1, br2, br3, br4, br5, br6, br7 = _wmma_i4_accumulate(
ab0, ab1, wr0, wr1, br0, br1, br2, br3, br4, br5, br6, br7
)
# Scale/bias may be FP16, BF16, or FP32 in the checkpoint. Load the
# original dtype and promote in registers instead of materializing FP32
# copies through PyTorch eager operations.
w_scale0 = tl.load(w_scale_ptr + row_w0, mask=mask_w0, other=0.0).to(tl.float32)
w_scale1 = tl.load(w_scale_ptr + row_w1, mask=mask_w1, other=0.0).to(tl.float32)
if HAS_BIAS:
bias0 = tl.load(bias_ptr + row_w0, mask=mask_w0, other=0.0).to(tl.float32)
bias1 = tl.load(bias_ptr + row_w1, mask=mask_w1, other=0.0).to(tl.float32)
else:
bias0 = tl.zeros((32,), tl.float32)
bias1 = tl.zeros((32,), tl.float32)
top_left = (tl0, tl1, tl2, tl3, tl4, tl5, tl6, tl7)
top_right = (tr0, tr1, tr2, tr3, tr4, tr5, tr6, tr7)
bottom_left = (bl0, bl1, bl2, bl3, bl4, bl5, bl6, bl7)
bottom_right = (br0, br1, br2, br3, br4, br5, br6, br7)
for element in tl.static_range(8):
row_in_tile = element * 2 + lane_group
row0 = tile_m + row_in_tile
row1 = tile_m + 16 + row_in_tile
col0 = tile_n + lane
col1 = tile_n + 16 + lane
a_scale0 = tl.load(a_scale_ptr + row0, mask=row0 < m_size, other=0.0).to(tl.float32)
a_scale1 = tl.load(a_scale_ptr + row1, mask=row1 < m_size, other=0.0).to(tl.float32)
value_tl = top_left[element].to(tl.float32) * a_scale0 * w_scale0 + bias0
value_tr = top_right[element].to(tl.float32) * a_scale0 * w_scale1 + bias1
value_bl = bottom_left[element].to(tl.float32) * a_scale1 * w_scale0 + bias0
value_br = bottom_right[element].to(tl.float32) * a_scale1 * w_scale1 + bias1
tl.store(out_ptr + row0 * stride_om + col0, value_tl, mask=(row0 < m_size) & (col0 < n_size))
tl.store(out_ptr + row0 * stride_om + col1, value_tr, mask=(row0 < m_size) & (col1 < n_size))
tl.store(out_ptr + row1 * stride_om + col0, value_bl, mask=(row1 < m_size) & (col0 < n_size))
tl.store(out_ptr + row1 * stride_om + col1, value_br, mask=(row1 < m_size) & (col1 < n_size))
@triton.jit
def _w4a4_dot_gemm_kernel(
a_ptr,
w_ptr,
a_scale_ptr,
w_scale_ptr,
bias_ptr,
out_ptr,
m_size,
n_size,
stride_am,
stride_wn,
stride_om,
K_SIZE: tl.constexpr,
HAS_BIAS: tl.constexpr,
BLOCK_M: tl.constexpr,
BLOCK_N: tl.constexpr,
BLOCK_K: tl.constexpr,
):
"""Multi-wave signed W4A4 GEMM lowered by Triton to RDNA3 IU4 WMMA."""
pid_m = tl.program_id(0)
pid_n = tl.program_id(1)
offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
offs_k_packed = tl.arange(0, BLOCK_K // 2)
accumulator = tl.zeros((BLOCK_M, BLOCK_N), tl.int32)
for k_base in range(0, K_SIZE, BLOCK_K):
packed_k = k_base // 2
a_bytes = tl.load(
a_ptr
+ offs_m[:, None] * stride_am
+ packed_k
+ offs_k_packed[None, :],
mask=offs_m[:, None] < m_size,
other=0,
)
w_bytes = tl.load(
w_ptr
+ offs_n[:, None] * stride_wn
+ packed_k
+ offs_k_packed[None, :],
mask=offs_n[:, None] < n_size,
other=0,
)
# tl.join adds a minor dimension; reshape restores low/high nibbles in
# their original K order without ever materializing an INT8 matrix.
a_i4 = tl.reshape(
tl.join((a_bytes & 0xF).to(tl.int4), (a_bytes >> 4).to(tl.int4)),
(BLOCK_M, BLOCK_K),
)
w_i4 = tl.reshape(
tl.join((w_bytes & 0xF).to(tl.int4), (w_bytes >> 4).to(tl.int4)),
(BLOCK_N, BLOCK_K),
)
accumulator += tl.dot(a_i4, tl.trans(w_i4), out_dtype=tl.int32)
a_scales = tl.load(
a_scale_ptr + offs_m, mask=offs_m < m_size, other=0.0
).to(tl.float32)
w_scales = tl.load(
w_scale_ptr + offs_n, mask=offs_n < n_size, other=0.0
).to(tl.float32)
values = accumulator.to(tl.float32) * a_scales[:, None] * w_scales[None, :]
if HAS_BIAS:
biases = tl.load(
bias_ptr + offs_n, mask=offs_n < n_size, other=0.0
).to(tl.float32)
values += biases[None, :]
output_offsets = offs_m[:, None] * stride_om + offs_n[None, :]
output_mask = (offs_m[:, None] < m_size) & (offs_n[None, :] < n_size)
tl.store(out_ptr + output_offsets, values, mask=output_mask)
def triton_w4a4_gemm(
activation: torch.Tensor,
weight: torch.Tensor,
activation_scale: torch.Tensor,
weight_scale: torch.Tensor,
bias: torch.Tensor | None = None,
output_dtype: torch.dtype = torch.float16,
) -> torch.Tensor:
if activation.ndim != 2 or weight.ndim != 2:
raise ValueError("W4A4 GEMM expects 2D packed activation and weight tensors")
if activation.dtype != torch.int8 or weight.dtype != torch.int8:
raise TypeError("W4A4 GEMM expects packed tensors stored as int8")
if activation.shape[1] != weight.shape[1]:
raise ValueError("W4A4 activation and weight K dimensions do not match")
m_size = activation.shape[0]
n_size = weight.shape[0]
k_size = activation.shape[1] * 2
if k_size % 64 != 0:
raise ValueError(f"W4A4 GEMM requires K divisible by 64, got {k_size}")
if not activation.is_contiguous() or not weight.is_contiguous():
raise ValueError("W4A4 GEMM expects contiguous packed tensors")
supported_metadata_dtypes = (torch.float16, torch.bfloat16, torch.float32)
if activation_scale.dtype not in supported_metadata_dtypes:
raise TypeError(f"Unsupported activation scale dtype {activation_scale.dtype}")
if weight_scale.dtype not in supported_metadata_dtypes:
raise TypeError(f"Unsupported weight scale dtype {weight_scale.dtype}")
if activation_scale.device != activation.device or weight_scale.device != activation.device:
raise ValueError("W4A4 scales must already be on the activation device")
if not activation_scale.is_contiguous() or not weight_scale.is_contiguous():
raise ValueError("W4A4 scales must be contiguous; the native hot path will not copy them")
if activation_scale.numel() != m_size or weight_scale.numel() != n_size:
raise ValueError("W4A4 scale shapes do not match activation or weight rows")
if bias is not None:
if bias.dtype not in supported_metadata_dtypes:
raise TypeError(f"Unsupported W4A4 bias dtype {bias.dtype}")
if bias.device != activation.device:
raise ValueError("W4A4 bias must already be on the activation device")
if not bias.is_contiguous():
raise ValueError("W4A4 bias must be contiguous; the native hot path will not copy it")
if bias.numel() != n_size:
raise ValueError("W4A4 bias shape does not match weight rows")
if output_dtype not in (torch.float16, torch.bfloat16):
raise TypeError(f"W4A4 output must be FP16 or BF16, got {output_dtype}")
output = torch.empty((m_size, n_size), device=activation.device, dtype=output_dtype)
block_m, block_n, block_k, num_warps = 64, 128, 64, 8
grid = (triton.cdiv(m_size, block_m), triton.cdiv(n_size, block_n))
_w4a4_dot_gemm_kernel[grid](
activation,
weight,
activation_scale,
weight_scale,
bias if bias is not None else weight_scale,
output,
m_size,
n_size,
stride_am=activation.stride(0),
stride_wn=weight.stride(0),
stride_om=output.stride(0),
K_SIZE=k_size,
HAS_BIAS=bias is not None,
BLOCK_M=block_m,
BLOCK_N=block_n,
BLOCK_K=block_k,
num_warps=num_warps,
num_stages=1,
)
return output
def _select_w4a4_tile(m_size: int, n_size: int) -> int:
"""Choose between the two precompiled WMMA layouts without autotuning.
A 32x32 program executes four 16x16 WMMA operations. It is faster for
dense tiles because it reuses A and W loads, but can do substantially more
work on narrow/tail-heavy matrices. Keep the larger tile when its WMMA
count is within 25% of the 16x16 layout.
"""
if min(m_size, n_size) <= 16:
return 16
wmma_16 = triton.cdiv(m_size, 16) * triton.cdiv(n_size, 16)
wmma_32 = 4 * triton.cdiv(m_size, 32) * triton.cdiv(n_size, 32)
return 32 if wmma_32 * 4 <= wmma_16 * 5 else 16
def _select_dispatch_n_chunk(m_size: int, n_size: int, k_size: int, tile_size: int) -> int:
"""Bound a single gfx1103 compute dispatch without changing IU4 math.
Profiling the 1024x1024 Krea2 workload showed individual native GEMM
dispatches as long as 1.96 seconds. The same iGPU also drives DWM, so
those monolithic dispatches make the desktop unresponsive and approach
Windows TDR limits. Split only genuinely large GEMMs; small/text paths
retain one launch. 1024 is divisible by both supported tile sizes.
"""
work = m_size * n_size * k_size
if work < 100_000_000_000 or n_size <= 1024:
return n_size
return min(n_size, 1024 - (1024 % tile_size))
def _profile_this_shape(key: tuple) -> bool:
if not _PROFILE_ENABLED or _PROFILE_FAILED:
return False
count = _PROFILE_COUNTS.get(key, 0)
_PROFILE_COUNTS[key] = count + 1
# The first call may include Triton compilation. Sample the second call
# once so the diagnostic reflects the steady-state device timeline.
return count == 1
def is_gfx1103_target() -> bool:
global _TARGET_CHECKED, _IS_GFX1103_TARGET, _TARGET_DESCRIPTION
if _TARGET_CHECKED:
return _IS_GFX1103_TARGET
try:
target = triton.runtime.driver.active.get_current_target()
except Exception:
# Triton's driver may not be initialized during module import. Do not
# cache that transient failure; the first real GPU forward may succeed.
return False
_IS_GFX1103_TARGET = (
target.backend == "hip" and target.arch == "gfx1103" and target.warp_size == 32
)
_TARGET_DESCRIPTION = (
f"backend={target.backend!r}, arch={target.arch!r}, warp_size={target.warp_size}"
)
_TARGET_CHECKED = True
return _IS_GFX1103_TARGET
def _require_gfx11_hip_target() -> None:
if not is_gfx1103_target():
raise RuntimeError(
"Native INT4 ConvRot requires the gfx1103 HIP Triton target with wave32; "
f"got {_TARGET_DESCRIPTION}"
)
def triton_rotate_activation(
x: torch.Tensor,
group_size: int,
output_dtype: torch.dtype | None = None,
) -> torch.Tensor:
_require_gfx11_hip_target()
if x.ndim != 2 or not x.is_contiguous():
raise ValueError("ConvRot expects a contiguous 2D activation tensor")
if x.dtype not in (torch.float16, torch.bfloat16, torch.float32):
raise TypeError(f"ConvRot expects FP16, BF16, or FP32 activation, got {x.dtype}")
if output_dtype is None:
output_dtype = x.dtype
if output_dtype not in (torch.float16, torch.bfloat16):
raise TypeError(f"ConvRot output must be FP16 or BF16, got {output_dtype}")
if group_size not in (16, 64, 256):
raise ValueError(f"ConvRot group size must be 16, 64, or 256, got {group_size}")
if x.shape[1] % group_size != 0:
raise ValueError(
f"Activation K={x.shape[1]} must be divisible by ConvRot group size {group_size}"
)
grouped = x.reshape(-1, group_size)
hadamard = _build_hadamard(group_size, device=x.device, dtype=x.dtype).contiguous()
output = torch.empty(grouped.shape, device=x.device, dtype=output_dtype)
block_n = 16 if group_size == 16 else 32
block_k = 16 if group_size == 16 else 32
grid = (triton.cdiv(grouped.shape[0], 16), triton.cdiv(group_size, block_n))
_convrot_matmul_kernel[grid](
grouped,
hadamard,
output,
grouped.shape[0],
GROUP_SIZE=group_size,
BLOCK_M=16,
BLOCK_N=block_n,
BLOCK_K=block_k,
num_warps=4,
)
return output.reshape_as(x)
def triton_quantize_int4_rowwise(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
_require_gfx11_hip_target()
if x.ndim != 2 or not x.is_contiguous():
raise ValueError("INT4 activation quantization expects a contiguous 2D tensor")
if x.dtype not in (torch.float16, torch.bfloat16):
raise TypeError(f"INT4 activation quantization expects FP16 or BF16, got {x.dtype}")
rows, k_size = x.shape
if k_size % 64 != 0:
raise ValueError(f"INT4 activation quantization requires K divisible by 64, got {k_size}")
block_size = triton.next_power_of_2(k_size)
packed = torch.empty((rows, k_size // 2), device=x.device, dtype=torch.int8)
scales = torch.empty((rows,), device=x.device, dtype=torch.float32)
num_warps = 8 if block_size >= 2048 else 4
_quantize_i4_rowwise_kernel[(rows,)](
x,
packed,
scales,
k_size,
BLOCK_SIZE=block_size,
num_warps=num_warps,
)
return packed, scales
def triton_convrot_w4a4_linear(
x: torch.Tensor,
weight: torch.Tensor,
weight_scale: torch.Tensor,
bias: torch.Tensor | None = None,
*,
convrot_groupsize: int = 256,
quant_group_size: int = 64,
linear_dtype: str = "int4",
output_dtype: torch.dtype | None = None,
) -> torch.Tensor:
global _LOGGED_NATIVE_PATH, _PROFILE_FAILED
_require_gfx11_hip_target()
if not _LOGGED_NATIVE_PATH:
logging.info(
"INT4 Fast: using native gfx1103 Triton ConvRot W4A4 "
"(v_wmma_i32_16x16x16_iu4, no eager fallback)."
)
_LOGGED_NATIVE_PATH = True
if linear_dtype != "int4":
raise ValueError(f"Native W4A4 kernel only supports linear_dtype='int4', got {linear_dtype!r}")
if quant_group_size != 64:
raise ValueError(f"Native W4A4 kernel requires quant_group_size=64, got {quant_group_size}")
if x.device != weight.device:
raise ValueError(f"Activation and packed weight must be on the same device, got {x.device} and {weight.device}")
if not x.is_contiguous():
raise ValueError("Native W4A4 activation must be contiguous; the hot path will not make an eager copy")
if not weight.is_contiguous():
raise ValueError("Native W4A4 packed weight must be contiguous; the hot path will not make an eager copy")
if output_dtype is None:
output_dtype = x.dtype
if output_dtype not in (torch.float16, torch.bfloat16):
raise TypeError(f"Native W4A4 output must be FP16 or BF16, got {output_dtype}")
original_shape = x.shape
x_2d = x.reshape(-1, original_shape[-1])
profile_key = (
x_2d.shape[0],
weight.shape[0],
x_2d.shape[1],
convrot_groupsize,
x.dtype,
output_dtype,
weight_scale.dtype,
None if bias is None else bias.dtype,
)
profile_this = _profile_this_shape(profile_key)
events = None
cpu_times = None
if profile_this:
try:
events = [torch.cuda.Event(enable_timing=True) for _ in range(4)]
events[0].record()
cpu_times = [time.perf_counter()]
except Exception as exc:
_PROFILE_FAILED = True
logging.warning("INT4 Fast profiler disabled after event setup failed: %s", exc)
rotated = triton_rotate_activation(x_2d, convrot_groupsize, output_dtype=output_dtype)
if events is not None:
cpu_times.append(time.perf_counter())
events[1].record()
packed_activation, activation_scale = triton_quantize_int4_rowwise(rotated)
if events is not None:
cpu_times.append(time.perf_counter())
events[2].record()
del rotated
output = triton_w4a4_gemm(
packed_activation,
weight,
activation_scale,
weight_scale,
bias,
output_dtype=output_dtype,
)
if events is not None:
cpu_times.append(time.perf_counter())
events[3].record()
try:
events[3].synchronize()
gpu_ms = [events[index].elapsed_time(events[index + 1]) for index in range(3)]
cpu_ms = [
(cpu_times[index + 1] - cpu_times[index]) * 1000.0
for index in range(3)
]
logging.info(
"INT4 Fast profile M=%d N=%d K=%d group=%d tile=64x128x64/8w: "
"GPU rotate=%.3fms quant=%.3fms gemm=%.3fms; "
"CPU submit rotate=%.3fms quant=%.3fms gemm=%.3fms",
x_2d.shape[0],
weight.shape[0],
x_2d.shape[1],
convrot_groupsize,
*gpu_ms,
*cpu_ms,
)
except Exception as exc:
_PROFILE_FAILED = True
logging.warning("INT4 Fast profiler disabled after timing failed: %s", exc)
return output.reshape(*original_shape[:-1], weight.shape[0])
Reference in New Issue View Git Blame Copy Permalink