gridbook 0.1.0__py3-none-any.whl

gridbook/__init__.py

@@ -0,0 +1,18 @@
+"""vLLM out-of-tree plugin for the NVFP4-CB / FP8-CB product-codebook formats.
+
+Registers a ``QuantizationConfig`` (``"gridbook"``, with ``"prismaquant"`` kept
+as a legacy alias) plus the linear and fused-MoE methods that serve
+codebook-quantized weights. The CUDA kernels ship as sources under
+``gridbook/csrc`` and are JIT-compiled on first use; without nvcc the plugin
+falls back to a correct-but-slow Triton path.
+
+``register`` is lazy so ``import gridbook.codec`` / ``import gridbook.kernels``
+(the format and correctness tests) work without vLLM installed.
+"""
+
+__version__ = "0.1.0"
+
+
+def register() -> None:
+    from .plugin import register as _register
+    _register()

gridbook/codec.py

@@ -0,0 +1,145 @@
+"""Self-contained (no `prismaquant` import at serve time) CB codec helpers:
+
+* load-time preprocessing that turns the shipped layout (LAYOUT.md) into the
+  small resident tensors the Triton kernel consumes — the flat codebook, the
+  pre-decoded fp4 scale plane, and an 8-byte-padded index stream. None of these
+  is a dense [N,K] weight (INV-1 holds);
+* activation QDQ that reproduces the emulation gate's served-activation buckets
+  (fp4 group-16 RTN / fp8 dynamic per-token) so served KL is comparable to the
+  emulated prediction.
+"""
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+
+VEC_DIM = 8
+SUPERBLOCK = 256
+FP4_GROUP = 16
+_E4M3 = torch.float8_e4m3fn
+NVFP4_GRID_MAX = 6.0
+FP8_ELEMENT_MAX = 448.0
+
+# E2M1 element grid (sorted ascending), for the fp4 activation RTN.
+_E2M1 = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0)
+
+# Two-tier v2 scale coding (docs/nvfp4-cb-plan/two-tier-scale-spec.md §1). Kept
+# in sync with prismaquant.nvfp4_cb_formats (the reference the kernel matches).
+SCALE_CODING_TWO_TIER = "two_tier"
+TWO_TIER_SUPER_BIAS = 127
+TWO_TIER_SUB_TABLE = (1.0, 1.125, 1.25, 1.375, 1.5, 1.625, 1.75, 1.875,
+                      2.0, 2.25, 2.5, 2.75, 3.0, 3.25, 3.5, 3.75)
+
+
+def type_size(k: int, is_fp4: bool) -> int:
+    return 4 * int(k) + (16 if is_fp4 else 0)
+
+
+def build_flat_codebook(sub_tables: list[torch.Tensor]) -> torch.Tensor:
+    """Concatenate product sub-tables (each (2^w, sub_dim)) into the flat layout
+    the kernel gathers from: block ``s`` = ``sub_tables[s].reshape(-1)`` (row
+    major, so entry (idx, local) sits at idx*sub_dim + local)."""
+    return torch.cat([t.reshape(-1).to(torch.bfloat16).contiguous()
+                      for t in sub_tables]).contiguous()
+
+
+def build_compose_table(sub_table) -> torch.Tensor:
+    """Two-tier v2 (docs/nvfp4-cb-plan/two-tier-scale-spec.md §1): the (256,16)
+    compose table ``T[c]·2^(E-127)`` flattened to (4096,) fp32, bit-exact to
+    ``nvfp4_cb_formats._two_tier_tables`` (float64 product -> fp32). The kernel
+    gathers ``compose[super_e*16 + code]`` per group — no resident fp32 plane
+    (spec §4/G4), just this 16 KiB constant table."""
+    T = torch.tensor(list(sub_table), dtype=torch.float64)
+    exps = torch.arange(256, dtype=torch.float64)
+    compose = (T[None, :] * torch.pow(2.0, exps[:, None] - 127.0)).to(
+        torch.float32)                                # (256, 16)
+    return compose.reshape(-1).contiguous()           # (4096,)
+
+
+def decode_fp4_scale_plane(qw: torch.Tensor, k: int) -> torch.Tensor:
+    """(N, n_sb*type_size) uint8 -> (N, n_sb*16) fp32 group-16 scales, decoded
+    from the E4M3 scale plane that follows each superblock's 4k index bytes."""
+    n, row_bytes = qw.shape
+    ts = type_size(k, is_fp4=True)
+    n_sb = row_bytes // ts
+    blk = qw.reshape(n, n_sb, ts)
+    plane = blk[:, :, 4 * k:4 * k + FP4_GROUP].contiguous()      # (N, n_sb, 16)
+    return plane.view(_E4M3).to(torch.float32).reshape(n, n_sb * FP4_GROUP)
+
+
+PAD_BYTES = 16
+
+
+def pad_qweight(qw: torch.Tensor) -> torch.Tensor:
+    """Right-pad each row by ``PAD_BYTES`` so the padded buffer satisfies BOTH
+    invariants every consumer of a padded row depends on:
+
+    1. **>= 8 bytes of read slack per row.** The decode/expand kernels extract a
+       codeword with an 8-byte window anchored at the codeword's first byte, so
+       the last codeword of the last superblock of the last row reads up to 7
+       bytes past the packed data. Without the slack that is an out-of-bounds
+       global read (illegal-memory-access, or silent garbage in the final
+       output rows).
+    2. **The padded row stride stays a 16-byte multiple** whenever the UNPADDED
+       one was. The fp8 CUTLASS entries (``cb_fused_prefill_mm_scaled``, the
+       persistent-TC prefill) take the row stride explicitly and TORCH_CHECK
+       ``stride(0) % 16 == 0`` — TMA needs 16-byte-aligned row starts. Every fp8
+       rung has ``type_size = 4k`` in {112,128,144,160,176,192}, so
+       ``row_bytes`` is 16-aligned and ``row_bytes + 16`` still is; the old
+       ``+ 8`` pad was NOT, which is why the padded buffer could not be shared
+       with the registered ``cb_qweight`` parameter and both had to stay
+       resident (see ``linear.process_weights_after_loading``). Pad width is
+       therefore load-bearing, not a spare-bytes choice: dropping it back to 8
+       re-breaks the fp8 prefill kernels' alignment check.
+
+    (fp4 rungs carry an odd ``type_size`` — ``4k+16`` v1, ``4k+9`` v2 — so their
+    row stride is not 16-aligned either way; they only ever hit kernels that
+    take the stride explicitly and require ``stride(1) == 1``.)
+
+    Consumers must read the row stride from the tensor (``.stride(0)``), never
+    derive it as ``size(1) - PAD_BYTES``.
+    """
+    return F.pad(qw.contiguous(), (0, PAD_BYTES), value=0).contiguous()
+
+
+# Signed E2M1 grid, cached per device: building it per call allocated a CPU
+# tensor and H2D-copied it on EVERY fp4 activation QDQ — a hidden sync in the
+# eager decode hot path, and a hard error under CUDA-graph capture (unpinned
+# CPU->CUDA copy). Warmup forwards populate the cache before any capture.
+_FP4_QDQ_GRID: dict = {}
+
+
+def _fp4_qdq_grid(device: torch.device) -> torch.Tensor:
+    grid = _FP4_QDQ_GRID.get(device)
+    if grid is None:
+        grid = torch.tensor(sorted({v for a in _E2M1 for v in (a, -a)}),
+                            dtype=torch.float32, device=device)
+        _FP4_QDQ_GRID[device] = grid
+    return grid
+
+
+def fp4_group16_act_qdq(x: torch.Tensor) -> torch.Tensor:
+    """W4A4 activation bucket: RTN to E2M1 at group-16 amax/6 scale (mirrors
+    format_registry `_make_rtn('fp4_e2m1', 16)`)."""
+    in_f = x.shape[-1]
+    grid = _fp4_qdq_grid(x.device)
+    w = x.reshape(-1, in_f).float().reshape(-1, in_f // FP4_GROUP, FP4_GROUP)
+    scale = w.abs().amax(dim=-1, keepdim=True).clamp_min(1e-8) / NVFP4_GRID_MAX
+    xg = w / scale
+    idx = torch.bucketize(xg.contiguous(), grid)
+    lo = grid[(idx - 1).clamp_min(0)]
+    hi = grid[idx.clamp_max(grid.numel() - 1)]
+    q = torch.where((hi - xg).abs() < (xg - lo).abs(), hi, lo)
+    return (q * scale).reshape(x.shape).to(x.dtype)
+
+
+def fp8_dynamic_act_qdq(x: torch.Tensor) -> torch.Tensor:
+    """W8A8 activation bucket: vLLM dynamic per-token E4M3 (mirrors
+    fp8_dynamic.fp8_dynamic_activation_qdq_vllm)."""
+    rows = x.reshape(-1, x.shape[-1]).float()
+    min_scale = 1.0 / (FP8_ELEMENT_MAX * 512.0)
+    scale = (rows.abs().amax(dim=-1, keepdim=True) / FP8_ELEMENT_MAX
+             ).clamp_min(min_scale)
+    q = (rows / scale).clamp(-FP8_ELEMENT_MAX, FP8_ELEMENT_MAX).to(_E4M3)
+    deq = q.to(torch.float32) * scale
+    return deq.reshape(x.shape).to(x.dtype)

gridbook/config.py

@@ -0,0 +1,425 @@
+"""``PrismaQuantConfig`` — the vLLM quantization config for the NVFP4-CB /
+FP8-CB out-of-tree lane (docs/nvfp4-cb-plan/serving-kernel.md §2, LAYOUT.md §4).
+
+vLLM auto-detects us from ``quant_method == "prismaquant"``. The exporter writes
+``config.json['quantization_config']`` as a *pointer* (``config_file`` ->
+``quant_config.json`` + ``codebook_file`` -> ``cb_codebooks.pqcb``); the full
+``config_groups`` / ``ignore`` live in ``quant_config.json``. We resolve that
+sidecar **lazily** (via ``get_current_vllm_config()``, the same handle
+``get_codebooks`` uses) since ``from_config`` runs before the model dir is
+plumbed. Inlined configs (``config_groups`` already present) are also accepted.
+
+**Mixed-container dispatch (serving-kernel.md §2).** A config group with a
+``"scheme"`` key is a CB group (our nvfp4_cb/fp8_cb vocabulary) -> our
+``PrismaQuantCBLinearMethod``. A group WITHOUT it uses the exact stock
+compressed-tensors vocabulary -> a real ``CompressedTensorsConfig`` we construct
+and delegate to (``CompressedTensorsW4A4Nvfp4`` for NVFP4 groups, the fp8 scheme
+for FP8_DYNAMIC). ``ignore`` -> ``UnquantizedLinearMethod``.
+"""
+from __future__ import annotations
+
+import json
+import os
+from typing import Any
+
+import torch
+from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
+from vllm.model_executor.layers.quantization.base_config import (
+    QuantizationConfig,
+    QuantizeMethodBase,
+)
+from vllm.model_executor.layers.vocab_parallel_embedding import (
+    UnquantizedEmbeddingMethod,
+    VocabParallelEmbedding,
+)
+
+try:
+    from vllm.model_executor.layers.fused_moe import RoutedExperts
+except Exception:  # pragma: no cover - older vLLM
+    RoutedExperts = None
+
+_MOE_LEAVES = ("gate_up_proj", "down_proj", "gate_proj", "up_proj")
+
+# vLLM fuses these siblings into one module; packed_modules_mapping is populated
+# by dispatch time, but we keep the standard mapping as a fallback.
+_FUSED_FALLBACK = {
+    "qkv_proj": ["q_proj", "k_proj", "v_proj"],
+    "gate_up_proj": ["gate_proj", "up_proj"],
+    "in_proj_qkvz": ["in_proj_qkv", "in_proj_z"],
+    "in_proj_ba": ["in_proj_b", "in_proj_a"],
+}
+
+
+def _canonical_prefix(prefix: str) -> str:
+    """vLLM serving prefix -> canonical target namespace. Some model classes
+    wrap the LM (`language_model.model.layers.*` on Qwen3.5-class VL) while
+    targets are canonical `model.layers.*`; strip the wrapper when the next
+    component is `model.` (measured via PRISMAQUANT_DEBUG_PREFIXES,
+    2026-07-22 — every LM layer resolved no-scheme without this)."""
+    if prefix.startswith("language_model.model."):
+        return prefix[len("language_model."):]
+    if prefix.startswith("language_model."):
+        return "model." + prefix[len("language_model."):]
+    # Pre-fix multimodal CHECKPOINT namespace (shipped 27B gridbook artifact):
+    # ``model.language_model.layers.*`` denotes the same Linear as canonical
+    # ``model.layers.*``. Normalising it here (as well as in
+    # ``_canonical_target``) keeps probe-side and target-side on one string.
+    if prefix.startswith("model.language_model."):
+        return "model." + prefix[len("model.language_model."):]
+    return prefix
+
+
+def _candidate_bases(name: str) -> list[str]:
+    """Every namespace vintage *name* can legitimately be matched against,
+    **most specific first** (the string as given, then its canonical form).
+
+    THE one place that answers "which namespace am I in?". A stored target /
+    serving prefix reaches us in one of three vintages — the old multimodal
+    CHECKPOINT form (``model.language_model.*``), the canonical form
+    (``model.*``), and the vLLM wrapper-class SERVING form
+    (``language_model.model.*``) — and ``apply_vllm_mapper`` can move the
+    stored keys into a *fourth*, the mapper's own namespace, AFTER
+    ``_ensure_resolved`` canonicalised them. Anything that matches a prefix
+    against ``target_scheme`` / ``ignore`` must therefore try both sides, and
+    must do so HERE: the dense fused path grew its own single-namespace copy of
+    this logic and silently mis-resolved for it (issue #1). A future fifth
+    vintage should mean editing this function and nothing else.
+    """
+    canonical = _canonical_prefix(name)
+    return [name] if canonical == name else [name, canonical]
+
+
+def _canonical_target(name: str) -> str:
+    """Stored ``config_groups[*].targets`` / ``ignore`` entry -> canonical
+    target namespace, so historical checkpoint-namespace artifacts resolve
+    against the canonicalised serving prefixes ``_canonical_prefix`` produces.
+
+    Rewrites (prefix-anchored only):
+      ``model.language_model.`` -> ``model.``          (old multimodal ckpt)
+      ``language_model.model.`` -> ``model.``          (serving wrapper form)
+      ``language_model.<rest>`` -> ``model.<rest>``
+    Everything else (``visual.*``, ``mtp.*``, plain ``model.layers.*``,
+    bare leaf names) passes through untouched."""
+    return _canonical_prefix(name)
+
+
+def _resolve_model_file(model_dir: str, fname: str) -> str:
+    """Local path for a sidecar file next to the model. When the model was
+    given as a Hub repo id (``vllm serve rdtand/...``) rather than a local
+    directory, fetch the sidecar from the Hub — vLLM's own loader handles the
+    weights that way, but OUR sidecars (quant_config.json, the .pqcb codebook
+    blob) were opened with a plain path join, which broke every serve-by-id
+    until 2026-07-22."""
+    if os.path.isdir(model_dir):
+        return os.path.join(model_dir, fname)
+    from huggingface_hub import hf_hub_download
+    return hf_hub_download(repo_id=model_dir, filename=fname)
+
+
+class PrismaQuantConfig(QuantizationConfig):
+    """Per-layer dispatch: CB decode / stock-CT delegation / unquantized."""
+
+    def __init__(self, raw_config: dict) -> None:
+        super().__init__()
+        self._raw_config = dict(raw_config or {})
+        self.codebook_file = self._raw_config.get("codebook_file",
+                                                  "cb_codebooks.pqcb")
+        # Resolved lazily (the sidecar quant_config.json needs the model dir).
+        self._resolved = False
+        self._full_config: dict = {}
+        self.config_groups: dict = {}
+        self.ignore: list[str] = []
+        self.target_scheme: dict[str, dict] = {}    # CB module -> scheme dict
+        self._cb_targets: set[str] = set()
+        self.ct_config = None                        # stock CompressedTensorsConfig
+        self._codebooks: dict[str, torch.Tensor] | None = None
+
+    # -- lazy resolution of the (possibly pointer) quant config --------------
+    def _ensure_resolved(self) -> None:
+        if self._resolved:
+            return
+        cfg = self._raw_config
+        if "config_groups" not in cfg:
+            cfg_file = cfg.get("config_file", "quant_config.json")
+            from vllm.config import get_current_vllm_config
+            model_dir = get_current_vllm_config().model_config.model
+            with open(_resolve_model_file(model_dir, cfg_file)) as fh:
+                cfg = json.load(fh)
+            self.codebook_file = cfg.get("codebook_file", self.codebook_file)
+        # Normalise stored namespaces ONCE, here, so all downstream resolution
+        # (ours and the delegated CT config's) sees canonical target names.
+        cfg = dict(cfg)
+        cfg["config_groups"] = {
+            name: {**g, "targets": [_canonical_target(t)
+                                    for t in g.get("targets", [])]}
+            for name, g in cfg["config_groups"].items()
+        }
+        cfg["ignore"] = [_canonical_target(i) for i in cfg.get("ignore", [])]
+        self._full_config = cfg
+        self.config_groups = cfg["config_groups"]
+        self.ignore = list(cfg["ignore"])
+        stock_groups: dict = {}
+        for name, g in self.config_groups.items():
+            if "scheme" in g:                        # CB group (our vocabulary)
+                for t in g["targets"]:
+                    self.target_scheme[t] = g["scheme"]
+                    self._cb_targets.add(t)
+            else:                                    # stock CT vocabulary
+                stock_groups[name] = g
+        self._alias_collapsed_shared_prefixes()
+        self.ct_config = (self._build_ct_config(stock_groups)
+                          if stock_groups else None)
+        self._resolved = True
+
+    def _alias_collapsed_shared_prefixes(self) -> None:
+        """HunYuan-V3-style shared-expert dispatch collapse. HYV3MoEFused builds
+        its shared MLP with ``prefix=f"{prefix}"`` — the ``.shared_mlp`` segment
+        never reaches ``get_quant_method``, which instead sees the PARENT-prefix
+        names ``…mlp.gate_up_proj`` / ``…mlp.down_proj``. Module paths (params,
+        checkpoint tensors) DO keep ``.shared_mlp.``, so only the dispatch key
+        collapses. Alias every ``….shared_mlp.<leaf>`` CB target and ignore
+        entry to its collapsed form so the CB method owns the shared expert
+        natively (packed decode in-kernel) instead of vLLM building plain bf16
+        Linears that the loader must fill by decode-at-load. Collision-safe: a
+        layer is either dense-MLP (real ``…mlp.<leaf>`` keys, no shared_mlp) or
+        MoE-with-shared (no real collapsed keys), and ``setdefault`` keeps any
+        real key authoritative. Archs that thread correct shared prefixes are
+        unaffected (their aliases match no module prefix). Runs before the
+        delegated-CT build so CT's ignore covers the aliases too."""
+        for t in [k for k in self.target_scheme if ".shared_mlp." in k]:
+            alias = t.replace(".shared_mlp.", ".")
+            self.target_scheme.setdefault(alias, self.target_scheme[t])
+            self._cb_targets.add(alias)
+        self.ignore.extend(ig.replace(".shared_mlp.", ".")
+                           for ig in list(self.ignore) if ".shared_mlp." in ig)
+
+    def _build_ct_config(self, stock_groups: dict):
+        """A stock CompressedTensorsConfig over the non-CB groups. They are
+        already CT vocabulary; we re-key quant_method, add our CB modules to
+        CT's ignore (so CT never owns them), and give it a valid top-level
+        format (our container's is a CB marker; stock groups carry per-group
+        formats that CT reads under "mixed-precision")."""
+        from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import (  # noqa: E501
+            CompressedTensorsConfig,
+        )
+        ct_dict = dict(self._full_config)
+        ct_dict["quant_method"] = "compressed-tensors"
+        ct_dict["config_groups"] = dict(stock_groups)
+        ct_dict["ignore"] = list(self.ignore) + sorted(self._cb_targets)
+        ct_dict.pop("codebook_file", None)
+        ct_dict.pop("provenance", None)
+        raw_fmt = str(self._full_config.get("format", ""))
+        if raw_fmt in ("", "nvfp4_cb", "fp8_cb", "cb", "mixed-precision"):
+            ct_dict["format"] = "mixed-precision"
+        return CompressedTensorsConfig.from_config(ct_dict)
+
+    def __repr__(self) -> str:
+        return (f"PrismaQuantConfig(resolved={self._resolved}, "
+                f"cb_targets={len(self.target_scheme)}, "
+                f"stock_ct={'yes' if self.ct_config is not None else 'no'})")
+
+    @classmethod
+    def get_name(cls):
+        return "gridbook"
+
+    def get_supported_act_dtypes(self) -> list[torch.dtype]:
+        return [torch.bfloat16, torch.float16]
+
+    @classmethod
+    def get_min_capability(cls) -> int:
+        return 80
+
+    @classmethod
+    def get_config_filenames(cls) -> list[str]:
+        return []
+
+    @classmethod
+    def from_config(cls, config: dict[str, Any]) -> "PrismaQuantConfig":
+        # Defer parsing: a pointer config resolves quant_config.json lazily.
+        return cls(config)
+
+    @classmethod
+    def override_quantization_method(cls, hf_quant_cfg, user_quant, **kwargs):
+        # "gridbook" is the registry key going forward; "prismaquant" is the
+        # legacy key older local artifacts carry — both dispatch here.
+        if user_quant in ("gridbook", "prismaquant"):
+            return "gridbook"
+        if hf_quant_cfg is not None and \
+                hf_quant_cfg.get("quant_method") in ("gridbook", "prismaquant"):
+            return "gridbook"
+        return None
+
+    # -- codebook sidecar (loaded once, shared across all layers) ------------
+    def get_codebooks(self) -> dict[str, torch.Tensor]:
+        if self._codebooks is None:
+            from safetensors.torch import load_file
+            from vllm.config import get_current_vllm_config
+            model_dir = get_current_vllm_config().model_config.model
+            self._codebooks = load_file(
+                _resolve_model_file(model_dir, self.codebook_file))
+        return self._codebooks
+
+    # -- per-prefix scheme resolution (handles vLLM fused qkv/gate_up) -------
+    def _is_ignored(self, prefix: str) -> bool:
+        return any(ig in base for base in _candidate_bases(prefix)
+                   for ig in self.ignore)
+
+    def shard_target_keys(self, prefix: str, *,
+                          unfused_fallback: bool = False) -> list[str]:
+        """``target_scheme`` keys naming the CB shards of (possibly fused)
+        *prefix*, in shard order — ``[]`` if none resolve.
+
+        THE single owner of fused-shard resolution: ``_scheme_for_prefix``
+        (which format does this module decode as?) and
+        ``PrismaQuantCBLinearMethod._shard_roles`` (which per-role codebooks
+        does it concatenate?) must agree module-for-module, and before issue #1
+        they were two hand-rolled copies that had already drifted — the copies
+        built their shard keys from the CANONICAL prefix only, so once
+        ``apply_vllm_mapper`` moved the stored keys into the mapper's namespace
+        a fused GDN ``in_proj_qkvz`` resolved to nothing (silent BF16
+        fall-through) and every dense ``_shard_roles`` returned ``[]`` (a
+        load-time width assert). Namespace choice is delegated wholesale to
+        ``_candidate_bases``.
+
+        Bases are tried in order and the FIRST base with any hit wins **whole**:
+        hits are never mixed across bases, because two vintages of one key can
+        name two different on-disk tensors and pairing shards across them would
+        silently fuse the wrong weights.
+
+        ``unfused_fallback`` reproduces ``_shard_roles``' extra ``or [leaf]``
+        rung — a plain Linear is its own single role. ``_scheme_for_prefix``
+        deliberately omits it (it has already tried the exact keys itself, and
+        a bare-leaf retry there would only re-ask the same question).
+        """
+        pmm = getattr(self, "packed_modules_mapping", {}) or {}
+        for base in _candidate_bases(prefix):
+            leaf = base.split(".")[-1]
+            shard_leaves = pmm.get(leaf) or _FUSED_FALLBACK.get(leaf)
+            if shard_leaves is None:
+                if not unfused_fallback:
+                    continue
+                shard_leaves = [leaf]
+            stem = base[: -len(leaf)]
+            hits = [stem + sl for sl in shard_leaves
+                    if stem + sl in self.target_scheme]
+            if hits:
+                return hits
+        return []
+
+    def _scheme_for_prefix(self, prefix: str) -> dict | None:
+        for base in _candidate_bases(prefix):
+            if base in self.target_scheme:
+                return self.target_scheme[base]
+        schemes = [self.target_scheme[k]
+                   for k in self.shard_target_keys(prefix)]
+        if not schemes:
+            return None
+        fmt_keys = ("grid", "mode", "k", "n_sub", "type_size")
+        sig = {kk: schemes[0][kk] for kk in fmt_keys}
+        for s in schemes[1:]:
+            if {kk: s[kk] for kk in fmt_keys} != sig:
+                raise ValueError(
+                    f"fused module {prefix} maps to mixed CB decode "
+                    "formats — export union-find should prevent this")
+        return schemes[0]
+
+    def get_quant_method(self, layer: torch.nn.Module,
+                         prefix: str) -> "QuantizeMethodBase | None":
+        self._ensure_resolved()
+        from .linear import PrismaQuantCBLinearMethod
+
+        # Keep the delegated CT config's fused-module mapping in lockstep.
+        if self.ct_config is not None:
+            self.ct_config.packed_modules_mapping = getattr(
+                self, "packed_modules_mapping", {}) or {}
+
+        if isinstance(layer, LinearBase):
+            # 1) CB target (has a "scheme") — ours (precise, fused-aware; ahead
+            #    of the substring ignore test).
+            scheme = self._scheme_for_prefix(prefix)
+            if os.environ.get("PRISMAQUANT_DEBUG_PREFIXES") == "1":
+                import sys
+                print(f"[pq-prefix] {prefix} -> "
+                      f"{'CB' if scheme is not None else 'no-scheme'}",
+                      file=sys.stderr, flush=True)
+            if scheme is not None:
+                return PrismaQuantCBLinearMethod(self, scheme, prefix)
+            # 2) explicitly-ignored -> BF16 passthrough.
+            if self._is_ignored(prefix):
+                return UnquantizedLinearMethod()
+            # 3) stock NVFP4 / FP8_DYNAMIC -> compressed-tensors delegation
+            #    (canonical prefix — CT targets are serving-namespace names).
+            if self.ct_config is not None:
+                return self.ct_config.get_quant_method(
+                    layer, _canonical_prefix(prefix))
+            return UnquantizedLinearMethod()
+
+        if isinstance(layer, VocabParallelEmbedding):
+            if self.ct_config is not None:
+                method = self.ct_config.get_quant_method(layer, prefix)
+                if method is not None:
+                    return method
+            return UnquantizedEmbeddingMethod()
+
+        # FusedMoE expert stacks (RoutedExperts): a CB expert group -> our MoE
+        # method; else delegate to the stock CT MoE path.
+        if RoutedExperts is not None and isinstance(layer, RoutedExperts):
+            scheme = self._moe_scheme_for_prefix(prefix)
+            if scheme is not None:
+                from .moe import PrismaQuantCBMoEMethod
+                return PrismaQuantCBMoEMethod(
+                    self, layer.moe_config, scheme, prefix)
+            if self.ct_config is not None:
+                return self.ct_config.get_quant_method(layer, prefix)
+            return None
+        return None
+
+    def _moe_scheme_for_prefix(self, prefix: str) -> dict | None:
+        """A CB expert stack (targets like ``…experts.gate_up_proj`` /
+        ``…experts.down_proj``) under this FusedMoE prefix — return its scheme
+        (uniform per layer, so any matching target's scheme is the layer's)."""
+        # Canonicalise BOTH sides, exactly as ``_scheme_for_prefix`` does for
+        # Linears. Without this the multimodal wrapper breaks experts ONLY:
+        # vLLM hands us the serving prefix ``language_model.model.layers.N.mlp.
+        # experts`` while the checkpoint-namespace targets read
+        # ``model.language_model.layers.N.mlp.experts.gate_up_proj``, so a raw
+        # ``startswith`` misses, no CB MoE method is created, no
+        # ``w13_cb_qweight``/``w2_cb_qweight`` params exist, and the arch's own
+        # expert mapping then derives ``experts.w2_weight.cb_qweight`` and
+        # AttributeErrors (35B CB serve boot). Dense Linears were unaffected
+        # because their lookup already canonicalised — that asymmetry WAS the bug.
+        #
+        # Structurally different from the dense lookup (the TARGET is longer
+        # than the prefix here, so this is a ``startswith``, not a key lookup),
+        # but the namespace question is the same one — so it comes from the same
+        # ``_candidate_bases``, on BOTH sides. Cross-vintage matches are safe
+        # here: ``_canonical_prefix`` only rewrites the ``language_model``
+        # wrapper, i.e. it renames the SAME module; it can never move a match to
+        # a different layer index or leaf.
+        bases = _candidate_bases(prefix)
+        for name, sch in self.target_scheme.items():
+            if name.split(".")[-1] not in _MOE_LEAVES:
+                continue
+            variants = _candidate_bases(name)
+            if any(v.startswith(b) for v in variants for b in bases):
+                return sch
+        return None
+
+    def apply_vllm_mapper(self, hf_to_vllm_mapper):
+        self._ensure_resolved()
+        # vLLM hands us the UNSTACKED mapper (get_unstacked_mapper()), so the
+        # q_proj->qkv_proj fusion is NOT rewritten (per-role leaf names survive
+        # for _scheme_for_prefix to re-fuse) — but genuine renames/prefixes ARE
+        # applied. For hybrid/VLM checkpoints that means the module-nesting
+        # prefix (e.g. Qwen3-VL: ``model.language_model.`` -> ``language_model.
+        # model.``) must be applied to the CB target keys too: _scheme_for_prefix
+        # matches serve-time prefixes EXACTLY (unlike the substring ignore test),
+        # so an un-remapped key silently falls through to unquantized and the
+        # cb_qweight load then fails ("no parameter named …cb_qweight"). Mirror
+        # exactly what the delegated stock-CT config does for its own targets.
+        self.ignore = hf_to_vllm_mapper.apply_list(self.ignore)
+        self.target_scheme = hf_to_vllm_mapper.apply_dict(self.target_scheme)
+        self._cb_targets = set(
+            hf_to_vllm_mapper.apply_list(sorted(self._cb_targets)))
+        if self.ct_config is not None:
+            self.ct_config.apply_vllm_mapper(hf_to_vllm_mapper)