mjolnir-kernel 0.1.0__py3-none-any.whl

mjolnir/banner.py

@@ -0,0 +1,127 @@
+"""The Mjolnir boot banner.
+
+Owned by the ``mjolnir`` package (not a vLLM patch): when the Mjolnir
+backend serves a model, ``emit()`` logs the hammer exactly once, with the
+info block (``Mjolnir v<version>``, author, repo, and the serving config
+name when the CLI set ``MJOLNIR_CONFIG``) placed to the art's right,
+vertically centred. It goes through vLLM's own logger, so it follows the
+server's formatter the same way vLLM's own boot logo does (ANSI colour
+only under a color formatter), and it can never break a server's
+start-up. ``python -m mjolnir.banner`` previews it without vLLM.
+"""
+
+from __future__ import annotations
+
+import os
+import sys
+
+__all__ = ["ART", "HEAD_START", "banner", "emit"]
+
+# The hammer. Committed art — do not re-flow.
+ART = [
+    "         ▓██",
+    "          █▒██",
+    "          █  ░█▒",
+    "          █ ▓█",
+    "          ██ █",
+    "          █  █░",
+    "         ██  ▒█",
+    "     █████▓  ██████",
+    "     █    ▓▒      ██",
+    "     ▒████    ████▓",
+    "          ░███",
+]
+# Rows at/below this index make up the head (brand orange); above it, the handle.
+HEAD_START = 7
+
+_HEAD = "\033[1;93m"  # bold brand orange (the hammer head)
+_WHITE = "\033[1m"  # bold default foreground (handle, author)
+_RESET = "\033[0m"
+_GAP = "   "
+_AUTHOR = "by Mieszko Syty (ms1design)"
+_URL = "https://github.com/thor-arena/mjolnir"
+
+
+def banner(version: str, color: bool = False, model: str | None = None) -> str:
+    """Render the banner: the art, with the info block (version, author,
+    repo, and ``Model: <config>`` when given) to its right, vertically
+    centred against the art. The message is one block of lines (no leading
+    or trailing empty line — ``emit()`` renders those as their own empty
+    log records, the way the vLLM boot logo's template does). ``color``
+    switches on the ANSI codes (head in the brand orange, the rest bold
+    white)."""
+    info = [f"Mjolnir v{version}", _AUTHOR, _URL]
+    if model:
+        info.append(f"Model: {model}")
+    n, k = len(ART), len(info)
+    start = (n - k) // 2
+    width = max(len(line) for line in ART)
+    rows = []
+    for i, line in enumerate(ART):
+        idx = i - start
+        if not color:
+            if 0 <= idx < k:
+                rows.append(line.ljust(width) + _GAP + info[idx])
+            else:
+                rows.append(line)
+            continue
+        art = line.ljust(width) + _GAP if 0 <= idx < k else line
+        segs = [(_HEAD if i >= HEAD_START else _WHITE, art)]
+        if 0 <= idx < k:
+            if idx == 0:
+                segs += [(_HEAD, "Mjolnir"), (_WHITE, f" v{version}")]
+            else:
+                segs += [(_WHITE, info[idx])]
+        rows.append("".join(f"{code}{text}{_RESET}" for code, text in segs))
+    return "\n".join(rows)
+
+
+_shown = False
+
+
+def emit() -> None:
+    """Log the banner once, through vLLM's logger when importable (its
+    formatter decides the ANSI colour, exactly like vLLM's own boot logo)
+    and a stdlib logger otherwise. The art is one multi-line record; the
+    empty lines on top and bottom are rendered as empty-string records
+    (vLLM's formatter turns them into bare prefixed lines, like the boot
+    logo's tail). The serving config name rides in ``MJOLNIR_CONFIG`` (set
+    by the ``mjolnir`` CLI's serve path). Never raises: a banner must not
+    be able to break a server's start."""
+    global _shown
+    if _shown:
+        return
+    _shown = True
+    try:
+        from mjolnir import __version__
+
+        model = os.environ.get("MJOLNIR_CONFIG", "").strip() or None
+        try:
+            from vllm.logger import current_formatter_type, init_logger
+
+            lgr = init_logger("vllm.mjolnir")
+            color = current_formatter_type(lgr) == "color"
+        except ImportError:
+            import logging
+
+            lgr = logging.getLogger("vllm.mjolnir")
+            color = sys.stderr.isatty()
+        lgr.info("")
+        lgr.info("%s", banner(__version__, color=color, model=model))
+        lgr.info("")
+    except Exception:  # noqa: BLE001 - the banner is decoration, not a dependency
+        pass
+
+
+if __name__ == "__main__":  # `python -m mjolnir.banner [--plain]`: preview without vLLM
+    import mjolnir
+
+    plain = "--plain" in sys.argv[1:]
+    msg = banner(
+        mjolnir.__version__,
+        color=sys.stdout.isatty() and not plain,
+        model=os.environ.get("MJOLNIR_CONFIG", "").strip() or None,
+    )
+    print("")
+    print(msg)
+    print("")

