mjolnir-kernel 0.1.1__tar.gz → 0.2.0__tar.gz

{mjolnir_kernel-0.1.1 → mjolnir_kernel-0.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: mjolnir-kernel
-Version: 0.1.1
+Version: 0.2.0
 Summary: Mjolnir: a paged decode-attention kernel for NVIDIA Jetson AGX Thor (sm_110) on tcgen05 / TMEM / TMA (CuTe DSL), with fp8 and 6-bit (e2m3) KV caches — packaged as the vLLM MJOLNIR attention backend.
 Project-URL: Homepage, https://github.com/thor-arena/mjolnir
 Project-URL: Documentation, https://github.com/thor-arena/mjolnir/tree/main/docs/kernel
@@ -141,7 +141,7 @@ runs on the fallback backend.
 | `MJOLNIR_ATTN_DRAFT_T5` | `1` | `0`: the drafter's layers (MTP / EAGLE) run on the fallback — isolates the draft-step cost (`mjolnir verify profile`). fp8 KV only: refused on fp6 KV, whose fallback cannot be CUDA-graphed |
 | `MJOLNIR_ATTN_NONCAUSAL` | `1` | `0`: a parallel drafter's non-causal block steps (DSpark / DFlash) go to the fallback instead of the decode kernel — fp8 KV only for the same reason as above |
 | `MJOLNIR_ATTN_FUSED_DRAFT` | `1` | `0`: tell vLLM the backend cannot advance draft metadata in place → one graph replay per draft step instead of the fused draft loop |
-| `MJOLNIR_LM_HEAD` | – | `fp8`: quantize the lm_head to e4m3 per row at load, Triton fp8-weight GEMM for the verify and every draft (`mjolnir.lmhead`; config `NVFP4_MJOLNIR_E2M3_LMH8`); the bf16 copy is freed automatically (untied heads only, +2.5 GB KV) |
+| `MJOLNIR_LM_HEAD` | – | `fp8`: quantize the lm_head to e4m3 per row at load, Triton fp8-weight GEMM for the verify and every draft (`mjolnir.lmhead`; config `NVFP4_MJOLNIR_E2M3_LMH8_VT64K`); the bf16 copy is freed automatically (untied heads only, +2.5 GB KV) |
 | `MJOLNIR_DRAFT_VOCAB` | – | path to a `draft_vocab.json` (`mjolnir calibrate --draft-vocab`): the drafter's logits only over those tokens (VocabTrim); `MJOLNIR_DRAFT_VOCAB_SIZE` keeps the first N (a shorter file is padded with the lowest unused ids, with a warning). The kept head stays in the drafter's native quant (NVFP4 4-bit in-kernel, FP8 row-slice) — never up-scaled to bf16; a bf16 head is the one exception (fp8 only when `MJOLNIR_LM_HEAD=fp8` is set). `MJOLNIR_DRAFT_VOCAB_NATIVE=0` forces the dequant→fp8 path |
  | `MJOLNIR_ATTN_VERBOSE` | – | `1`: print the kernel's compile-time notes (ring stages, TMEM columns) — the standalone harness always does |
  | `MJOLNIR_ATTN_BANNER` | `1` | `0`: skip the Mjolnir hammer banner in the vLLM boot log (`mjolnir.banner.emit`, once per start, coloured with the server's formatter) |
@@ -166,8 +166,8 @@ expands the pages a prefill / mixed step reads to e4m3 for the FA4 fallback
 (channel scales applied, one descale per head). The scales come from a calibration run on fp8 KV:
 
 ```bash
-python -m mjolnir.calibrate --config <serve.yaml> --out kv_scales.json   # in the image
-MJOLNIR_ATTN_KV_SCALES=kv_scales.json vllm serve ... --kv-cache-dtype fp6_e2m3 --attention-backend MJOLNIR
+python -m mjolnir.calibrate --config <serve.yaml> --out kv_scales/e2m3.json   # in the image
+MJOLNIR_ATTN_KV_SCALES=kv_scales/e2m3.json vllm serve ... --kv-cache-dtype fp6_e2m3 --attention-backend MJOLNIR
 ```
 
 (`mjolnir calibrate` / `mjolnir serve` do both for you —
@@ -185,6 +185,31 @@ How they compare in speed and accuracy:
 [`quant`](https://github.com/thor-arena/mjolnir/blob/main/docs/kernel/quant.md).
 Wiring a vLLM cache-write path for fp6 is an open contribution.
 
+## Platforms and architectures
+
+The kernel is written for one architecture today — Thor (`sm_110`) — but the
+package is organized by platform: `mjolnir/arch/` is the implemented-kernel
+registry (`mjolnir.arch.thor` re-exports the exact runner surface the backend
+consumes), and `mjolnir/platforms.py` is the torch-free leaf that maps a
+compute capability to it. Thor is implemented; DGX Spark (`sm_121`) and Jetson
+Orin (`sm_87`) are declared research — named in the matrix, no kernel yet.
+
+`kernel_for((11, 0))` returns `mjolnir.arch.thor`; any other compute
+capability raises `PlatformNotSupported` naming the requested `sm_XX` and the
+shipped platforms. The vLLM backend never sees that exception: its metadata
+builder consults the pure `unsupported_reason()` and takes the normal
+ineligible-config path — that layer's decode steps run on the fallback
+backend, the serve continues. On a host without a usable CUDA device,
+`MJOLNIR_ARCH` (`sm_110` or `11.0`) pins the architecture instead of the
+torch read (the standalone harness).
+
+There is no remembered platform state anywhere in the stack: the CLI side
+resolves its platform per invocation (flag > `$MJOLNIR_PLATFORM` > host
+detection — four source classes with provenance, the `/tmp/nv_jetson_model`
+model copy corroboration-only), so the kernel is always resolved from the
+device the run actually happens on. See
+[docs/faq/platforms.md](../../docs/faq/platforms.md).
+
 ## Tests
 
 From a checkout, on a Thor:
@@ -213,6 +238,8 @@ is printed as information.
 | `mjolnir/t5_decode.py` | the CuTe-DSL kernel, the split-KV merge, the harness and the emulator |
 | `mjolnir/runner.py` | plan buffers, compile cache, launch, KV views (no vLLM import) |
 | `mjolnir/backend.py` | the vLLM backend, metadata builder, impl |
+| `mjolnir/platforms.py` | arch → kernel resolution (`kernel_for`, `PlatformNotSupported`), `MJOLNIR_ARCH` detection — torch-free |
+| `mjolnir/arch/thor.py` | the Thor (`sm_110`) kernel surface: re-exports of the runner |
 | `mjolnir/__init__.py` | `register()` — the vLLM plugin entry point |
 
 Apache-2.0, except `mjolnir/t5_decode.py`, which is derived from an NVIDIA

{mjolnir_kernel-0.1.1 → mjolnir_kernel-0.2.0}/README.md

@@ -113,7 +113,7 @@ runs on the fallback backend.
 | `MJOLNIR_ATTN_DRAFT_T5` | `1` | `0`: the drafter's layers (MTP / EAGLE) run on the fallback — isolates the draft-step cost (`mjolnir verify profile`). fp8 KV only: refused on fp6 KV, whose fallback cannot be CUDA-graphed |
 | `MJOLNIR_ATTN_NONCAUSAL` | `1` | `0`: a parallel drafter's non-causal block steps (DSpark / DFlash) go to the fallback instead of the decode kernel — fp8 KV only for the same reason as above |
 | `MJOLNIR_ATTN_FUSED_DRAFT` | `1` | `0`: tell vLLM the backend cannot advance draft metadata in place → one graph replay per draft step instead of the fused draft loop |
-| `MJOLNIR_LM_HEAD` | – | `fp8`: quantize the lm_head to e4m3 per row at load, Triton fp8-weight GEMM for the verify and every draft (`mjolnir.lmhead`; config `NVFP4_MJOLNIR_E2M3_LMH8`); the bf16 copy is freed automatically (untied heads only, +2.5 GB KV) |
+| `MJOLNIR_LM_HEAD` | – | `fp8`: quantize the lm_head to e4m3 per row at load, Triton fp8-weight GEMM for the verify and every draft (`mjolnir.lmhead`; config `NVFP4_MJOLNIR_E2M3_LMH8_VT64K`); the bf16 copy is freed automatically (untied heads only, +2.5 GB KV) |
 | `MJOLNIR_DRAFT_VOCAB` | – | path to a `draft_vocab.json` (`mjolnir calibrate --draft-vocab`): the drafter's logits only over those tokens (VocabTrim); `MJOLNIR_DRAFT_VOCAB_SIZE` keeps the first N (a shorter file is padded with the lowest unused ids, with a warning). The kept head stays in the drafter's native quant (NVFP4 4-bit in-kernel, FP8 row-slice) — never up-scaled to bf16; a bf16 head is the one exception (fp8 only when `MJOLNIR_LM_HEAD=fp8` is set). `MJOLNIR_DRAFT_VOCAB_NATIVE=0` forces the dequant→fp8 path |
  | `MJOLNIR_ATTN_VERBOSE` | – | `1`: print the kernel's compile-time notes (ring stages, TMEM columns) — the standalone harness always does |
  | `MJOLNIR_ATTN_BANNER` | `1` | `0`: skip the Mjolnir hammer banner in the vLLM boot log (`mjolnir.banner.emit`, once per start, coloured with the server's formatter) |
@@ -138,8 +138,8 @@ expands the pages a prefill / mixed step reads to e4m3 for the FA4 fallback
 (channel scales applied, one descale per head). The scales come from a calibration run on fp8 KV:
 
 ```bash
-python -m mjolnir.calibrate --config <serve.yaml> --out kv_scales.json   # in the image
-MJOLNIR_ATTN_KV_SCALES=kv_scales.json vllm serve ... --kv-cache-dtype fp6_e2m3 --attention-backend MJOLNIR
+python -m mjolnir.calibrate --config <serve.yaml> --out kv_scales/e2m3.json   # in the image
+MJOLNIR_ATTN_KV_SCALES=kv_scales/e2m3.json vllm serve ... --kv-cache-dtype fp6_e2m3 --attention-backend MJOLNIR
 ```
 
 (`mjolnir calibrate` / `mjolnir serve` do both for you —
@@ -157,6 +157,31 @@ How they compare in speed and accuracy:
 [`quant`](https://github.com/thor-arena/mjolnir/blob/main/docs/kernel/quant.md).
 Wiring a vLLM cache-write path for fp6 is an open contribution.
 
+## Platforms and architectures
+
+The kernel is written for one architecture today — Thor (`sm_110`) — but the
+package is organized by platform: `mjolnir/arch/` is the implemented-kernel
+registry (`mjolnir.arch.thor` re-exports the exact runner surface the backend
+consumes), and `mjolnir/platforms.py` is the torch-free leaf that maps a
+compute capability to it. Thor is implemented; DGX Spark (`sm_121`) and Jetson
+Orin (`sm_87`) are declared research — named in the matrix, no kernel yet.
+
+`kernel_for((11, 0))` returns `mjolnir.arch.thor`; any other compute
+capability raises `PlatformNotSupported` naming the requested `sm_XX` and the
+shipped platforms. The vLLM backend never sees that exception: its metadata
+builder consults the pure `unsupported_reason()` and takes the normal
+ineligible-config path — that layer's decode steps run on the fallback
+backend, the serve continues. On a host without a usable CUDA device,
+`MJOLNIR_ARCH` (`sm_110` or `11.0`) pins the architecture instead of the
+torch read (the standalone harness).
+
+There is no remembered platform state anywhere in the stack: the CLI side
+resolves its platform per invocation (flag > `$MJOLNIR_PLATFORM` > host
+detection — four source classes with provenance, the `/tmp/nv_jetson_model`
+model copy corroboration-only), so the kernel is always resolved from the
+device the run actually happens on. See
+[docs/faq/platforms.md](../../docs/faq/platforms.md).
+
 ## Tests
 
 From a checkout, on a Thor:
@@ -185,6 +210,8 @@ is printed as information.
 | `mjolnir/t5_decode.py` | the CuTe-DSL kernel, the split-KV merge, the harness and the emulator |
 | `mjolnir/runner.py` | plan buffers, compile cache, launch, KV views (no vLLM import) |
 | `mjolnir/backend.py` | the vLLM backend, metadata builder, impl |
+| `mjolnir/platforms.py` | arch → kernel resolution (`kernel_for`, `PlatformNotSupported`), `MJOLNIR_ARCH` detection — torch-free |
+| `mjolnir/arch/thor.py` | the Thor (`sm_110`) kernel surface: re-exports of the runner |
 | `mjolnir/__init__.py` | `register()` — the vLLM plugin entry point |
 
 Apache-2.0, except `mjolnir/t5_decode.py`, which is derived from an NVIDIA

{mjolnir_kernel-0.1.1 → mjolnir_kernel-0.2.0}/mjolnir/__init__.py

@@ -24,7 +24,7 @@ from __future__ import annotations
 import os
 
 __all__ = ["BACKEND_NAME", "BACKEND_PATH", "__version__", "register"]
-__version__ = "0.1.1"
+__version__ = "0.2.0"
 
 BACKEND_NAME = "MJOLNIR"
 BACKEND_PATH = "mjolnir.backend.MjolnirAttentionBackend"

mjolnir_kernel-0.2.0/mjolnir/arch/__init__.py

@@ -0,0 +1,7 @@
+"""Implemented-kernel registry: one subpackage per supported architecture
+(``mjolnir.arch.thor``).
+
+``mjolnir.platforms.kernel_for`` maps a compute capability to a dotted module
+path here. Importing a subpackage pulls the CUTLASS DSL (heavy) — only do it
+when a device is real.
+"""

mjolnir_kernel-0.2.0/mjolnir/arch/thor.py

@@ -0,0 +1,42 @@
+"""The Thor (sm_110) kernel: re-exports of the exact runner surface the vLLM
+backend consumes. Adding an architecture = adding a sibling subpackage with
+the same surface and a ``_SHIPPED`` row in ``mjolnir.platforms``:
+
+* ``Runner`` — the decode-runner class the backend instantiates; the one
+  name a family subpackage must define itself (all others are re-exports of
+  shared, family-independent machinery from :mod:`mjolnir.runner`);
+* the shared helpers every family consumes unchanged: ``StepPlan``,
+  ``SUPPORTED_PAGE_SIZES``, ``detect_kv_layout``, ``env_flag``,
+  ``split_kv_cache``.
+
+The backend resolves the class through :func:`mjolnir.platforms.load_kernel`
+at build time — it never imports a family module directly, so landing a
+family edits nothing but ``platforms._SHIPPED`` and its own subpackage.
+Importing this module pulls the CUTLASS DSL (heavy): only do it when a
+device is real. The compiled-kernel cache fingerprint
+(``runner._cache_fingerprint``) hashes the kernel module the shared runner
+imports — family-blind until the runner itself is split per family: a
+second family must extend the fingerprint (or key its own cache) or it
+would reuse this family's objects.
+"""
+
+from mjolnir.runner import (
+    SUPPORTED_PAGE_SIZES,
+    StepPlan,
+    T5DecodeRunner,
+    detect_kv_layout,
+    env_flag,
+    split_kv_cache,
+)
+
+Runner = T5DecodeRunner
+
+__all__ = [
+    "SUPPORTED_PAGE_SIZES",
+    "Runner",
+    "StepPlan",
+    "T5DecodeRunner",
+    "detect_kv_layout",
+    "env_flag",
+    "split_kv_cache",
+]

{mjolnir_kernel-0.1.1 → mjolnir_kernel-0.2.0}/mjolnir/backend.py

@@ -35,20 +35,27 @@ from __future__ import annotations
 
 import copy
 import logging
-from typing import Any, ClassVar
+from typing import TYPE_CHECKING, Any, ClassVar
 
 import torch
 
 from . import banner
+from .platforms import detect, load_kernel, unsupported_reason
 from .runner import (
     SUPPORTED_PAGE_SIZES,
     StepPlan,
-    T5DecodeRunner,
     detect_kv_layout,
     env_flag,
     split_kv_cache,
 )
 
+if TYPE_CHECKING:
+    # The annotations name the Thor runner; a second family's runner duck-types
+    # the same surface and arrives through its own arch subpackage. At runtime
+    # the class is resolved via platforms.load_kernel — never by importing a
+    # family module here (see mjolnir.arch).
+    from .runner import T5DecodeRunner
+
 logger = logging.getLogger("vllm.mjolnir")  # under vLLM's handler
 
 # --------------------------------------------------------------------------
@@ -238,6 +245,13 @@ def ineligible_reason(kv_cache_spec, layer_names, vllm_config) -> str | None:
     return None
 
 
+def _unsupported_arch_reason(device) -> str | None:
+    """None on a supported compute capability; else the platform registry's
+    verdict (torch-free leaf), routed into the same t5=None fallback as every
+    ineligible config."""
+    return unsupported_reason(torch.cuda.get_device_capability(device))
+
+
 # --------------------------------------------------------------------------
 # fp6 (e2m3) KV: preconditions and per-channel scales
 # --------------------------------------------------------------------------
@@ -395,6 +409,10 @@ class MjolnirMetadataBuilder(_BaseBuilder):  # type: ignore[misc, valid-type]
         self._fp6 = self._kv_format == "e2m3"
         self._page = kv_cache_spec.block_size
         self._t5_reason = ineligible_reason(kv_cache_spec, layer_names, vllm_config)
+        if self._t5_reason is None:
+            # platform guard (surgical): a GPU without a Mjolnir kernel
+            # degrades exactly like any other ineligible config
+            self._t5_reason = _unsupported_arch_reason(device)
         # Diagnosis knobs for the draft path (MTP / EAGLE drafter layers):
         # MJOLNIR_ATTN_DRAFT_T5=0 runs every drafter step on the fallback;
         # MJOLNIR_ATTN_FUSED_DRAFT=0 makes vLLM rebuild the draft metadata
@@ -438,7 +456,7 @@ class MjolnirMetadataBuilder(_BaseBuilder):  # type: ignore[misc, valid-type]
         q_mode = {"bf16": "bf16", "bfloat16": "bf16", "fp8": "fp8", "e4m3": "fp8"}[q_mode]
         if self._fp6 and q_mode != "fp8":
             raise ValueError("fp6_e2m3 KV runs on the fp8-Q path only: unset MJOLNIR_ATTN_Q_DTYPE")
-        self.t5 = T5DecodeRunner(
+        self.t5 = load_kernel(detect()).Runner(
             num_heads=num_heads,
             num_kv_heads=kv_cache_spec.num_kv_heads,
             head_size=kv_cache_spec.head_size,

{mjolnir_kernel-0.1.1 → mjolnir_kernel-0.2.0}/mjolnir/calibrate.py

@@ -9,14 +9,14 @@ 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
+        --out /out/e2m3.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``.
+with its own drafter to ``kv_scales/<method>.kv_scales.json``.
 """
 
 from __future__ import annotations

mjolnir_kernel-0.2.0/mjolnir/platforms.py

@@ -0,0 +1,103 @@
+"""Kernel-side platform leaf: which architecture is this GPU, and which kernel
+implements it.
+
+Import-safe everywhere (torch is imported inside :func:`detect` only, never at
+module level), so the vLLM backend can consult it and any host can read the
+registry without pulling the GPU stack.
+
+``kernel_for`` resolves a compute capability to a dotted module path under
+:mod:`mjolnir.arch`; :func:`load_kernel` imports it — heavy (the CUTLASS DSL
+comes with it), so only call it when a device is real.
+"""
+
+from __future__ import annotations
+
+import os
+from types import ModuleType
+
+# The platforms Mjolnir ships support for: (name, sm, implemented). A row
+# without a kernel (spark / orin) exists so the unsupported-arch error can
+# name the whole matrix; its kernel arrives as mjolnir/arch/<name>.py. The
+# kernel package cannot import the CLI's registry, so the sm rows are
+# mirrored here — the two sides are pinned in sync by
+# tests/cli/test_kernel_platforms.py::test_kernel_and_cli_registries_agree_on_the_sm_rows.
+_SHIPPED: tuple[tuple[str, tuple[int, int], bool], ...] = (
+    ("thor", (11, 0), True),
+    ("spark", (12, 1), False),
+    ("orin", (8, 7), False),
+)
+
+
+class PlatformNotSupported(RuntimeError):
+    """The GPU's compute capability has no Mjolnir kernel."""
+
+
+def detect() -> tuple[int, int]:
+    """The device's compute capability: ``$MJOLNIR_ARCH`` (``sm_110`` or
+    ``11.0``) wins, else torch reads it off the CUDA device."""
+    raw = os.environ.get("MJOLNIR_ARCH", "").strip()
+    if raw:
+        return _parse_arch(raw)
+    try:
+        import torch
+    except ImportError as e:
+        raise RuntimeError(
+            "Mjolnir cannot detect the GPU compute capability (torch is not importable): "
+            "set MJOLNIR_ARCH, e.g. MJOLNIR_ARCH=sm_110"
+        ) from e
+    if not torch.cuda.is_available():
+        raise RuntimeError(
+            "Mjolnir cannot detect the GPU compute capability (no CUDA device): "
+            "set MJOLNIR_ARCH, e.g. MJOLNIR_ARCH=sm_110"
+        )
+    cap = torch.cuda.get_device_capability()
+    return int(cap[0]), int(cap[1])
+
+
+def _parse_arch(raw: str) -> tuple[int, int]:
+    """``sm_110`` / ``11.0`` / ``8.7`` -> (major, minor)."""
+    s = raw.strip().lower().removeprefix("sm_")
+    if s.isdigit():  # the sm_ form packs the minor into the last digit: sm_110 = 11.0
+        return int(s[:-1]), int(s[-1])
+    major, sep, minor = s.partition(".")
+    if sep and major.isdigit() and minor.isdigit():
+        return int(major), int(minor)
+    raise RuntimeError(f"MJOLNIR_ARCH={raw!r}: expected 'sm_<major><minor>' (e.g. sm_110) or '<major>.<minor>'")
+
+
+def kernel_for(cc: tuple[int, int]) -> str:
+    """The dotted module path of the kernel implementing ``cc`` — currently
+    only ``mjolnir.arch.thor``; anything else raises :class:`PlatformNotSupported`."""
+    for name, sm, implemented in _SHIPPED:
+        if sm == tuple(cc):
+            if implemented:
+                return f"mjolnir.arch.{name}"
+            break  # shipped, but the kernel has not landed: unsupported for now
+    raise PlatformNotSupported(_unsupported_msg(cc))
+
+
+def load_kernel(cc: tuple[int, int]) -> ModuleType:
+    """Import and return the kernel module for ``cc`` (heavy: pulls the
+    CUTLASS DSL — only call this when a device is real)."""
+    import importlib
+
+    return importlib.import_module(kernel_for(cc))
+
+
+def unsupported_reason(cc: tuple[int, int]) -> str | None:
+    """None when ``cc`` has a kernel; else the reason the vLLM backend routes
+    into its ineligible-config degrade path (names the requested sm and the
+    fallback that takes over). Pure — no device access."""
+    for _name, sm, implemented in _SHIPPED:
+        if sm == tuple(cc) and implemented:
+            return None
+    return _unsupported_msg(cc) + " Decode steps run on the fallback backend."
+
+
+def _unsupported_msg(cc: tuple[int, int]) -> str:
+    impl = [f"{n} (sm_{maj}{minr})" for n, (maj, minr), ok in _SHIPPED if ok]
+    research = [f"{n} (sm_{maj}{minr})" for n, (maj, minr), ok in _SHIPPED if not ok]
+    msg = f"Mjolnir has no kernel for sm_{cc[0]}{cc[1]} — implemented: {', '.join(impl)}"
+    if research:
+        msg += f"; research (no kernel yet): {', '.join(research)}"
+    return msg