"""Cheaper MTP drafter for Qwen3.8-Flash-Next (Qwen4ExpMTP) on vLLM 0.31 Model Runner V2 - lossless.

Same recipe as autotrust/Qwen3.8-27B-NVFP4-DGX-SPARK, adapted to the Flash-Next MTP layer. After
MTPSpeculator.load_draft_model:
  1. the drafter gets its own lm_head restricted to a draft vocabulary (64k rows gathered from the BF16 target
     lm_head, special tokens always kept); greedy drafting takes the argmax over the kept rows and maps it back
     to target ids (speculative_config.use_local_argmax_reduction=True routes drafting through get_top_tokens);
  2. that head is re-quantized to NVFP4 W4A16 (Marlin GEMV);
  3. the BF16 linear layers of the MTP block (fc_embedding, fc_hidden, attention qkv/o, shared expert) are
     re-quantized to NVFP4 W4A16 too. The MTP routed experts are already NVFP4 in the checkpoint; the router,
     the hyper-connection mixers and the QSA indexer stay BF16.
Bytes read per draft step at batch 1: ~1.45 GB -> ~0.16 GB (BF16 248k x 2560 head = 1.27 GB alone).
The target model and the lm_head it verifies with are never touched, so accepted tokens are exactly the
target's own (greedy: identical output; sampling: same distribution).

In-process use: install() before vllm.LLM(...) with VLLM_ENABLE_V1_MULTIPROCESSING=0.
vllm serve: set Q38_DRAFT_VOCAB (and optionally Q38_DRAFT_HEAD_Q / Q38_DRAFT_MTP_Q) and put the directory that
holds sitecustomize.py (or the plugin) on PYTHONPATH; see README.
"""
import os
import sys
import types

import torch
import torch.nn as nn
import torch.nn.functional as F

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

MTP_QUANT_KEYS = ("fc_embedding", "fc_hidden", "qkv_proj", "o_proj", "shared_expert")


class _Bf16(nn.Module):
    def __init__(self, W):
        super().__init__()
        self.register_buffer("w", W.contiguous(), persistent=False)

    def forward(self, x):
        return F.linear(x, self.w)


def make_linear(W, kind):
    if kind in (None, "bf16"):
        return _Bf16(W)
    from nvfp4_ops import NVFP4Linear
    return NVFP4Linear(W, kind)


class _QuantMethod:
    """quant_method replacement for a vLLM linear layer (keeps the module, swaps the GEMM)."""

    def __init__(self, layer, kind):
        layer.dq = make_linear(layer.weight.data, kind)

    def apply(self, layer, x, bias=None):
        y = layer.dq(x)
        return y + bias if bias is not None else y


def install(vocab_path=None, head_q="marlin", mtp_q="marlin"):
    import vllm.v1.worker.gpu.spec_decode.mtp.speculator as ms

    if getattr(ms.MTPSpeculator.load_draft_model, "_q38", False):
        return
    keep = torch.load(vocab_path) if vocab_path else None
    orig = ms.MTPSpeculator.load_draft_model

    def load_draft_model(self, target_model, target_attn_layer_names):
        model = orig(self, target_model, target_attn_layer_names)
        W = model.lm_head.weight.data
        assert W.dtype in (torch.bfloat16, torch.float16), W.dtype
        V, dev = W.shape[0], W.device
        idx = keep.nonzero()[:, 0].to(dev) if keep is not None else torch.arange(V, device=dev)
        pad = (-idx.numel()) % 64            # NVFP4 GEMMs want N % 64 == 0: pad with copies of a kept id
        if pad:
            idx = torch.cat([idx, idx[:1].expand(pad)])
        model.register_buffer("dh_idx", idx, persistent=False)
        model.draft_head = make_linear(W[idx].contiguous(), head_q)

        def get_top_tokens(self_, h):
            return self_.dh_idx[self_.draft_head(h).argmax(-1)]

        def compute_logits(self_, h, spec_step_idx=0):
            lk = self_.draft_head(h).to(torch.float32)
            out = lk.new_full((h.shape[0], V), float("-inf"))
            out.index_copy_(1, self_.dh_idx, lk)
            return out

        model.get_top_tokens = types.MethodType(get_top_tokens, model)
        model.compute_logits = types.MethodType(compute_logits, model)

        done = []
        if mtp_q not in (None, "bf16"):
            for name, mod in model.model.named_modules():
                w = getattr(mod, "weight", None)
                if (hasattr(mod, "quant_method") and isinstance(w, torch.Tensor) and w.dim() == 2
                        and w.dtype == torch.bfloat16 and any(k in name for k in MTP_QUANT_KEYS)
                        and w.shape[0] % 64 == 0 and w.shape[1] % 16 == 0):
                    mod.quant_method = _QuantMethod(mod, mtp_q)
                    done.append(f"{name}{tuple(w.shape)}")
        print(f"[q38-draft] vocab rows={idx.numel()} head_q={head_q} mtp_q={mtp_q} "
              f"quantized MTP linears ({len(done)}): {done}", flush=True)
        return model

    load_draft_model._q38 = True
    ms.MTPSpeculator.load_draft_model = load_draft_model


def install_from_env():
    vocab = os.environ.get("Q38_DRAFT_VOCAB")
    if vocab:
        install(vocab, os.environ.get("Q38_DRAFT_HEAD_Q", "marlin"), os.environ.get("Q38_DRAFT_MTP_Q", "marlin"))