"""YuE2-specific decode kernels; packed weight arithmetic lives in OrbitQuant."""
from pathlib import Path
import types
from functools import cache
import torch
from torch.nn import functional as F
from fuse_packed import project_qkv as original_project_qkv, mlp_forward as original_mlp

_KERNEL = None
_W8A8_KERNEL = [False]


def kernel_dir(name):
    """Per-architecture kernel package directory: <name>-sm<major><minor> when shipped, else <name>."""
    base = Path(__file__).parent / "kernels"
    major, minor = torch.cuda.get_device_capability()
    candidate = base / f"{name}-sm{major}{minor}"
    return candidate if candidate.is_dir() else base / name


def kernel():
    global _KERNEL
    if _KERNEL is None:
        from kernels import get_local_kernel
        _KERNEL = get_local_kernel(kernel_dir("fast"), "yue2_qkv_fused")
    return _KERNEL


@cache
def rms_kernel():
    from kernels import get_local_kernel
    return get_local_kernel(kernel_dir("rms"), "yue2_qkv_fused")


@cache
def fused_kernel():
    from kernels import get_local_kernel
    return get_local_kernel(kernel_dir("fused"), "yue2_qkv_fused")


def _fused_args(linear, device):
    constants = linear._activation_kernel_constant_tensors(device)
    ac, asc, wc, wsc = linear._int8_surrogate_constants(device)
    return constants, ac, asc, wc, wsc


def rms_project_fused(x, norm, linear):
    """Single kernel: RMS-folded RPBH quantization in the prologue, packed W4A4 DP4A GEMV body."""
    constants, ac, asc, wc, wsc = _fused_args(linear, x.device)
    out = fused_kernel().fused_gemv(x, norm.weight, constants['permutation'], constants['signs'],
        constants['boundaries'], norm.eps, linear.activation_eps, linear.rotation.normalization,
        linear.packed_weight_indices, linear.row_norms, ac, wc, asc, wsc)
    return out.reshape(*x.shape[:-1], linear.out_features)


def plain_project_fused(x, linear, virtual_threads):
    """Single kernel: RPBH quantization (no RMS) in the prologue + packed GEMV; rows 1..8 only."""
    constants, ac, asc, wc, wsc = _fused_args(linear, x.device)
    out = fused_kernel().fused_gemv(x, None, constants['permutation'], constants['signs'],
        constants['boundaries'], 0.0, linear.activation_eps, linear.rotation.normalization,
        linear.packed_weight_indices, linear.row_norms, ac, wc, asc, wsc, virtual_threads=virtual_threads)
    return out.reshape(*x.shape[:-1], linear.out_features)


def rms_project2(x, h, norm, linear):
    """v2 warp-local fused kernel: returns (projection, x_new) with x_new = x + h (or None)."""
    if int(getattr(linear, "_fused_route", lambda: 0)()) == 3:
        return rms_project_add(x, h, norm, linear)
    constants, ac, asc, wc, wsc = _fused_args(linear, x.device)
    out, x_new = fused_kernel().fused_gemv2(x, norm.weight, constants['permutation'], constants['signs'],
        constants['boundaries'], norm.eps, linear.activation_eps, linear.rotation.normalization,
        linear.packed_weight_indices, linear.row_norms, ac, wc, asc, wsc, h=h)
    return out.reshape(*x.shape[:-1], linear.out_features), x_new


def rms_project_add(x, h, norm, linear):
    """rmsquant with folded residual add (single CTA, bit-exact vs add + rmsquant) + library GEMV."""
    from orbitquant.kernels.native_packed_matmul import matmul_packed_w4a4_int8_with_native_kernel
    constants = linear._activation_kernel_constant_tensors(x.device)
    if h is None:
        packed, norms = rms_kernel().rmsquant(x, norm.weight, constants['permutation'], constants['signs'], constants['boundaries'], norm.eps, linear.activation_eps)
        x_new = None
    else:
        packed, norms, x_new = fused_kernel().rmsquant_add(x, h, norm.weight, constants['permutation'], constants['signs'], constants['boundaries'], norm.eps, linear.activation_eps)
    ac, asc, wc, wsc = linear._int8_surrogate_constants(x.device)
    out = matmul_packed_w4a4_int8_with_native_kernel(packed, linear.packed_weight_indices, norms, linear.row_norms, ac, wc,
        activation_scale=asc, weight_scale=wsc, in_features=linear.in_features, out_features=linear.out_features, output_dtype=x.dtype)
    return out.reshape(*x.shape[:-1], linear.out_features), x_new


