recipes/docker-images/glm52-nf3-hybrid/overlay/hybrid_loader.py

"""
hybrid_loader.py — Path A: TP4-native NVFP4/NF3 hybrid loader for vLLM (modelopt path).
Deploy via .pth (`import hybrid_loader`). Fires in every worker.

GLM-5.2 hybrid checkpoint: routed experts are per-layer mixed — top-K (=64) NVFP4 (crisp) +
the rest NF3 (3-bit, group-32). Non-experts bf16 (excluded). Stock modelopt would allocate
uniform NVFP4 for all 256 experts (~420 GiB → OOM on 384). This loader allocates COMPACT
two-group storage (~326 GiB) and does a reference two-pass forward.

Interceptions (armed by an import hook on vllm...modelopt):
  1. ModelOptNvFp4Config.from_config -> stash `hybrid_bit_map` (stock strips unknown keys).
  2. FusedMoEMethodCls -> HybridNvFp4MoE.

NO mapping patch: stock RoutedExperts.make_expert_params_mapping is PREFIX-based, so
`...gate_proj.weight_packed` -> `...routed_experts.w13_weight_packed`, `...weight_scale` ->
`...w13_weight_scale`, etc. We register compact params under those exact names; the per-layer
weight_loader demuxes NVFP4 vs NF3 by the expert's group (this layer's remap) and TP-shards.

Facts baked in (verified against the eldritch image):
  * apply(layer, x, topk_weights, topk_ids, shared_experts, shared_experts_input): routing is
    upstream (sigmoid/noaux_tc/norm/2.5x already applied); shared experts run by the runner
    separately -> apply returns ROUTED-ONLY output.
  * moe_kernel stays None for hybrid layers -> is_monolithic False (via experts_cls) -> apply
    dispatched; maybe_make_prepare_finalize overridden to None so MK-init doesn't raise.
  * on-disk NVFP4 block scales are LINEAR (2D) -> dequantize_to_dtype(..., swizzle=False).
  * TP=4, moe_intermediate=2048 -> clean chunk, no padding. gate/up shard dim0, down shard dim1.
"""
import os, sys, importlib.abc, importlib.util

_NF3_VALS = [-1.0, -0.6047, -0.3563, -0.1275, 0.1275, 0.3563, 0.6047, 1.0]
_HYBRID_DEBUG = os.environ.get("HYBRID_DEBUG") not in (None, "", "0")  # eager-only: per-tier norms
# triton | b12x | b12x_nf3 | ref.  b12x_nf3 = BOTH tiers through the b12x W4A16
# CuteDSL kernel: kept-64 NVFP4 as weight_layout="packed", NF3-192 as the new
# weight_layout="nf3_2p1" (e4m3_k32 scales).  HYBRID_NF3 is ignored in that mode.
_HYBRID_KEPT = os.environ.get("HYBRID_KEPT", "triton")
_HYBRID_TIER = os.environ.get("HYBRID_TIER", "both")  # both|a|b (b12x_nf3 isolation)
_HYBRID_ACT_CAPTURE = os.environ.get("HYBRID_ACT_CAPTURE", "")  # dir: save MoE-layer inputs for GPTQ Hessians
_act_store = {"n": 0, "buf": {}, "flushed": {}}
_HYBRID_NF3 = os.environ.get("HYBRID_NF3", "fast")    # fast | ref  (NF3 3-bit tier)
_HYBRID_PROFILE = os.environ.get("HYBRID_PROFILE", "0") == "1"  # per-tier CUDA-event timing at prefill M>64
# TC-decode fast path for m<=8 (2026-07-05): compile a SECOND launch per tier at
# size_m=max_m (small-size_m compiles are the broken family) with block-8 direct
# top-k + fused sum. Measured standalone: 0.27ms vs prod 0.36-0.38ms/launch at
# decode m, numerics rel<0.012. Gated OFF by default; flag-off = today's bytes.
_HYBRID_TC_DECODE = os.environ.get("HYBRID_TC_DECODE", "0") == "1"
# v2-lite fusion: tiers A+B share ONE fc2_out staging (disjoint (token,k)
# route slots via the -1-dropping expert maps); tier B compiles with
# zero_fc2_output=False so A's staged routes survive, and a SINGLE topk_sum
# over the union replaces (sum_a, sum_b, add). Packed path only — tc-decode
# steps keep the per-tier fused-sum + add.
_HYBRID_FUSE_V2 = os.environ.get("HYBRID_FUSE_V2", "0") == "1"
_prof_store = {"n": 0, "a_ms": 0.0, "b_ms": 0.0, "rows": 0, "pend": []}
_dbg_first_build = [True]   # keep kept-tier originals for ONLY the first b12x layer (VRAM)
_HBM = None  # cached hybrid_bit_map (per-process); read once from the checkpoint config.json

# ---- b12x_nf3 backend constants / shared runtime ----
# Pinned CTA tiles (fc1_tile_k, fc1_tile_n, fc2_tile_k, fc2_tile_n): the NF3
# flat-span weight layout is packed for a SPECIFIC tile_n, but the kernel's
# auto tile selection is m-dependent (fc1_tile_n flips 128<->256 across m).
# (64,256,64,256) validated for our shapes (fc1 N=1024 K=6144 / fc2 N=6144
# K=512) at BOTH moe_block_size 8 (decode) and 64 (prefill), both scale
# formats: smem fits (45-76KB <= 100.9KB) and the SM121 reg-count table has
# entries for (256,{1,4},16,4). It is also exactly what auto-selection picks
# for the max-m prefill, so prefill throughput is the natural one.
_B12X_NF3_TILES = (64, 256, 64, 256)
_B12X_NF3_DECODE_M = 8          # <=8 -> preplanned TC-decode launch (fused topk sum)
_B12X_NF3_MAX_TOKENS = int(os.environ.get("HYBRID_B12X_MAX_TOKENS", "8192"))
# v1.5 stream-fusion: run tier-A (NVFP4) and tier-B (NF3) launches on CONCURRENT
# streams instead of sequentially. Tier B gets its OWN scratch set (buffers_b) so
# nothing races; fork/join via pre-created events (graph-capture legal). Failure
# mode if the kernel host ignores torch stream ctx = launches serialize on one
# stream = no overlap, still correct. Default OFF.
_HYBRID_FUSE_STREAMS = os.environ.get("HYBRID_FUSE_STREAMS", "0") == "1"
_b12x_nf3_rt = {                # module-level, shared across ALL layers (one scratch set)
    "max_m": None,              # fixed at first apply: max(env, first-call m)
    "topk": None,
    "launches": {},             # (E, layout, scale_fmt, topk, max_m) -> (dec, pre)
    "buffers": None,            # W4A16PackedBuffers planned at max_m/route_E=256
    "buffers_b": None,          # v1.5: tier-B scratch set (only when FUSE_STREAMS)
    "stream_b": None,           # v1.5: side stream + fork/join events
    "ev_fork": None,
    "ev_join": None,
    "out_a": None,              # [max_m, H] bf16 per-tier outputs (fully overwritten
    "out_b": None,              #  by every run_w4a16_moe call that uses them)
}


def _load_hbm(quant_config=None):
    """Reliable in-worker hybrid_bit_map: config-object attr, else read config.json off disk.
    (from_config stash on the config object does NOT survive pickling to TP workers.)"""
    global _HBM
    if _HBM is not None:
        return _HBM
    hbm = getattr(quant_config, "hybrid_bit_map", None) if quant_config is not None else None
    if hbm is None:
        try:
            import json, os
            from vllm.config import get_current_vllm_config
            mp = get_current_vllm_config().model_config.model
            cfgp = os.path.join(mp, "config.json")
            qc = json.load(open(cfgp)).get("quantization_config", {})
            hbm = qc.get("hybrid_bit_map")
            print(f"[hybrid_loader] hbm read from {cfgp}: {len(hbm) if hbm else 0} layers", flush=True)
        except Exception as e:
            print("[hybrid_loader] hbm config.json read failed:", e, flush=True)
    _HBM = hbm
    return hbm


