import torch
import deepinfer  # noqa: sets vendor path
from deepinfer.model_executor.models.v4_modules import mhc_ops
import tile_kernels.torch as ref
torch.manual_seed(0)
t, h, M = 6, 256, 4
m3 = M * (M + 2)
res = torch.randn(1, t, M, h, dtype=torch.bfloat16)
fn = torch.randn(m3, M * h) * 0.02
w = torch.rand(M * h) + 0.5
print('n_splits', mhc_ops.default_n_splits())
for ns in (1, 16):
    try:
        out = mhc_ops.mhc_pre_norm_fn(res, fn, w, 1e-6, n_splits=ns)
        r = ref.mhc_pre_norm_fn_ref(res, fn, w, 1e-6)
        print(ns, 'max diff', (out - r).abs().max().item())
    except Exception as e:
        print(ns, 'ERR', type(e).__name__, str(e)[:300])
x = torch.randn(t, 2, 128, dtype=torch.bfloat16)
o = mhc_ops.expand_to_mhc(x, 4); print('expand ok', o.shape)