def plain_project2(x, linear, virtual_threads):
    constants, ac, asc, wc, wsc = _fused_args(linear, x.device)
    out, _ = fused_kernel().fused_gemv2(x, None, constants['permutation'], constants['signs'],
        constants['boundaries'], 0.0, linear.activation_eps, linear.rotation.normalization,
        linear.packed_weight_indices, linear.row_norms, ac, wc, asc, wsc, virtual_threads=virtual_threads)
    return out.reshape(*x.shape[:-1], linear.out_features)


@cache
def _compiled_pointwise():
    import torch as _torch
    def rms(x, weight, eps):
        return x * _torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + eps).to(x.dtype) * weight
    def rotary(x, cos, sin):
        half = x.shape[-1] // 2
        x1, x2 = x[..., :half], x[..., half:]
        cos, sin = cos.to(x.dtype), sin.to(x.dtype)
        return _torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], dim=-1)
    def silu_mul(x):
        gate, up = x.chunk(2, dim=-1)
        return F.silu(gate) * up
    opts = dict(dynamic=True)
    return (_torch.compile(rms, **opts), _torch.compile(rotary, **opts), _torch.compile(silu_mul, **opts))


def install_nar_compile(model, enabled=True):
    """Compile the NAR RMSNorm / RoPE / SwiGLU pointwise chains (same math, fewer kernels)."""
    rms_c, rotary_c, silu_mul_c = _compiled_pointwise() if enabled else (None, None, None)
    for layer in model.model.layers:
        for norm in (layer.nar_input_layernorm, layer.nar_pre_mlp_layernorm):
            if not hasattr(norm, "_eager_forward"):
                norm._eager_forward = norm.forward
            norm.forward = (types.MethodType(lambda self, x: rms_c(x, self.weight, self.eps), norm) if enabled else norm._eager_forward)
        attn = layer.nar_self_attn
        if not hasattr(attn, "_eager_project_qkv"):
            attn._eager_project_qkv = attn.project_qkv
        if enabled:
            def project_qkv(self, x, cos, sin, _rot=rotary_c):
                b, t, _ = x.shape
                q, k, v = self.packed_qkv(x).split((self.num_heads * self.head_dim, self.num_kv_heads * self.head_dim, self.num_kv_heads * self.head_dim), dim=-1)
                q = self.q_norm(q.reshape(b, t, self.num_heads, self.head_dim))
                k = self.k_norm(k.reshape(b, t, self.num_kv_heads, self.head_dim))
                return _rot(q, cos.unsqueeze(2), sin.unsqueeze(2)), _rot(k, cos.unsqueeze(2), sin.unsqueeze(2)), v.reshape(b, t, self.num_kv_heads, self.head_dim)
            attn.project_qkv = types.MethodType(project_qkv, attn)
            for norm in (attn.q_norm, attn.k_norm):
                if not hasattr(norm, "_eager_forward"):
                    norm._eager_forward = norm.forward
                norm.forward = types.MethodType(lambda self, x: rms_c(x, self.weight, self.eps), norm)
        else:
            attn.project_qkv = attn._eager_project_qkv
            for norm in (attn.q_norm, attn.k_norm):
                if hasattr(norm, "_eager_forward"):
                    norm.forward = norm._eager_forward
        mlp = layer.nar_mlp
        if not hasattr(mlp, "_eager_forward"):
            mlp._eager_forward = mlp.forward
        mlp.forward = (types.MethodType(lambda self, x, _sm=silu_mul_c: self.down_proj(_sm(self.packed_gate_up(x))), mlp) if enabled else mlp._eager_forward)
    model._nar_compiled = bool(enabled)