def _unpack_nf3(packed, scale, out_cols, blk=32):
    """packed:[R, out_cols//8*3] uint8, scale:[R, out_cols//blk] fp8 -> [R, out_cols] bf16."""
    import torch
    nf = torch.tensor(_NF3_VALS, device=packed.device, dtype=torch.float32)
    R = packed.shape[0]
    p = packed.reshape(R, out_cols // 8, 3).to(torch.int32)
    w24 = p[..., 0] | (p[..., 1] << 8) | (p[..., 2] << 16)
    codes = torch.stack([(w24 >> (3 * i)) & 0x7 for i in range(8)], -1).reshape(R, out_cols)
    return (nf[codes.long()] * scale.float().repeat_interleave(blk, 1)).to(torch.bfloat16)


def _fp8_ne_transform(weights):
    """Dequant fp8 non-expert weights -> bf16 on the fly. Speaks BOTH scale
    dialects: `.weight_scale_fp8` (rev-2, bf16 per-channel) and `.weight_scale`
    U8 (rev-3 MXFP8, e8m0 per-32 groups, [out, in//32]). Expert tensors never
    match (uint8-packed weights / F8 scales) and pass through untouched."""
    import torch
    scales, pend = {}, {}

    def _deq(w, s):
        kind, st = s
        if kind == "chan":
            return (w.to(torch.float32) * st.to(torch.float32).unsqueeze(1)).to(torch.bfloat16)
        sc = torch.pow(2.0, st.to(torch.float32) - 127.0)
        return (w.to(torch.float32) * sc.repeat_interleave(32, 1)).to(torch.bfloat16)

    def _stage(t):
        # fastsafetensors yields CUDA tensors; holding tier pairs in VRAM
        # mid-shard OOMs the load. CPU-stage ONLY the F8 tier (16GB of 327) —
        # experts keep streaming GPU-direct.
        return t.cpu() if getattr(t, "is_cuda", False) else t

    for name, t in weights:
        if name.endswith(".weight_scale_fp8"):
            wn = name[:-len(".weight_scale_fp8")] + ".weight"
            w = pend.pop(wn, None)
            if w is not None:
                yield wn, _deq(w, ("chan", _stage(t)))
            else:
                scales[wn] = ("chan", _stage(t))
        elif (name.endswith(".weight_scale")
              and getattr(t, "dtype", None) == torch.uint8 and t.dim() == 2):
            wn = name[:-len(".weight_scale")] + ".weight"
            w = pend.pop(wn, None)
            if w is not None:
                yield wn, _deq(w, ("mx", _stage(t)))
            else:
                scales[wn] = ("mx", _stage(t))
        elif name.endswith(".weight") and getattr(t, "dtype", None) == torch.float8_e4m3fn:
            s = scales.pop(name, None)
            if s is not None:
                yield name, _deq(_stage(t), s)
            else:
                pend[name] = _stage(t)
        else:
            yield name, t
    for wn, t in pend.items():   # unmatched fp8 weight (shouldn't happen) -> best-effort
        yield wn, t.to(torch.bfloat16)


def _deq_nv(w, s, g, gs, deq):
    """NVFP4 dequant (linear on-disk scales). g is [2] for fused w13 or scalar for w2."""
    import torch
    if g.numel() == 2:  # fused w13: gate rows use g[0], up rows use g[1]
        half = w.shape[0] // 2
        a = deq(w[:half], s[:half], g[0], torch.bfloat16, gs, swizzle=False)
        b = deq(w[half:], s[half:], g[1], torch.bfloat16, gs, swizzle=False)
        return torch.cat([a, b], 0)
    return deq(w, s, g.reshape(()), torch.bfloat16, gs, swizzle=False)


def _patch(mod):
    import re, torch
    import torch.nn.functional as F
    from vllm.model_executor.utils import set_weight_attrs
    from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import dequantize_to_dtype
    from vllm.distributed import (get_tensor_model_parallel_rank,
                                  get_tensor_model_parallel_world_size)
    Base = mod.ModelOptNvFp4FusedMoE
    Cfg = mod.ModelOptNvFp4Config

    # ---- 1. keep hybrid_bit_map alive through from_config ----
    _orig_fc = getattr(Cfg.from_config, "__func__", Cfg.from_config)
    def _from_config(cls, config):
        obj = _orig_fc(cls, config)
        hbm = None
        kf = None
        if isinstance(config, dict):
            hbm = config.get("hybrid_bit_map")
            kf = config.get("kept_format")
            if isinstance(config.get("quantization"), dict):
                hbm = hbm or config["quantization"].get("hybrid_bit_map")
                kf = kf or config["quantization"].get("kept_format")
        try:
            obj.hybrid_bit_map = hbm
            obj.kept_format = kf
            print(f"[hybrid_loader] hybrid_bit_map stashed: "
                  f"{len(hbm) if hbm else 0} layers"
                  f"{' kept_format=' + kf if kf else ''}", flush=True)
        except Exception as e:
            print("[hybrid_loader] stash failed:", e, flush=True)
        return obj
    Cfg.from_config = classmethod(_from_config)

    class HybridNvFp4MoE(Base):
        def maybe_make_prepare_finalize(self, *a, **k):
            return None  # we own the forward via apply(); no external prepare/finalize

        def _bits(self, layer):
            hbm = _load_hbm(self.quant_config)
            pfx = getattr(layer, "prefix", None)
            if pfx is None:
                pfx = getattr(layer, "layer_name", "") or ""
            m = re.search(r"layers\.(\d+)\b", str(pfx))
            b = hbm.get(str(int(m.group(1)))) if (hbm and m) else None
            if b is None:
                print(f"[hybrid][MISS] pfx={pfx!r} hbm_len={len(hbm) if hbm else 0} "
                      f"idx={m.group(1) if m else None}", flush=True)
            return b

        def create_weights(self, layer, num_experts, hidden_size,
                           intermediate_size_per_partition, params_dtype, **extra):
            bits = self._bits(layer)
            # v3: kept tier is MXFP4 (e2m1 + ue8m0/32) ONLY for hybrid-mapped
            # layers; the MTP/uniform layer (bits None) stays NVFP4.
            kept_mx = (bits is not None and
                       getattr(self.quant_config, "kept_format", None)
                       == "mxfp4_e8m0k32")
            if bits is None:
                # Non-hybrid MoE (the MTP/nextn layer) — its experts are uniform NVFP4.
                # Route it through OUR path as all-kept (bits=4) so it uses our Triton NVFP4
                # kernel + our weight_loader. (Falling to super() uses the stock MoE loader,
                # whose weight-scale quant_method check at routed_experts.py:906 rejects it
                # under our method -> the MTP boot crash.)
                bits = [4] * num_experts
            H, I = hidden_size, intermediate_size_per_partition
            gs = self.quant_config.group_size
            tp_rank = get_tensor_model_parallel_rank()
            tp_size = get_tensor_model_parallel_world_size()
            kept = [e for e, b in enumerate(bits) if b == 4]
            dem = [e for e, b in enumerate(bits) if b == 3]
            Kn, Km = len(kept), len(dem)
            layer.hyb = {"remap": {**{e: (0, i) for i, e in enumerate(kept)},
                                   **{e: (1, i) for i, e in enumerate(dem)}},
                         "H": H, "I": I, "gs": gs, "E": num_experts,
                         "kept_mx": kept_mx}
            _pfx = getattr(layer, "prefix", None) or getattr(layer, "layer_name", "?")
            layer.hyb["lname"] = str(_pfx)
            print(f"[hybrid] {_pfx}: {Kn}{'MXFP4' if kept_mx else 'NVFP4'} + "
                  f"{Km}NF3 (tp{tp_rank}/{tp_size} I={I})", flush=True)

            def wl(param, loaded, name_mapped=None, *, shard_id=None,
                   expert_id=None, return_success=False, **kw):
                nm = name_mapped or ""
                if "input_scale" in nm:            # W4A16 reference -> unused
                    return True
                grp, li = layer.hyb["remap"][int(expert_id)]
                fam = "w13" if "w13_" in nm else "w2"
                sh = shard_id
                if "weight_scale_2" in nm:         # NVFP4 per-tensor global (kept only)
                    tgt = getattr(layer, f"{fam}_weight_scale_2")
                    if fam == "w13":
                        tgt.data[li, 0 if sh == "w1" else 1] = loaded.reshape(()).to(tgt.dtype)
                    else:
                        tgt.data[li] = loaded.reshape(()).to(tgt.dtype)
                    return True
                # TP shard the block-quantized 2D tensor (gate/up -> dim0, down -> dim1)
                if tp_size > 1 and loaded.ndim >= 2:
                    if sh in ("w1", "w3"):
                        loaded = loaded.chunk(tp_size, 0)[tp_rank]
                    elif sh == "w2":
                        loaded = loaded.chunk(tp_size, 1)[tp_rank]
                if "weight_scale" in nm:           # block scale -> real storage, demux by group
                    tgt = getattr(layer, f"{fam}_nv_s" if grp == 0 else f"{fam}_nf_s")
                elif "weight_packed" in nm:        # NF3 packed weight
                    tgt = getattr(layer, f"{fam}_weight_packed")
                else:                              # plain NVFP4 weight
                    tgt = getattr(layer, f"{fam}_weight")
                d = tgt.data[li]
                if fam == "w13" and sh in ("w1", "w3"):   # gate->top half, up->bottom half
                    half = d.shape[0] // 2
                    d = d[:half] if sh == "w1" else d[half:]
                d.copy_(loaded.reshape(d.shape).to(d.dtype))
                return True

            def P(name, shape, dt=torch.uint8):
                p = torch.nn.Parameter(torch.zeros(shape, dtype=dt,
                                       device=torch.cuda.current_device()), requires_grad=False)
                set_weight_attrs(p, {"weight_loader": wl})
                layer.register_parameter(name, p)

            mk = lambda n: max(n, 1)
            # --- names the stock (prefix-based) mapping produces (routing needs no patch) ---
            P("w13_weight",         (mk(Kn), 2 * I, H // 2))         # NVFP4 weight (kept)
            P("w13_weight_packed",  (mk(Km), 2 * I, H // 8 * 3))     # NF3 packed weight (demoted)
            P("w13_weight_scale",   (1,))                            # dispatcher (routes to *_s)
            P("w13_weight_scale_2", (mk(Kn), 2), torch.float32)      # NVFP4 global (kept)
            P("w13_input_scale",    (1,), torch.float32)             # dispatcher (ignored)
            P("w2_weight",          (mk(Kn), H, I // 2))
            P("w2_weight_packed",   (mk(Km), H, I // 8 * 3))
            P("w2_weight_scale",    (1,))
            P("w2_weight_scale_2",  (mk(Kn),), torch.float32)
            P("w2_input_scale",     (1,), torch.float32)
            # --- real block-scale storage (filled by dispatcher; not mapping-routed) ---
            # mx kept tier: ue8m0 scales at group 32 (uint8) instead of e4m3/gs.
            nv_gs = 32 if kept_mx else gs
            nv_dt = torch.uint8 if kept_mx else torch.float8_e4m3fn
            for nm, sh, dt_ in [("w13_nv_s", (mk(Kn), 2 * I, H // nv_gs), nv_dt),
                                ("w13_nf_s", (mk(Km), 2 * I, H // 32), torch.float8_e4m3fn),
                                ("w2_nv_s",  (mk(Kn), H, I // nv_gs), nv_dt),
                                ("w2_nf_s",  (mk(Km), H, I // 32), torch.float8_e4m3fn)]:
                layer.register_parameter(nm, torch.nn.Parameter(
                    torch.zeros(sh, dtype=dt_,
                                device=torch.cuda.current_device()), requires_grad=False))

        def _build_kept_b12x(self, layer):
            """Build a REAL b12x NVFP4 fused-MoE kernel over just the Kn kept experts.
            Reuses the production kernel (graph-safe, sm120) via a cloned num_experts=Kn
            FusedMoEConfig. Validated standalone: correct + skips out-of-range ids + cudagraph-safe.
            apply() remaps topk so kept->[0,Kn), non-kept->Kn (the kernel drops the sentinel)."""
            import dataclasses
            import torch.nn as _nn
            from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import (
                select_nvfp4_moe_backend, convert_to_nvfp4_moe_kernel_format,
                make_nvfp4_moe_quant_config, make_nvfp4_moe_kernel)
            from vllm.model_executor.layers.fused_moe.config import FusedMoEParallelConfig
            from vllm.model_executor.layers.quantization.utils.quant_utils import kNvfp4Static
            from vllm.model_executor.layers.fused_moe.activation import MoEActivation
            E = layer.hyb["E"]
            dev = layer.w13_weight.device
            remap = layer.hyb["remap"]
            Kn = sum(1 for (grp, li) in remap.values() if grp == 0)
            # Build the kept kernel as a NON-PARALLEL (tp=1) MoE over the PER-RANK intermediate:
            # the weights are already TP-sharded by the weight_loader (ispp per rank), so the
            # kernel must see intermediate_size == ispp and tp=1 (else it mis-indexes full vs
            # sharded -> per-element garbage). The layer's post-apply all-reduce handles TP.
            ispp = self.moe.intermediate_size_per_partition
            kept_moe = dataclasses.replace(
                self.moe, num_experts=Kn, num_local_experts=Kn, num_logical_experts=Kn,
                intermediate_size=ispp,
                moe_parallel_config=FusedMoEParallelConfig.make_no_parallel())
            backend, experts_cls = select_nvfp4_moe_backend(
                config=kept_moe, weight_key=kNvfp4Static, activation_key=None)
            KL = _nn.Module()
            KL.activation = getattr(layer, "activation", MoEActivation.SILU)
            KL.moe_config = kept_moe
            KL.local_num_experts = Kn
            a13 = torch.ones(Kn, 2, device=dev, dtype=torch.float32)
            a2 = torch.ones(Kn, device=dev, dtype=torch.float32)
            w13ks2 = layer.w13_weight_scale_2[:, 0].contiguous()
            conv = convert_to_nvfp4_moe_kernel_format(
                nvfp4_backend=backend, layer=KL,
                w13=layer.w13_weight, w13_scale=layer.w13_nv_s, w13_scale_2=w13ks2, a13_scale=a13,
                w2=layer.w2_weight, w2_scale=layer.w2_nv_s, w2_scale_2=layer.w2_weight_scale_2,
                a2_scale=a2, is_act_and_mul=True, use_a16=True)
            (cw13, cw13s, cw13s2, ca13, cw2, cw2s, cw2s2, ca2) = conv
            for k, v in [("w13_weight", cw13), ("w13_weight_scale", cw13s),
                         ("w13_weight_scale_2", cw13s2), ("w13_input_scale", ca13),
                         ("w2_weight", cw2), ("w2_weight_scale", cw2s),
                         ("w2_weight_scale_2", cw2s2), ("w2_input_scale", ca2)]:
                setattr(KL, k, v)
            qconf = make_nvfp4_moe_quant_config(backend, cw13s, cw2s, cw13s2, cw2s2,
                                                ca13, ca2, use_a16=True)
            kk = make_nvfp4_moe_kernel(qconf, kept_moe, experts_cls, routing_tables=None)
            kk.fused_experts.process_weights_after_loading(KL)
            # Mark this method as owning its modular kernel (supports_internal_mk -> True) so
            # vLLM's post-load maybe_init_modular_kernel() returns early instead of rebuilding
            # the parent kernel from the (freed) standard weight attrs.
            self.moe_kernel = kk
            # remap: global -> kept-local, else sentinel Kn (b12x skips ids >= num_experts)
            nvfp4_remap = torch.full((E,), Kn, dtype=torch.int32, device=dev)
            for g, (grp, li) in remap.items():
                if grp == 0:
                    nvfp4_remap[g] = li
            if _HYBRID_DEBUG and not _b12x_nf3_rt.get("_selftested"):
                _b12x_nf3_rt["_selftested"] = True
                try:
                    from vllm.distributed import get_tensor_model_parallel_rank as _gr
                    _rk = _gr()
                except Exception:
                    _rk = -1
                xt = torch.randn(4, layer.hyb["H"], device=dev, dtype=torch.bfloat16) * 0.03
                ids = (torch.arange(32, device=dev, dtype=torch.int32) % Kn).reshape(4, 8)
                tww = torch.full((4, 8), 0.125, device=dev, dtype=torch.float32)
                o = kk.apply(xt, KL.w13_weight, KL.w2_weight, tww, ids,
                             activation=KL.activation, global_num_experts=Kn,
                             expert_map=None, apply_router_weight_on_input=False,
                             shared_experts=None, shared_experts_input=None)
                torch.cuda.synchronize()
                print(f"[selftest] r{_rk} {getattr(layer, 'prefix', '?')}: "
                      f"out={o.float().norm():.4f} inf={bool(torch.isinf(o).any())} "
                      f"nan={bool(torch.isnan(o).any())} "
                      f"cw13={KL.w13_weight.view(torch.uint8).float().mean():.3f} "
                      f"cw13s2_max={KL.w13_weight_scale_2.float().abs().max():.3e} "
                      f"cw2s2_max={KL.w2_weight_scale_2.float().abs().max():.3e}", flush=True)
            layer.hyb["kept_kernel"] = kk
            layer.hyb["kept_KL"] = KL
            layer.hyb["Kn"] = Kn
            layer.hyb["nvfp4_remap"] = nvfp4_remap
            # CRITICAL: keep the converted tensors alive past this function. b12x's
            # process_weights compacts the source weight to (0,) but its prepared weights
            # VIEW the converted `cw13`/`cw2`; if those locals are freed on return the views
            # dangle -> per-layer garbage (works in a script where they stay in scope).
            layer.hyb["_keepalive"] = conv
            # free the BIG compact kept originals (KL holds the converted copies) -> flat VRAM.
            # keep the tiny *_weight_scale* params (get_fused_moe_quant_config may still read them).
            keep_dbg = _HYBRID_DEBUG and _dbg_first_build[0]   # keep originals for ONE layer only
            if keep_dbg:
                _dbg_first_build[0] = False
            for nm in ("w13_weight", "w2_weight", "w13_nv_s", "w2_nv_s"):
                if not keep_dbg and hasattr(layer, nm):
                    try:
                        delattr(layer, nm)
                    except Exception:
                        setattr(layer, nm, None)
            print(f"[hybrid] {getattr(layer, 'prefix', '?')}: kept b12x kernel built "
                  f"(Kn={Kn}, backend={backend}, kept_ispp={kept_moe.intermediate_size_per_partition}, "
                  f"kept_tp={kept_moe.moe_parallel_config.tp_size}, "
                  f"w13={tuple(KL.w13_weight.shape)})", flush=True)

        def _build_kept_mxfp4(self, layer):
            """v3 kept tier: MXFP4 (e2m1 + ue8m0/32) over the Kn kept experts
            via the stock mxfp4 oracle chain. The B12X backend's convert is a
            passthrough — repacking happens inside process_weights_after_loading.
            A16 by default; A8 opt-in via HYBRID_MX_A8=1 (KLD-gate before daily).
            Mirrors _build_kept_b12x's no-parallel-clone + keepalive shape."""
            import dataclasses
            import os as _os
            import torch.nn as _nn
            from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
                select_mxfp4_moe_backend, convert_weight_to_mxfp4_moe_kernel_format,
                make_mxfp4_moe_quant_config, make_mxfp4_moe_kernel)
            from vllm.model_executor.layers.fused_moe.config import FusedMoEParallelConfig
            from vllm.model_executor.layers.fused_moe.activation import MoEActivation
            E = layer.hyb["E"]
            dev = layer.w13_weight.device
            remap = layer.hyb["remap"]
            Kn = sum(1 for (grp, li) in remap.values() if grp == 0)
            ispp = self.moe.intermediate_size_per_partition
            kept_moe = dataclasses.replace(
                self.moe, num_experts=Kn, num_local_experts=Kn, num_logical_experts=Kn,
                intermediate_size=ispp,
                moe_parallel_config=FusedMoEParallelConfig.make_no_parallel())
            act_key = None
            if _os.environ.get("HYBRID_MX_A8", "0") == "1":
                from vllm.model_executor.layers.quantization.utils.quant_utils import (
                    kFp8StaticTensorSym)
                act_key = kFp8StaticTensorSym
            backend, experts_cls = select_mxfp4_moe_backend(
                kept_moe, activation_key=act_key)
            KL = _nn.Module()
            KL.activation = getattr(layer, "activation", MoEActivation.SILU)
            KL.moe_config = kept_moe
            KL.local_num_experts = Kn
            cw13, cw2, cw13s, cw2s, _b13, _b2 = (
                convert_weight_to_mxfp4_moe_kernel_format(
                    backend, KL, layer.w13_weight, layer.w2_weight,
                    layer.w13_nv_s, layer.w2_nv_s))
            for k, v in [("w13_weight", cw13), ("w13_weight_scale", cw13s),
                         ("w2_weight", cw2), ("w2_weight_scale", cw2s)]:
                setattr(KL, k, v)
            qconf = make_mxfp4_moe_quant_config(backend, cw13s, cw2s, layer=KL)
            kk = make_mxfp4_moe_kernel(qconf, kept_moe, experts_cls,
                                       routing_tables=None)
            kk.fused_experts.process_weights_after_loading(KL)
            self.moe_kernel = kk
            mx_remap = torch.full((E,), Kn, dtype=torch.int32, device=dev)
            for g, (grp, li) in remap.items():
                if grp == 0:
                    mx_remap[g] = li
            layer.hyb["kept_kernel"] = kk
            layer.hyb["kept_KL"] = KL
            layer.hyb["nvfp4_remap"] = mx_remap
            layer.hyb["_keepalive_mx"] = (cw13, cw2, cw13s, cw2s)
            for nm in ("w13_weight", "w2_weight", "w13_nv_s", "w2_nv_s"):
                if hasattr(layer, nm):
                    try:
                        delattr(layer, nm)
                    except Exception:
                        setattr(layer, nm, None)
            print(f"[hybrid] {getattr(layer, 'prefix', '?')}: kept MXFP4 kernel "
                  f"built (Kn={Kn}, backend={backend}, a8={act_key is not None})",
                  flush=True)

        def _build_b12x_nf3(self, layer):
            """HYBRID_KEPT=b12x_nf3: drive BOTH tiers through the b12x W4A16
            CuteDSL MoE kernel.

            Object A (kept NVFP4)  -> the PRODUCTION vLLM chain via
              _build_kept_b12x (select_nvfp4_moe_backend -> convert_to_nvfp4_
              moe_kernel_format -> make_nvfp4_moe_kernel, no-parallel clone,
              keepalive). Manual composition through the prepare building
              blocks is numerically WRONG for varying real-range scales
              (harness-proven 2026-07-01: only the production convert chain
              passes vs the dequant reference; probe test_b12x.py TESTS A-D).
            Object B (NF3 192)     -> weight_layout="nf3_2p1"/e4m3_k32 packed
              from the 837 checkpoint planes (chunked over experts) through
              our NF3 kernel format (GPU unit test + in-model norms PASS).
            Launches/buffers are built lazily at first apply() (topk + real max
            m known there; the first forward is vLLM's EAGER profile run, so
            nothing compiles inside CUDA-graph capture)."""
            import torch
            from b12x.moe.fused.w4a16.prepare import (
                PreparedNF3MoeWeights,
                W4A16PackedWeights,
                _make_workspace,
                _nf3_pack_code_experts,
                _nf3_pack_scale_experts,
                _permute_nvfp4_scales,
                _repack_weight,
            )
            import nf3_replan
            hyb = layer.hyb
            E, H, I = hyb["E"], hyb["H"], hyb["I"]
            remap = hyb["remap"]
            dev = layer.w13_weight.device
            Kn = sum(1 for (grp, _li) in remap.values() if grp == 0)
            Km = sum(1 for (grp, _li) in remap.values() if grp == 1)
            emap_a = torch.full((E,), -1, dtype=torch.int32, device=dev)
            emap_b = torch.full((E,), -1, dtype=torch.int32, device=dev)
            for g, (grp, li) in remap.items():
                (emap_a if grp == 0 else emap_b)[g] = li
            fc1_tn, fc2_tn = _B12X_NF3_TILES[1], _B12X_NF3_TILES[3]

            # ---- object B: NF3 -> "nf3_2p1" (chunked over experts) ----
            # (built FIRST: object A's production build frees the kept originals)
            prep_b = None
            if Km > 0:
                chunk = 16   # bound transient VRAM (codes int32 = ~400MB/16 w13 experts)
                w13_planes, w2_planes = [], []
                for e0 in range(0, Km, chunk):
                    codes = nf3_replan.unpack_837_codes(
                        layer.w13_weight_packed[e0:e0 + chunk], H)
                    w13_planes.append(_nf3_pack_code_experts(
                        codes, size_k=H, size_n=2 * I, tile_n=fc1_tn))
                    del codes
                for e0 in range(0, Km, chunk):
                    codes = nf3_replan.unpack_837_codes(
                        layer.w2_weight_packed[e0:e0 + chunk], I)
                    w2_planes.append(_nf3_pack_code_experts(
                        codes, size_k=I, size_n=H, tile_n=fc2_tn))
                    del codes
                w13_nf3 = torch.cat(w13_planes, 0).contiguous(); del w13_planes
                w2_nf3 = torch.cat(w2_planes, 0).contiguous(); del w2_planes
                w13_ns = _nf3_pack_scale_experts(
                    layer.w13_nf_s.float(), size_k=H, size_n=2 * I)
                w2_ns = _nf3_pack_scale_experts(
                    layer.w2_nf_s.float(), size_k=I, size_n=H)
                nf3_global = torch.full(
                    (Km,), 2.0 ** 116, dtype=torch.float32, device=dev)
                prep_b = PreparedNF3MoeWeights(
                    w13=w13_nf3, w13_scale=w13_ns, w13_global_scale=nf3_global,
                    w2=w2_nf3, w2_scale=w2_ns,
                    w2_global_scale=nf3_global.clone(),
                    workspace=_make_workspace(dev),
                    hidden_size=H, intermediate_size=I, num_experts=Km,
                    is_gated=True, params_dtype=torch.bfloat16,
                    fc1_tile_n=fc1_tn, fc2_tile_n=fc2_tn)

            hyb["prepB"] = prep_b
            hyb["prepA"] = None  # filled by the kept build below
            hyb["emap_a"], hyb["emap_b"] = emap_a, emap_b
            hyb["Kn"], hyb["Km"] = Kn, Km
            keep_dbg = _HYBRID_DEBUG and _dbg_first_build[0]
            if keep_dbg:
                _dbg_first_build[0] = False
            # ---- object A: kept NVFP4 -> manual "packed" composition.
            # REHABILITATED 2026-07-01: byte-identical to the stock prepare entry
            # (verified) and PASSES the FIXED (swizzle=False) dequant reference on
            # ranks 0/1/2 at rel 0.006. The earlier condemnation used a broken
            # reference. Produces weight_layout="packed" -> TC-decode launches
            # compile; NO modular kernel / workspace manager in the forward.
            if Kn > 0:
                g13 = layer.w13_weight_scale_2[:Kn, 0].contiguous()
                g2 = layer.w2_weight_scale_2[:Kn].contiguous()
                w13_packed = _repack_weight(
                    layer.w13_weight.contiguous(), size_k=H, size_n=2 * I)
                w2_packed = _repack_weight(
                    layer.w2_weight.contiguous(), size_k=I, size_n=H)
                w13_ps, w13_pg = _permute_nvfp4_scales(
                    layer.w13_nv_s, g13, size_k=H, size_n=2 * I,
                    a_dtype=torch.bfloat16)
                w2_ps, w2_pg = _permute_nvfp4_scales(
                    layer.w2_nv_s, g2, size_k=I, size_n=H,
                    a_dtype=torch.bfloat16)
                hyb["prepA"] = W4A16PackedWeights(
                    w13=w13_packed, w13_scale=w13_ps, w13_global_scale=w13_pg,
                    w2=w2_packed, w2_scale=w2_ps, w2_global_scale=w2_pg,
                    workspace=_make_workspace(dev),
                    hidden_size=H, intermediate_size=I, num_experts=Kn,
                    is_gated=True, params_dtype=torch.bfloat16,
                    source_format="modelopt_nvfp4", w13_layout="w13",
                    weight_layout="packed", scale_format="e4m3_k16")
                if not keep_dbg:
                    for nm in ("w13_weight", "w2_weight", "w13_nv_s", "w2_nv_s"):
                        p_ = getattr(layer, nm, None)
                        if p_ is not None and getattr(p_, "data", None) is not None:
                            p_.data = p_.data.new_empty((0,))
            if not keep_dbg:
                for nm in ("w13_weight_packed", "w2_weight_packed",
                           "w13_nf_s", "w2_nf_s"):
                    p = getattr(layer, nm, None)
                    if p is not None and getattr(p, "data", None) is not None:
                        p.data = p.data.new_empty((0,))
            print(f"[hybrid] {getattr(layer, 'prefix', '?')}: b12x_nf3 built "
                  f"(Kn={Kn} via production chain + Km={Km} nf3_2p1, "
                  f"tiles={_B12X_NF3_TILES})", flush=True)

        def _ensure_b12x_nf3_runtime(self, layer, m, topk):
            """First-apply init: pinned-tile preplanned launches (per object) +
            ONE module-level scratch/buffer set. First apply = vLLM's eager
            profile run at max_num_batched_tokens, so max_m sizes itself to the
            real serving ceiling and nothing compiles during graph capture."""
            import dataclasses
            import torch
            from b12x.moe.fused.w4a16.kernel import compile_w4a16_fused_moe
            from b12x.moe.fused.w4a16.host import (
                make_w4a16_packed_buffers, max_packed_route_slots)
            hyb = layer.hyb
            st = _b12x_nf3_rt
            E, H, I = hyb["E"], hyb["H"], hyb["I"]
            dev = hyb["emap_a"].device
            if st["max_m"] is None:
                st["max_m"] = max(_B12X_NF3_MAX_TOKENS, int(m))
                st["topk"] = int(topk)
            if int(topk) != st["topk"]:
                raise RuntimeError(
                    f"b12x_nf3: topk changed {st['topk']} -> {topk}")
            props = torch.cuda.get_device_properties(dev)
            sms = int(props.multi_processor_count)
            max_shared_mem = int(getattr(
                props, "shared_memory_per_block_optin", 101_376))

            def launches(prepared, zero_out=True):
                key = (prepared.num_experts, prepared.weight_layout,
                       prepared.scale_format, st["topk"], st["max_m"], H, I,
                       zero_out)
                got = st["launches"].get(key)
                if got is not None:
                    return got
                common = dict(
                    hidden_size=H, intermediate_size=I,
                    num_experts=prepared.num_experts, top_k=st["topk"],
                    activation="silu", apply_router_weight_on_input=False,
                    element_dtype="bf16", fast_math=True, sms=sms,
                    max_shared_mem=max_shared_mem,
                    weight_layout=prepared.weight_layout,
                    scale_format=prepared.scale_format,
                    force_tile_config=_B12X_NF3_TILES)
                # ONE launch per object for ALL m (1..max_m): spec-0 (size_m>=2)
                # block-64 packed-route + expert_map + zero_fc2_output=True.
                # VALIDATED on real ckpt data at m=1/8/33 (rel 0.02-0.05).
                # TC-decode/block-8 and the size_m=1 compile are BROKEN for these
                # shapes+tiles (rel 1.4 / 197) — never use them. Decode padding
                # cost ~0: blocks exist only for ACTIVE experts; pad slots waste
                # ALU, not HBM bytes, and decode is bandwidth-bound.
                cap_slots = max_packed_route_slots(st["max_m"] * st["topk"], 64, E)
                pre = compile_w4a16_fused_moe(
                    size_m=st["max_m"], zero_fc2_output=zero_out,
                    moe_block_size=64, max_m_blocks=(cap_slots + 63) // 64,
                    direct_topk_routes=False, tc_decode_fused_sum=False,
                    **common)
                assert (int(pre.fc1_tile_n), int(pre.fc2_tile_n)) == (
                    _B12X_NF3_TILES[1], _B12X_NF3_TILES[3]), "tile pin failed"
                dec = pre
                if _HYBRID_TC_DECODE:
                    # dec launch: SMALL size_m (8) direct-topk + fused-sum at FORCED
                    # PIN tiles. Needs the op-boundary tile-config passthrough fix
                    # (b12x_nf3_*_fix/kernel.py): the stock op re-resolves its kernel
                    # WITHOUT force_tile_config, silently swapping PIN->auto so the
                    # 256-pack was read with wrong geometry (garbage rel 1.3-1.5 at
                    # m=1,6,8). Fixed kernel honors the forced PIN -> reads the SAME
                    # production 256-pack (NO 2nd pack, NO pool cost). Opus-validated
                    # rel 0.006-0.011 @ m=1..8, single-compile reuse worst 0.0099.
                    # 0.277ms vs packed 0.365ms = ~14ms/step decode bounty.
                    # Run side: LUT-mapped local ids, NO expert_map (direct-topk).
                    try:
                        cand = compile_w4a16_fused_moe(
                            size_m=_B12X_NF3_DECODE_M, zero_fc2_output=False,
                            moe_block_size=8,
                            max_m_blocks=_B12X_NF3_DECODE_M * st["topk"],
                            direct_topk_routes=True, tc_decode_fused_sum=True,
                            **common)  # common carries force_tile_config=_B12X_NF3_TILES
                        assert (int(cand.fc1_tile_n), int(cand.fc2_tile_n)) == (
                            _B12X_NF3_TILES[1], _B12X_NF3_TILES[3]), "tc tile pin failed"
                        dec = cand
                        print(f"[hybrid] tc-decode launch armed (force PIN) "
                              f"E={prepared.num_experts} tiles="
                              f"{int(cand.fc1_tile_n)},{int(cand.fc2_tile_n)}",
                              flush=True)
                    except Exception as _te:
                        print(f"[hybrid] tc-decode compile failed, packed fallback: "
                              f"{_te}", flush=True)
                st["launches"][key] = (dec, pre)
                return st["launches"][key]

            if hyb.get("prepA") is not None:
                hyb["launchA"] = launches(hyb["prepA"])
            if hyb.get("prepB") is not None:
                # v2-lite: B must NOT zero the shared staging (A staged first).
                # A-only / B-only debug tiers are incompatible with that.
                if _HYBRID_FUSE_V2 and _HYBRID_TIER in ("a", "b"):
                    raise RuntimeError("HYBRID_FUSE_V2 is incompatible with "
                                       "HYBRID_TIER=a/b")
                hyb["launchB"] = launches(hyb["prepB"],
                                          zero_out=not _HYBRID_FUSE_V2)
            if st["buffers"] is None:
                prep_any = hyb.get("prepA") or hyb.get("prepB")
                if prep_any is None:
                    # uniform-NVFP4 layer (MTP) first: no NF3 object anywhere yet;
                    # the kept kernel manages its own workspace -> no buffers needed.
                    hyb["_rt_ready"] = True
                    return
                buf = make_w4a16_packed_buffers(
                    prep_any, m=st["max_m"], topk=st["topk"],
                    dtype=torch.bfloat16, device=dev, route_num_experts=E)
                # the preplanned prefill launch validates route capacity at
                # moe_block_size=64; the plan's own block choice can be smaller
                # for small max_m -> upsize the route buffers if needed.
                need_slots = max_packed_route_slots(st["max_m"] * st["topk"], 64, E)
                need_blocks = (need_slots + 63) // 64
                if (buf.packed_route_indices.numel() < need_slots
                        or buf.block_expert_ids.numel() < need_blocks):
                    buf = dataclasses.replace(
                        buf,
                        packed_route_indices=torch.empty(
                            (need_slots,), dtype=torch.int32, device=dev),
                        block_expert_ids=torch.empty(
                            (need_blocks,), dtype=torch.int32, device=dev))
                st["buffers"] = buf
                st["out_a"] = buf.output          # [max_m, H], fully overwritten per call
                st["out_b"] = torch.empty_like(buf.output)
                if _HYBRID_FUSE_STREAMS:
                    # v1.5: tier-B gets its OWN scratch (no cross-stream races) +
                    # a dedicated side stream + pre-created fork/join events
                    # (event creation during graph capture is illegal -> here).
                    buf_b = make_w4a16_packed_buffers(
                        prep_any, m=st["max_m"], topk=st["topk"],
                        dtype=torch.bfloat16, device=dev, route_num_experts=E)
                    if (buf_b.packed_route_indices.numel() < need_slots
                            or buf_b.block_expert_ids.numel() < need_blocks):
                        buf_b = dataclasses.replace(
                            buf_b,
                            packed_route_indices=torch.empty(
                                (need_slots,), dtype=torch.int32, device=dev),
                            block_expert_ids=torch.empty(
                                (need_blocks,), dtype=torch.int32, device=dev))
                    st["buffers_b"] = buf_b
                    st["out_b"] = buf_b.output
                    st["stream_b"] = torch.cuda.Stream(device=dev)
                    st["ev_fork"] = torch.cuda.Event()
                    st["ev_join"] = torch.cuda.Event()
                    print("[hybrid] v1.5 stream-fusion armed (tier-B scratch + "
                          "side stream)", flush=True)
                if _HYBRID_FUSE_V2:
                    # dedicated union staging: fc2 results from BOTH tiers
                    # live here across the A->B sequence. Must NOT alias the
                    # fc1 workspace (kernel default) — B's FC1 would clobber
                    # A's staged rows (learned the hard way: /Optics/ loops).
                    st["fc2_union"] = torch.empty(
                        (st["max_m"] * st["topk"] * H,),
                        dtype=torch.bfloat16, device=dev)
                    print("[hybrid] v2-lite fusion armed (dedicated union "
                          "staging + single topk_sum)", flush=True)
                print(f"[hybrid] b12x_nf3 runtime ready: max_m={st['max_m']} "
                      f"topk={st['topk']} sms={sms}", flush=True)
            hyb["_rt_ready"] = True

        def _apply_b12x_nf3(self, layer, x, topk_weights, topk_ids):
            import torch
            from b12x.moe.fused.w4a16.kernel import run_w4a16_moe
            hyb = layer.hyb
            st = _b12x_nf3_rt
            m = int(x.shape[0])
            topk = int(topk_ids.shape[1])
            if not hyb.get("_rt_ready"):
                self._ensure_b12x_nf3_runtime(layer, m, topk)
            if m > st["max_m"]:
                raise RuntimeError(
                    f"b12x_nf3: m={m} exceeds planned capacity {st['max_m']}; "
                    "set HYBRID_B12X_MAX_TOKENS >= max_num_batched_tokens")
            decode = m <= _B12X_NF3_DECODE_M
            if _HYBRID_ACT_CAPTURE and m > 8:
                try:
                    from vllm.distributed import get_tensor_model_parallel_rank
                    if get_tensor_model_parallel_rank() == 0:
                        pfx = str(hyb.get("lname") or getattr(layer, "prefix", None)
                                  or id(layer)).replace("/", "_").replace(".", "_")
                        b = _act_store["buf"].setdefault(pfx, {"x": [], "ids": []})
                        if _act_store["flushed"].get(pfx, 0) < 32768:  # cap tokens/layer
                            b["x"].append(x.detach().to(torch.float16).cpu())
                            b["ids"].append(topk_ids.detach().to(torch.int32).cpu())
                            if sum(t.shape[0] for t in b["x"]) >= 8192:
                                import os as _os
                                _os.makedirs(_HYBRID_ACT_CAPTURE, exist_ok=True)
                                fn = f"{_HYBRID_ACT_CAPTURE}/{pfx}.pt"
                                prev = torch.load(fn) if _os.path.exists(fn) else {"x": [], "ids": []}
                                prev["x"].append(torch.cat(b["x"]))
                                prev["ids"].append(torch.cat(b["ids"]))
                                torch.save(prev, fn)
                                _act_store["flushed"][pfx] = _act_store["flushed"].get(pfx, 0) + sum(t.shape[0] for t in b["x"])
                                b["x"], b["ids"] = [], []
                except Exception as _ce:
                    if not _act_store.get("_warned"):
                        _act_store["_warned"] = True
                        print("[hybrid][act-capture] failed:", repr(_ce), flush=True)
            tw = (topk_weights if topk_weights.dtype == torch.float32
                  else topk_weights.float())
            if not tw.is_contiguous():
                tw = tw.contiguous()
            buf = st["buffers"]

            def run(prepared, launch_pair, emap, out, rbuf=None,
                    skip_sum=False, force_pre=False, fc2u=None):
                rbuf = buf if rbuf is None else rbuf
                use_dec = (_HYBRID_TC_DECODE and decode and not force_pre
                           and launch_pair[0] is not launch_pair[1])
                launch = launch_pair[0] if use_dec else launch_pair[1]
                ids = (topk_ids if topk_ids.dtype == torch.int32
                       else topk_ids.to(torch.int32))
                if not ids.is_contiguous():
                    ids = ids.contiguous()
                if use_dec:
                    # direct-topk path: kernel reads FLAT local ids and skips
                    # negatives itself; expert_map must NOT be used. emap doubles
                    # as the global->local LUT (graph-safe gather).
                    ids = emap[ids.long()].to(torch.int32).contiguous()
                    em = None
                else:
                    em = emap   # kernel translates global->local + drops -1 (zero=True pairs)
                return run_w4a16_moe(
                    x, prepared, tw, ids,
                    activation="silu",
                    intermediate_cache13=rbuf.intermediate_cache13,
                    intermediate_cache2=rbuf.intermediate_cache2,
                    output=out,
                    fc1_c_tmp=rbuf.fc1_c_tmp, fc2_c_tmp=rbuf.fc2_c_tmp,
                    packed_route_indices=rbuf.packed_route_indices,
                    block_expert_ids=rbuf.block_expert_ids,
                    packed_route_count=rbuf.packed_route_count,
                    expert_offsets=rbuf.expert_offsets,
                    expert_map=em,
                    fused_launch=launch,
                    # v2-lite kwargs exist only in our tcfix kernel; pass
                    # them only when engaged so stock kernel dirs stay
                    # boot-compatible with this loader.
                    **({"skip_topk_sum": True} if skip_sum else {}),
                    **({"fc2_out_override": fc2u} if fc2u is not None else {}))

            if hyb["Km"] == 0:
                # uniform-NVFP4 layer (MTP/nextn): single tier through OUR launcher
                out = torch.empty((m, hyb["H"]), dtype=x.dtype, device=x.device)
                out = run(hyb["prepA"], hyb["launchA"], hyb["emap_a"], out)
                if _HYBRID_DEBUG and not hyb.get("_dbg"):
                    hyb["_dbg"] = True
                    print(f"[dbg-nf3] {getattr(layer, 'prefix', '?')} m={m} "
                          f"A-only={out.float().norm():.1f}", flush=True)
                return out
            out_a = out_b = None
            if _HYBRID_TIER == "a":
                out = run(hyb["prepA"], hyb["launchA"], hyb["emap_a"],
                          st["out_a"][:m]).clone()
            elif _HYBRID_TIER == "b":
                out = run(hyb["prepB"], hyb["launchB"], hyb["emap_b"],
                          st["out_b"][:m]).clone()
            elif _HYBRID_FUSE_V2 and not (
                    _HYBRID_TC_DECODE and decode
                    and hyb["launchA"][0] is not hyb["launchA"][1]
                    and hyb["launchB"][0] is not hyb["launchB"][1]):
                # v2-lite: A zeros the shared fc2 staging and stages its
                # routes; B (compiled zero_fc2_output=False) stages its
                # disjoint routes; ONE topk_sum over the union -> out_b is
                # the complete result. No out_a+out_b add. force_pre pins
                # both tiers to the packed path (a lone tc-armed tier would
                # bypass staging and double-count).
                run(hyb["prepA"], hyb["launchA"], hyb["emap_a"],
                    st["out_a"][:m], skip_sum=True, force_pre=True,
                    fc2u=st["fc2_union"])
                out = run(hyb["prepB"], hyb["launchB"], hyb["emap_b"],
                          st["out_b"][:m], force_pre=True,
                          fc2u=st["fc2_union"])
            elif _HYBRID_FUSE_STREAMS and st["stream_b"] is not None:
                # v1.5: fork tier-B onto the side stream so both grouped GEMMs
                # co-schedule on the SMs; join before the sum. Events pre-created
                # (graph-capture legal); tier-B uses its own scratch set.
                s = torch.cuda.current_stream()
                sb = st["stream_b"]
                st["ev_fork"].record(s)
                sb.wait_event(st["ev_fork"])
                out_a = run(hyb["prepA"], hyb["launchA"], hyb["emap_a"],
                            st["out_a"][:m])
                with torch.cuda.stream(sb):
                    out_b = run(hyb["prepB"], hyb["launchB"], hyb["emap_b"],
                                st["out_b"][:m], rbuf=st["buffers_b"])
                st["ev_join"].record(sb)
                s.wait_event(st["ev_join"])
                out = out_a + out_b
            else:
                prof = _HYBRID_PROFILE and m > 64
                if prof:
                    import time as _time
                    _w0 = _time.perf_counter()
                    ev = [torch.cuda.Event(enable_timing=True) for _ in range(3)]
                    ev[0].record()
                out_a = run(hyb["prepA"], hyb["launchA"], hyb["emap_a"],
                            st["out_a"][:m])
                if prof:
                    ev[1].record()
                out_b = run(hyb["prepB"], hyb["launchB"], hyb["emap_b"],
                            st["out_b"][:m])
                if prof:
                    ev[2].record()
                    _prof_store["wall_ms"] = _prof_store.get("wall_ms", 0.0) + \
                        (_time.perf_counter() - _w0) * 1000.0
                    _prof_store["pend"].append((m, ev))
                    _prof_store["n"] += 1
                    if _prof_store["n"] % 624 == 0:  # ~8 full 78-layer prefill passes
                        try:
                            torch.cuda.synchronize()
                            for pm, pe in _prof_store["pend"]:
                                _prof_store["a_ms"] += pe[0].elapsed_time(pe[1])
                                _prof_store["b_ms"] += pe[1].elapsed_time(pe[2])
                                _prof_store["rows"] += pm
                            _prof_store["pend"] = []
                            from vllm.distributed import get_tensor_model_parallel_rank
                            if get_tensor_model_parallel_rank() == 0:
                                n = _prof_store["n"]
                                a = _prof_store["a_ms"] / n
                                b = _prof_store["b_ms"] / n
                                mavg = _prof_store["rows"] / n
                                w = _prof_store.get("wall_ms", 0.0) / n
                                print(f"[hybrid-prof] n={n} avg_m={mavg:.0f} "
                                      f"A(nvfp4)={a:.3f}ms B(nf3)={b:.3f}ms "
                                      f"wall={w:.3f}ms /layer-chunk "
                                      f"-> MoE-gpu {78*(a+b):.0f}ms MoE-wall {78*w:.0f}ms "
                                      f"per {mavg:.0f}-tok chunk",
                                      flush=True)
                        except Exception as _pe:
                            print("[hybrid-prof] harvest failed:", repr(_pe), flush=True)
                out = out_a + out_b
            if _HYBRID_DEBUG and not hyb.get("_dbg"):
                hyb["_dbg"] = True
                na = "-" if out_a is None else f"{out_a.float().norm():.1f}"
                nb = "-" if out_b is None else f"{out_b.float().norm():.1f}"
                print(f"[dbg-nf3] {getattr(layer, 'prefix', '?')} m={m} "
                      f"decode={decode} tier={_HYBRID_TIER} x={x.float().norm():.1f} "
                      f"A={na} B={nb} out={out.float().norm():.1f} "
                      f"nan={bool(torch.isnan(out).any())}", flush=True)
                if getattr(layer, "w13_weight", None) is not None and \
                        layer.w13_weight.numel() > 0:
                    try:
                        ra = self._apply_ref(layer, x, topk_weights, topk_ids, 0)
                        rb = self._apply_ref(layer, x, topk_weights, topk_ids, 1)
                        ea = float((out_a.float() - ra.float()).abs().sum()
                                   / (ra.float().abs().sum() + 1e-9))
                        eb = float((out_b.float() - rb.float()).abs().sum()
                                   / (rb.float().abs().sum() + 1e-9))
                        print(f"[dbgcmp-nf3] A_rel={ea:.4f} B_rel={eb:.4f}",
                              flush=True)
                    except Exception as _e:
                        print("[dbgcmp-nf3] ref failed:", repr(_e), flush=True)
            return out

        def process_weights_after_loading(self, layer):
            if not hasattr(layer, "hyb"):
                return super().process_weights_after_loading(layer)
            if _HYBRID_KEPT == "b12x_nf3":
                try:
                    self._build_b12x_nf3(layer)
                    layer.hyb["_b12x_nf3"] = True
                    layer.hyb["_b12x"] = False
                    layer.hyb["_kept_triton"] = False
                    layer.hyb["_fast"] = False
                    return
                except Exception as e:
                    import traceback
                    print("[hybrid] b12x_nf3 build FAILED -> triton/ref "
                          "fallback:", e, flush=True)
                    traceback.print_exc()
                    layer.hyb["_b12x_nf3"] = False
                    # originals are freed only on success -> the stock flow
                    # below still has everything it needs.
            E = layer.hyb["E"]
            dev = layer.w13_weight.device
            # --- NF3 tier: expert map (demoted global id -> local NF3 index, else -1) ---
            emap = torch.full((E,), -1, dtype=torch.int32, device=dev)
            for g, (grp, li) in layer.hyb["remap"].items():
                if grp == 1:
                    emap[g] = li
            layer.hyb["emap_nf3"] = emap
            if _HYBRID_NF3 == "fast":
                try:
                    import nf3_kernel  # noqa: F401 (mounted alongside loader; warm import)
                    layer.hyb["_fast"] = True
                except Exception as e:
                    print("[hybrid] nf3_kernel import failed -> reference:", e, flush=True)
                    layer.hyb["_fast"] = False
            else:
                layer.hyb["_fast"] = False
            # --- NVFP4 kept tier ---
            layer.hyb["_b12x"] = False
            layer.hyb["_kept_triton"] = False
            if _HYBRID_KEPT in ("triton", "b12x_nf3"):
                # Custom Triton NVFP4 grouped-GEMM: reads the checkpoint weights DIRECTLY
                # (no convert/prepare/copy), graph-safe. emap: kept global -> local, else -1.
                # ("b12x_nf3" reaches here ONLY after a failed b12x_nf3 build ->
                #  fall back to the proven triton kept tier; originals intact.)
                kmap = torch.full((E,), -1, dtype=torch.int32, device=dev)
                for gid, (grp, li) in layer.hyb["remap"].items():
                    if grp == 0:
                        kmap[gid] = li
                layer.hyb["emap_nvfp4"] = kmap
                layer.hyb["g13"] = layer.w13_weight_scale_2[:, 0].contiguous()
                layer.hyb["g2"] = layer.w2_weight_scale_2.contiguous()
                try:
                    import nvfp4_kernel  # noqa: F401
                    layer.hyb["_kept_triton"] = True
                except Exception as e:
                    print("[hybrid] nvfp4_kernel import failed -> reference:", e, flush=True)
            elif _HYBRID_KEPT == "b12x":
                try:
                    self._build_kept_b12x(layer)
                    layer.hyb["_b12x"] = True
                except Exception as e:
                    import traceback
                    print("[hybrid] kept b12x build FAILED -> reference NVFP4:", e, flush=True)
                    traceback.print_exc()

        def _apply_ref(self, layer, x, topk_weights, topk_ids, only_grp):
            H, I, gs = layer.hyb["H"], layer.hyb["I"], layer.hyb["gs"]
            remap = layer.hyb["remap"]
            out = torch.zeros_like(x)
            for e in torch.unique(topk_ids).tolist():
                if e < 0 or e not in remap:
                    continue
                grp, li = remap[e]
                if grp != only_grp:
                    continue
                if grp == 0:
                    w13 = _deq_nv(layer.w13_weight[li], layer.w13_nv_s[li], layer.w13_weight_scale_2[li], gs, dequantize_to_dtype)
                    w2 = _deq_nv(layer.w2_weight[li], layer.w2_nv_s[li], layer.w2_weight_scale_2[li], gs, dequantize_to_dtype)
                else:
                    w13 = _unpack_nf3(layer.w13_weight_packed[li], layer.w13_nf_s[li], H)
                    w2 = _unpack_nf3(layer.w2_weight_packed[li], layer.w2_nf_s[li], I)
                sel = (topk_ids == e)
                tok = sel.any(-1)
                if not tok.any():
                    continue
                g, u = (x[tok] @ w13.t()).chunk(2, -1)
                y = (F.silu(g) * u) @ w2.t()
                wgt = (topk_weights * sel).sum(-1)[tok].unsqueeze(-1).to(y.dtype)
                out[tok] += wgt * y
            return out

        def apply(self, layer, x, topk_weights, topk_ids,
                  shared_experts=None, shared_experts_input=None):
            # routed experts only; shared experts + routing handled by the runner.
            if not hasattr(layer, "hyb"):
                return super().apply(layer, x, topk_weights, topk_ids,
                                     shared_experts, shared_experts_input)
            if layer.hyb.get("_b12x_nf3"):
                return self._apply_b12x_nf3(layer, x, topk_weights, topk_ids)
            H, I, E = layer.hyb["H"], layer.hyb["I"], layer.hyb["E"]
            # --- NVFP4 kept tier ---
            if layer.hyb.get("_kept_triton"):
                import nvfp4_kernel
                kept_out = nvfp4_kernel.nvfp4_moe_layer(
                    x, layer.w13_weight, layer.w13_nv_s, layer.hyb["g13"],
                    layer.w2_weight, layer.w2_nv_s, layer.hyb["g2"],
                    topk_ids, topk_weights, layer.hyb["emap_nvfp4"], E, H, I, 16)
            elif layer.hyb.get("_b12x"):
                kk = layer.hyb["kept_kernel"]
                KL = layer.hyb["kept_KL"]
                Kn = layer.hyb["Kn"]
                kept_ids = layer.hyb["nvfp4_remap"][topk_ids]   # kept->[0,Kn), else Kn (b12x skips)
                kept_out = kk.apply(
                    x, KL.w13_weight, KL.w2_weight, topk_weights, kept_ids,
                    activation=KL.activation, global_num_experts=Kn, expert_map=None,
                    apply_router_weight_on_input=False,
                    shared_experts=None, shared_experts_input=None)
            else:
                kept_out = self._apply_ref(layer, x, topk_weights, topk_ids, 0)
            # --- NF3 tier (192 experts) ---
            if layer.hyb.get("_fast"):
                import nf3_kernel
                nf3_out = nf3_kernel.nf3_moe_layer(
                    x, layer.w13_weight_packed, layer.w13_nf_s,
                    layer.w2_weight_packed, layer.w2_nf_s,
                    topk_ids, topk_weights, layer.hyb["emap_nf3"], E, 32, H, I)
            else:
                nf3_out = self._apply_ref(layer, x, topk_weights, topk_ids, 1)
            out = kept_out + nf3_out
            if _HYBRID_DEBUG and not layer.hyb.get("_dbg"):
                layer.hyb["_dbg"] = True
                print(f"[dbg] {getattr(layer, 'prefix', '?')} "
                      f"x={x.float().norm():.1f} kept={kept_out.float().norm():.1f} "
                      f"nf3={nf3_out.float().norm():.1f} out={out.float().norm():.1f} "
                      f"nan={bool(torch.isnan(out).any())}", flush=True)
                if layer.hyb.get("_b12x") and hasattr(layer, "w13_weight"):
                    try:
                        rk = self._apply_ref(layer, x, topk_weights, topk_ids, 0)
                        ratio = float(kept_out.float().norm() / (rk.float().norm() + 1e-9))
                        erel = float((kept_out.float() - rk.float()).abs().sum()
                                     / (rk.float().abs().sum() + 1e-9))
                        print(f"[dbgcmp] {getattr(layer, 'prefix', '?')} "
                              f"b12x={kept_out.float().norm():.3f} ref={rk.float().norm():.3f} "
                              f"ratio={ratio:.3f} elem_rel={erel:.3f}", flush=True)
                    except Exception as _e:
                        print("[dbgcmp] ref failed:", repr(_e), flush=True)
            return out

    mod.ModelOptNvFp4Config.FusedMoEMethodCls = HybridNvFp4MoE
    print("[hybrid_loader] HybridNvFp4MoE installed", flush=True)

    # ---- fp8 non-expert dequant-on-load (weight_scale_fp8 modules are in `ignore` -> bf16) ----
    try:
        from vllm.model_executor.model_loader.default_loader import DefaultModelLoader
        if not getattr(DefaultModelLoader, "_hyb_wrapped", False):
            _oga = DefaultModelLoader.get_all_weights
            def get_all_weights(self, model_config, model):
                return _fp8_ne_transform(_oga(self, model_config, model))
            DefaultModelLoader.get_all_weights = get_all_weights
            DefaultModelLoader._hyb_wrapped = True
            print("[hybrid_loader] fp8 non-expert dequant-on-load installed", flush=True)
    except Exception as e:
        print("[hybrid_loader] wrap get_all_weights failed:", e, flush=True)

    # ---- 3. HYBRID_MXFP8_NATIVE: serve the F8 ne-tier via online-mxfp8 ----
    # Instead of dequant-to-bf16 (2 B/param resident), excluded LinearBase
    # modules whose disk tensors are F8 get Mxfp8OnlineLinearMethod: loader
    # still feeds them our dequanted bf16, the method re-quantizes e8m0/32 at
    # load (bit-exact round trip vs rev-3 disk) and serves through the B12X
    # fp8 GEMM (VLLM_USE_B12X_FP8_GEMM=1 forces B12xMxfp8LinearKernel).
    # Net: ~-3.65 GiB/GPU weights. Activations on these linears become
    # dynamic-A8 (Festr serving convention) — KLD re-gate before publishing.
    if os.environ.get("HYBRID_MXFP8_NATIVE", "0") == "1":
        try:
            import json as _json
            from vllm.model_executor.layers.quantization.online.mxfp8 import (
                Mxfp8OnlineLinearMethod)
            from vllm.model_executor.layers.linear import (
                LinearBase, UnquantizedLinearMethod)
            from vllm.model_executor.layers.vocab_parallel_embedding import (
                ParallelLMHead)
            _tier_path = os.environ.get(
                "HYBRID_MXFP8_TIER_JSON",
                "/opt/venv/lib/python3.12/site-packages/mxfp8_tier.json")
            _allow = set(_json.load(open(_tier_path))["module_prefixes"])

            def _mx_norm(p):
                i = p.find("layers.")
                return p[i:] if i >= 0 else p

            _base = mod.ModelOptQuantConfigBase
            if not getattr(_base, "_hyb_mxfp8_overlay", False):
                _orig_gqm = _base.get_quant_method
                _n_over = [0]

                def get_quant_method(self, layer, prefix):
                    m = _orig_gqm(self, layer, prefix)
                    if (type(m) is UnquantizedLinearMethod
                            and isinstance(layer, LinearBase)
                            and not isinstance(layer, ParallelLMHead)
                            and _mx_norm(prefix) in _allow):
                        _n_over[0] += 1
                        if _n_over[0] <= 4 or _n_over[0] % 128 == 0:
                            print(f"[hybrid_loader] mxfp8 overlay #{_n_over[0]}: "
                                  f"{prefix}", flush=True)
                        return Mxfp8OnlineLinearMethod()
                    return m

                _base.get_quant_method = get_quant_method
                _base._hyb_mxfp8_overlay = True
                print(f"[hybrid_loader] mxfp8 ne-tier overlay armed "
                      f"({len(_allow)} module prefixes)", flush=True)
        except Exception as e:
            import traceback
            print("[hybrid_loader] mxfp8 overlay FAILED:", e, flush=True)
            traceback.print_exc()


class _Hook(importlib.abc.MetaPathFinder):
    T = "vllm.model_executor.layers.quantization.modelopt"
    def find_spec(self, name, path, target=None):
        if name == self.T and not getattr(self, "_d", False):
            self._d = True
            spec = importlib.util.find_spec(name)
            if spec:
                real_exec = spec.loader.exec_module
                def ex(module):
                    real_exec(module)
                    try:
                        _patch(module)
                    except Exception as e:
                        import traceback
                        print("[hybrid_loader] patch FAILED:", e, flush=True)
                        traceback.print_exc()
                spec.loader.exec_module = ex
            return spec
        return None


if not any(isinstance(f, _Hook) for f in sys.meta_path):
    sys.meta_path.insert(0, _Hook())
    print("[hybrid_loader] import hook armed", flush=True)