""" 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)