def install_fused_gemv(model, enabled=True, virtual_threads=None):
    """Route decode-time o_proj/down_proj through the fused kernel; rms_project switches via the flag."""
    import torch as _torch
    if virtual_threads is None:
        virtual_threads = 512 if _torch.cuda.get_device_capability() == (8, 9) else 256
    model._fused_gemv = int(enabled)   # 0 off, 1 = v1 (CTA prologue), 2 = v2 (warp-local prologue, folded residual), 3 = rmsquant_add + library GEMV
    model._fused_vt = int(virtual_threads)
    if enabled:
        fused_kernel()
    for layer in model.model.layers:
        for linear in (layer.self_attn.packed_qkv, layer.mlp.packed_gate_up):
            linear._fused_route = (lambda m=model: int(getattr(m, "_fused_gemv", 0)))
    for layer in model.model.layers:
        for linear in (layer.self_attn.o_proj, layer.mlp.down_proj):
            if not hasattr(linear, "_library_forward"):
                linear._library_forward = linear.forward
            if enabled and int(enabled) != 3:
                def forward(self, x, _vt=virtual_threads, _v=int(enabled)):
                    rows = x.numel() // x.shape[-1]
                    if x.is_cuda and x.dtype == torch.bfloat16 and 1 <= rows <= 8 and self.in_features in (2048, 6144) and self.bias is None:
                        return plain_project2(x, self, _vt) if _v == 2 else plain_project_fused(x, self, _vt)
                    return self._library_forward(x)
                linear.forward = types.MethodType(forward, linear)
            else:
                linear.forward = linear._library_forward

# Single-warp quantizers (kernels/fused: wq_rms / wq_plain / wq_swiglu). Off by default.
_WQ = {"rms": False, "oproj": False, "down": False}


def install_wq(model, rms=False, oproj=False, down=False):
    """Enable the single-warp decode quantizers: RMS (both norms), o_proj input, SwiGLU+down input."""
    _WQ.update(rms=bool(rms), oproj=bool(oproj), down=bool(down))
    if oproj or down:
        for layer in model.model.layers:
            for lin in (layer.self_attn.o_proj, layer.mlp.down_proj):
                if getattr(lin, "bias", None) is not None:
                    raise ValueError("warp quantizer paths assume bias-free o_proj/down_proj")
    model._wq = dict(_WQ)
    return dict(_WQ)


def _wq_gemv(packed, norms, linear, out_dtype):
    from orbitquant.kernels.native_packed_matmul import matmul_packed_w4a4_int8_with_native_kernel
    ac, asc, wc, wsc = linear._int8_surrogate_constants(packed.device)
    return matmul_packed_w4a4_int8_with_native_kernel(packed, linear.packed_weight_indices,
        norms, linear.row_norms, ac, wc, activation_scale=asc, weight_scale=wsc,
        in_features=linear.in_features, out_features=linear.out_features, output_dtype=out_dtype)


_EPI_ADD = [False]


def install_epi_add(model, enabled):
    """Fold the residual add of o_proj/down_proj into the native GEMV epilogue (its per-column bias slot)
    for single-row decode: out = fmaf(sum, scale, residual) with one BF16 rounding instead of two."""
    _EPI_ADD[0] = bool(enabled)
    model._epi_add = bool(enabled)


def _epi_gemv(packed, norms, linear, out_dtype, residual):
    from orbitquant.kernels.native_packed_matmul import matmul_packed_w4a4_int8_with_native_kernel
    ac, asc, wc, wsc = linear._int8_surrogate_constants(packed.device)
    return matmul_packed_w4a4_int8_with_native_kernel(packed, linear.packed_weight_indices,
        norms, linear.row_norms, ac, wc, activation_scale=asc, weight_scale=wsc,
        in_features=linear.in_features, out_features=linear.out_features,
        bias=residual.reshape(-1), output_dtype=out_dtype)


