[CI] Prune redundant CPU test overhead (#34309)
Co-authored-by: hnyls2002 <lsyincs@gmail.com>
This commit is contained in:
co-authored by
hnyls2002
parent
6ad3f2d8fd
commit
85cdf1178d
@@ -9603,7 +9603,7 @@ def m3_fp8_attn_gemm_enabled(args) -> bool:
|
||||
# (sglang.srt.runtime_context). The two functions below are LEGACY shims kept
|
||||
# for the existing call-sites; they publish/read the same live object by
|
||||
# reference. Do not add new call-sites — the counts are ratcheted
|
||||
# (decrease-only) by test/registered/unit/test_legacy_global_ratchet.py.
|
||||
# (decrease-only) by scripts/lint/check_legacy_global_ratchet.py.
|
||||
# Imports are in-function so the two modules stay cycle-free at import time.
|
||||
def set_global_server_args_for_scheduler(server_args: ServerArgs):
|
||||
"""Legacy publish shim (role=scheduler) — prefer
|
||||
|
||||
@@ -0,0 +1,659 @@
|
||||
import itertools
|
||||
import math
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
precision = {
|
||||
torch.bfloat16: 1e-2,
|
||||
torch.float16: 1e-3,
|
||||
torch.float32: 1e-5,
|
||||
}
|
||||
|
||||
|
||||
BLOCK_N, BLOCK_K = 64, 128
|
||||
factor_for_scale = 1e-3
|
||||
fp8_max, fp8_min = 400, -400
|
||||
|
||||
|
||||
def parametrize(**params):
|
||||
def decorator(func):
|
||||
def wrapper(self):
|
||||
for combo in itertools.product(*params.values()):
|
||||
kwargs = dict(zip(params.keys(), combo))
|
||||
with self.subTest(**kwargs):
|
||||
func(self, **kwargs)
|
||||
|
||||
return wrapper
|
||||
|
||||
return decorator
|
||||
|
||||
|
||||
def SiluAndMul(x: torch.Tensor) -> torch.Tensor:
|
||||
d = x.shape[-1] // 2
|
||||
return F.silu(x[..., :d]) * x[..., d:]
|
||||
|
||||
|
||||
def GeluAndMul(x: torch.Tensor, approximate="tanh") -> torch.Tensor:
|
||||
d = x.shape[-1] // 2
|
||||
return F.gelu(x[..., :d], approximate=approximate) * x[..., d:]
|
||||
|
||||
|
||||
def per_token_quant_int8(x):
|
||||
x = x.float()
|
||||
absmax = x.abs().max(dim=-1).values
|
||||
absmax = absmax.clamp_min(1e-10).unsqueeze(-1)
|
||||
scale_x = absmax / 127
|
||||
x_q = x.mul(127 / absmax)
|
||||
x_q = torch.round(x_q).to(torch.int8)
|
||||
|
||||
return x_q, scale_x
|
||||
|
||||
|
||||
def convert_weight(weight, scale_block_size, A_dtype):
|
||||
N, K = weight.size()
|
||||
fp8_max = 448.0
|
||||
scale_block_size_N, scale_block_size_K = scale_block_size # (128, 128)
|
||||
|
||||
pad_N = (scale_block_size_N - (N % scale_block_size_N)) % scale_block_size_N
|
||||
pad_K = (scale_block_size_K - (K % scale_block_size_K)) % scale_block_size_K
|
||||
|
||||
if pad_N > 0 or pad_K > 0:
|
||||
weight = torch.nn.functional.pad(weight, (0, pad_K, 0, pad_N))
|
||||
|
||||
weight_blocks = weight.view(
|
||||
math.ceil(N / scale_block_size_N),
|
||||
scale_block_size_N,
|
||||
math.ceil(K / scale_block_size_K),
|
||||
scale_block_size_K,
|
||||
) # (8, 128, 8, 128)
|
||||
weight_blocks = weight_blocks.permute(0, 2, 1, 3).contiguous() # (8, 8, 128, 128)
|
||||
|
||||
# Step 2: compute per-block max abs values → scale
|
||||
abs_max = weight_blocks.abs().amax(dim=(-2, -1), keepdim=True) # (8, 8, 1, 1)
|
||||
scales = abs_max / fp8_max
|
||||
scales = torch.where(
|
||||
scales == 0, torch.ones_like(scales), scales
|
||||
) # avoid division by zero
|
||||
|
||||
q_fp8 = (weight_blocks / scales).to(torch.float8_e4m3fn)
|
||||
q_fp8_reshape = q_fp8.permute(0, 2, 1, 3).contiguous()
|
||||
|
||||
if pad_N > 0 or pad_K > 0:
|
||||
q_fp8_reshape = q_fp8_reshape.view(N + pad_N, K + pad_K)
|
||||
q_fp8_reshape = q_fp8_reshape[:N, :K].contiguous()
|
||||
else:
|
||||
q_fp8_reshape = q_fp8_reshape.view(N, K)
|
||||
|
||||
dq_weight = q_fp8.float() * scales
|
||||
dq_weight = dq_weight.permute(0, 2, 1, 3).contiguous() # (8, 128, 8, 128)
|
||||
|
||||
if pad_N > 0 or pad_K > 0:
|
||||
w_dq = dq_weight.view(N + pad_N, K + pad_K).to(A_dtype)
|
||||
w_dq = w_dq[:N, :K].contiguous()
|
||||
else:
|
||||
w_dq = dq_weight.view(N, K).to(A_dtype)
|
||||
|
||||
scales = scales.view(
|
||||
math.ceil(N / scale_block_size_N), math.ceil(K / scale_block_size_K)
|
||||
)
|
||||
|
||||
return q_fp8_reshape, scales, w_dq
|
||||
|
||||
|
||||
def native_w8a8_per_token_matmul(A, B, As, Bs, bias, output_dtype=torch.bfloat16):
|
||||
"""Matrix multiplication function that supports per-token input quantization and per-column weight quantization"""
|
||||
A = A.to(torch.float32)
|
||||
B = B.to(torch.float32)
|
||||
|
||||
assert A.shape[-1] == B.shape[-1], "Dimension mismatch"
|
||||
assert B.ndim == 2 and B.is_contiguous(), "B must be a 2D contiguous tensor"
|
||||
|
||||
# Reshape input
|
||||
M = A.numel() // A.shape[-1]
|
||||
B = B.t() # Transpose weight matrix
|
||||
N, K = B.shape
|
||||
origin_C_shape = A.shape[:-1] + (K,)
|
||||
A = A.reshape(M, N)
|
||||
|
||||
# As is per-token [M, 1], Bs is per-column [1, K]
|
||||
C = torch.matmul(A, B) # [M, K]
|
||||
C = As * C * Bs.view(1, -1) # Broadcast per-column scale
|
||||
|
||||
if bias is not None:
|
||||
C.add_(bias.view(1, -1))
|
||||
|
||||
return C.reshape(origin_C_shape).to(output_dtype)
|
||||
|
||||
|
||||
def torch_naive_moe(a, w1, w2, b, routed_scaling_factor, output_dtype=torch.bfloat16):
|
||||
|
||||
a = a.to(torch.float32)
|
||||
w1 = w1.to(torch.float32)
|
||||
w2 = w2.to(torch.float32)
|
||||
b = b.to(torch.float32) if b is not None else None
|
||||
|
||||
ic1 = torch.matmul(a, w1.transpose(0, 1))
|
||||
ic2 = SiluAndMul(ic1)
|
||||
ic3 = torch.matmul(ic2, w2.transpose(0, 1))
|
||||
|
||||
out = ic3 if b is None else ic3 + b * routed_scaling_factor
|
||||
|
||||
return out.to(output_dtype)
|
||||
|
||||
|
||||
def torch_w8a8_per_column_moe(
|
||||
a, w1_q, w2_q, w1_s, w2_s, b, routed_scaling_factor, output_dtype=torch.bfloat16
|
||||
):
|
||||
|
||||
a = a.to(torch.float32)
|
||||
b = b.to(torch.float32) if b is not None else None
|
||||
|
||||
# Perform per-token quantization
|
||||
a_q, a_s = per_token_quant_int8(a)
|
||||
|
||||
ic1 = native_w8a8_per_token_matmul(
|
||||
a_q, w1_q, a_s, w1_s, bias=None, output_dtype=torch.float32
|
||||
)
|
||||
ic2 = SiluAndMul(ic1)
|
||||
|
||||
a1_q, a1_s = per_token_quant_int8(ic2)
|
||||
ic3 = native_w8a8_per_token_matmul(
|
||||
a1_q, w2_q, a1_s, w2_s, bias=None, output_dtype=torch.float32
|
||||
)
|
||||
|
||||
out = ic3 if b is None else ic3 + b * routed_scaling_factor
|
||||
|
||||
return out.to(output_dtype)
|
||||
|
||||
|
||||
def scaled_weight(weight, scales):
|
||||
E, N, K = weight.shape
|
||||
pad_N = (BLOCK_N - (N % BLOCK_N)) % BLOCK_N
|
||||
pad_K = (BLOCK_K - (K % BLOCK_K)) % BLOCK_K
|
||||
|
||||
if pad_N > 0 or pad_K > 0:
|
||||
weight = torch.nn.functional.pad(weight, (0, pad_K, 0, pad_N))
|
||||
|
||||
weight_block = (
|
||||
weight.view(E, math.ceil(N / BLOCK_N), BLOCK_N, math.ceil(K / BLOCK_K), BLOCK_K)
|
||||
.permute(0, 1, 3, 2, 4)
|
||||
.float()
|
||||
.contiguous()
|
||||
)
|
||||
|
||||
weight_scaled = (
|
||||
(
|
||||
weight_block
|
||||
* scales.view(E, math.ceil(N / BLOCK_N), math.ceil(K / BLOCK_K), 1, 1)
|
||||
)
|
||||
.permute(0, 1, 3, 2, 4)
|
||||
.contiguous()
|
||||
)
|
||||
if pad_N > 0 or pad_K > 0:
|
||||
weight_scaled = weight_scaled.view(E, N + pad_N, K + pad_K)
|
||||
weight_scaled = weight_scaled[..., :N, :K].contiguous()
|
||||
else:
|
||||
weight_scaled = weight_scaled.view(E, N, K)
|
||||
return weight_scaled
|
||||
|
||||
|
||||
def torch_naive_fused_moe(a, w1, w2, score, topk, renormalize):
|
||||
B, D = a.shape
|
||||
a = a.view(B, -1, D).repeat(1, topk, 1).reshape(-1, D)
|
||||
out = torch.zeros(B * topk, w2.shape[1], dtype=a.dtype, device=a.device)
|
||||
score = torch.softmax(score, dim=-1, dtype=torch.float32)
|
||||
topk_weight, topk_ids = torch.topk(score, topk)
|
||||
|
||||
if renormalize:
|
||||
topk_weight = topk_weight / topk_weight.sum(dim=-1, keepdim=True)
|
||||
|
||||
topk_weight = topk_weight.view(-1)
|
||||
topk_ids = topk_ids.view(-1)
|
||||
for i in range(w1.shape[0]):
|
||||
mask = topk_ids == i
|
||||
if mask.sum():
|
||||
out[mask] = SiluAndMul(a[mask] @ w1[i].transpose(0, 1)) @ w2[i].transpose(
|
||||
0, 1
|
||||
)
|
||||
return (
|
||||
out.view(B, -1, w2.shape[1]) * topk_weight.view(B, -1, 1).to(out.dtype)
|
||||
).sum(dim=1)
|
||||
|
||||
|
||||
def moe_gptoss_act(x, alpha: float = 1.702, limit: float = 7.0):
|
||||
x_glu, x_linear = x[..., ::2], x[..., 1::2]
|
||||
# Clamp the input values
|
||||
x_glu = x_glu.clamp(min=None, max=limit)
|
||||
x_linear = x_linear.clamp(min=-limit, max=limit)
|
||||
out_glu = x_glu * torch.sigmoid(alpha * x_glu)
|
||||
# Note we add an extra bias of 1 to the linear layer
|
||||
return out_glu * (x_linear + 1.0)
|
||||
|
||||
|
||||
def torch_naive_gptoss_fused_moe(
|
||||
x,
|
||||
w1,
|
||||
w2,
|
||||
w1_bias,
|
||||
w2_bias,
|
||||
topk_weights,
|
||||
topk_ids,
|
||||
activation_alpha,
|
||||
swiglu_limit,
|
||||
len_experts,
|
||||
) -> torch.Tensor:
|
||||
|
||||
# Ref code from https://huggingface.co/deepseek-ai/DeepSeek-V2/blob/e0828e3cc0a03408724b80c3cc92c8e072db8d01/modeling_deepseek.py#L589
|
||||
cnts = topk_ids.new_zeros((topk_ids.shape[0], len_experts))
|
||||
cnts.scatter_(1, topk_ids.to(torch.int64), 1)
|
||||
tokens_per_expert = cnts.sum(dim=0)
|
||||
idxs = topk_ids.view(-1).argsort()
|
||||
|
||||
sorted_tokens = x[idxs // topk_ids.shape[1]]
|
||||
tokens_per_expert = tokens_per_expert.cpu().numpy()
|
||||
|
||||
outputs = []
|
||||
start_idx = 0
|
||||
for i, num_tokens in enumerate(tokens_per_expert):
|
||||
end_idx = start_idx + num_tokens
|
||||
if num_tokens == 0:
|
||||
continue
|
||||
tokens_for_this_expert = sorted_tokens[start_idx:end_idx]
|
||||
|
||||
layer_w13_weight = w1[i]
|
||||
layer_w13_weight_bias = w1_bias[i]
|
||||
layer_w2_weight_bias = w2_bias[i]
|
||||
layer_w2_weight = w2[i]
|
||||
|
||||
gate_up = F.linear(
|
||||
tokens_for_this_expert,
|
||||
layer_w13_weight,
|
||||
bias=layer_w13_weight_bias.to(torch.bfloat16),
|
||||
)
|
||||
gate_up = moe_gptoss_act(gate_up, activation_alpha, swiglu_limit)
|
||||
expert_out = F.linear(
|
||||
gate_up, layer_w2_weight, bias=layer_w2_weight_bias.to(torch.bfloat16)
|
||||
)
|
||||
outputs.append(expert_out)
|
||||
start_idx = end_idx
|
||||
|
||||
outs = torch.cat(outputs, dim=0) if len(outputs) else sorted_tokens.new_empty(0)
|
||||
new_x = torch.empty_like(outs)
|
||||
|
||||
new_x[idxs] = outs
|
||||
final_out = (
|
||||
new_x.view(*topk_ids.shape, -1)
|
||||
.type(topk_weights.dtype)
|
||||
.mul_(topk_weights.unsqueeze(dim=-1))
|
||||
.sum(dim=1)
|
||||
.type(new_x.dtype)
|
||||
)
|
||||
return final_out
|
||||
|
||||
|
||||
def torch_naive_fused_moe_gptoss(
|
||||
a,
|
||||
w1,
|
||||
w2,
|
||||
w1_bias,
|
||||
w2_bias,
|
||||
topk_weight,
|
||||
topk_ids,
|
||||
renormalize,
|
||||
activation_alpha,
|
||||
swiglu_limit,
|
||||
len_experts,
|
||||
):
|
||||
if renormalize:
|
||||
topk_weight = topk_weight / topk_weight.sum(dim=-1, keepdim=True)
|
||||
|
||||
return torch_naive_gptoss_fused_moe(
|
||||
a,
|
||||
w1,
|
||||
w2,
|
||||
w1_bias,
|
||||
w2_bias,
|
||||
topk_weight,
|
||||
topk_ids,
|
||||
activation_alpha,
|
||||
swiglu_limit,
|
||||
len_experts,
|
||||
)
|
||||
|
||||
|
||||
def torch_w8a8_per_column_fused_moe(a, w1, w2, w1_s, w2_s, topk_weight, topk_ids, topk):
|
||||
"""This function performs fused moe with per-column int8 quantization using native torch."""
|
||||
|
||||
B, D = a.shape
|
||||
# Perform per-token quantization
|
||||
a_q, a_s = per_token_quant_int8(a)
|
||||
# Repeat tokens to match topk
|
||||
a_q = a_q.view(B, -1, D).repeat(1, topk, 1).reshape(-1, D)
|
||||
# Also repeat the scale
|
||||
a_s = a_s.view(B, -1, 1).repeat(1, topk, 1).reshape(-1, 1) # [B*topk, 1]
|
||||
|
||||
out = torch.zeros(B * topk, w2.shape[1], dtype=torch.float32, device=a.device)
|
||||
|
||||
# Calculate routing
|
||||
topk_weight = topk_weight.view(-1)
|
||||
topk_ids = topk_ids.view(-1)
|
||||
# Process each expert
|
||||
for i in range(w1.shape[0]):
|
||||
mask = topk_ids == i
|
||||
if mask.sum():
|
||||
# First MLP layer: note that a_s is now per-token
|
||||
inter_out = native_w8a8_per_token_matmul(
|
||||
a_q[mask],
|
||||
w1[i],
|
||||
a_s[mask],
|
||||
w1_s[i],
|
||||
bias=None,
|
||||
output_dtype=torch.float32,
|
||||
)
|
||||
# Activation function
|
||||
act_out = SiluAndMul(inter_out)
|
||||
# Quantize activation output with per-token
|
||||
act_out_q, act_out_s = per_token_quant_int8(act_out)
|
||||
# Second MLP layer
|
||||
out[mask] = native_w8a8_per_token_matmul(
|
||||
act_out_q,
|
||||
w2[i],
|
||||
act_out_s,
|
||||
w2_s[i],
|
||||
bias=None,
|
||||
output_dtype=torch.float32,
|
||||
)
|
||||
# Apply routing weights and sum
|
||||
return (
|
||||
(out.view(B, -1, w2.shape[1]) * topk_weight.view(B, -1, 1).to(out.dtype))
|
||||
.sum(dim=1)
|
||||
.to(a.dtype)
|
||||
)
|
||||
|
||||
|
||||
def native_fp8_fused_moe(a, w1, w2, topk_weight, topk_ids, topk):
|
||||
B, D = a.shape
|
||||
a = a.view(B, -1, D).repeat(1, topk, 1).reshape(-1, D).float()
|
||||
out = torch.zeros(B * topk, w2.shape[1], dtype=torch.float32, device=a.device)
|
||||
|
||||
# Calculate routing
|
||||
topk_weight = topk_weight.view(-1)
|
||||
topk_ids = topk_ids.view(-1)
|
||||
|
||||
for i in range(w1.shape[0]):
|
||||
mask = topk_ids == i
|
||||
if mask.sum():
|
||||
ic0 = torch.matmul(a[mask], w1[i].transpose(0, 1))
|
||||
ic1 = SiluAndMul(ic0)
|
||||
out[mask] = torch.matmul(ic1, w2[i].transpose(0, 1))
|
||||
|
||||
return (
|
||||
(out.view(B, -1, w2.shape[1]) * topk_weight.view(B, -1, 1).to(out.dtype))
|
||||
.sum(dim=1)
|
||||
.to(a.dtype)
|
||||
)
|
||||
|
||||
|
||||
# https://github.com/NVIDIA/TensorRT-Model-Optimizer/blob/main/modelopt/torch/quantization/qtensor/mxfp4_tensor.py
|
||||
class MXFP4QuantizeUtil:
|
||||
E2M1_max = 6.0
|
||||
|
||||
E2M1_values = [0, 0.5, 1, 1.5, 2, 3, 4, 6]
|
||||
E2M1_bounds = torch.tensor([0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5])
|
||||
|
||||
block_size = 32
|
||||
|
||||
@classmethod
|
||||
def quantize(cls, input: torch.Tensor) -> tuple:
|
||||
"""Converting a tensor to a quantized format based on MXFP4 quantization. Only E4M3 is supported.
|
||||
Args:
|
||||
input (torch.Tensor): The input tensor to be quantized.
|
||||
"""
|
||||
|
||||
def cast_fp4(x):
|
||||
sign = torch.sign(x)
|
||||
sign_bit = (2 - sign) // 2
|
||||
ord_ = torch.sum(
|
||||
(x.abs().unsqueeze(-1) - cls.E2M1_bounds.to(x.device)) > 0, dim=-1
|
||||
)
|
||||
fp4_val = (sign_bit * 0b1000 + ord_).to(torch.uint8)
|
||||
return fp4_val
|
||||
|
||||
def fuse_uint4_to_uint8(x):
|
||||
# If the last dimension is odd, pad with zeros
|
||||
# If this behavior is not desired, please modify the code accordingly
|
||||
left_side = x[..., 0::2] # Even indices (0, 2, 4...)
|
||||
right_side = x[..., 1::2] # Odd indices (1, 3, 5...)
|
||||
new_data = (
|
||||
right_side.clone() << 4
|
||||
) # Put odd indices (higher addresses) in high bits
|
||||
new_data[
|
||||
..., : left_side.shape[-1]
|
||||
] += left_side # Put even indices in low bits
|
||||
return new_data
|
||||
|
||||
original_shape = input.shape
|
||||
original_dtype = input.dtype
|
||||
input = input.view(-1, cls.block_size)
|
||||
# get scales
|
||||
input_amax = input.abs().max(dim=-1, keepdim=True).values
|
||||
descale = input_amax / cls.E2M1_max
|
||||
min_value = torch.tensor(-127.0, device=descale.device)
|
||||
e8m0_scale = torch.ceil(torch.maximum(torch.log2(descale), min_value))
|
||||
|
||||
input = (input / torch.exp2(e8m0_scale)).view(original_shape)
|
||||
input_q = cast_fp4(input)
|
||||
input_q = fuse_uint4_to_uint8(input_q)
|
||||
e8m0_scale = (e8m0_scale + 127).to(torch.uint8)
|
||||
return input_q, e8m0_scale
|
||||
|
||||
@classmethod
|
||||
def dequantize(cls, quantized_data, dtype: torch.dtype, scale):
|
||||
"""Dequantze MXFP4 packed tensor to a target dtype."""
|
||||
|
||||
def unfuse_uint8_to_uint4(x):
|
||||
"""Unfuse uint8 values back to uint4 values.
|
||||
This is the inverse operation of fuse_uint4_to_uint8.
|
||||
"""
|
||||
# Extract the lower 4 bits (even indices)
|
||||
left_side = x & 0x0F
|
||||
|
||||
# Extract the upper 4 bits (odd indices)
|
||||
right_side = (x >> 4) & 0x0F
|
||||
|
||||
# Create a new tensor with alternating values
|
||||
shape = list(x.shape)
|
||||
shape[-1] = shape[-1] * 2
|
||||
result = torch.zeros(shape, dtype=torch.uint8, device=x.device)
|
||||
|
||||
# Fill in the values - even indices get low bits, odd indices get high bits
|
||||
result[..., 0::2] = left_side # Even indices from low bits
|
||||
result[..., 1::2] = right_side # Odd indices from high bits
|
||||
|
||||
return result
|
||||
|
||||
e8m0_scale = scale
|
||||
|
||||
# Unfuse the uint8 values back to uint4
|
||||
x_unfused = unfuse_uint8_to_uint4(quantized_data)
|
||||
# print("@@@ x_unfused: ", x_unfused)
|
||||
# Extract sign and magnitude
|
||||
sign = 1 - 2 * ((x_unfused & 0b1000) >> 3).to(
|
||||
torch.float32
|
||||
) # Extract sign bit and convert to +1/-1
|
||||
magnitude = x_unfused & 0b0111 # Extract magnitude bits
|
||||
magnitude = magnitude.to(torch.long)
|
||||
|
||||
# Create a tensor with the E2M1 values
|
||||
values = torch.tensor(cls.E2M1_values, device=quantized_data.device)
|
||||
|
||||
# Use gather to index the values tensor properly
|
||||
# We need to reshape magnitude to match the dimensions we want to gather along
|
||||
original_shape = magnitude.shape
|
||||
x_float = values[magnitude.reshape(-1)].reshape(original_shape)
|
||||
|
||||
# Apply sign and scale
|
||||
x_float = sign.float() * x_float
|
||||
|
||||
# Reshape to apply block-wise scaling
|
||||
x_float = x_float.reshape(-1, cls.block_size)
|
||||
|
||||
# Apply the E8M0 scale
|
||||
scale_factor = torch.exp2(e8m0_scale.float() - 127)
|
||||
scale_factor = scale_factor.reshape(-1, 1) # Reshape for proper broadcasting
|
||||
|
||||
# Apply scaling and reshape back to original shape
|
||||
x_float = x_float * scale_factor
|
||||
|
||||
# Reshape back to the original shape
|
||||
return x_float.reshape(original_shape).to(dtype)
|
||||
|
||||
|
||||
def make_non_contiguous(x: torch.Tensor) -> torch.Tensor:
|
||||
# Make a tensor non-contiguous without changing shape.
|
||||
if not x.is_contiguous():
|
||||
return x
|
||||
|
||||
last_dim = x.shape[-1]
|
||||
expanded = torch.empty(*x.shape[:-1], last_dim + 32, dtype=x.dtype, device=x.device)
|
||||
expanded[..., :last_dim].copy_(x)
|
||||
return expanded.narrow(-1, 0, last_dim)
|
||||
|
||||
|
||||
def awq_reverse_reorder_int_tensor(int_tensor, bits: int):
|
||||
assert bits == 4
|
||||
|
||||
int_tensor = int_tensor.T.contiguous()
|
||||
compress_ratio = 32 // bits
|
||||
assert int_tensor.shape[-1] % compress_ratio == 0
|
||||
|
||||
order_map = [0, 2, 4, 6, 1, 3, 5, 7]
|
||||
order_tensor = torch.tensor(
|
||||
order_map, dtype=torch.int32, device=int_tensor.device
|
||||
).reshape(1, -1)
|
||||
order_tensor = order_tensor.repeat(int_tensor.shape[1] // compress_ratio, 1)
|
||||
order_tensor = order_tensor + torch.arange(
|
||||
0,
|
||||
int_tensor.shape[1],
|
||||
compress_ratio,
|
||||
dtype=torch.int32,
|
||||
device=int_tensor.device,
|
||||
).reshape(-1, 1)
|
||||
order_tensor = order_tensor.reshape(-1)
|
||||
|
||||
reverse_order_tensor = torch.arange(order_tensor.shape[0])[order_tensor]
|
||||
reverse_order_tensor = reverse_order_tensor[order_tensor]
|
||||
int_tensor = int_tensor[:, reverse_order_tensor]
|
||||
return int_tensor
|
||||
|
||||
|
||||
def unpack_and_dequant_awq(
|
||||
awq_qweight: torch.Tensor,
|
||||
awq_qzeros: torch.Tensor,
|
||||
awq_scales: torch.Tensor,
|
||||
bits: int,
|
||||
group_size: int,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
awq_qweight (`torch.LongTensor`):
|
||||
Expected shape: (in_features, out_features // (32 // bits))
|
||||
awq_qzeros (`torch.LongTensor`):
|
||||
Expected shape: (in_features // group_size, out_features // (32 // bits))
|
||||
awq_scales (`torch.LongTensor`):
|
||||
Expected shape: (in_features // group_size, out_features)
|
||||
|
||||
Returns:
|
||||
fp16_weight (`torch.LongTensor`):
|
||||
With shape (in_features, out_features).
|
||||
zeros (`torch.LongTensor`):
|
||||
With shape (in_features // group_size, out_features).
|
||||
"""
|
||||
assert bits == 4
|
||||
|
||||
qzeros = awq_qzeros
|
||||
qweight = awq_qweight
|
||||
qweight = qweight.T.contiguous()
|
||||
|
||||
scales = awq_scales
|
||||
scales = scales.reshape(-1, 1, scales.shape[-1])
|
||||
|
||||
infeatures = awq_qweight.shape[0]
|
||||
|
||||
wf = torch.tensor(
|
||||
list(range(0, 32, bits)), dtype=torch.int32, device=qzeros.device
|
||||
).unsqueeze(0)
|
||||
zeros = torch.bitwise_right_shift(torch.unsqueeze(qzeros, 2), wf.unsqueeze(0)).to(
|
||||
torch.int16 if bits == 8 else torch.int8
|
||||
)
|
||||
|
||||
torch.bitwise_and(zeros, (2**bits) - 1, out=zeros)
|
||||
|
||||
zeros = zeros.reshape(-1, 1, zeros.shape[1] * zeros.shape[2])
|
||||
|
||||
weight = torch.bitwise_right_shift(
|
||||
torch.unsqueeze(qweight, 1), wf.unsqueeze(-1)
|
||||
).to(torch.int16 if bits == 8 else torch.int8)
|
||||
torch.bitwise_and(weight, (2**bits) - 1, out=weight)
|
||||
weight = weight.reshape(-1, group_size, weight.shape[2])
|
||||
|
||||
weight = weight.view(-1, weight.shape[-1])
|
||||
zeros = zeros.view(-1, zeros.shape[-1])
|
||||
|
||||
zeros = zeros.T.contiguous()
|
||||
zeros = awq_reverse_reorder_int_tensor(zeros, bits)
|
||||
weight = awq_reverse_reorder_int_tensor(weight, bits)
|
||||
|
||||
# Dequantize weights.
|
||||
scales = awq_scales
|
||||
zeros = zeros.contiguous()
|
||||
scale_zeros = zeros * scales
|
||||
|
||||
g_idx = torch.tensor(
|
||||
[i // group_size for i in range(infeatures)], dtype=torch.int32
|
||||
)
|
||||
scale_mat = scales[g_idx]
|
||||
scale_zeros_mat = scale_zeros[g_idx].to(torch.bfloat16)
|
||||
|
||||
qdq_weight_T = weight * scale_mat - scale_zeros_mat.to(torch.bfloat16)
|
||||
|
||||
fp16_weight = qdq_weight_T.T
|
||||
|
||||
return fp16_weight, zeros
|
||||
|
||||
|
||||
def unpack_4bit_to_32bit_signed(qweight, qzeros):
|
||||
# Unpack 4-bit values and interpret them as signed integers
|
||||
unpacked_weights = torch.zeros(
|
||||
(qweight.shape[0] * 8, qweight.shape[1]),
|
||||
dtype=torch.int8,
|
||||
device=qweight.device,
|
||||
requires_grad=False,
|
||||
)
|
||||
unpacked_zeros = torch.zeros(
|
||||
(qzeros.shape[0], qzeros.shape[1] * 8),
|
||||
dtype=torch.int8,
|
||||
device=qzeros.device,
|
||||
requires_grad=False,
|
||||
)
|
||||
|
||||
for row in range(unpacked_weights.shape[0]):
|
||||
i = row % 8
|
||||
unpacked_weights[row, :] = (qweight[row // 8, :] >> (4 * i)) & 0xF
|
||||
|
||||
for col in range(unpacked_zeros.shape[1]):
|
||||
i = col % 8
|
||||
unpacked_zeros[:, col] = (qzeros[:, col // 8] >> (4 * i)) & 0xF
|
||||
|
||||
return unpacked_weights, unpacked_zeros + 1
|
||||
|
||||
|
||||
def unpack_and_dequant_gptq(qweight, qzeros, scales):
|
||||
unpacked_qweight, unpacked_qzeros = unpack_4bit_to_32bit_signed(qweight, qzeros)
|
||||
group_size = unpacked_qweight.shape[0] // scales.shape[0]
|
||||
scales = scales.repeat_interleave(group_size, dim=0)
|
||||
unpacked_qzeros = unpacked_qzeros.repeat_interleave(group_size, dim=0)
|
||||
unpacked_qweight = (unpacked_qweight - unpacked_qzeros) * scales
|
||||
|
||||
return unpacked_qweight.T
|
||||
@@ -0,0 +1,22 @@
|
||||
"""Deterministic batch-shape coverage for the input-logprob sweeps."""
|
||||
|
||||
import itertools
|
||||
|
||||
|
||||
def coverage_cases(menu, max_seqs):
|
||||
"""Every singleton, every ordered pair, and wider heterogeneous cases.
|
||||
|
||||
`menu` order is load-bearing: width >= 3 walks it cyclically, so reordering
|
||||
it changes which wide combinations run.
|
||||
"""
|
||||
yield from ((item,) for item in menu)
|
||||
yield from itertools.product(menu, repeat=2)
|
||||
for width in range(3, max_seqs + 1):
|
||||
for offset in range(len(menu)):
|
||||
yield tuple(menu[(offset + step) % len(menu)] for step in range(width))
|
||||
yield tuple(menu[(offset - step) % len(menu)] for step in range(width))
|
||||
# Cyclic windows never repeat an item; adjacent duplicates need their own.
|
||||
for index, item in enumerate(menu):
|
||||
other = menu[(index + 1) % len(menu)]
|
||||
yield (item,) * (width - 1) + (other,)
|
||||
yield (other,) + (item,) * (width - 1)
|
||||
Reference in New Issue
Block a user