93 lines
3.1 KiB
Python
93 lines
3.1 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Debug script to understand the scale computation issue.
|
|
"""
|
|
|
|
import torch
|
|
import triton
|
|
import triton.language as tl
|
|
from nvfp4_utils import _compute_quant_and_scale, _compute_dequant
|
|
|
|
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
|
|
@triton.jit
|
|
def debug_scale_kernel(
|
|
src_ptr,
|
|
scale_ptr,
|
|
s_dec_ptr,
|
|
BLOCK_SIZE_OUT_DIM: tl.constexpr,
|
|
BLOCK_SIZE_QUANT_DIM: tl.constexpr,
|
|
stride_src_outer,
|
|
stride_src_quant,
|
|
outer_dim,
|
|
quant_dim,
|
|
):
|
|
"""Debug kernel to print scale values."""
|
|
outer_idx = tl.program_id(0)
|
|
quant_idx = tl.program_id(1)
|
|
|
|
start_outer = outer_idx * BLOCK_SIZE_OUT_DIM
|
|
start_quant = quant_idx * BLOCK_SIZE_QUANT_DIM
|
|
|
|
offs_outer = tl.arange(0, BLOCK_SIZE_OUT_DIM)[:, None]
|
|
offs_quant = tl.arange(0, BLOCK_SIZE_QUANT_DIM)[None, :]
|
|
|
|
mask_outer = (start_outer + offs_outer) < outer_dim
|
|
mask_quant = (start_quant + offs_quant) < quant_dim
|
|
full_mask = mask_outer & mask_quant
|
|
|
|
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)
|
|
|
|
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=tl.float8e4nv
|
|
)
|
|
|
|
# Store scale and s_dec for inspection
|
|
# scale_tensor has shape [BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_MX_SCALE] where BLOCK_SIZE_QUANT_MX_SCALE = BLOCK_SIZE_QUANT_DIM // 16
|
|
BLOCK_SIZE_QUANT_MX_SCALE: tl.constexpr = BLOCK_SIZE_QUANT_DIM // 16
|
|
scale_offsets = (start_outer + offs_outer) * stride_src_outer + (start_quant // 16) * stride_src_quant
|
|
tl.store(scale_ptr + scale_offsets, scale_tensor.reshape([BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_MX_SCALE])[:, 0].to(tl.float32), mask=mask_outer)
|
|
tl.store(s_dec_ptr + outer_idx, s_dec)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
if not torch.cuda.is_available():
|
|
print("CUDA not available")
|
|
exit(1)
|
|
|
|
# Test with all ones
|
|
x = torch.ones(128, 64, device=DEVICE, dtype=torch.float16)
|
|
x_2d = x.view(-1, x.shape[-1])
|
|
outer_dim, quant_dim = x_2d.shape
|
|
|
|
scale_output = torch.zeros(outer_dim, quant_dim // 16, device=DEVICE, dtype=torch.float32)
|
|
s_dec_output = torch.zeros(outer_dim // 128, device=DEVICE, dtype=torch.float32)
|
|
|
|
grid = (
|
|
triton.cdiv(outer_dim, 128),
|
|
triton.cdiv(quant_dim, 64),
|
|
)
|
|
|
|
debug_scale_kernel[grid](
|
|
src_ptr=x_2d,
|
|
scale_ptr=scale_output,
|
|
s_dec_ptr=s_dec_output,
|
|
BLOCK_SIZE_OUT_DIM=128,
|
|
BLOCK_SIZE_QUANT_DIM=64,
|
|
stride_src_outer=x_2d.stride(0),
|
|
stride_src_quant=x_2d.stride(1),
|
|
outer_dim=outer_dim,
|
|
quant_dim=quant_dim,
|
|
)
|
|
|
|
print(f"Scale values (first block): {scale_output[0, :4]}")
|
|
print(f"s_dec value: {s_dec_output[0]}")
|
|
print(f"Expected s_dec: {1.0 / (6 * 448 / 1.0)} = {1.0 / 2688}")
|
|
print(f"Expected scale: {(1.0 / 6) * 2688} = {2688 / 6} = {448}")
|
|
|
|
|