def _quantize_like_linear(x2, linear):
    """The activation quantizer OrbitQuantLinear.forward itself picks for this linear at decode."""
    constants = linear._activation_kernel_constant_tensors(x2.device)
    if linear._native_w4_activation_available(x2):
        from orbitquant.kernels.native_packed_matmul import quantize_activations_packed_w4_with_native_kernel
        threads = 512 if torch.cuda.get_device_capability(x2.device) == (8, 9) else 256
        return quantize_activations_packed_w4_with_native_kernel(x2, constants['permutation'], constants['signs'],
            constants['boundaries'], eps=linear.activation_eps, inv_sqrt_block=linear.rotation.normalization, threads=threads)
    from orbitquant.kernels.triton_cuda import quantize_activations_packed_w4_with_triton
    return quantize_activations_packed_w4_with_triton(x2, rotation=linear.rotation, codebook=linear.activation_codebook,
                                                      eps=linear.activation_eps, constant_tensors=constants)


def oproj_project(h, linear, residual=None):
    """o_proj for decode rows (+ optional residual): warp quantizer (wq_oproj) and/or epilogue residual (epi_add)."""
    x2 = h.reshape(-1, h.shape[-1])
    fuse_add = residual is not None and _EPI_ADD[0] and x2.shape[0] == 1
    if not _WQ["oproj"] and not fuse_add:
        out = linear(h)
        return out if residual is None else residual + out
    constants = linear._activation_kernel_constant_tensors(h.device)
    if _WQ["oproj"]:
        vt = 512 if torch.cuda.get_device_capability(h.device) == (8, 9) else 256
        packed, norms = fused_kernel().wq_plain(x2, constants['permutation'], constants['signs'], constants['boundaries'],
                                                linear.activation_eps, linear.rotation.normalization, vt)
    else:
        packed, norms = _quantize_like_linear(x2, linear)
    if fuse_add:
        out = _epi_gemv(packed, norms, linear, h.dtype, residual)
        return out.reshape(*residual.shape)
    out = _wq_gemv(packed, norms, linear, h.dtype).reshape(*h.shape[:-1], linear.out_features)
    return out if residual is None else residual + out


def down_project(projected, linear, residual=None):
    """down_proj from gate|up (+ optional residual): SwiGLU warp quantizer (wq_down) and/or epilogue residual (epi_add)."""
    x2 = projected.reshape(-1, projected.shape[-1])
    fuse_add = residual is not None and _EPI_ADD[0] and x2.shape[0] == 1
    if not _WQ["down"] and not fuse_add:
        out = linear(kernel().swiglu(projected))
        return out if residual is None else residual + out
    constants = linear._activation_kernel_constant_tensors(projected.device)
    if _WQ["down"]:
        packed, norms = fused_kernel().wq_swiglu(x2, constants['permutation'], constants['signs'], constants['boundaries'],
                                                 linear.activation_eps, linear.rotation.normalization)
    else:
        packed, norms = _quantize_like_linear(kernel().swiglu(x2), linear)
    if fuse_add:
        out = _epi_gemv(packed, norms, linear, projected.dtype, residual)
        return out.reshape(*residual.shape)
    out = _wq_gemv(packed, norms, linear, projected.dtype).reshape(*projected.shape[:-1], linear.out_features)
    return out if residual is None else residual + out



def rms_project(x, norm, linear):
    """Keep x*gamma in FP32 and fold the positive RMS scale into token norms."""
    route = getattr(linear, "_fused_route", None)
    if route is not None and route():
        if route() == 2:
            return rms_project2(x, None, norm, linear)[0]
        return rms_project_fused(x, norm, linear)
    from orbitquant.kernels.native_packed_matmul import matmul_packed_w4a4_int8_with_native_kernel
    constants = linear._activation_kernel_constant_tensors(x.device)
    if _WQ["rms"]:
        packed, norms = fused_kernel().wq_rms(x.reshape(-1, x.shape[-1]), norm.weight, constants['permutation'],
            constants['signs'], constants['boundaries'], norm.eps, linear.activation_eps)
    else:
        packed, norms = rms_kernel().rmsquant(x, norm.weight, constants['permutation'],
            constants['signs'], constants['boundaries'], norm.eps, linear.activation_eps)
    ac, asc, wc, wsc = linear._int8_surrogate_constants(x.device)
    out = matmul_packed_w4a4_int8_with_native_kernel(packed, linear.packed_weight_indices,
        norms, linear.row_norms, ac, wc, activation_scale=asc, weight_scale=wsc,
        in_features=linear.in_features, out_features=linear.out_features, output_dtype=x.dtype)
    return out.reshape(*x.shape[:-1], linear.out_features)


