453 lines
16 KiB
Python
453 lines
16 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Precision test for fake_quantize_q and fake_quantize_kv kernels.
|
|
|
|
Tests that the kernels correctly:
|
|
1. Quantize and dequantize Q, K, V tensors
|
|
2. Preserve approximate values within quantization error
|
|
3. Handle edge cases (zeros, small/large values, padding)
|
|
4. Work with different shapes and dtypes
|
|
"""
|
|
|
|
import torch
|
|
import triton
|
|
import triton.language as tl
|
|
import numpy as np
|
|
from typing import Tuple, Dict, Any
|
|
import unittest
|
|
|
|
from quant_utils import fake_quantize_q, fake_quantize_kv
|
|
|
|
|
|
def cosine_similarity(tensor1: torch.Tensor, tensor2: torch.Tensor) -> float:
|
|
"""Compute cosine similarity between two tensors."""
|
|
t1_flat = tensor1.flatten().float()
|
|
t2_flat = tensor2.flatten().float()
|
|
|
|
dot_product = torch.dot(t1_flat, t2_flat)
|
|
norm1 = torch.norm(t1_flat)
|
|
norm2 = torch.norm(t2_flat)
|
|
|
|
if norm1 == 0 or norm2 == 0:
|
|
return 0.0
|
|
|
|
return (dot_product / (norm1 * norm2)).item()
|
|
|
|
|
|
def compute_metrics(original: torch.Tensor, quantized: torch.Tensor, mask: torch.Tensor = None) -> Dict[str, float]:
|
|
"""Compute precision metrics between original and quantized tensors."""
|
|
if mask is not None:
|
|
# Only compute metrics on valid (non-padding) elements
|
|
original = original[mask]
|
|
quantized = quantized[mask]
|
|
|
|
diff = (original.float() - quantized.float()).abs()
|
|
|
|
metrics = {
|
|
'max_diff': diff.max().item(),
|
|
'mean_diff': diff.mean().item(),
|
|
'std_diff': diff.std().item(),
|
|
'cosine_similarity': cosine_similarity(original, quantized),
|
|
}
|
|
|
|
# Relative error (avoid division by zero)
|
|
abs_original = original.float().abs()
|
|
rel_error = diff / (abs_original + 1e-8)
|
|
metrics['max_relative_error'] = rel_error.max().item()
|
|
metrics['mean_relative_error'] = rel_error.mean().item()
|
|
|
|
return metrics
|
|
|
|
|
|
def test_fake_quantize_q(
|
|
q: torch.Tensor,
|
|
BLOCK_M: int = 128,
|
|
HEAD_DIM: int = 128,
|
|
verbose: bool = False
|
|
) -> Tuple[torch.Tensor, Dict[str, Any]]:
|
|
"""
|
|
Test fake_quantize_q kernel.
|
|
|
|
Args:
|
|
q: Input Q tensor with shape (Z, H, N_CTX_Q, HEAD_DIM)
|
|
BLOCK_M: Block size for sequence dimension
|
|
HEAD_DIM: Head dimension (must be multiple of 16)
|
|
verbose: Whether to print detailed information
|
|
|
|
Returns:
|
|
Tuple of (fake_q, metrics_dict)
|
|
"""
|
|
assert q.is_cuda, "Input must be on CUDA"
|
|
assert q.dim() == 4, f"Q must be 4D, got shape {q.shape}"
|
|
assert HEAD_DIM % 16 == 0, f"HEAD_DIM must be multiple of 16, got {HEAD_DIM}"
|
|
|
|
Z, H, N_CTX_Q, HEAD_DIM_actual = q.shape
|
|
assert HEAD_DIM_actual == HEAD_DIM, f"HEAD_DIM mismatch: {HEAD_DIM_actual} != {HEAD_DIM}"
|
|
|
|
# Create output tensor
|
|
fake_q = torch.empty_like(q)
|
|
|
|
# Launch kernel
|
|
grid = (triton.cdiv(N_CTX_Q, BLOCK_M), Z * H, 1)
|
|
|
|
fake_quantize_q[grid](
|
|
q, fake_q,
|
|
q.stride(0), q.stride(1), # stride_z_q, stride_h_q
|
|
q.stride(2), q.stride(3), # stride_tok_q, stride_d_q
|
|
fake_q.stride(0), fake_q.stride(1), # fake_stride_z_q, fake_stride_h_q
|
|
fake_q.stride(2), fake_q.stride(3), # fake_stride_tok_q, fake_stride_d_q
|
|
H, N_CTX_Q,
|
|
BLOCK_M=BLOCK_M, HEAD_DIM=HEAD_DIM
|
|
)
|
|
|
|
# Compute metrics
|
|
metrics = compute_metrics(q, fake_q)
|
|
metrics['shape'] = list(q.shape)
|
|
metrics['dtype'] = str(q.dtype)
|
|
|
|
if verbose:
|
|
print(f"fake_quantize_q test:")
|
|
print(f" Shape: {q.shape}")
|
|
print(f" Dtype: {q.dtype}")
|
|
print(f" Max diff: {metrics['max_diff']:.6f}")
|
|
print(f" Mean diff: {metrics['mean_diff']:.6f}")
|
|
print(f" Cosine similarity: {metrics['cosine_similarity']:.6f}")
|
|
print(f" Max relative error: {metrics['max_relative_error']:.6f}")
|
|
|
|
return fake_q, metrics
|
|
|
|
|
|
def test_fake_quantize_kv(
|
|
k: torch.Tensor,
|
|
v: torch.Tensor,
|
|
BLOCK_N: int = 128,
|
|
HEAD_DIM: int = 128,
|
|
verbose: bool = False
|
|
) -> Tuple[torch.Tensor, torch.Tensor, Dict[str, Any]]:
|
|
"""
|
|
Test fake_quantize_kv kernel.
|
|
|
|
Args:
|
|
k: Input K tensor with shape (Z, H, N_CTX_KV, HEAD_DIM)
|
|
v: Input V tensor with shape (Z, H, N_CTX_KV, HEAD_DIM)
|
|
BLOCK_N: Block size for sequence dimension
|
|
HEAD_DIM: Head dimension (must be multiple of 16)
|
|
verbose: Whether to print detailed information
|
|
|
|
Returns:
|
|
Tuple of (fake_k, fake_v, metrics_dict)
|
|
"""
|
|
assert k.is_cuda and v.is_cuda, "Inputs must be on CUDA"
|
|
assert k.shape == v.shape, f"K and V must have same shape: {k.shape} != {v.shape}"
|
|
assert k.dim() == 4, f"K must be 4D, got shape {k.shape}"
|
|
assert HEAD_DIM % 16 == 0, f"HEAD_DIM must be multiple of 16, got {HEAD_DIM}"
|
|
|
|
Z, H, N_CTX_KV, HEAD_DIM_actual = k.shape
|
|
assert HEAD_DIM_actual == HEAD_DIM, f"HEAD_DIM mismatch: {HEAD_DIM_actual} != {HEAD_DIM}"
|
|
|
|
# Create output tensors
|
|
fake_k = torch.empty_like(k)
|
|
fake_v = torch.empty_like(v)
|
|
|
|
# Launch kernel
|
|
grid = (triton.cdiv(N_CTX_KV, BLOCK_N), Z * H, 1)
|
|
|
|
fake_quantize_kv[grid](
|
|
k, v, fake_k, fake_v,
|
|
k.stride(0), k.stride(1), # stride_z_kv, stride_h_kv
|
|
k.stride(2), k.stride(3), # stride_tok_kv, stride_d_kv
|
|
fake_k.stride(0), fake_k.stride(1), # fake_stride_z_kv, fake_stride_h_kv
|
|
fake_k.stride(2), fake_k.stride(3), # fake_stride_tok_kv, fake_stride_d_kv
|
|
H, N_CTX_KV,
|
|
BLOCK_N=BLOCK_N, HEAD_DIM=HEAD_DIM
|
|
)
|
|
|
|
# Compute metrics for both K and V
|
|
metrics_k = compute_metrics(k, fake_k)
|
|
metrics_v = compute_metrics(v, fake_v)
|
|
|
|
metrics = {
|
|
'k': metrics_k,
|
|
'v': metrics_v,
|
|
'shape': list(k.shape),
|
|
'dtype': str(k.dtype),
|
|
}
|
|
|
|
if verbose:
|
|
print(f"fake_quantize_kv test:")
|
|
print(f" Shape: {k.shape}")
|
|
print(f" Dtype: {k.dtype}")
|
|
print(f" K - Max diff: {metrics_k['max_diff']:.6f}, Cosine sim: {metrics_k['cosine_similarity']:.6f}")
|
|
print(f" V - Max diff: {metrics_v['max_diff']:.6f}, Cosine sim: {metrics_v['cosine_similarity']:.6f}")
|
|
|
|
return fake_k, fake_v, metrics
|
|
|
|
|
|
class FakeQuantizeKernelsTest:
|
|
"""Test suite for fake_quantize_q and fake_quantize_kv kernels."""
|
|
|
|
def __init__(self, device: str = "cuda", dtype: torch.dtype = torch.bfloat16):
|
|
self.device = device
|
|
self.dtype = dtype
|
|
self.HEAD_DIM = 128
|
|
self.BLOCK_M = 128
|
|
self.BLOCK_N = 128
|
|
|
|
def test_basic_q(self, verbose: bool = False) -> Dict[str, Any]:
|
|
"""Test fake_quantize_q with basic random input."""
|
|
print("\n" + "="*60)
|
|
print("Test: Basic Q quantization")
|
|
print("="*60)
|
|
|
|
Z, H, N_CTX_Q = 2, 8, 256
|
|
q = torch.randn(Z, H, N_CTX_Q, self.HEAD_DIM, device=self.device, dtype=self.dtype)
|
|
|
|
fake_q, metrics = test_fake_quantize_q(q, BLOCK_M=self.BLOCK_M, HEAD_DIM=self.HEAD_DIM, verbose=verbose)
|
|
|
|
# Assertions
|
|
assert fake_q.shape == q.shape, f"Shape mismatch: {fake_q.shape} != {q.shape}"
|
|
assert fake_q.dtype == q.dtype, f"Dtype mismatch: {fake_q.dtype} != {q.dtype}"
|
|
assert metrics['cosine_similarity'] > 0.9, f"Cosine similarity too low: {metrics['cosine_similarity']}"
|
|
|
|
print("✓ Basic Q test passed")
|
|
return metrics
|
|
|
|
def test_basic_kv(self, verbose: bool = False) -> Dict[str, Any]:
|
|
"""Test fake_quantize_kv with basic random inputs."""
|
|
print("\n" + "="*60)
|
|
print("Test: Basic KV quantization")
|
|
print("="*60)
|
|
|
|
Z, H, N_CTX_KV = 2, 8, 512
|
|
k = torch.randn(Z, H, N_CTX_KV, self.HEAD_DIM, device=self.device, dtype=self.dtype)
|
|
v = torch.randn(Z, H, N_CTX_KV, self.HEAD_DIM, device=self.device, dtype=self.dtype)
|
|
|
|
fake_k, fake_v, metrics = test_fake_quantize_kv(
|
|
k, v, BLOCK_N=self.BLOCK_N, HEAD_DIM=self.HEAD_DIM, verbose=verbose
|
|
)
|
|
|
|
# Assertions
|
|
assert fake_k.shape == k.shape, f"K shape mismatch: {fake_k.shape} != {k.shape}"
|
|
assert fake_v.shape == v.shape, f"V shape mismatch: {fake_v.shape} != {v.shape}"
|
|
assert metrics['k']['cosine_similarity'] > 0.9, f"K cosine similarity too low: {metrics['k']['cosine_similarity']}"
|
|
assert metrics['v']['cosine_similarity'] > 0.9, f"V cosine similarity too low: {metrics['v']['cosine_similarity']}"
|
|
|
|
print("✓ Basic KV test passed")
|
|
return metrics
|
|
|
|
def test_small_values(self, verbose: bool = False) -> Dict[str, Any]:
|
|
"""Test with small input values."""
|
|
print("\n" + "="*60)
|
|
print("Test: Small values")
|
|
print("="*60)
|
|
|
|
Z, H, N_CTX_Q = 1, 4, 128
|
|
q = torch.randn(Z, H, N_CTX_Q, self.HEAD_DIM, device=self.device, dtype=self.dtype) * 0.01
|
|
|
|
fake_q, metrics = test_fake_quantize_q(q, BLOCK_M=self.BLOCK_M, HEAD_DIM=self.HEAD_DIM, verbose=verbose)
|
|
|
|
print("✓ Small values test passed")
|
|
return metrics
|
|
|
|
def test_large_values(self, verbose: bool = False) -> Dict[str, Any]:
|
|
"""Test with large input values."""
|
|
print("\n" + "="*60)
|
|
print("Test: Large values")
|
|
print("="*60)
|
|
|
|
Z, H, N_CTX_Q = 1, 4, 128
|
|
q = torch.randn(Z, H, N_CTX_Q, self.HEAD_DIM, device=self.device, dtype=self.dtype) * 10.0
|
|
|
|
fake_q, metrics = test_fake_quantize_q(q, BLOCK_M=self.BLOCK_M, HEAD_DIM=self.HEAD_DIM, verbose=verbose)
|
|
|
|
print("✓ Large values test passed")
|
|
return metrics
|
|
|
|
def test_zeros(self, verbose: bool = False) -> Dict[str, Any]:
|
|
"""Test with zero input."""
|
|
print("\n" + "="*60)
|
|
print("Test: Zeros")
|
|
print("="*60)
|
|
|
|
Z, H, N_CTX_Q = 1, 4, 128
|
|
q = torch.zeros(Z, H, N_CTX_Q, self.HEAD_DIM, device=self.device, dtype=self.dtype)
|
|
|
|
fake_q, metrics = test_fake_quantize_q(q, BLOCK_M=self.BLOCK_M, HEAD_DIM=self.HEAD_DIM, verbose=verbose)
|
|
|
|
# Zeros should remain approximately zero after quantization
|
|
max_val = fake_q.abs().max().item()
|
|
assert max_val < 1e-3, f"Zeros should remain near zero, got max: {max_val}"
|
|
|
|
print("✓ Zeros test passed")
|
|
return metrics
|
|
|
|
def test_different_shapes(self, verbose: bool = False) -> Dict[str, Any]:
|
|
"""Test with different tensor shapes."""
|
|
print("\n" + "="*60)
|
|
print("Test: Different shapes")
|
|
print("="*60)
|
|
|
|
test_shapes = [
|
|
(1, 1, 64, self.HEAD_DIM), # Single batch, single head, small seq
|
|
(1, 8, 128, self.HEAD_DIM), # Single batch, multiple heads
|
|
(4, 16, 256, self.HEAD_DIM), # Multiple batches, many heads
|
|
(2, 8, 512, self.HEAD_DIM), # Medium size
|
|
]
|
|
|
|
all_metrics = []
|
|
for shape in test_shapes:
|
|
Z, H, N_CTX_Q, HEAD_DIM = shape
|
|
q = torch.randn(*shape, device=self.device, dtype=self.dtype)
|
|
|
|
fake_q, metrics = test_fake_quantize_q(
|
|
q, BLOCK_M=self.BLOCK_M, HEAD_DIM=self.HEAD_DIM, verbose=verbose
|
|
)
|
|
|
|
assert fake_q.shape == q.shape
|
|
assert metrics['cosine_similarity'] > 0.85
|
|
all_metrics.append(metrics)
|
|
|
|
if verbose:
|
|
print(f" Shape {shape}: cosine_sim={metrics['cosine_similarity']:.4f}")
|
|
|
|
print("✓ Different shapes test passed")
|
|
return all_metrics
|
|
|
|
def test_different_dtypes(self, verbose: bool = False) -> Dict[str, Any]:
|
|
"""Test with different dtypes."""
|
|
print("\n" + "="*60)
|
|
print("Test: Different dtypes")
|
|
print("="*60)
|
|
|
|
Z, H, N_CTX_Q = 1, 4, 128
|
|
dtypes = [torch.bfloat16, torch.float16, torch.float32]
|
|
|
|
all_metrics = []
|
|
for dtype in dtypes:
|
|
q = torch.randn(Z, H, N_CTX_Q, self.HEAD_DIM, device=self.device, dtype=dtype)
|
|
|
|
fake_q, metrics = test_fake_quantize_q(
|
|
q, BLOCK_M=self.BLOCK_M, HEAD_DIM=self.HEAD_DIM, verbose=verbose
|
|
)
|
|
|
|
assert fake_q.dtype == q.dtype
|
|
assert metrics['cosine_similarity'] > 0.85
|
|
all_metrics.append(metrics)
|
|
|
|
if verbose:
|
|
print(f" {dtype}: cosine_sim={metrics['cosine_similarity']:.4f}")
|
|
|
|
print("✓ Different dtypes test passed")
|
|
return all_metrics
|
|
|
|
def test_non_divisible_sequence_length(self, verbose: bool = False) -> Dict[str, Any]:
|
|
"""Test with sequence lengths that don't divide evenly by block size."""
|
|
print("\n" + "="*60)
|
|
print("Test: Non-divisible sequence length")
|
|
print("="*60)
|
|
|
|
Z, H = 1, 4
|
|
# Use sequence length that doesn't divide evenly by BLOCK_M
|
|
N_CTX_Q = 150 # Not divisible by 128
|
|
|
|
q = torch.randn(Z, H, N_CTX_Q, self.HEAD_DIM, device=self.device, dtype=self.dtype)
|
|
|
|
fake_q, metrics = test_fake_quantize_q(q, BLOCK_M=self.BLOCK_M, HEAD_DIM=self.HEAD_DIM, verbose=verbose)
|
|
|
|
# Check that padding is handled correctly (last block should have zeros for padding)
|
|
assert fake_q.shape == q.shape
|
|
assert metrics['cosine_similarity'] > 0.85
|
|
|
|
print("✓ Non-divisible sequence length test passed")
|
|
return metrics
|
|
|
|
def test_kv_different_shapes(self, verbose: bool = False) -> Dict[str, Any]:
|
|
"""Test fake_quantize_kv with different shapes."""
|
|
print("\n" + "="*60)
|
|
print("Test: KV different shapes")
|
|
print("="*60)
|
|
|
|
test_shapes = [
|
|
(1, 1, 64, self.HEAD_DIM),
|
|
(1, 8, 128, self.HEAD_DIM),
|
|
(2, 16, 256, self.HEAD_DIM),
|
|
]
|
|
|
|
all_metrics = []
|
|
for shape in test_shapes:
|
|
Z, H, N_CTX_KV, HEAD_DIM = shape
|
|
k = torch.randn(*shape, device=self.device, dtype=self.dtype)
|
|
v = torch.randn(*shape, device=self.device, dtype=self.dtype)
|
|
|
|
fake_k, fake_v, metrics = test_fake_quantize_kv(
|
|
k, v, BLOCK_N=self.BLOCK_N, HEAD_DIM=self.HEAD_DIM, verbose=verbose
|
|
)
|
|
|
|
assert fake_k.shape == k.shape
|
|
assert fake_v.shape == v.shape
|
|
assert metrics['k']['cosine_similarity'] > 0.85
|
|
assert metrics['v']['cosine_similarity'] > 0.85
|
|
all_metrics.append(metrics)
|
|
|
|
print("✓ KV different shapes test passed")
|
|
return all_metrics
|
|
|
|
def run_all_tests(self, verbose: bool = False) -> Dict[str, Any]:
|
|
"""Run all tests."""
|
|
print("\n" + "="*60)
|
|
print("Running Fake Quantize Kernels Precision Tests")
|
|
print("="*60)
|
|
|
|
if not torch.cuda.is_available():
|
|
print("CUDA not available, skipping tests")
|
|
return {}
|
|
|
|
results = {}
|
|
|
|
try:
|
|
results['basic_q'] = self.test_basic_q(verbose=verbose)
|
|
results['basic_kv'] = self.test_basic_kv(verbose=verbose)
|
|
results['small_values'] = self.test_small_values(verbose=verbose)
|
|
results['large_values'] = self.test_large_values(verbose=verbose)
|
|
results['zeros'] = self.test_zeros(verbose=verbose)
|
|
results['different_shapes'] = self.test_different_shapes(verbose=verbose)
|
|
results['different_dtypes'] = self.test_different_dtypes(verbose=verbose)
|
|
results['non_divisible_seq'] = self.test_non_divisible_sequence_length(verbose=verbose)
|
|
results['kv_different_shapes'] = self.test_kv_different_shapes(verbose=verbose)
|
|
|
|
print("\n" + "="*60)
|
|
print("✓ All tests passed!")
|
|
print("="*60)
|
|
|
|
except Exception as e:
|
|
print(f"\n✗ Test failed with error: {e}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
raise
|
|
|
|
return results
|
|
|
|
|
|
def main():
|
|
"""Main function to run the precision tests."""
|
|
tester = FakeQuantizeKernelsTest(device="cuda", dtype=torch.bfloat16)
|
|
results = tester.run_all_tests(verbose=True)
|
|
|
|
# Print summary
|
|
print("\n" + "="*60)
|
|
print("Test Summary")
|
|
print("="*60)
|
|
|
|
if 'basic_q' in results:
|
|
print(f"Basic Q - Cosine similarity: {results['basic_q']['cosine_similarity']:.6f}")
|
|
print(f"Basic Q - Max diff: {results['basic_q']['max_diff']:.6f}")
|
|
|
|
if 'basic_kv' in results:
|
|
print(f"Basic K - Cosine similarity: {results['basic_kv']['k']['cosine_similarity']:.6f}")
|
|
print(f"Basic V - Cosine similarity: {results['basic_kv']['v']['cosine_similarity']:.6f}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|
|
|