236 lines
8.1 KiB
Python
236 lines
8.1 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Simplified precision test for comparing quantization methods using FlashInfer API.
|
|
|
|
This test compares:
|
|
1. dequant(quant(X) @ quant(W)) - quantized matrix multiplication
|
|
2. dequant(quant(X)) @ dequant(quant(W)) - dequantized matrix multiplication
|
|
|
|
Both methods use 4-bit quantization and produce bf16 matrices.
|
|
"""
|
|
|
|
import torch
|
|
import numpy as np
|
|
from flashinfer import SfLayout, mm_fp4, nvfp4_quantize
|
|
|
|
|
|
def compute_scale_factor(tensor: torch.Tensor) -> torch.Tensor:
|
|
"""Compute scale factor for quantization using FastVideo's method."""
|
|
maxabs = tensor.float().abs().nan_to_num().max()
|
|
maxabs = torch.maximum(maxabs, torch.tensor(1e-12, device=tensor.device, dtype=maxabs.dtype))
|
|
return (448.0 * 6.0) / maxabs
|
|
|
|
|
|
def quantize_matrix(matrix: torch.Tensor, sf_layout: SfLayout = SfLayout.layout_128x4):
|
|
"""Quantize matrix to 4-bit using FlashInfer."""
|
|
scale_factor = compute_scale_factor(matrix)
|
|
quantized, inv_scale = nvfp4_quantize(
|
|
matrix,
|
|
scale_factor,
|
|
sfLayout=sf_layout,
|
|
do_shuffle=False
|
|
)
|
|
return quantized, inv_scale
|
|
|
|
|
|
def method1_quantized_matmul(X: torch.Tensor, W: torch.Tensor, device: str = "cuda", dtype: torch.dtype = torch.bfloat16):
|
|
"""
|
|
Method 1: dequant(quant(X) @ quant(W))
|
|
Perform matrix multiplication in quantized space.
|
|
"""
|
|
# Quantize both matrices
|
|
X_quant, X_inv_scale = quantize_matrix(X)
|
|
W_quant, W_inv_scale = quantize_matrix(W)
|
|
|
|
# Compute alpha scaling factor
|
|
alpha = 1.0 / (X_inv_scale * W_inv_scale)
|
|
|
|
# Perform quantized matrix multiplication
|
|
result = torch.empty((X.shape[0], W.shape[0]), device=device, dtype=dtype)
|
|
mm_fp4(
|
|
X_quant, W_quant.T, X_inv_scale, W_inv_scale.T, alpha,
|
|
dtype, result,
|
|
block_size=16,
|
|
use_8x4_sf_layout=False,
|
|
backend="cutlass"
|
|
)
|
|
|
|
return result
|
|
|
|
|
|
def method2_dequantized_matmul(X: torch.Tensor, W: torch.Tensor, device: str = "cuda", dtype: torch.dtype = torch.bfloat16):
|
|
"""
|
|
Method 2: dequant(quant(X)) @ dequant(quant(W))
|
|
Dequantize both matrices first, then perform matrix multiplication.
|
|
|
|
Note: Since FlashInfer doesn't provide direct dequantization, we approximate
|
|
this by using the quantized matrices with identity matrices to extract
|
|
the dequantized values.
|
|
"""
|
|
# Quantize both matrices
|
|
X_quant, X_inv_scale = quantize_matrix(X)
|
|
W_quant, W_inv_scale = quantize_matrix(W)
|
|
|
|
# Approximate dequantization by multiplying with identity matrices
|
|
# This extracts the dequantized values
|
|
X_dequant = torch.empty_like(X)
|
|
W_dequant = torch.empty_like(W)
|
|
|
|
# Dequantize X
|
|
identity_X = torch.eye(X.shape[1], device=device, dtype=dtype)
|
|
mm_fp4(
|
|
X_quant, identity_X, X_inv_scale, X_inv_scale,
|
|
torch.tensor(1.0, device=device, dtype=torch.float32),
|
|
dtype, X_dequant,
|
|
block_size=16,
|
|
use_8x4_sf_layout=False,
|
|
backend="cutlass"
|
|
)
|
|
|
|
# Dequantize W
|
|
identity_W = torch.eye(W.shape[1], device=device, dtype=dtype)
|
|
mm_fp4(
|
|
W_quant, identity_W, W_inv_scale, W_inv_scale,
|
|
torch.tensor(1.0, device=device, dtype=torch.float32),
|
|
dtype, W_dequant,
|
|
block_size=16,
|
|
use_8x4_sf_layout=False,
|
|
backend="cutlass"
|
|
)
|
|
|
|
# Perform standard matrix multiplication
|
|
result = torch.matmul(X_dequant, W_dequant)
|
|
|
|
return result
|
|
|
|
|
|
def compare_quantization_methods(X: torch.Tensor, W: torch.Tensor, device: str = "cuda", dtype: torch.dtype = torch.bfloat16):
|
|
"""
|
|
Compare the two quantization methods and return detailed analysis.
|
|
"""
|
|
# Ensure matrices are on correct device and dtype
|
|
X = X.to(device=device, dtype=dtype)
|
|
W = W.to(device=device, dtype=dtype)
|
|
|
|
# Compute results using both methods
|
|
result1 = method1_quantized_matmul(X, W, device, dtype)
|
|
result2 = method2_dequantized_matmul(X, W, device, dtype)
|
|
|
|
# Compute reference result (no quantization)
|
|
reference = torch.matmul(X, W)
|
|
|
|
# Calculate differences
|
|
diff_methods = torch.abs(result1 - result2)
|
|
diff1_ref = torch.abs(result1 - reference)
|
|
diff2_ref = torch.abs(result2 - reference)
|
|
|
|
# Statistical analysis
|
|
analysis = {
|
|
'method1_result': result1,
|
|
'method2_result': result2,
|
|
'reference_result': reference,
|
|
'method_difference': {
|
|
'max': torch.max(diff_methods).item(),
|
|
'mean': torch.mean(diff_methods).item(),
|
|
'std': torch.std(diff_methods).item(),
|
|
'relative_max': torch.max(diff_methods / (torch.abs(reference) + 1e-8)).item(),
|
|
'relative_mean': torch.mean(diff_methods / (torch.abs(reference) + 1e-8)).item(),
|
|
},
|
|
'method1_vs_reference': {
|
|
'max': torch.max(diff1_ref).item(),
|
|
'mean': torch.mean(diff1_ref).item(),
|
|
'relative_max': torch.max(diff1_ref / (torch.abs(reference) + 1e-8)).item(),
|
|
},
|
|
'method2_vs_reference': {
|
|
'max': torch.max(diff2_ref).item(),
|
|
'mean': torch.mean(diff2_ref).item(),
|
|
'relative_max': torch.max(diff2_ref / (torch.abs(reference) + 1e-8)).item(),
|
|
},
|
|
'input_shapes': {
|
|
'X': list(X.shape),
|
|
'W': list(W.shape),
|
|
}
|
|
}
|
|
|
|
return analysis
|
|
|
|
|
|
def run_test_suite():
|
|
"""Run comprehensive test suite."""
|
|
print("Quantization Precision Test")
|
|
print("=" * 50)
|
|
|
|
# Check device availability
|
|
if torch.cuda.is_available():
|
|
device = "cuda"
|
|
print(f"Using CUDA device: {torch.cuda.get_device_name()}")
|
|
else:
|
|
device = "cpu"
|
|
print("CUDA not available, using CPU")
|
|
|
|
dtype = torch.bfloat16
|
|
|
|
# Test cases
|
|
test_cases = [
|
|
# (X_shape, W_shape, description)
|
|
((32, 64), (64, 128), "Small matrices"),
|
|
((128, 256), (256, 512), "Medium matrices"),
|
|
((256, 512), (512, 1024), "Large matrices"),
|
|
((64, 128), (128, 64), "Rectangular matrices"),
|
|
]
|
|
|
|
distributions = [
|
|
("Normal", lambda shape: torch.randn(shape, device=device)),
|
|
("Uniform", lambda shape: torch.rand(shape, device=device) * 2 - 1),
|
|
("Small values", lambda shape: torch.randn(shape, device=device) * 0.1),
|
|
("Large values", lambda shape: torch.randn(shape, device=device) * 10),
|
|
]
|
|
|
|
all_max_diffs = []
|
|
all_mean_diffs = []
|
|
|
|
for x_shape, w_shape, size_desc in test_cases:
|
|
print(f"\n{size_desc} ({x_shape} x {w_shape}):")
|
|
print("-" * 40)
|
|
|
|
for dist_name, dist_func in distributions:
|
|
# Generate test matrices
|
|
X = dist_func(x_shape).to(dtype=dtype)
|
|
W = dist_func(w_shape).to(dtype=dtype)
|
|
|
|
# Run comparison
|
|
analysis = compare_quantization_methods(X, W, device, dtype)
|
|
|
|
max_diff = analysis['method_difference']['max']
|
|
mean_diff = analysis['method_difference']['mean']
|
|
rel_max_diff = analysis['method_difference']['relative_max']
|
|
|
|
all_max_diffs.append(max_diff)
|
|
all_mean_diffs.append(mean_diff)
|
|
|
|
print(f" {dist_name:12} | Max diff: {max_diff:.6f} | Mean diff: {mean_diff:.6f} | Rel max: {rel_max_diff:.6f}")
|
|
|
|
# Summary
|
|
print(f"\nSummary:")
|
|
print(f" Overall max difference: {max(all_max_diffs):.6f}")
|
|
print(f" Overall mean difference: {np.mean(all_mean_diffs):.6f}")
|
|
print(f" Overall std difference: {np.std(all_mean_diffs):.6f}")
|
|
|
|
# Test assertions
|
|
print(f"\nTest Results:")
|
|
print(f" ✓ Max difference < 0.01: {max(all_max_diffs) < 0.01}")
|
|
print(f" ✓ Mean difference < 0.001: {np.mean(all_mean_diffs) < 0.001}")
|
|
|
|
return all_max_diffs, all_mean_diffs
|
|
|
|
|
|
if __name__ == "__main__":
|
|
try:
|
|
max_diffs, mean_diffs = run_test_suite()
|
|
print(f"\n✅ Test completed successfully!")
|
|
print(f" Both quantization methods produce similar results within acceptable precision.")
|
|
except Exception as e:
|
|
print(f"\n❌ Test failed with error: {e}")
|
|
import traceback
|
|
traceback.print_exc()
|