def norm_forward(self, x):
    if x.is_cuda and x.dtype == torch.bfloat16 and x.numel() // x.shape[-1] <= 32:
        return kernel().rmsnorm(x, self.weight, self.eps)
    return self._reference_forward(self, x)


def project_qkv(self, x, cos, sin):
    if not x.is_cuda or x.dtype != torch.bfloat16 or x.shape[1] != 1:
        return original_project_qkv(self, x, cos, sin)
    b, t, _ = x.shape
    sizes = (self.num_heads * self.head_dim, self.num_kv_heads * self.head_dim,
             self.num_kv_heads * self.head_dim)
    q, k, v = self.packed_qkv(x).split(sizes, -1)
    q = self.q_norm(q.reshape(b, t, self.num_heads, self.head_dim))
    k = self.k_norm(k.reshape(b, t, self.num_kv_heads, self.head_dim))
    return (kernel().rope(q, cos.unsqueeze(2), sin.unsqueeze(2)),
            kernel().rope(k, cos.unsqueeze(2), sin.unsqueeze(2)),
            v.reshape(b, t, self.num_kv_heads, self.head_dim))


def project_qkv_cache(self, x, cos, sin, positions, keys, values):
    return kernel().qkv_decode(self.packed_qkv(x), self.q_norm.weight, self.k_norm.weight,
                               cos, sin, positions, keys, values, self.num_heads,
                               self.q_norm.eps)


def mlp_forward(self, x):
    if x.is_cuda and x.dtype == torch.bfloat16 and x.shape[1] == 1:
        return self.down_proj(kernel().swiglu(self.packed_gate_up(x)))
    return original_mlp(self, x)


def _int8_rows(weight, chunk=8192):
    """Per-row absmax INT8 rows, zero-padded to a multiple of 8 rows for cuBLASLt INT8 GEMM.

    ``weight`` may be one [N, K] tensor or a sequence of row blocks (concatenated logically). Rows are
    quantized ``chunk`` at a time straight into the output, so the FP32 transient stays at
    ``chunk * K * 4`` bytes instead of two full copies of the head (2.5 GiB for the ABC head)."""
    blocks = list(weight) if isinstance(weight, (list, tuple)) else [weight]
    total = sum(int(b.shape[0]) for b in blocks)
    k = int(blocks[0].shape[1]); device = blocks[0].device
    padded = total + (-total) % 8
    rows = torch.zeros((padded, k), dtype=torch.int8, device=device)
    scale = torch.ones((padded,), dtype=torch.float32, device=device)
    out = 0
    for block in blocks:
        for start in range(0, int(block.shape[0]), chunk):
            part = block[start:start + chunk].float()
            sc = part.abs().amax(-1).clamp_min(1e-12) / 127.0
            rows[out:out + part.shape[0]] = (part / sc[:, None]).round().clamp(-127, 127).to(torch.int8)
            scale[out:out + part.shape[0]] = sc
            out += part.shape[0]
            del part
    return rows, scale


def install_int8_heads(model, *, drop_bf16=True):
    """Replace the BF16 output head by INT8 row subsets used by ABC and semantic sampling."""
    from yue2.protocol import EOD, ABC_END, MUSIC_END, CODEC_OFFSET, CODEC_SIZE
    head = model.lm_head
    if getattr(head, "_int8_abc", None) is not None:
        return
    w = head.weight.data
    head._int8_abc = _int8_rows([w[:EOD], w[ABC_END:ABC_END + 1]])      # no 620 MB BF16 concat copy
    head._int8_sem = _int8_rows(w[MUSIC_END:CODEC_OFFSET + CODEC_SIZE])
    if drop_bf16:
        # Keep a 1-row placeholder parameter so module bookkeeping stays valid.
        head.weight = torch.nn.Parameter(w[:1].clone(), requires_grad=False)
        del w
        torch.cuda.empty_cache()
    head._int8_heads = True


