"""Tune vLLM's Triton fused-MoE kernel for this GPU on prain's slot tensors.

usage: python scripts/tune_moe.py OUT_DIR [FORMAT...]   (scripts/env.sh points
                                    VLLM_TUNED_CONFIG_FOLDER at build/moe-configs)
Times candidate configs on one prain round (16 slots of 128 routed experts, expert_map;
up to 32 tokens route into the slots like decode, more tokens spread like prefill),
Gemma 4 expert shapes, bf16/int8/int4 slots, and writes the fastest per batch size as vLLM
config files for every slot count E (the best config hardly depends on E). Every config in
the files compiled, ran and matched the default config's output on this GPU.
"""

import itertools
import json
import os
import sys
import time

import torch
from vllm.config import VllmConfig, set_current_vllm_config
from vllm.model_executor.layers.fused_moe import override_config
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
from vllm.model_executor.layers.fused_moe.config import (int4_w4a16_moe_quant_config,
                                                         int8_w8a16_moe_quant_config)
from vllm.model_executor.layers.fused_moe.fused_moe import (fused_experts, get_config_file_name,
                                                            get_default_config)

from prain.quant import GROUP, quantize

H, I, TOPK, EXPERTS, SLOTS = 2816, 704, 8, 128, 16
SIZES = (1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048)
FORMATS = {"bf16": (None, None), "int8": ("int8_w8a16", None), "int4": ("int4_w4a16", [0, GROUP])}


def candidates(tokens):
    # ponytail: ~8 s per Triton compile on this host, so a small space (~70 compiles per
    # format); widen it when a GPU compiles faster.
    if tokens <= 32:
        space = itertools.product((16,), (32, 64, 128), (64, 128, 256), (1,), (2, 4), (2,))
    else:
        space = itertools.product((32, 64), (64, 128), (32, 64), (8,), (4, 8), (2,))
    for bm, bn, bk, group, warps, stages in space:
        yield {"BLOCK_SIZE_M": bm, "BLOCK_SIZE_N": bn, "BLOCK_SIZE_K": bk, "GROUP_SIZE_M": group,
               "SPLIT_K": 1, "num_warps": warps, "num_stages": stages}  # int4 needs SPLIT_K


def timed(fn, repeat):
    fn()
    torch.cuda.synchronize()
    start = time.perf_counter()
    for _ in range(repeat):
        fn()
    torch.cuda.synchronize()
    return (time.perf_counter() - start) * 1000 / repeat


def main(out_dir, formats):
    os.makedirs(out_dir, exist_ok=True)
    torch.manual_seed(0)
    w13 = torch.randn(SLOTS, 2 * I, H, device="cuda").mul_(0.02).bfloat16()
    w2 = torch.randn(SLOTS, H, I, device="cuda").mul_(0.02).bfloat16()
    table = torch.full((EXPERTS + 1,), -1, dtype=torch.int32, device="cuda")
    table[:SLOTS] = torch.arange(SLOTS, dtype=torch.int32, device="cuda")
    for name in formats:
        dtype, block_shape = FORMATS[name]
        if name == "bf16":
            a, b, quant_config = w13, w2, None
        else:
            (a, s13), (b, s2) = quantize(w13, name), quantize(w2, name)
            make = int8_w8a16_moe_quant_config if name == "int8" else int4_w4a16_moe_quant_config
            quant_config = make(w1_scale=s13, w2_scale=s2, w1_zp=None, w2_zp=None,
                                block_shape=block_shape)
        best = {}
        for tokens in SIZES:
            x = torch.randn(tokens, H, device="cuda").mul_(0.5).bfloat16()
            pool = SLOTS if tokens <= 32 else EXPERTS
            ids = torch.stack([torch.randperm(pool, device="cuda")[:TOPK]
                               for _ in range(tokens)]).to(torch.int32)
            weights = torch.softmax(torch.randn(tokens, TOPK, device="cuda"), -1)

            def run():
                return fused_experts(x, a, b, weights, ids, activation=MoEActivation.GELU_TANH,
                                     global_num_experts=EXPERTS + 1, expert_map=table,
                                     quant_config=quant_config)

            default = get_default_config(tokens, SLOTS, I, a.shape[2], TOPK, dtype, block_shape)
            with override_config(default):  # vLLM's own choice when no config file exists
                default_ms = timed(run, 10 if tokens <= 32 else 3)
                reference = run()
            results = [(default_ms, json.dumps(default, sort_keys=True))]  # never worse than it
            for config in candidates(tokens):
                # catch inside: vLLM's override_config skips its restore when the body raises
                with override_config(config):
                    try:
                        ms = timed(run, 10 if tokens <= 32 else 3)
                        out = run()
                    except Exception:  # noqa: BLE001 - configs that do not compile or fit
                        continue
                if (out.float() - reference.float()).norm() <= 1e-2 * reference.float().norm() + 1e-6:
                    results.append((ms, json.dumps(config, sort_keys=True)))
            ms, config = min(results)
            best[str(tokens)] = json.loads(config)
            print(f"{name} tokens={tokens:4d}: default {default_ms:8.3f} ms, best {ms:8.3f} ms "
                  f"({len(results)} configs) {config}", flush=True)
        for e in range(1, EXPERTS + 1):  # any slot count a placement can give a layer
            with open(os.path.join(out_dir, get_config_file_name(e, I, dtype, block_shape)), "w") as f:
                json.dump(best, f, indent=1)


if __name__ == "__main__":
    with set_current_vllm_config(VllmConfig()):
        main(sys.argv[1], sys.argv[2:] or list(FORMATS))