"""Cross-check: tile_kernels.quant.per_token_cast output feeds deep_gemm's FP8 GEMM.

This is the exact ReplicatedLinear.apply_weight_fp8 path
(deepinfer/model_executor/layers/linear.py) with CPU tensors.
"""

import os
import sys

sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "..", "..", "deepinfer", "neuron_ops", "vendor"))

import torch

import deep_gemm
import tile_kernels

CAST_PARAMS = dict(fmt="e4m3", num_per_channels=32, round_sf=True,
                   use_packed_ue8m0=True, use_tma_aligned_col_major_sf=True)


def calc_diff(x, y):
    x, y = x.double(), y.double()
    denom = (x * x + y * y).sum().item()
    return 0.0 if denom == 0 else 1 - 2 * (x * y).sum().item() / denom


m, n, k = 48, 128, 256
x = torch.randn(m, k, dtype=torch.bfloat16)
w = torch.randn(n, k, dtype=torch.bfloat16)

xq, x_sf = tile_kernels.quant.per_token_cast(x, **CAST_PARAMS)
wq, w_sf = tile_kernels.quant.per_token_cast(w, **CAST_PARAMS)
print(f"{xq.dtype=} {x_sf.dtype=} {x_sf.shape=} {x_sf.stride()=}")
print(f"{wq.dtype=} {w_sf.dtype=} {w_sf.shape=} {w_sf.stride()=}")

out = torch.empty(m, n, dtype=torch.bfloat16)
deep_gemm.fp8_fp4_gemm_nt((xq, x_sf), (wq, w_sf), out, None, recipe=(1, 1, 32))

ref = x.float() @ w.float().T
diff = calc_diff(out.float(), ref)
print(f"diff vs bf16 reference: {diff:.3e}")
assert diff < 2e-3, diff

# The dequantized operands must agree with the GEMM result exactly.
from deep_gemm._quant import dequantize  # noqa: E402

x_deq = dequantize(xq, deep_gemm.transform_sf_into_required_layout(x_sf, m, k, (1, 1, 32), is_sfa=True))
w_deq = dequantize(wq, deep_gemm.transform_sf_into_required_layout(w_sf, n, k, (1, 1, 32), is_sfa=False))
exact = (x_deq @ w_deq.T).to(torch.bfloat16)
print("exact match with dequantized matmul:", torch.equal(out, exact))
print(f"quantization error (dequant vs bf16): {calc_diff(exact.float(), ref):.3e}")
print("OK")