def _int8_gemv(x, rows, scale):
    """x [B,K] BF16 -> logits [B,N] BF16 via per-token INT8 activations; W8A8 DP4A GEMV kernel or cuBLASLt INT8 GEMM."""
    b, k = x.shape
    if _W8A8_KERNEL[0] and b <= 8:
        fk = fused_kernel()
        xi, xs = fk.quantize_rows_int8(x)
        return fk.gemv_w8a8(xi, xs, rows, scale)
    xf = x.float()
    xs = xf.abs().amax(-1, keepdim=True).clamp_min(1e-12) / 127.0
    xi = torch.zeros((32, k), device=x.device, dtype=torch.int8)
    xi[:b] = (xf / xs).round().clamp(-127, 127).to(torch.int8)
    acc = torch._int_mm(xi, rows.t())[:b]
    return (acc.float() * (xs * scale[None, :])).to(x.dtype)


def use_w8a8_head_kernel(enabled=True):
    if enabled:
        fk = fused_kernel()
        if not hasattr(fk, "gemv_w8a8"):
            raise RuntimeError("fused kernel package lacks gemv_w8a8")
    _W8A8_KERNEL[0] = bool(enabled)


def install_w4_heads(model, quant_config, *, drop_bf16=True):
    """Quantize the ABC and semantic output-head row subsets with OrbitQuant W4A4 (same recipe as the
    transformer projections); the native packed GEMV serves them at decode. Irreversible in-process."""
    from yue2.protocol import EOD, ABC_END, MUSIC_END, CODEC_OFFSET, CODEC_SIZE
    from orbitquant import OrbitQuantConfig
    from orbitquant.layers import OrbitQuantLinear
    head = model.lm_head
    if getattr(head, "_w4_abc", None) is not None:
        return
    quant = OrbitQuantConfig.from_dict(dict(quant_config))
    w = head.weight.data
    abc_rows = torch.cat([w[:EOD], w[ABC_END:ABC_END + 1]], 0).contiguous()
    head._w4_abc = OrbitQuantLinear.from_weight(abc_rows, bias=None, in_features=w.shape[1], out_features=abc_rows.shape[0],
                                                source_weight_layout="out_in", config=quant, module_name="lm_head.abc").eval().to(w.device)
    del abc_rows
    sem_rows = w[MUSIC_END:CODEC_OFFSET + CODEC_SIZE].contiguous()
    head._w4_sem = OrbitQuantLinear.from_weight(sem_rows, bias=None, in_features=w.shape[1], out_features=sem_rows.shape[0],
                                                source_weight_layout="out_in", config=quant, module_name="lm_head.semantic").eval().to(w.device)
    del sem_rows
    if drop_bf16:
        head.weight = torch.nn.Parameter(w[:1].clone(), requires_grad=False)
        del w
        torch.cuda.empty_cache()
    head._w4_heads = True


# vae_core=256: half the compiled-decoder activation peak (-0.44 GiB) at equal speed; lazy_handoff: the NAR keeps the
# INT8 semantic cache and dequantizes per layer (-0.19 GiB NAR peak). Both bit-identical to the previous outputs.
_BASE = ("lookahead", "resident", "abc_crop", "vae_compile", "vae_mode=default", "vae_core=256")
MODES = {
    "off": (),
    "exact": _BASE,
    # nar_compile (torch.compile of the NAR pointwise chains) is not in any profile any more: it saved 0.2 s per warm song
    # on RTX 4090 but costs 1.1 s of cache load per process and changes the latents (fused fp32 math), so outputs would
    # depend on whether the process was warm. The flag still works for experiments.
    "fast": _BASE + ("kv_handoff", "lazy_handoff", "int8_heads", "w8a8"),
    "lowmem": _BASE + ("kv_handoff", "lazy_handoff", "int8_heads", "w8a8", "int8_embed"),
    # turbo: OrbitQuant W4A4 output heads (same recipe as the projections); larger logit change (KL ~0.016) for ~10% faster ABC
    "turbo": _BASE + ("kv_handoff", "lazy_handoff", "w4_heads", "int8_embed"),
}


