#!/usr/bin/env python3 """ Basic tests for the fake_quantize function. Tests core functionality with simple inputs and edge cases. """ import torch import triton import triton.language as tl from typing import Tuple from qat_attn import fake_quantize from nvfp4_utils import _downcast_to_mxfp, _upcast_from_mxfp, MXFP_BLOCK_SIZE, _compute_quant_and_scale, _compute_dequant # Import get_fake_quant from test_quantization_precision try: from test_quantization_precision import get_fake_quant except ImportError: # Fallback: define get_fake_quant here if import fails from flashinfer import SfLayout, nvfp4_quantize, e2m1_and_ufp8sf_scale_to_float def get_fake_quant(x: torch.Tensor): orig_shape = x.shape orig_dtype = x.dtype device = x.device x = x.view(-1, x.shape[-1]) x_global_sf = (448 * 6) / x.float().abs().nan_to_num().max() x_fp4, x_scale = nvfp4_quantize(x, x_global_sf, sfLayout=SfLayout.layout_128x4, do_shuffle=False) x_dequant = e2m1_and_ufp8sf_scale_to_float(x_fp4, x_scale, 1 / x_global_sf) val = x_dequant.view(orig_shape).to(orig_dtype).to(device) return val DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") def cosine_similarity(tensor1, tensor2): """ Compute cosine similarity between two tensors. Args: tensor1: First tensor tensor2: Second tensor (same shape as tensor1) Returns: Cosine similarity value (scalar) """ # Flatten tensors for computation t1_flat = tensor1.flatten().float() t2_flat = tensor2.flatten().float() # Compute cosine similarity: (A · B) / (||A|| * ||B||) dot_product = torch.dot(t1_flat, t2_flat) norm1 = torch.norm(t1_flat) norm2 = torch.norm(t2_flat) # Avoid division by zero if norm1 == 0 or norm2 == 0: return 0.0 cos_sim = dot_product / (norm1 * norm2) return cos_sim.item() @triton.jit def fake_quantize_kernel( src_ptr, dst_ptr, high_prec_ptr, BLOCK_SIZE_OUT_DIM: tl.constexpr, BLOCK_SIZE_QUANT_DIM: tl.constexpr, dst_dtype: tl.constexpr, mx_tensor_dtype: tl.constexpr, stride_src_outer, stride_src_quant, stride_dst_outer, stride_dst_quant, stride_high_prec_outer, stride_high_prec_quant, outer_dim, quant_dim, ): """ Kernel wrapper to test fake_quantize function. """ outer_idx = tl.program_id(0) quant_idx = tl.program_id(1) # Compute offsets start_outer = outer_idx * BLOCK_SIZE_OUT_DIM start_quant = quant_idx * BLOCK_SIZE_QUANT_DIM # Create offset arrays offs_outer = tl.arange(0, BLOCK_SIZE_OUT_DIM)[:, None] offs_quant = tl.arange(0, BLOCK_SIZE_QUANT_DIM)[None, :] # Create masks for valid elements mask_outer = (start_outer + offs_outer) < outer_dim mask_quant = (start_quant + offs_quant) < quant_dim full_mask = mask_outer & mask_quant # Load source tensor src_offsets = (start_outer + offs_outer) * stride_src_outer + (start_quant + offs_quant) * stride_src_quant src_tensor = tl.load(src_ptr + src_offsets, mask=full_mask, other=0.0) # Call fake_quantize quantized_tensor, high_prec_tensor = fake_quantize( src_tensor=src_tensor, BLOCK_SIZE_OUT_DIM=BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_DIM=BLOCK_SIZE_QUANT_DIM, dst_dtype=dst_dtype, mx_tensor_dtype=mx_tensor_dtype ) # Store results dst_offsets = (start_outer + offs_outer) * stride_dst_outer + (start_quant + offs_quant) * stride_dst_quant high_prec_offsets = (start_outer + offs_outer) * stride_high_prec_outer + (start_quant + offs_quant) * stride_high_prec_quant tl.store(dst_ptr + dst_offsets, quantized_tensor, mask=full_mask) tl.store(high_prec_ptr + high_prec_offsets, high_prec_tensor, mask=full_mask) def call_fake_quantize( x: torch.Tensor, BLOCK_SIZE_OUT_DIM: int = 128, BLOCK_SIZE_QUANT_DIM: int = 128, dst_dtype: torch.dtype = None, mx_tensor_dtype = tl.float8e4nv ) -> Tuple[torch.Tensor, torch.Tensor]: """ Call fake_quantize on a tensor. Args: x: Input tensor (2D or can be reshaped to 2D) BLOCK_SIZE_OUT_DIM: Block size for outer dimension BLOCK_SIZE_QUANT_DIM: Block size for quantization dimension (must be multiple of 16) dst_dtype: Output dtype (defaults to input dtype) mx_tensor_dtype: Quantization dtype (tl.float8e4nv or tl.uint8 for fp4) Returns: Tuple of (quantized_tensor, high_precision_tensor) """ assert x.is_cuda, "Input must be on CUDA" assert BLOCK_SIZE_QUANT_DIM % 16 == 0, f"BLOCK_SIZE_QUANT_DIM must be multiple of 16, got {BLOCK_SIZE_QUANT_DIM}" orig_shape = x.shape orig_dtype = x.dtype # Reshape to 2D x_2d = x.view(-1, x.shape[-1]) outer_dim, quant_dim = x_2d.shape if dst_dtype is None: dst_dtype = orig_dtype # Map torch dtype to triton dtype dtype_map = { torch.float32: tl.float32, torch.float16: tl.float16, torch.bfloat16: tl.bfloat16, } triton_dst_dtype = dtype_map.get(dst_dtype, tl.float16) # Allocate outputs quantized_output = torch.empty_like(x_2d, dtype=dst_dtype) high_prec_output = torch.empty_like(x_2d, dtype=torch.float32) # Launch kernel grid = ( triton.cdiv(outer_dim, BLOCK_SIZE_OUT_DIM), triton.cdiv(quant_dim, BLOCK_SIZE_QUANT_DIM), ) fake_quantize_kernel[grid]( src_ptr=x_2d, dst_ptr=quantized_output, high_prec_ptr=high_prec_output, BLOCK_SIZE_OUT_DIM=BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_DIM=BLOCK_SIZE_QUANT_DIM, dst_dtype=triton_dst_dtype, mx_tensor_dtype=mx_tensor_dtype, stride_src_outer=x_2d.stride(0), stride_src_quant=x_2d.stride(1), stride_dst_outer=quantized_output.stride(0), stride_dst_quant=quantized_output.stride(1), stride_high_prec_outer=high_prec_output.stride(0), stride_high_prec_quant=high_prec_output.stride(1), outer_dim=outer_dim, quant_dim=quant_dim, ) return quantized_output.view(orig_shape).to(dst_dtype), high_prec_output.view(orig_shape) def test_basic_functionality(): """Test basic functionality of fake_quantize.""" print("Test 1: Basic functionality") torch.manual_seed(42) x = torch.randn(128, 64, device=DEVICE, dtype=torch.float16) quantized, high_prec = call_fake_quantize(x, BLOCK_SIZE_OUT_DIM=128, BLOCK_SIZE_QUANT_DIM=64) # Check shapes assert quantized.shape == x.shape, f"Shape mismatch: {quantized.shape} vs {x.shape}" assert high_prec.shape == x.shape, f"Shape mismatch: {high_prec.shape} vs {x.shape}" # Check that quantization happened (values should be different) diff = (quantized.float() - x.float()).abs() max_diff = diff.max().item() mean_diff = diff.mean().item() assert max_diff > 0, "Quantization should change values" # Compute cosine similarity between quantized and original cos_sim = cosine_similarity(quantized, x) # Check that high_prec matches original high_prec_diff = (high_prec - x.float()).abs().max().item() assert high_prec_diff < 1e-5, f"High precision tensor should match original, got diff: {high_prec_diff}" print(f" ✓ Shape check passed") print(f" ✓ Quantization occurred (max diff: {max_diff:.6f}, mean diff: {mean_diff:.6f}, cosine sim: {cos_sim:.6f})") print(f" ✓ High precision tensor matches original (diff: {high_prec_diff:.6e})") def test_different_shapes(): """Test fake_quantize with different tensor shapes.""" print("\nTest 2: Different shapes") torch.manual_seed(42) test_shapes = [ (64, 64), (128, 128), (256, 128), (128, 256), (32, 64), ] for shape in test_shapes: x = torch.randn(shape, device=DEVICE, dtype=torch.float16) # Use block sizes that divide evenly block_out = min(128, shape[0]) block_quant = min(128, shape[1]) # Ensure block_quant is multiple of 16 block_quant = (block_quant // 16) * 16 if block_quant == 0: block_quant = 16 quantized, high_prec = call_fake_quantize( x, BLOCK_SIZE_OUT_DIM=block_out, BLOCK_SIZE_QUANT_DIM=block_quant ) assert quantized.shape == x.shape, f"Shape mismatch for {shape}" assert high_prec.shape == x.shape, f"Shape mismatch for {shape}" print(f" ✓ Shape {shape} passed") def test_different_dtypes(): """Test fake_quantize with different input dtypes.""" print("\nTest 3: Different dtypes") torch.manual_seed(42) shape = (128, 64) dtypes = [torch.float16, torch.bfloat16, torch.float32] for dtype in dtypes: x = torch.randn(shape, device=DEVICE, dtype=dtype) quantized, high_prec = call_fake_quantize( x, BLOCK_SIZE_OUT_DIM=128, BLOCK_SIZE_QUANT_DIM=64, dst_dtype=dtype ) assert quantized.dtype == dtype, f"Dtype mismatch: {quantized.dtype} vs {dtype}" assert quantized.shape == x.shape, f"Shape mismatch for dtype {dtype}" print(f" ✓ Dtype {dtype} passed") def test_edge_cases(): """Test fake_quantize with edge cases.""" print("\nTest 4: Edge cases") # Test with zeros print(" Testing zeros...") x_zeros = torch.zeros(128, 64, device=DEVICE, dtype=torch.float16) quantized, high_prec = call_fake_quantize(x_zeros, BLOCK_SIZE_OUT_DIM=128, BLOCK_SIZE_QUANT_DIM=64) assert quantized.shape == x_zeros.shape # Quantized zeros should still be close to zero assert quantized.abs().max().item() < 1e-3, "Quantized zeros should be close to zero" print(" ✓ Zeros passed") # Test with ones (FAIL) print(" Testing ones...") x_ones = torch.ones(128, 64, device=DEVICE, dtype=torch.float16) print(f"x_ones: {x_ones}") quantized, high_prec = call_fake_quantize(x_ones, BLOCK_SIZE_OUT_DIM=128, BLOCK_SIZE_QUANT_DIM=64) print(f"quantized: {quantized}") print(f"high_prec: {high_prec}") assert quantized.shape == x_ones.shape # Quantized ones should be close to 1.0 diff = (quantized - x_ones).abs().max().item() assert diff < 0.1, f"Quantized ones should be close to 1.0, got max diff: {diff}" print(f" ✓ Ones passed (max diff: {diff:.6f})") # Test with small values print(" Testing small values...") x_small = torch.randn(128, 64, device=DEVICE, dtype=torch.float16) * 0.001 quantized, high_prec = call_fake_quantize(x_small, BLOCK_SIZE_OUT_DIM=128, BLOCK_SIZE_QUANT_DIM=64) assert quantized.shape == x_small.shape print(" ✓ Small values passed") # Test with large values print(" Testing large values...") x_large = torch.randn(128, 64, device=DEVICE, dtype=torch.float16) * 100.0 quantized, high_prec = call_fake_quantize(x_large, BLOCK_SIZE_OUT_DIM=128, BLOCK_SIZE_QUANT_DIM=64) assert quantized.shape == x_large.shape print(" ✓ Large values passed") def test_fp4_quantization(): """Test fake_quantize with FP4 quantization (uint8).""" print("\nTest 5: FP4 quantization (uint8)") torch.manual_seed(42) x = torch.randn(128, 64, device=DEVICE, dtype=torch.float16) quantized, high_prec = call_fake_quantize( x, BLOCK_SIZE_OUT_DIM=128, BLOCK_SIZE_QUANT_DIM=64, mx_tensor_dtype=tl.uint8 # FP4 uses uint8 ) assert quantized.shape == x.shape assert high_prec.shape == x.shape # Check that quantization happened diff = (quantized.float() - x.float()).abs() max_diff = diff.max().item() mean_diff = diff.mean().item() assert max_diff > 0, "Quantization should change values" # Compute cosine similarity between quantized and original cos_sim = cosine_similarity(quantized, x) print(f" ✓ FP4 quantization passed (max diff: {max_diff:.6f}, mean diff: {mean_diff:.6f}, cosine sim: {cos_sim:.6f})") def test_finite_values(): """Test that fake_quantize produces finite values.""" print("\nTest 7: Finite values check") torch.manual_seed(42) x = torch.randn(128, 64, device=DEVICE, dtype=torch.float16) quantized, high_prec = call_fake_quantize(x, BLOCK_SIZE_OUT_DIM=128, BLOCK_SIZE_QUANT_DIM=64) # Check for NaN and Inf assert not torch.isnan(quantized).any(), "Quantized tensor contains NaN" assert not torch.isinf(quantized).any(), "Quantized tensor contains Inf" assert not torch.isnan(high_prec).any(), "High precision tensor contains NaN" assert not torch.isinf(high_prec).any(), "High precision tensor contains Inf" print(" ✓ All values are finite") def test_different_block_sizes(): """Test fake_quantize with different block sizes.""" print("\nTest 8: Different block sizes") torch.manual_seed(42) x = torch.randn(256, 128, device=DEVICE, dtype=torch.float16) block_configs = [ (64, 64), (128, 64), (64, 128), (128, 128), ] for block_out, block_quant in block_configs: quantized, high_prec = call_fake_quantize( x, BLOCK_SIZE_OUT_DIM=block_out, BLOCK_SIZE_QUANT_DIM=block_quant ) assert quantized.shape == x.shape assert high_prec.shape == x.shape print(f" ✓ Block size ({block_out}, {block_quant}) passed") # ============================================================================ # MXFP Kernel Tests # ============================================================================ @triton.jit def mxfp_fake_quantize_kernel( src_ptr, dst_ptr, high_prec_ptr, BLOCK_SIZE_OUT_DIM: tl.constexpr, BLOCK_SIZE_QUANT_DIM: tl.constexpr, dst_dtype: tl.constexpr, mx_tensor_dtype: tl.constexpr, stride_src_outer, stride_src_quant, stride_dst_outer, stride_dst_quant, stride_high_prec_outer, stride_high_prec_quant, outer_dim, quant_dim, ): """ Kernel wrapper to test MXFP fake_quantize using _downcast_to_mxfp and _upcast_from_mxfp. """ outer_idx = tl.program_id(0) quant_idx = tl.program_id(1) # Determine if we are dealing with fp4 types is_fp4: tl.constexpr = mx_tensor_dtype == tl.uint8 K_DIVISOR: tl.constexpr = 2 if is_fp4 else 1 BLOCK_SIZE_QUANT_MX_SCALE: tl.constexpr = BLOCK_SIZE_QUANT_DIM // MXFP_BLOCK_SIZE BLOCK_SIZE_QUANT_MX_TENSOR: tl.constexpr = BLOCK_SIZE_QUANT_DIM // K_DIVISOR # Allocate temporary storage for quantized tensor and scale # These need to be stored in shared memory or we need to allocate them # For simplicity, we'll compute quantize and dequantize in the same kernel # Compute offsets start_outer = outer_idx * BLOCK_SIZE_OUT_DIM start_quant = quant_idx * BLOCK_SIZE_QUANT_DIM start_mx_scale_quant = quant_idx * BLOCK_SIZE_QUANT_MX_SCALE start_mx_quant = quant_idx * BLOCK_SIZE_QUANT_MX_TENSOR # Create offset arrays offs_outer = tl.arange(0, BLOCK_SIZE_OUT_DIM)[:, None] offs_quant = tl.arange(0, BLOCK_SIZE_QUANT_DIM)[None, :] offs_mx_scale_quant = tl.arange(0, BLOCK_SIZE_QUANT_MX_SCALE)[None, :] offs_mx_quant = tl.arange(0, BLOCK_SIZE_QUANT_MX_TENSOR)[None, :] # Create masks for valid elements mask_outer = (start_outer + offs_outer) < outer_dim mask_quant = (start_quant + offs_quant) < quant_dim full_mask = mask_outer & mask_quant # Load source tensor src_offsets = (start_outer + offs_outer) * stride_src_outer + (start_quant + offs_quant) * stride_src_quant src_tensor = tl.load(src_ptr + src_offsets, mask=full_mask, other=0.0) # Store high precision tensor (original) high_prec_offsets = (start_outer + offs_outer) * stride_high_prec_outer + (start_quant + offs_quant) * stride_high_prec_quant tl.store(high_prec_ptr + high_prec_offsets, src_tensor.to(tl.float32), mask=full_mask) # Quantize using _compute_quant_and_scale (same as fake_quantize uses internally) valid_mask = tl.full(src_tensor.shape, 1, tl.int1) quantized_tensor, scale_tensor, s_dec = _compute_quant_and_scale( src_tensor=src_tensor, valid_src_mask=valid_mask, mx_tensor_dtype=mx_tensor_dtype ) # Dequantize using _compute_dequant (same as fake_quantize uses internally) dequantized_tensor = _compute_dequant( mx_tensor=quantized_tensor, scale=scale_tensor, s_dec=s_dec, BLOCK_SIZE_OUT_DIM=BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_DIM=BLOCK_SIZE_QUANT_DIM, dst_dtype=dst_dtype ) # Store result dst_offsets = (start_outer + offs_outer) * stride_dst_outer + (start_quant + offs_quant) * stride_dst_quant tl.store(dst_ptr + dst_offsets, dequantized_tensor, mask=full_mask) def call_mxfp_fake_quantize( x: torch.Tensor, BLOCK_SIZE_OUT_DIM: int = 128, BLOCK_SIZE_QUANT_DIM: int = 128, dst_dtype: torch.dtype = None, mx_tensor_dtype = tl.uint8 # Default to FP4 (uint8) ) -> Tuple[torch.Tensor, torch.Tensor]: """ Call MXFP fake_quantize on a tensor using _downcast_to_mxfp and _upcast_from_mxfp logic. Args: x: Input tensor (2D or can be reshaped to 2D) BLOCK_SIZE_OUT_DIM: Block size for outer dimension BLOCK_SIZE_QUANT_DIM: Block size for quantization dimension (must be multiple of 16) dst_dtype: Output dtype (defaults to input dtype) mx_tensor_dtype: Quantization dtype (tl.uint8 for fp4 or tl.float8e4nv for fp8) Returns: Tuple of (quantized_tensor, high_precision_tensor) """ assert x.is_cuda, "Input must be on CUDA" assert BLOCK_SIZE_QUANT_DIM % 16 == 0, f"BLOCK_SIZE_QUANT_DIM must be multiple of 16, got {BLOCK_SIZE_QUANT_DIM}" orig_shape = x.shape orig_dtype = x.dtype # Reshape to 2D x_2d = x.view(-1, x.shape[-1]) outer_dim, quant_dim = x_2d.shape if dst_dtype is None: dst_dtype = orig_dtype # Map torch dtype to triton dtype dtype_map = { torch.float32: tl.float32, torch.float16: tl.float16, torch.bfloat16: tl.bfloat16, } triton_dst_dtype = dtype_map.get(dst_dtype, tl.float16) # Allocate outputs quantized_output = torch.empty_like(x_2d, dtype=dst_dtype) high_prec_output = torch.empty_like(x_2d, dtype=torch.float32) # Launch kernel grid = ( triton.cdiv(outer_dim, BLOCK_SIZE_OUT_DIM), triton.cdiv(quant_dim, BLOCK_SIZE_QUANT_DIM), ) mxfp_fake_quantize_kernel[grid]( src_ptr=x_2d, dst_ptr=quantized_output, high_prec_ptr=high_prec_output, BLOCK_SIZE_OUT_DIM=BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_DIM=BLOCK_SIZE_QUANT_DIM, dst_dtype=triton_dst_dtype, mx_tensor_dtype=mx_tensor_dtype, stride_src_outer=x_2d.stride(0), stride_src_quant=x_2d.stride(1), stride_dst_outer=quantized_output.stride(0), stride_dst_quant=quantized_output.stride(1), stride_high_prec_outer=high_prec_output.stride(0), stride_high_prec_quant=high_prec_output.stride(1), outer_dim=outer_dim, quant_dim=quant_dim, ) return quantized_output.view(orig_shape).to(dst_dtype), high_prec_output.view(orig_shape) def test_mxfp_basic_functionality(): """Test basic functionality of MXFP fake_quantize.""" print("\n" + "=" * 60) print("MXFP Kernel Tests") print("=" * 60) print("\nTest 1: Basic functionality (MXFP)") torch.manual_seed(42) x = torch.randn(128, 64, device=DEVICE, dtype=torch.float16) quantized, high_prec = call_mxfp_fake_quantize(x, BLOCK_SIZE_OUT_DIM=128, BLOCK_SIZE_QUANT_DIM=64) # Check shapes assert quantized.shape == x.shape, f"Shape mismatch: {quantized.shape} vs {x.shape}" assert high_prec.shape == x.shape, f"Shape mismatch: {high_prec.shape} vs {x.shape}" # Check that quantization happened (values should be different) diff = (quantized.float() - x.float()).abs() max_diff = diff.max().item() mean_diff = diff.mean().item() assert max_diff > 0, "Quantization should change values" # Compute cosine similarity between quantized and original cos_sim = cosine_similarity(quantized, x) # Check that high_prec matches original high_prec_diff = (high_prec - x.float()).abs().max().item() assert high_prec_diff < 1e-5, f"High precision tensor should match original, got diff: {high_prec_diff}" print(f" ✓ Shape check passed") print(f" ✓ Quantization occurred (max diff: {max_diff:.6f}, mean diff: {mean_diff:.6f}, cosine sim: {cos_sim:.6f})") print(f" ✓ High precision tensor matches original (diff: {high_prec_diff:.6e})") def test_mxfp_different_shapes(): """Test MXFP fake_quantize with different tensor shapes.""" print("\nTest 2: Different shapes (MXFP)") torch.manual_seed(42) test_shapes = [ (64, 64), (128, 128), (256, 128), (128, 256), (32, 64), ] for shape in test_shapes: x = torch.randn(shape, device=DEVICE, dtype=torch.float16) # Use block sizes that divide evenly block_out = min(128, shape[0]) block_quant = min(128, shape[1]) # Ensure block_quant is multiple of 16 block_quant = (block_quant // 16) * 16 if block_quant == 0: block_quant = 16 quantized, high_prec = call_mxfp_fake_quantize( x, BLOCK_SIZE_OUT_DIM=block_out, BLOCK_SIZE_QUANT_DIM=block_quant ) assert quantized.shape == x.shape, f"Shape mismatch for {shape}" assert high_prec.shape == x.shape, f"Shape mismatch for {shape}" print(f" ✓ Shape {shape} passed") def test_mxfp_different_dtypes(): """Test MXFP fake_quantize with different input dtypes.""" print("\nTest 3: Different dtypes (MXFP)") torch.manual_seed(42) shape = (128, 64) dtypes = [torch.float16, torch.bfloat16, torch.float32] for dtype in dtypes: x = torch.randn(shape, device=DEVICE, dtype=dtype) quantized, high_prec = call_mxfp_fake_quantize( x, BLOCK_SIZE_OUT_DIM=128, BLOCK_SIZE_QUANT_DIM=64, dst_dtype=dtype ) assert quantized.dtype == dtype, f"Dtype mismatch: {quantized.dtype} vs {dtype}" assert quantized.shape == x.shape, f"Shape mismatch for dtype {dtype}" print(f" ✓ Dtype {dtype} passed") def test_mxfp_edge_cases(): """Test MXFP fake_quantize with edge cases.""" print("\nTest 4: Edge cases (MXFP)") # Test with zeros print(" Testing zeros...") x_zeros = torch.zeros(128, 64, device=DEVICE, dtype=torch.float16) quantized, high_prec = call_mxfp_fake_quantize(x_zeros, BLOCK_SIZE_OUT_DIM=128, BLOCK_SIZE_QUANT_DIM=64) assert quantized.shape == x_zeros.shape # Quantized zeros should still be close to zero assert quantized.abs().max().item() < 1e-3, "Quantized zeros should be close to zero" print(" ✓ Zeros passed") # Test with ones print(" Testing ones...") x_ones = torch.ones(128, 64, device=DEVICE, dtype=torch.float16) quantized, high_prec = call_mxfp_fake_quantize(x_ones, BLOCK_SIZE_OUT_DIM=128, BLOCK_SIZE_QUANT_DIM=64) assert quantized.shape == x_ones.shape # Quantized ones should be close to 1.0 diff = (quantized - x_ones).abs().max().item() assert diff < 0.1, f"Quantized ones should be close to 1.0, got max diff: {diff}" print(f" ✓ Ones passed (max diff: {diff:.6f})") # Test with small values print(" Testing small values...") x_small = torch.randn(128, 64, device=DEVICE, dtype=torch.float16) * 0.001 quantized, high_prec = call_mxfp_fake_quantize(x_small, BLOCK_SIZE_OUT_DIM=128, BLOCK_SIZE_QUANT_DIM=64) assert quantized.shape == x_small.shape print(" ✓ Small values passed") # Test with large values print(" Testing large values...") x_large = torch.randn(128, 64, device=DEVICE, dtype=torch.float16) * 100.0 quantized, high_prec = call_mxfp_fake_quantize(x_large, BLOCK_SIZE_OUT_DIM=128, BLOCK_SIZE_QUANT_DIM=64) assert quantized.shape == x_large.shape print(" ✓ Large values passed") def test_mxfp_fp4_quantization(): """Test MXFP fake_quantize with FP4 quantization (uint8).""" print("\nTest 5: FP4 quantization (MXFP, uint8)") torch.manual_seed(42) x = torch.randn(128, 64, device=DEVICE, dtype=torch.float16) quantized, high_prec = call_mxfp_fake_quantize( x, BLOCK_SIZE_OUT_DIM=128, BLOCK_SIZE_QUANT_DIM=64, mx_tensor_dtype=tl.uint8 # FP4 uses uint8 ) assert quantized.shape == x.shape assert high_prec.shape == x.shape # Check that quantization happened diff = (quantized.float() - x.float()).abs() max_diff = diff.max().item() mean_diff = diff.mean().item() assert max_diff > 0, "Quantization should change values" # Compute cosine similarity between quantized and original cos_sim = cosine_similarity(quantized, x) print(f" ✓ FP4 quantization passed (max diff: {max_diff:.6f}, mean diff: {mean_diff:.6f}, cosine sim: {cos_sim:.6f})") def test_mxfp_fp8_quantization(): """Test MXFP fake_quantize with FP8 quantization (float8e4nv).""" print("\nTest 6: FP8 quantization (MXFP, float8e4nv)") torch.manual_seed(42) x = torch.randn(128, 64, device=DEVICE, dtype=torch.float16) quantized, high_prec = call_mxfp_fake_quantize( x, BLOCK_SIZE_OUT_DIM=128, BLOCK_SIZE_QUANT_DIM=64, mx_tensor_dtype=tl.float8e4nv # FP8 uses float8e4nv ) assert quantized.shape == x.shape assert high_prec.shape == x.shape # Check that quantization happened diff = (quantized.float() - x.float()).abs() max_diff = diff.max().item() mean_diff = diff.mean().item() assert max_diff > 0, "Quantization should change values" # Compute cosine similarity between quantized and original cos_sim = cosine_similarity(quantized, x) print(f" ✓ FP8 quantization passed (max diff: {max_diff:.6f}, mean diff: {mean_diff:.6f}, cosine sim: {cos_sim:.6f})") def test_mxfp_finite_values(): """Test that MXFP fake_quantize produces finite values.""" print("\nTest 7: Finite values check (MXFP)") torch.manual_seed(42) x = torch.randn(128, 64, device=DEVICE, dtype=torch.float16) quantized, high_prec = call_mxfp_fake_quantize(x, BLOCK_SIZE_OUT_DIM=128, BLOCK_SIZE_QUANT_DIM=64) # Check for NaN and Inf assert not torch.isnan(quantized).any(), "Quantized tensor contains NaN" assert not torch.isinf(quantized).any(), "Quantized tensor contains Inf" assert not torch.isnan(high_prec).any(), "High precision tensor contains NaN" assert not torch.isinf(high_prec).any(), "High precision tensor contains Inf" print(" ✓ All values are finite") def test_mxfp_different_block_sizes(): """Test MXFP fake_quantize with different block sizes.""" print("\nTest 8: Different block sizes (MXFP)") torch.manual_seed(42) x = torch.randn(256, 128, device=DEVICE, dtype=torch.float16) block_configs = [ (64, 64), (128, 64), (64, 128), (128, 128), ] for block_out, block_quant in block_configs: quantized, high_prec = call_mxfp_fake_quantize( x, BLOCK_SIZE_OUT_DIM=block_out, BLOCK_SIZE_QUANT_DIM=block_quant ) assert quantized.shape == x.shape assert high_prec.shape == x.shape print(f" ✓ Block size ({block_out}, {block_quant}) passed") def test_get_fake_quant_metrics(): """Test get_fake_quant function with cosine similarity, mean_diff, and max_diff metrics.""" print("\n" + "=" * 60) print("Testing get_fake_quant with cosine similarity, mean_diff, and max_diff") print("=" * 60) if not torch.cuda.is_available(): print("CUDA not available, skipping test") return device = torch.device("cuda") torch.manual_seed(42) # Test with different shapes and dtypes test_cases = [ ((128, 64), torch.bfloat16, "2D bfloat16"), ((256, 128), torch.bfloat16, "2D bfloat16 large"), ((2, 8, 128), torch.bfloat16, "3D bfloat16"), ((128, 64), torch.float16, "2D float16"), ((1, 8, 256, 128), torch.bfloat16, "4D bfloat16"), ] for shape, dtype, name in test_cases: print(f"\nTest: {name} - Shape: {shape}") x = torch.randn(*shape, device=device, dtype=dtype) x_fq = get_fake_quant(x) # Check shapes match assert x_fq.shape == x.shape, f"Shape mismatch: {x_fq.shape} vs {x.shape}" assert x_fq.dtype == x.dtype, f"Dtype mismatch: {x_fq.dtype} vs {x.dtype}" # Compute metrics diff = (x_fq.float() - x.float()).abs() max_diff = diff.max().item() mean_diff = diff.mean().item() cos_sim = cosine_similarity(x_fq, x) print(f" Max diff: {max_diff:.6f}") print(f" Mean diff: {mean_diff:.6f}") print(f" Cosine similarity: {cos_sim:.6f}") # Basic sanity checks assert max_diff < 10.0, f"Max difference too large: {max_diff}" assert cos_sim > 0.5, f"Cosine similarity too low: {cos_sim}" print("\n✓ All metric tests passed!") if __name__ == "__main__": if not torch.cuda.is_available(): print("CUDA not available. These tests require CUDA.") exit(1) print("=" * 60) print("Running basic tests for fake_quantize function") print("=" * 60) print(f"Device: {DEVICE}") print() try: test_basic_functionality() test_different_shapes() test_different_dtypes() test_edge_cases() test_fp4_quantization() test_finite_values() # Test get_fake_quant with metrics test_get_fake_quant_metrics() # Run MXFP kernel tests # test_mxfp_basic_functionality() # test_mxfp_different_shapes() # test_mxfp_different_dtypes() # test_mxfp_edge_cases() # test_mxfp_fp4_quantization() # test_mxfp_fp8_quantization() # test_mxfp_finite_values() # test_mxfp_different_block_sizes() 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() exit(1)