mjolnir/calibrate.py

@@ -0,0 +1,182 @@
+"""Calibrate per-(layer, KV head, channel) K/V scales for the fp6 (e2m3) KV cache.
+
+Runs the model in-process on vLLM with the Mjolnir backend on **fp8** KV and
+``MJOLNIR_ATTN_KV_CALIBRATE=1``: every K/V row the backend writes to the
+cache is folded into a running per-channel amax (``mjolnir.kvscales``). After
+the engine's own warm-up and profiling runs the accumulators are zeroed, a
+varied prompt set is generated (prose, code, math, multilingual, long
+contexts), and the result is written as a ``mjolnir.kv_scales/2`` file —
+what ``MJOLNIR_ATTN_KV_SCALES`` points ``--kv-cache-dtype fp6_e2m3`` at.
+
+    python -m mjolnir.calibrate --config /configs/Qwen/Qwen3.8-27B/NVFP4_MJOLNIR/config.yaml \\
+        --out /out/kv_scales.json --download-dir /data/models/huggingface
+    python -m mjolnir.calibrate --model Qwen/Qwen3-0.6B --out /tmp/scales.json   # small models / tests
+
+The scales cover every attention layer that runs on the Mjolnir backend,
+drafters included (drafting writes its KV too): the built-in MTP head, and an
+external drafter whose ``speculative-config.attention_backend`` is MJOLNIR
+(DSpark: its 5 layers come out as layers 64-68). ``mjolnir calibrate`` runs this in the image and writes a config
+with its own drafter to ``<config>.kv_scales.json``.
+"""
+
+from __future__ import annotations
+
+import argparse
+import dataclasses
+import os
+import sys
+import time
+
+TOPICS = [
+    "Explain how a transformer's attention mechanism works, step by step, with a small numeric example.",
+    "Write a Python function that parses an ISO-8601 timestamp without external libraries, then test it.",
+    "def quicksort(xs):\n    # complete this implementation and explain its complexity\n",
+    "Summarize the causes and consequences of the 1929 stock market crash in five paragraphs.",
+    "Solve: a train leaves at 9:40 at 72 km/h, another at 10:05 at 90 km/h on the same track. When do they meet?",
+    "Translate into German, French and Polish: 'The quick brown fox jumps over the lazy dog near the river bank.'",
+    "Write a SQL query that returns the top three customers by revenue per region, with window functions.",
+    "Compose a short story about a lighthouse keeper who finds a message in a bottle.",
+    "List the differences between TCP and UDP and when you would use each, as a table.",
+    "Prove that the square root of two is irrational.",
+    "Describe the architecture of the NVIDIA Blackwell GPU memory hierarchy for a systems engineer.",
+    "Rust: implement a thread-safe bounded queue with condition variables; explain the invariants.",
+    "What are the main ideas of Kant's Critique of Pure Reason? Answer for a first-year student.",
+    "日本の四季について、それぞれの季節の特徴を説明してください。",
+    "Напиши краткое руководство по выращиванию томатов на балконе.",
+    'Given the JSON {"a": [1, 2, {"b": null}], "c": "x"}, write a jq expression that extracts b and explain it.',
+]
+
+LONG_SOURCE = (
+    "In the long history of computing, memory bandwidth has repeatedly set the pace of progress. "
+    "Each generation of accelerators added arithmetic faster than it added bytes per second, and "
+    "every workload that streams data once — decoding a language model token by token is one — "
+    "became a question of how few bytes it could read. Caches help only when data is reused; a "
+    "key-value cache of a long conversation is read in full for every new token, so its format "
+    "decides the speed. "
+)
+
+
+def long_prompt(tokens: int, i: int) -> str:
+    """A long context (~`tokens` words) ending in a question about it."""
+    words_per_rep = len(LONG_SOURCE.split())
+    reps = max(1, tokens // words_per_rep)
+    body = "".join(f"[section {k + i}] {LONG_SOURCE}" for k in range(reps))
+    return body + "\nQuestion: summarize the argument above in three sentences and name its weakest step."
+
+
+def _engine_parser():
+    from vllm.engine.arg_utils import EngineArgs
+
+    try:
+        from vllm.utils.argparse_utils import FlexibleArgumentParser
+    except ImportError:  # older vLLM
+        from vllm.utils import FlexibleArgumentParser
+    return EngineArgs.add_cli_args(FlexibleArgumentParser()), EngineArgs
+
+
+def engine_kwargs_from_config(path: str) -> dict:
+    """A `vllm serve --config` YAML as LLM(...) kwargs, parsed by vLLM's own
+    argument parser — so every value gets the type `vllm serve` gives it
+    (``16K`` -> 16384, ``kv-cache-dtype-skip-layers: sliding_window`` -> a
+    list, nested dicts -> configs). Serve-only keys (port, served-model-name,
+    tool parsers, ...) are dropped; only values that differ from the engine
+    defaults are returned."""
+    import tempfile
+
+    import yaml
+
+    parser, engine_args = _engine_parser()
+    data = yaml.safe_load(open(path)) or {}
+    known = {k: v for k, v in data.items() if "--" + str(k).replace("_", "-") in parser._option_string_actions}
+    with tempfile.NamedTemporaryFile("w", suffix=".yaml", delete=False) as f:
+        yaml.safe_dump(known, f)
+    try:
+        ns = vars(parser.parse_args(["--config", f.name]))
+    finally:
+        os.unlink(f.name)
+    base = vars(parser.parse_args([]))
+    fields = {fld.name for fld in dataclasses.fields(engine_args)}
+    kw = {}
+    for k, v in ns.items():
+        if k not in fields:
+            continue
+        try:
+            same = v == base.get(k)
+        except Exception:  # noqa: BLE001 - uncomparable config objects: keep the value
+            same = False
+        if not same:
+            kw[k] = v
+    return kw
+
+
+def main(argv=None) -> int:
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    src = ap.add_mutually_exclusive_group(required=True)
+    src.add_argument("--config", help="vllm serve YAML (the serving config to calibrate)")
+    src.add_argument("--model", help="a model id (no config: defaults for everything else)")
+    ap.add_argument("--out", required=True, help="where to write the mjolnir.kv_scales/2 JSON")
+    ap.add_argument("--download-dir", default=None)
+    ap.add_argument("--max-model-len", type=int, default=32768)
+    ap.add_argument("--gpu-mem", type=float, default=None, help="gpu_memory_utilization (default: the config's)")
+    ap.add_argument("--prompts", type=int, default=48, help="short prompts (cycled through the topic list)")
+    ap.add_argument("--long", type=int, default=6, help="long-context prompts")
+    ap.add_argument("--long-tokens", type=int, default=6000, help="approx. words per long prompt")
+    ap.add_argument("--max-tokens", type=int, default=96, help="tokens generated per prompt")
+    args = ap.parse_args(argv)
+
+    os.environ["MJOLNIR_ATTN_KV_CALIBRATE"] = "1"
+    os.environ.setdefault("VLLM_ENABLE_V1_MULTIPROCESSING", "0")  # the accumulators live in this process
+    import mjolnir
+    from mjolnir import kvscales
+
+    mjolnir.register()
+    from vllm import LLM, SamplingParams
+
+    kw = engine_kwargs_from_config(args.config) if args.config else {"model": args.model, "dtype": "bfloat16"}
+    for k in ("served_model_name",):
+        kw.pop(k, None)
+    kw.update(
+        kv_cache_dtype="fp8",  # calibrate on the fp8 path; the K/V are observed before quantization
+        enforce_eager=True,
+        max_model_len=min(int(kw.get("max_model_len") or args.max_model_len), args.max_model_len),
+        seed=0,
+    )
+    kw["attention_backend"] = "MJOLNIR" if mjolnir.has_native_member() else "CUSTOM"
+    kw.pop("attention_config", None)
+    if args.gpu_mem is not None:
+        kw["gpu_memory_utilization"] = args.gpu_mem
+    if args.download_dir:
+        kw["download_dir"] = args.download_dir
+    print(f"calibrate: LLM({', '.join(f'{k}={v!r}' for k, v in sorted(kw.items()))})", flush=True)
+    try:
+        llm = LLM(**kw)
+    except TypeError:
+        backend = kw.pop("attention_backend")
+        llm = LLM(attention_config={"backend": backend}, **kw)
+
+    kvscales.reset()  # drop what the warm-up / profiling runs wrote
+    prompts = [TOPICS[i % len(TOPICS)] for i in range(args.prompts)]
+    budget = max(256, kw["max_model_len"] - args.max_tokens - 64)
+    prompts += [long_prompt(min(args.long_tokens, budget * 3 // 4), i) for i in range(args.long)]
+    sp = SamplingParams(temperature=0.8, top_p=0.95, max_tokens=args.max_tokens, seed=0)
+    t0 = time.time()
+    outs = llm.generate(prompts, sp)
+    toks = sum(len(o.prompt_token_ids) + len(o.outputs[0].token_ids) for o in outs)
+    doc = kvscales.dump(args.out, model=str(kw.get("model", "")))
+    if not doc["amax"]:
+        print("calibrate: no K/V was observed — did the Mjolnir backend run?", file=sys.stderr)
+        return 1
+    print(f"calibrate: {len(prompts)} prompts, {toks} tokens in {time.time() - t0:.0f}s", flush=True)
+    for name, e in doc["amax"].items():
+        k_max, v_max = max(max(r) for r in e["k"]), max(max(r) for r in e["v"])
+        k_med = sorted(x for r in e["k"] for x in r)[len(e["k"]) * len(e["k"][0]) // 2]
+        print(
+            f"  {name}: K amax {k_max:8.3f} (median channel {k_med:7.3f})  V amax {v_max:8.3f}"
+            f"  ({len(e['k'])} heads x {len(e['k'][0])} channels)"
+        )
+    print(f"wrote {args.out}  ({len(doc['amax'])} layers)")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

mjolnir/draftvocab.py

@@ -0,0 +1,150 @@
+"""Count which tokens the target model generates — the drafter's vocabulary
+for VocabTrim (``MJOLNIR_DRAFT_VOCAB``, see ``mjolnir.lmhead``).
+
+Runs the serving config in-process (vLLM ``LLM``), samples a varied prompt
+set (several completions per prompt, temperature 0.7), counts the generated
+token ids and writes them most frequent first, with the coverage — the share
+of generated tokens inside the top-K — for a few K. Coverage is measured
+held out (ranking from half the prompts, scored on the other half), since the
+in-sample number is 100% as soon as K passes the distinct-token count. Ids the
+run never generated follow in id order (BPE merge rank), so a top-K always has
+K tokens. VocabTrim found the target's own generations the best source
+(arXiv 2506.22694).
+
+    python -m mjolnir.draftvocab --config /configs/Qwen/Qwen3.8-27B/NVFP4_MJOLNIR_E2M3/config.yaml \\
+        --out /out/draft_vocab.json --download-dir /data/models/huggingface
+
+File: ``{"format": "mjolnir.draft_vocab/1", "model": ..., "generated_tokens": N,
+"coverage": {"32768": 0.997, ...}, "tokens": [ids, most frequent first]}``.
+Special tokens (EOS, the chat / think markers) are always included.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import sys
+import time
+from collections import Counter
+
+FORMAT = "mjolnir.draft_vocab/1"
+COVERAGE_K = (4096, 8192, 16384, 32768, 49152, 65536)
+
+EXTRA_TOPICS = [
+    "Write a detailed README for a command-line tool that benchmarks GPUs.",
+    "Explain gradient descent to a high-school student, then give the math.",
+    "Write TypeScript code for a debounced search box in React, with tests.",
+    "Summarize the plot of a classic novel of your choice and analyze its themes.",
+    "Give a step-by-step recipe for sourdough bread with timings.",
+    "Compare PostgreSQL and SQLite for an embedded device; answer as a table and a recommendation.",
+    "Write a formal email declining a meeting and proposing two alternative times.",
+    "Debug this Python: `for i in range(10): print(i/0)` — explain the error and fix it.",
+    "Describe the water cycle in Spanish, then in English.",
+    "Write a C function that reverses a linked list in place; explain each line.",
+    "List ten interview questions for a senior backend engineer with model answers.",
+    "Explain how TLS 1.3 handshakes work, with a sequence of messages.",
+]
+
+
+def rank(counts: Counter, special: set, vocab: int) -> list[int]:
+    """Special tokens, then generated ids by count, then the rest in id order."""
+    out = [t for t in sorted(special) if t not in counts] + [t for t, _ in counts.most_common()]
+    seen = set(out)
+    return out + [t for t in range(vocab) if t not in seen]
+
+
+def held_out_coverage(train: Counter, test: Counter, special: set, vocab: int) -> dict[str, float]:
+    """Share of ``test``'s tokens inside the top-K of a ranking built from ``train``."""
+    total = sum(test.values())
+    if not total:
+        return {str(k): 0.0 for k in COVERAGE_K}
+    order = rank(train, special, vocab)
+    pos = {t: i for i, t in enumerate(order)}
+    out = {}
+    for k in COVERAGE_K:
+        out[str(k)] = round(sum(c for t, c in test.items() if pos.get(t, vocab) < k) / total, 5)
+    return out
+
+
+def main(argv=None) -> int:
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    src = ap.add_mutually_exclusive_group(required=True)
+    src.add_argument("--config", help="vllm serve YAML")
+    src.add_argument("--model", help="a model id")
+    ap.add_argument("--out", required=True)
+    ap.add_argument("--download-dir", default=None)
+    ap.add_argument("--max-model-len", type=int, default=16384)
+    ap.add_argument("--prompts", type=int, default=192)
+    ap.add_argument("--samples", type=int, default=4, help="completions per prompt")
+    ap.add_argument("--max-tokens", type=int, default=768)
+    ap.add_argument("--long", type=int, default=4, help="long-context prompts")
+    args = ap.parse_args(argv)
+
+    os.environ.setdefault("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
+    for k in ("MJOLNIR_DRAFT_VOCAB", "MJOLNIR_LM_HEAD"):  # count with the plain heads
+        os.environ.pop(k, None)
+    import mjolnir
+    from mjolnir.calibrate import TOPICS, engine_kwargs_from_config, long_prompt
+
+    mjolnir.register()
+    from vllm import LLM, SamplingParams
+
+    kw = engine_kwargs_from_config(args.config) if args.config else {"model": args.model, "dtype": "bfloat16"}
+    kw.pop("served_model_name", None)
+    kw["max_model_len"] = min(int(kw.get("max_model_len") or args.max_model_len), args.max_model_len)
+    if args.download_dir:
+        kw["download_dir"] = args.download_dir
+    print(f"draftvocab: LLM({', '.join(f'{k}={v!r}' for k, v in sorted(kw.items()))})", flush=True)
+    llm = LLM(**kw)
+    topics = TOPICS + EXTRA_TOPICS
+    prompts = [topics[i % len(topics)] for i in range(args.prompts)]
+    prompts += [long_prompt(min(4000, kw["max_model_len"] // 2), i) for i in range(args.long)]
+    msgs = [[{"role": "user", "content": p}] for p in prompts]
+    sp = SamplingParams(temperature=0.7, top_p=0.95, max_tokens=args.max_tokens, n=args.samples, seed=0)
+    t0 = time.time()
+    outs = llm.chat(msgs, sp)
+    counts: Counter = Counter()
+    halves: tuple[Counter, Counter] = (Counter(), Counter())
+    for i, o in enumerate(outs):
+        for c in o.outputs:
+            counts.update(c.token_ids)
+            halves[i % 2].update(c.token_ids)
+    total = sum(counts.values())
+    if not total:
+        print("draftvocab: nothing generated", file=sys.stderr)
+        return 1
+    tok = llm.get_tokenizer()
+    special = set(getattr(tok, "all_special_ids", []) or [])
+    try:
+        special |= set(tok.get_added_vocab().values())
+    except Exception:  # noqa: BLE001
+        pass
+    vocab = len(tok)
+    ranked = rank(counts, special, vocab)
+    coverage = held_out_coverage(halves[0], halves[1], special, vocab)
+    doc = {
+        "format": FORMAT,
+        "model": str(kw.get("model", "")),
+        "generated_tokens": total,
+        "distinct_tokens": len(counts),
+        "coverage": coverage,
+        "coverage_kind": "held-out",
+        "measured_tokens": len(counts),
+        "tokens": ranked,
+    }
+    os.makedirs(os.path.dirname(os.path.abspath(args.out)), exist_ok=True)
+    with open(args.out, "w") as f:
+        json.dump(doc, f)
+    print(
+        f"draftvocab: {len(prompts)} prompts x {args.samples}, {total} tokens generated in {time.time() - t0:.0f}s, "
+        f"{len(counts)} distinct"
+    )
+    for k in COVERAGE_K:
+        print(f"  top {k:>6}: {coverage[str(k)] * 100:6.2f}% of held-out generated tokens")
+    print(f"wrote {args.out}")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())