def install_triton_quant(model, names=("self_attn.o_proj",), enabled=True):
    """Route decode-time activation quantization of the named OrbitQuant linears through the fused
    Triton kernel instead of the native pack kernel. Measured on RTX 4090: no gain and it changes
    logits (different rounding path), kept only as an experiment flag. ``enabled=False`` restores
    the class method so interleaved runs in one process do not inherit the override."""
    import types
    count = 0
    for layer in model.model.layers:
        for name in names:
            module = layer
            for part in name.split("."):
                module = getattr(module, part)
            if enabled:
                module._native_w4_activation_available = types.MethodType(lambda self, x: False, module)
            else:
                module.__dict__.pop("_native_w4_activation_available", None)
            count += 1
    model._triton_quant = tuple(names) if enabled else ()
    return count


def apply_mode(pipe, model, name, warm=False):
    """Apply a named runtime mode (see MODES) to a loaded pipeline/model.

    ``warm=False`` (one request per process, the CLI default) drops the torch.compile paths: their
    per-process cache load (about 2.4 s for the VAE decoder and 1.1 s for the NAR pointwise chains on
    RTX 4090) exceeds what they save on a single song (0.5 s and 0.15 s). ``warm=True`` keeps them for
    long-lived processes that serve many requests."""
    if name not in MODES:
        raise ValueError(f"unknown mode {name!r}; choose from {sorted(MODES)}")
    flags = set(MODES[name])
    if not warm:
        # one request per process: the compiled VAE decoder costs ~2 s of cache load for a 0.5 s gain
        flags -= {"vae_compile"}
        flags = {f for f in flags if not f.startswith("vae_mode=")}
    kv = dict(f.split("=", 1) for f in flags if "=" in f)
    model._lookahead = "lookahead" in flags
    model._unroll2 = False
    pipe.keep_resident = "resident" in flags
    pipe.vae_cudnn_benchmark = "vae_compile" in flags
    pipe.vae_compile = "vae_compile" in flags
    pipe.vae_compile_mode = kv.get("vae_mode", "reduce-overhead")
    pipe.vae_graph = False
    model.lm_head._abc_crop = "abc_crop" in flags
    model._kv_handoff_enabled = "kv_handoff" in flags
    install_nar_compile(model, enabled="nar_compile" in flags)
    install_fused_gemv(model, enabled=0)
    use_w8a8_head_kernel("w8a8" in flags)
    if "int8_heads" in flags:
        install_int8_heads(model)
    if "w4_heads" in flags:
        install_w4_heads(model, pipe.orbit_manifest["config"])
    if "int8_embed" in flags:
        install_int8_embed(model)
    install_triton_quant(model, enabled="tq_oproj" in flags)
    install_wq(model, rms="wq_rms" in flags, oproj="wq_oproj" in flags, down="wq_down" in flags)
    install_epi_add(model, "epi_add" in flags)
    if "vae_core" in kv:
        pipe.vae_core_frames = int(kv["vae_core"])
    model._nar_graph = "nar_graph" in flags
    model._lazy_handoff = "lazy_handoff" in flags
    model._mode = name
    return sorted(flags)


def install_int8_embed(model):
    """INT8 per-row token embeddings with a dequantizing gather (memory: 756 MB -> 378 MB)."""
    emb = model.model.embed_tokens
    if getattr(emb, "_int8_rows", None) is not None:
        return
    rows, scale = _int8_rows(emb.weight.data)
    emb._int8_rows, emb._int8_scale, emb._dtype = rows, scale, emb.weight.dtype
    emb.weight = torch.nn.Parameter(emb.weight.data[:1].clone(), requires_grad=False)
    torch.cuda.empty_cache()
    def forward(self, ids):
        return (self._int8_rows[ids].to(torch.float32) * self._int8_scale[ids][..., None]).to(self._dtype)
    emb.forward = types.MethodType(forward, emb)


