mjolnir-kernel 0.1.0__py3-none-any.whl

mjolnir/lmhead.py

@@ -0,0 +1,379 @@
+"""Opt-in lm_head experiments for decode (vLLM GPU model runner).
+
+On Qwen3.8-27B NVFP4 the bf16 lm_head (248,320 x 5,120, 2.54 GB) is read by
+the target's verify and by every MTP draft step — ~40% of a decode step's
+bytes. Two experiments, both off unless their env var is set (a config opts
+in through the ``.env`` sidecar inside its config dir, see
+``mjolnir_cli.dockerctl``):
+
+* ``MJOLNIR_LM_HEAD=fp8`` — the lm_head weight quantized to e4m3 per row at
+  load (scale = row amax / 448) and multiplied by a Triton fp8-weight GEMM
+  (bf16 activations, fp32 accumulate): half the bytes for the verify and the
+  drafts. The fp8 head makes the bf16 copy dead weight, so the bf16 weight is
+  freed automatically (an untied head only), returning 2.5 GB to the KV cache.
+* ``MJOLNIR_DRAFT_VOCAB=<json>`` (+ ``MJOLNIR_DRAFT_VOCAB_SIZE``) — VocabTrim
+  for the drafter: its logits are computed only over the most frequent
+  target tokens (``mjolnir calibrate --draft-vocab``), the rest are -inf.
+  Verification still uses the full vocabulary, so the output distribution is
+  unchanged (rejection sampling stays exact); only acceptance can drop.
+
+The hooks wrap ``GPUModelRunner.load_model`` (installed by
+``mjolnir.register`` when either variable is set) and run after the target
+and the drafter are loaded, before CUDA-graph capture.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import os
+
+import torch
+
+try:
+    import triton
+    import triton.language as tl
+except ImportError:  # pragma: no cover
+    triton = None
+    tl = None
+
+_log = logging.getLogger("vllm.mjolnir")
+
+ENV_LM_HEAD = "MJOLNIR_LM_HEAD"
+ENV_DRAFT_VOCAB = "MJOLNIR_DRAFT_VOCAB"
+ENV_DRAFT_VOCAB_SIZE = "MJOLNIR_DRAFT_VOCAB_SIZE"
+FP8_MAX = 448.0
+SMALL_M = 64  # beyond this many rows (prompt logprobs) the dequantized chunks win
+
+
+def wanted() -> bool:
+    return bool(os.environ.get(ENV_LM_HEAD, "").strip() or os.environ.get(ENV_DRAFT_VOCAB, "").strip())
+
+
+# ── fp8-weight GEMM ─────────────────────────────────────────────────────────
+
+if triton is not None:
+
+    @triton.jit
+    def _fp8_head_kernel(
+        x_ptr,
+        w_ptr,
+        s_ptr,
+        o_ptr,
+        M,
+        N,
+        K,
+        stride_xm,
+        stride_om,
+        BLOCK_M: tl.constexpr,
+        BLOCK_N: tl.constexpr,
+        BLOCK_K: tl.constexpr,
+    ):
+        pid_n = tl.program_id(0)
+        pid_m = tl.program_id(1)
+        rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+        rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+        rk = tl.arange(0, BLOCK_K)
+        acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+        for k0 in range(0, K, BLOCK_K):
+            kk = k0 + rk
+            x = tl.load(
+                x_ptr + rm[:, None] * stride_xm + kk[None, :],
+                mask=(rm[:, None] < M) & (kk[None, :] < K),
+                other=0.0,
+            )
+            w = tl.load(
+                w_ptr + rn[:, None] * K + kk[None, :],
+                mask=(rn[:, None] < N) & (kk[None, :] < K),
+                other=0.0,
+            )
+            acc += tl.dot(x.to(tl.bfloat16), tl.trans(w.to(tl.bfloat16)))
+        s = tl.load(s_ptr + rn, mask=rn < N, other=0.0)
+        acc = acc * s[None, :]
+        tl.store(
+            o_ptr + rm[:, None] * stride_om + rn[None, :],
+            acc.to(o_ptr.dtype.element_ty),
+            mask=(rm[:, None] < M) & (rn[None, :] < N),
+        )
+
+
+ROW_CHUNK = 4096  # rows per step when quantizing / testing a head (bounds the temporaries)
+
+
+def quantize_rows(w: torch.Tensor, chunk: int = ROW_CHUNK) -> tuple[torch.Tensor, torch.Tensor]:
+    """(N, K) float -> (N, K) e4m3 and (N,) fp32 row scales (amax / 448).
+
+    Row chunks of ``chunk``: a whole-head fp32 copy of the 248K x 5120 lm_head
+    is 5 GiB, and the abs / divide / clamp temporaries of one expression are
+    several more. Freed or not, the caching allocator keeps them reserved,
+    and vLLM's startup memory profile (on Thor: one unified pool) counts them
+    as used — they came straight out of the KV cache. Chunked, the peak is a
+    few hundred MiB."""
+    N = w.shape[0]
+    q = torch.empty(w.shape, dtype=torch.float8_e4m3fn, device=w.device)
+    scale = torch.empty(N, dtype=torch.float32, device=w.device)
+    for r0 in range(0, N, chunk):
+        w32 = w[r0 : r0 + chunk].float()
+        sc = (w32.abs().amax(dim=1) / FP8_MAX).clamp_min(1e-12)
+        q[r0 : r0 + chunk] = (w32.div_(sc[:, None])).clamp_(-FP8_MAX, FP8_MAX).to(torch.float8_e4m3fn)
+        scale[r0 : r0 + chunk] = sc
+        del w32
+    return q, scale
+
+
+def fp8_linear_reference(x: torch.Tensor, w8: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
+    return (x.float() @ (w8.float() * scale[:, None]).t()).to(x.dtype)
+
+
+def fp8_linear(x: torch.Tensor, w8: torch.Tensor, scale: torch.Tensor, chunk: int = 16384) -> torch.Tensor:
+    """x (..., K) bf16/fp16 @ (w8 * scale)^T -> (..., N), x's dtype."""
+    shape = x.shape
+    x2 = x.reshape(-1, shape[-1])
+    if x2.stride(-1) != 1:
+        x2 = x2.contiguous()
+    M, K = x2.shape
+    N = w8.shape[0]
+    out = torch.empty(M, N, dtype=x.dtype, device=x.device)
+    if M <= SMALL_M and triton is not None:
+        BM, BN, BK = 16, 64, 128
+        grid = (triton.cdiv(N, BN), triton.cdiv(M, BM))
+        _fp8_head_kernel[grid](
+            x2, w8, scale, out, M, N, K, x2.stride(0), out.stride(0), BLOCK_M=BM, BLOCK_N=BN, BLOCK_K=BK,
+            num_warps=4, num_stages=4,
+        )  # fmt: skip
+    else:  # many rows: dequantize a slab of the head at a time
+        for n0 in range(0, N, chunk):
+            w = w8[n0 : n0 + chunk].to(x.dtype) * scale[n0 : n0 + chunk, None].to(x.dtype)
+            out[:, n0 : n0 + chunk] = x2 @ w.t()
+    return out.reshape(*shape[:-1], N)
+
+
+class Fp8HeadMethod:
+    """Stands in for the lm_head's quant method: ``apply`` runs the fp8
+    GEMM; everything else (``embedding`` for a tied head, attributes) goes
+    to the original method."""
+
+    def __init__(self, base, w8: torch.Tensor, scale: torch.Tensor):
+        self._base = base
+        self.w8 = w8
+        self.scale = scale
+
+    def apply(self, layer, x, bias=None):
+        out = fp8_linear(x, self.w8, self.scale)
+        return out + bias if bias is not None else out
+
+    def __getattr__(self, name):
+        return getattr(self._base, name)
+
+
+def _self_test(w: torch.Tensor, w8: torch.Tensor, scale: torch.Tensor) -> float:
+    """Relative error of the fp8 head vs bf16 on random activations."""
+    g = torch.Generator(device=w.device).manual_seed(0)
+    x = torch.randn(4, w.shape[1], device=w.device, dtype=w.dtype, generator=g)
+    rows = torch.arange(0, w.shape[0], max(1, w.shape[0] // ROW_CHUNK), device=w.device)[:ROW_CHUNK]
+    ref = x.float() @ w.index_select(0, rows).float().t()  # a strided row sample: no full fp32 copy
+    w8s = w8.view(torch.uint8).index_select(0, rows).view(torch.float8_e4m3fn)  # fp8 has no index_select kernel
+    got = fp8_linear(x, w8s, scale.index_select(0, rows)).float()
+    return float((got - ref).norm() / ref.norm().clamp_min(1e-12))
+
+
+def _is_tied(lm_head, model) -> bool:
+    ptr = lm_head.weight.data_ptr()
+    for name, mod in model.named_modules():
+        if mod is not lm_head and name.endswith("embed_tokens") and hasattr(mod, "weight"):
+            if mod.weight.data_ptr() == ptr:
+                return True
+    return False
+
+
+def install_fp8_head(lm_head, model=None) -> bool:
+    if getattr(lm_head, "_mjolnir_fp8", False):
+        return True
+    qm = getattr(lm_head, "quant_method", None)
+    if qm is None or "Unquantized" not in type(qm).__name__:
+        _log.warning("Mjolnir lm_head fp8: %s is not an unquantized head; left as is", type(qm).__name__)
+        return False
+    w = lm_head.weight.data
+    w8, scale = quantize_rows(w)
+    err = _self_test(w, w8, scale)
+    if not err < 0.05:
+        _log.warning("Mjolnir lm_head fp8: self-test relative error %.4f; left in bf16", err)
+        return False
+    lm_head.quant_method = Fp8HeadMethod(qm, w8, scale)
+    lm_head._mjolnir_fp8 = True
+    freed = ""
+    if model is not None:  # the fp8 head is now the only consumer: the bf16 copy is dead weight
+        if _is_tied(lm_head, model):
+            freed = " (tied to the embedding: bf16 weight kept)"
+        else:
+            lm_head.weight.data = torch.empty(0, dtype=w.dtype, device=w.device)
+            torch.cuda.empty_cache()
+            freed = f", bf16 weight freed ({w.numel() * w.element_size() / 2**30:.2f} GiB)"
+    _log.info(
+        "Mjolnir lm_head fp8: %s x %s e4m3 per-row, self-test rel. error %.4f%s",
+        w.shape[0],
+        w.shape[1],
+        err,
+        freed,
+    )
+    return True
+
+
+# ── VocabTrim for the drafter ───────────────────────────────────────────────
+
+
+def load_vocab(path: str, size: int | None = None) -> list[int]:
+    """Token ids, most frequent first, from a ``mjolnir.draft_vocab/1`` file
+    (``{"format": ..., "tokens": [ids], ...}``) or a plain JSON list."""
+    with open(path) as f:
+        doc = json.load(f)
+    ids = doc["tokens"] if isinstance(doc, dict) else doc
+    ids = [int(t) for t in ids]
+    return ids[:size] if size else ids
+
+
+def pad_vocab(ids: list[int], size: int | None, vocab: int) -> list[int]:
+    """Fill a short id list up to ``size`` with the lowest unused ids.
+
+    A calibration that saw fewer distinct tokens than the requested size
+    (6K distinct from ~50K generated tokens is typical) would otherwise trim
+    the drafter far below the size the config names. Low ids come first
+    because BPE vocabularies are ordered by merge rank, a fair frequency
+    prior for tokens the calibration never generated."""
+    seen = {i for i in ids if 0 <= i < vocab}
+    out = [i for i in dict.fromkeys(ids) if 0 <= i < vocab]
+    if size is None or len(out) >= size:
+        return out[:size] if size else out
+    need = min(size, vocab) - len(out)
+    for i in range(vocab):
+        if need <= 0:
+            break
+        if i not in seen:
+            out.append(i)
+            need -= 1
+    return out
+
+
+def install_draft_vocab(draft_model, ids: list[int], size: int | None = None) -> bool:
+    lm_head = getattr(draft_model, "lm_head", None)
+    lp = getattr(draft_model, "logits_processor", None)
+    if lm_head is None or lp is None or lm_head.weight.numel() == 0:
+        _log.warning("Mjolnir draft vocab: the drafter has no usable lm_head; left as is")
+        return False
+    if getattr(lp, "scale", 1.0) != 1.0 or getattr(lp, "soft_cap", None) is not None:
+        _log.warning("Mjolnir draft vocab: logit scale / soft cap not supported; left as is")
+        return False
+    w = lm_head.weight.data
+    vocab = int(getattr(lp, "org_vocab_size", w.shape[0]))
+    calibrated = len({i for i in ids if 0 <= i < vocab})
+    ids = pad_vocab(ids, size, vocab)
+    if size and calibrated < size:
+        _log.warning(
+            "Mjolnir draft vocab: the file has %d tokens, padded to %d with the lowest unused ids "
+            "(re-run `mjolnir calibrate --draft-vocab` with more prompts for a measured ranking)",
+            calibrated,
+            len(ids),
+        )
+    idx = torch.tensor(sorted(set(ids)), dtype=torch.long, device=w.device)
+    fp8 = getattr(lm_head, "_mjolnir_fp8", False) or os.environ.get(ENV_LM_HEAD, "").strip() == "fp8"
+    sub = w.index_select(0, idx).contiguous()
+    if fp8:
+        sub8, sub_scale = quantize_rows(sub)
+
+        def head(h):
+            return fp8_linear(h, sub8, sub_scale)
+    else:
+
+        def head(h):
+            return torch.nn.functional.linear(h, sub)
+
+    # row chunks bound the temporaries when a caller (vLLM's startup memory
+    # profile) passes thousands of rows at once
+    def compute_logits(hidden_states, *args, **kwargs):
+        h = hidden_states.reshape(-1, hidden_states.shape[-1])
+        out = torch.full((h.shape[0], vocab), float("-inf"), dtype=h.dtype, device=h.device)
+        for r0 in range(0, h.shape[0], ROW_CHUNK):
+            out[r0 : r0 + ROW_CHUNK].index_copy_(1, idx, head(h[r0 : r0 + ROW_CHUNK]).to(out.dtype))
+        return out.reshape(*hidden_states.shape[:-1], vocab)
+
+    def get_top_tokens(hidden_states, *args, **kwargs):
+        h = hidden_states.reshape(-1, hidden_states.shape[-1])
+        top = torch.cat([head(h[r0 : r0 + ROW_CHUNK]).argmax(dim=-1) for r0 in range(0, max(1, h.shape[0]), ROW_CHUNK)])
+        return idx[top].reshape(hidden_states.shape[:-1])
+
+    draft_model.compute_logits = compute_logits
+    if hasattr(draft_model, "get_top_tokens"):
+        draft_model.get_top_tokens = get_top_tokens
+    draft_model._mjolnir_draft_vocab = int(idx.numel())
+    _log.info(
+        "Mjolnir draft vocab: drafter logits over %d of %d tokens (%s head, %.0f MB)",
+        idx.numel(),
+        vocab,
+        "fp8" if fp8 else "bf16",
+        sub.numel() * (1 if fp8 else sub.element_size()) / 2**20,
+    )
+    return True
+
+
+# ── the hook ────────────────────────────────────────────────────────────────
+
+
+def _lm_heads(model) -> list:
+    seen, out = set(), []
+    for name, mod in model.named_modules():
+        if name.split(".")[-1] == "lm_head" and hasattr(mod, "quant_method") and id(mod) not in seen:
+            seen.add(id(mod))
+            out.append(mod)
+    return out
+
+
+def after_load(runner) -> None:
+    """Apply the opted-in experiments to a loaded GPU model runner."""
+    target = getattr(runner, "model", None)
+    spec = getattr(runner, "speculator", None)
+    draft = getattr(spec, "model", None) if spec is not None else None
+    vocab_path = os.environ.get(ENV_DRAFT_VOCAB, "").strip()
+    if vocab_path and draft is not None:  # before fp8 may free the bf16 weight
+        try:
+            size = int(os.environ.get(ENV_DRAFT_VOCAB_SIZE, "0") or 0) or None
+            install_draft_vocab(draft, load_vocab(vocab_path), size)
+        except Exception as e:  # noqa: BLE001 - an experiment never breaks serving
+            _log.warning("Mjolnir draft vocab: not installed (%r)", e)
+    elif vocab_path:
+        _log.warning("Mjolnir draft vocab: %s set but this config has no drafter", ENV_DRAFT_VOCAB)
+    if os.environ.get(ENV_LM_HEAD, "").strip().lower() == "fp8":
+        for model in (target, draft):
+            if model is None:
+                continue
+            for head in _lm_heads(model):
+                try:
+                    install_fp8_head(head, target)
+                except Exception as e:  # noqa: BLE001
+                    _log.warning("Mjolnir lm_head fp8: not installed (%r)", e)
+    if torch.cuda.is_available():
+        # hand the quantization temporaries back before vLLM profiles memory
+        # for the KV cache (a cached-but-free block still counts as used)
+        torch.cuda.empty_cache()
+
+
+def install_hooks() -> None:
+    """Wrap GPUModelRunner.load_model (idempotent; no-op unless opted in)."""
+    if not wanted():
+        return
+    try:
+        from vllm.v1.worker.gpu import model_runner as mr
+    except Exception as e:  # noqa: BLE001
+        _log.warning("Mjolnir lm_head hooks: vLLM's GPU model runner not importable (%r)", e)
+        return
+    cls = mr.GPUModelRunner
+    if getattr(cls, "_mjolnir_lmhead_hooked", False):
+        return
+    orig = cls.load_model
+
+    def load_model(self, *args, **kwargs):
+        out = orig(self, *args, **kwargs)
+        after_load(self)
+        return out
+
+    cls.load_model = load_model
+    cls._mjolnir_lmhead_hooked = True

mjolnir/qpack.py

@@ -0,0 +1,132 @@
+"""Fused dynamic fp8 Q quantization + packing for the Mjolnir decode kernel.
+
+The kernel reads Q as ``(b, h_k, n_q, d)`` e4m3 with one descale per
+``(sequence, KV head)``, where row ``n = qi * h_r + r`` holds query token
+``qi`` of KV-head group member ``r``. vLLM hands Q over as ``(tokens, h_q, d)``
+bf16/fp16, any token and head stride (often a view of the fused QKV buffer).
+
+Doing that with PyTorch ops costs ~10 launches per layer (cast, abs, amax,
+clamp, divide, copy, divide, cast, permute-copy, copy) — at short contexts
+more GPU time than the attention kernel saves. ``pack_q_fp8`` is one Triton
+launch: one program per (sequence, KV head) loads its ``q_len * h_r * d``
+values once, takes the amax, writes the descale and the packed e4m3 rows.
+"""
+
+from __future__ import annotations
+
+import torch
+
+try:  # Triton ships with vLLM / torch on CUDA; the harness may run without it
+    import triton
+    import triton.language as tl
+except ImportError:  # pragma: no cover - no Triton: the torch path is used
+    triton = None
+    tl = None
+
+FP8_MAX = 448.0
+
+
+def available() -> bool:
+    return triton is not None
+
+
+if triton is not None:
+
+    @triton.jit
+    def _pack_q_fp8_kernel(
+        q_ptr,
+        out_ptr,
+        qd_ptr,
+        mul_ptr,
+        stride_t,
+        stride_h,
+        q_len: tl.constexpr,
+        h_k: tl.constexpr,
+        h_r: tl.constexpr,
+        d: tl.constexpr,
+        BLOCK_N: tl.constexpr,
+        BLOCK_D: tl.constexpr,
+        FP8_MAX_C: tl.constexpr,
+        HAS_MUL: tl.constexpr,
+    ):
+        b = tl.program_id(0)
+        k = tl.program_id(1)
+        n_q = q_len * h_r
+        rows = tl.arange(0, BLOCK_N)  # n = qi * h_r + r
+        cols = tl.arange(0, BLOCK_D)
+        qi = rows // h_r
+        r = rows % h_r
+        token = b * q_len + qi
+        head = k * h_r + r
+        mask = (rows[:, None] < n_q) & (cols[None, :] < d)
+        src = q_ptr + token[:, None] * stride_t + head[:, None] * stride_h + cols[None, :]
+        x = tl.load(src, mask=mask, other=0.0).to(tl.float32)
+        if HAS_MUL:  # per-(q head, channel) factor: fp6 K channel scales folded into Q
+            x = x * tl.load(mul_ptr + head[:, None] * d + cols[None, :], mask=mask, other=0.0)
+        amax = tl.maximum(tl.max(tl.max(tl.abs(x), axis=1), axis=0), 1e-6)
+        scale = amax / FP8_MAX_C
+        tl.store(qd_ptr + b * h_k + k, scale)
+        y = tl.minimum(tl.maximum(x / scale, -FP8_MAX_C), FP8_MAX_C)
+        dst = out_ptr + ((b * h_k + k) * n_q + rows[:, None]) * d + cols[None, :]
+        tl.store(dst, y.to(out_ptr.dtype.element_ty), mask=mask)
+
+
+def pack_q_fp8(
+    q3: torch.Tensor,
+    b: int,
+    q_len: int,
+    h_k: int,
+    h_r: int,
+    out: torch.Tensor,
+    qd: torch.Tensor,
+    mul: torch.Tensor | None = None,
+) -> None:
+    """Quantize ``q3`` (``(b * q_len, h_k * h_r, d)``, unit stride on ``d``)
+    into ``out`` (contiguous ``(b, h_k, q_len * h_r, d)`` e4m3) and ``qd``
+    (contiguous ``(b, h_k)`` fp32 descales: amax / 448). ``mul`` (contiguous
+    ``(h_k * h_r, d)`` fp32) multiplies Q before quantization. One launch."""
+    d = q3.shape[-1]
+    assert q3.stride(-1) == 1, "Q needs a unit stride on head_dim"
+    assert out.is_contiguous() and qd.is_contiguous()
+    if mul is not None:
+        assert mul.is_contiguous() and mul.dtype == torch.float32 and mul.numel() == h_k * h_r * d
+    n_q = q_len * h_r
+    block_n = triton.next_power_of_2(n_q)
+    block_d = triton.next_power_of_2(d)
+    warps = 8 if block_n * block_d >= 4096 else 4
+    _pack_q_fp8_kernel[(b, h_k)](
+        q3,
+        out,
+        qd,
+        mul if mul is not None else qd,
+        q3.stride(0),
+        q3.stride(1),
+        q_len=q_len,
+        h_k=h_k,
+        h_r=h_r,
+        d=d,
+        BLOCK_N=block_n,
+        BLOCK_D=block_d,
+        FP8_MAX_C=FP8_MAX,
+        HAS_MUL=mul is not None,
+        num_warps=warps,
+    )
+
+
+def pack_q_fp8_reference(q3: torch.Tensor, b: int, q_len: int, h_k: int, h_r: int, mul: torch.Tensor | None = None):
+    """The PyTorch formulation (the runner's fallback, and the test oracle)."""
+    d = q3.shape[-1]
+    qg = q3[: b * q_len].float()
+    if mul is not None:
+        qg = qg * mul.view(h_k * h_r, d)
+    qg = qg.reshape(b, q_len, h_k, h_r, d)
+    amax = qg.abs().amax(dim=(1, 3, 4)).clamp_min(1e-6)  # (b, h_k)
+    scale = amax / FP8_MAX
+    packed = (
+        (qg / scale[:, None, :, None, None])
+        .clamp(-FP8_MAX, FP8_MAX)
+        .to(torch.float8_e4m3fn)
+        .permute(0, 2, 1, 3, 4)
+        .reshape(b, h_k, q_len * h_r, d)
+    )
+    return packed, scale