import torch  # must precede qtip_kernels so libc10 is loaded (the .so links it)
try:
    import qtip_kernels
except Exception as _qtip_err:  # ABI mismatch (torch upgrade) / not built / no CUDA
    qtip_kernels = None
    import warnings
    warnings.warn(
        f"qtip_kernels unavailable ({_qtip_err}); CUDA trellis kernels disabled — "
        f"falling back to torch decode. has_kernel() is forced False, so quantize/"
        f"hfize/eval run the (slower) manifest path. Rebuild with nvcc to re-enable.")

kernels = [
    (5120, 1, 5120, 4),
    (5120, 1, 5120, 3),
    (5120, 1, 5120, 2),
    (1024, 1, 5120, 4),
    (1024, 1, 5120, 3),
    (1024, 1, 5120, 2),
    (8192, 1, 5120, 4),
    (8192, 1, 5120, 3),
    (8192, 1, 5120, 2),
    (5120, 1, 8192, 4),
    (5120, 1, 8192, 3),
    (5120, 1, 8192, 2),
    (1024, 1, 3072, 4),
    (8192, 1, 3072, 4),
    (3072, 1, 8192, 4),
    (3072, 1, 3072, 4),
    (53248, 1, 16384, 2),
    (53248, 1, 16384, 3),
    (53248, 1, 16384, 4),
    (16384, 1, 53248, 2),
    (16384, 1, 53248, 3),
    (16384, 1, 53248, 4),
    (1024, 1, 16384, 2),
    (1024, 1, 16384, 3),
    (1024, 1, 16384, 4),
    (16384, 1, 16384, 2),
    (16384, 1, 16384, 3),
    (16384, 1, 16384, 4),
    (4096, 1, 14336, 2),
    (4096, 1, 14336, 3),
    (4096, 1, 14336, 4),
    (14336, 1, 4096, 2),
    (14336, 1, 4096, 3),
    (14336, 1, 4096, 4),
    (1024, 1, 4096, 2),
    (1024, 1, 4096, 3),
    (1024, 1, 4096, 4),
    (4096, 1, 4096, 2),
    (4096, 1, 11008, 2),
    (4096, 1, 12288, 2),
    (11008, 1, 4096, 2),
    (12288, 1, 4096, 2),
    (22016, 1, 4096, 2),
    (8192, 1, 8192, 2),
    (10240, 1, 8192, 2),
    (10240, 1, 8192, 3),
    (10240, 1, 8192, 4),
    (57344, 1, 8192, 2),
    (57344, 1, 8192, 3),
    (57344, 1, 8192, 4),
    (8192, 1, 1024, 2),
    (8192, 1, 28672, 2),
    (28672, 1, 8192, 2),
    (1024, 1, 8192, 2),
    (4096, 1, 4096, 3),
    (4096, 1, 11008, 3),
    (4096, 1, 12288, 3),
    (11008, 1, 4096, 3),
    (12288, 1, 4096, 3),
    (22016, 1, 4096, 3),
    (8192, 1, 8192, 3),
    (8192, 1, 1024, 3),
    (8192, 1, 28672, 3),
    (28672, 1, 8192, 3),
    (1024, 1, 8192, 3),
    (4096, 1, 4096, 4),
    (4096, 1, 11008, 4),
    (4096, 1, 12288, 4),
    (11008, 1, 4096, 4),
    (12288, 1, 4096, 4),
    (22016, 1, 4096, 4),
    (8192, 1, 8192, 4),
    (8192, 1, 1024, 4),
    (8192, 1, 28672, 4),
    (28672, 1, 8192, 4),
    (1024, 1, 8192, 4),
    (27648, 1, 5120, 2),
    (27648, 1, 5120, 3),
    (27648, 1, 5120, 4),
    (5120, 1, 27648, 2),
    (5120, 1, 27648, 3),
    (5120, 1, 27648, 4),
    # Qwen3.5-27B text-only projections (k4_hf: R=4, k2_hf: R=2)
    (5120, 1, 6144, 4),
    (5120, 1, 6144, 2),
    (5120, 1, 17408, 4),
    (5120, 1, 17408, 2),
    (6144, 1, 5120, 4),
    (6144, 1, 5120, 2),
    (10240, 1, 5120, 4),
    (10240, 1, 5120, 2),
    (12288, 1, 5120, 4),
    (12288, 1, 5120, 2),
    (17408, 1, 5120, 4),
    (17408, 1, 5120, 2),
    # Qwen3.5-35B-A3B MoE (body K=4, routed experts K=2)
    (512, 1, 2048, 4),
    (1024, 1, 2048, 2),
    (2048, 1, 512, 2),
    (2048, 1, 512, 4),
    (2048, 1, 4096, 4),
    (4096, 1, 2048, 4),
    (8192, 1, 2048, 4),
    # Qwen3.5-122B-A10B MoE (hidden=3072, moe_inter=1024; body K=4, routed experts K=2)
    (2048, 1, 3072, 2),    # routed experts gate_up_proj
    (3072, 1, 1024, 2),    # routed experts down_proj
    (512, 1, 3072, 4),     # self_attn k_proj / v_proj
    (3072, 1, 1024, 4),    # shared_expert down_proj
    (12288, 1, 3072, 4),   # linear_attn in_proj_qkv
    (16384, 1, 3072, 4),   # self_attn q_proj
]

kdict = {}

for m, n, k, bitrate in kernels:
    torch.library.define(
        f"quip_lib::decompress_matvec_qtip_{m}_{n}_{k}_{bitrate}",
        "(Tensor compressed, Tensor x, Tensor codebook) -> Tensor")

    name = f"decompress_matvec_qtip_{m}_{n}_{k}_{bitrate}"
    kernel_name = f"qtip_kernels.decompress_matvec_16_9_{bitrate}_1_{m}_{n}_{k}"
    exec(f"""\
@torch.library.register_fake("quip_lib::{name}")
def {name}_abstract(
        compressed: torch.Tensor,
        x: torch.Tensor,
        codebook: torch.Tensor) -> torch.Tensor:
    return torch.zeros(1, {m}, dtype=torch.float32, device=x.device)

@torch.library.impl("quip_lib::{name}", "cuda")
def {name}_cuda(
        compressed: torch.Tensor,
        x: torch.Tensor,
        codebook: torch.Tensor) -> torch.Tensor:
    out = torch.zeros(({m}, 1), dtype=torch.float32, device=x.device)
    {kernel_name}(out, compressed.reshape(-1).view(torch.int32), x.to(torch.float16).T, codebook.reshape(-1))
    return out.T
    """)