def head_forward(self, x):
    if getattr(self, "_w4_heads", False):
        from yue2.protocol import EOD, ABC_END, MUSIC_END, CODEC_OFFSET, CODEC_SIZE
        lead = x.shape[:-1]
        x2 = x.reshape(-1, x.shape[-1])
        out = torch.zeros((x2.shape[0], self._vocab), device=x.device, dtype=x.dtype)
        if self._semantic_phase:
            out[:, MUSIC_END:CODEC_OFFSET + CODEC_SIZE] = self._w4_sem(x2)
        else:
            logits = self._w4_abc(x2)
            out[:, :EOD] = logits[:, :EOD]
            out[:, ABC_END] = logits[:, EOD]
        return out.reshape(*lead, self._vocab)
    if getattr(self, "_int8_heads", False):
        from yue2.protocol import EOD, ABC_END, MUSIC_END, CODEC_OFFSET, CODEC_SIZE
        lead = x.shape[:-1]
        x2 = x.reshape(-1, x.shape[-1])
        out = torch.zeros((x2.shape[0], self._vocab), device=x.device, dtype=x.dtype)
        if self._semantic_phase:
            out[:, MUSIC_END:CODEC_OFFSET + CODEC_SIZE] = _int8_gemv(x2, *self._int8_sem)[:, :CODEC_OFFSET + CODEC_SIZE - MUSIC_END]
        else:
            logits = _int8_gemv(x2, *self._int8_abc)
            out[:, :EOD] = logits[:, :EOD]
            out[:, ABC_END] = logits[:, EOD]
        return out.reshape(*lead, self._vocab)
    if not self._semantic_phase:
        if not getattr(self, "_abc_crop", False):
            return F.linear(x, self.weight, self.bias)
        # ABC sampling only reads ordinary text classes [:EOD] plus ABC_END, so
        # computing only those rows is exact (the rest is masked to -inf anyway).
        from yue2.protocol import EOD, ABC_END
        out = torch.zeros((*x.shape[:-1], self.weight.shape[0]), device=x.device, dtype=x.dtype)
        out[..., :EOD] = F.linear(x, self.weight[:EOD])
        out[..., ABC_END:ABC_END + 1] = F.linear(x, self.weight[ABC_END:ABC_END + 1])
        return out
    from yue2.protocol import MUSIC_END, CODEC_OFFSET, CODEC_SIZE
    end = CODEC_OFFSET + CODEC_SIZE
    small = F.linear(x, self.weight[MUSIC_END:end])
    # Keep the public vocabulary shape and finite masked values for BF16 CFG.
    out = torch.zeros((*x.shape[:-1], self.weight.shape[0]), device=x.device, dtype=x.dtype)
    out[..., MUSIC_END:end] = small
    return out


def install(model, profile, *, fuse_rms_quant=False):
    if fuse_rms_quant and profile not in {"fast", "memory"}:
        raise ValueError("RMS/activation fusion requires fast or memory profile")
    if profile == "reference" or (getattr(model, "_fast_profile", None) == profile and
                                  getattr(model, "_fuse_rms_quant", False) == fuse_rms_quant):
        return
    if not getattr(model, "_orbitquant_fused", False):
        raise ValueError("Fast profiles require grouped packed projections")
    kernel()
    if fuse_rms_quant:
        for layer in model.model.layers:
            for linear in (layer.self_attn.packed_qkv, layer.mlp.packed_gate_up):
                if not (linear.in_features == 2048 and linear.activation_bits == 4 and
                        linear.weight_bits == 4 and linear.rotation.block_size == 2048 and
                        linear.bias is None):
                    raise ValueError("RMS fusion requires bias-free 2048-wide W4A4 RPBH projections")
        rms_kernel()
    model._fuse_rms_quant = bool(fuse_rms_quant)
    norms = [model.model.norm]
    for layer in model.model.layers:
        attn = layer.self_attn
        norms.extend((layer.input_layernorm, layer.post_attention_layernorm,
                      attn.q_norm, attn.k_norm))
        attn.project_qkv = types.MethodType(project_qkv, attn)
        if profile in {"fast", "memory"}:
            attn.project_qkv_cache = types.MethodType(project_qkv_cache, attn)
        layer.mlp.forward = types.MethodType(mlp_forward, layer.mlp)
    for norm in norms:
        norm._reference_forward = type(norm).forward
        norm.forward = types.MethodType(norm_forward, norm)
    model.lm_head._semantic_phase = False
    model.lm_head._vocab = model.lm_head.weight.shape[0]
    model.lm_head.forward = types.MethodType(head_forward, model.lm_head)
    model._compact_sampling = profile in {"fast", "memory"}
    model._fast_profile = profile