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

{mjolnir_kernel-0.1.0 → mjolnir_kernel-0.2.0}/.gitignore

@@ -3,6 +3,12 @@ __pycache__/
 .venv/
 .pytest_cache/
 .ruff_cache/
+.env
+# coverage artifacts (the CI test stage writes these; they regenerate)
+.coverage
+.coverage.*
+coverage.xml
+.coverage_html/
 dist/
 build/
 *.egg-info/

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: mjolnir-kernel
-Version: 0.1.0
+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
@@ -18,12 +18,10 @@ Classifier: Intended Audience :: Developers
 Classifier: Intended Audience :: Science/Research
 Classifier: Operating System :: POSIX :: Linux
 Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.10
-Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
-Requires-Python: >=3.10
+Requires-Python: >=3.12
 Provides-Extra: dsl
 Requires-Dist: nvidia-cutlass-dsl>=4.2; extra == 'dsl'
 Description-Content-Type: text/markdown
@@ -36,7 +34,7 @@ packaged as a vLLM attention backend.
 
 ```bash
 vllm serve <model> --attention-backend MJOLNIR      # the mjolnir-vllm image
-vllm serve <model> --attention-backend CUSTOM       # stock vLLM 0.31.0 + this package
+vllm serve <model> --attention-backend CUSTOM       # stock vLLM 0.32.0 + this package
 ```
 
 This package is the kernel only. Serving, calibration and benchmarking are the
@@ -89,10 +87,10 @@ fused QKV buffer.
 
 ## Install
 
-**On the image** (`narandill/mjolnir-vllm:v0.31.0`) it is pre-installed and the
+**On the image** (`narandill/mjolnir-vllm:v0.32.0`) it is pre-installed and the
 `MJOLNIR` backend name is patched into vLLM — nothing to do.
 
-**On a stock vLLM 0.31.0** (inside a container that has CUDA 13 and `sm_110`):
+**On a stock vLLM 0.32.0** (inside a container that has CUDA 13 and `sm_110`):
 
 ```bash
 pip install "mjolnir-kernel[dsl]"          # [dsl] pulls nvidia-cutlass-dsl
@@ -143,8 +141,8 @@ 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_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) |
+| `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) |
 | `MJOLNIR_ATTN_KV_LAYOUT` | auto | force the KV view: `BHN2D BNH2D B2NHD B2HND 2BNHD 2BHND` |
@@ -168,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 —
@@ -187,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:
@@ -215,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.0 → mjolnir_kernel-0.2.0}/README.md

@@ -6,7 +6,7 @@ packaged as a vLLM attention backend.
 
 ```bash
 vllm serve <model> --attention-backend MJOLNIR      # the mjolnir-vllm image
-vllm serve <model> --attention-backend CUSTOM       # stock vLLM 0.31.0 + this package
+vllm serve <model> --attention-backend CUSTOM       # stock vLLM 0.32.0 + this package
 ```
 
 This package is the kernel only. Serving, calibration and benchmarking are the
@@ -59,10 +59,10 @@ fused QKV buffer.
 
 ## Install
 
-**On the image** (`narandill/mjolnir-vllm:v0.31.0`) it is pre-installed and the
+**On the image** (`narandill/mjolnir-vllm:v0.32.0`) it is pre-installed and the
 `MJOLNIR` backend name is patched into vLLM — nothing to do.
 
-**On a stock vLLM 0.31.0** (inside a container that has CUDA 13 and `sm_110`):
+**On a stock vLLM 0.32.0** (inside a container that has CUDA 13 and `sm_110`):
 
 ```bash
 pip install "mjolnir-kernel[dsl]"          # [dsl] pulls nvidia-cutlass-dsl
@@ -113,8 +113,8 @@ 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_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) |
+| `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) |
 | `MJOLNIR_ATTN_KV_LAYOUT` | auto | force the KV view: `BHN2D BNH2D B2NHD B2HND 2BNHD 2BHND` |
@@ -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.0 → mjolnir_kernel-0.2.0}/mjolnir/__init__.py

@@ -23,8 +23,8 @@ from __future__ import annotations
 
 import os
 
-__all__ = ["register", "BACKEND_PATH", "BACKEND_NAME", "__version__"]
-__version__ = "0.1.0"
+__all__ = ["BACKEND_NAME", "BACKEND_PATH", "__version__", "register"]
+__version__ = "0.2.0"
 
 BACKEND_NAME = "MJOLNIR"
 BACKEND_PATH = "mjolnir.backend.MjolnirAttentionBackend"
@@ -41,7 +41,7 @@ def _registry():
     try:
         from vllm.v1.attention.backends.registry import AttentionBackendEnum, register_backend
     except ImportError:  # older vLLM layout
-        from vllm.attention.backends.registry import (  # type: ignore
+        from vllm.attention.backends.registry import (
             AttentionBackendEnum,
             register_backend,
         )

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.0 → 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
 
 # --------------------------------------------------------------------------
@@ -57,12 +64,12 @@ logger = logging.getLogger("vllm.mjolnir")  # under vLLM's handler
 try:
     from vllm.v1.attention.backends.registry import AttentionBackendEnum
 except ImportError:  # older layout
-    from vllm.attention.backends.registry import AttentionBackendEnum  # type: ignore
+    from vllm.attention.backends.registry import AttentionBackendEnum
 
 try:
     from vllm.v1.attention.backend import AttentionCGSupport
 except ImportError:
-    from vllm.v1.attention.backends.utils import AttentionCGSupport  # type: ignore
+    from vllm.v1.attention.backends.utils import AttentionCGSupport
 
 FP8_CACHE_DTYPES = ("fp8", "fp8_e4m3")
 FP6_CACHE_DTYPES = ("fp6_e2m3",)
@@ -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,
@@ -718,13 +736,6 @@ class MjolnirImpl(_BaseImpl):  # type: ignore[misc, valid-type]
         row = runner.kv_row_bytes
         if self._kv_layout is None:
             self._kv_layout = detect_kv_layout(kv_cache.shape, runner.h_k, runner.page_size, row)
-            # logger.info(
-            #     "Mjolnir: KV cache %s %s -> layout %s (%s)",
-            #     tuple(kv_cache.shape),
-            #     tuple(kv_cache.stride()),
-            #     self._kv_layout,
-            #     runner.kv_format,
-            # )
         return split_kv_cache(kv_cache, self._kv_layout, row)
 
     def _write_kv(self, layer, key, value, kv_cache, k_cache, v_cache, md):

{mjolnir_kernel-0.1.0 → mjolnir_kernel-0.2.0}/mjolnir/banner.py

@@ -122,6 +122,6 @@ if __name__ == "__main__":  # `python -m mjolnir.banner [--plain]`: preview with
         color=sys.stdout.isatty() and not plain,
         model=os.environ.get("MJOLNIR_CONFIG", "").strip() or None,
     )
-    print("")
+    print()
     print(msg)
-    print("")
+    print()

{mjolnir_kernel-0.1.0 → 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
@@ -86,7 +86,8 @@ def engine_kwargs_from_config(path: str) -> dict:
     import yaml
 
     parser, engine_args = _engine_parser()
-    data = yaml.safe_load(open(path)) or {}
+    with open(path) as f:
+        data = yaml.safe_load(f) 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)

{mjolnir_kernel-0.1.0 → mjolnir_kernel-0.2.0}/mjolnir/draftvocab.py

@@ -22,6 +22,7 @@ Special tokens (EOS, the chat / think markers) are always included.
 from __future__ import annotations
 
 import argparse
+import contextlib
 import json
 import os
 import sys
@@ -116,10 +117,8 @@ def main(argv=None) -> int:
         return 1
     tok = llm.get_tokenizer()
     special = set(getattr(tok, "all_special_ids", []) or [])
-    try:
+    with contextlib.suppress(Exception):
         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)

{mjolnir_kernel-0.1.0 → mjolnir_kernel-0.2.0}/mjolnir/kvfp6.py

@@ -47,7 +47,7 @@ def e2m3_values(device) -> torch.Tensor:
     """(64,) fp32: the value of every e2m3 code."""
     key = str(device)
     if key not in _LUTF:
-        from .t5_decode import minifloat_table
+        from .quant import minifloat_table
 
         _LUTF[key] = minifloat_table("e2m3").to(device)
     return _LUTF[key]
@@ -57,7 +57,7 @@ def e2m3_to_e4m3_lut(device) -> torch.Tensor:
     """(64,) uint8: e4m3fn byte of every e2m3 code (exact)."""
     key = str(device)
     if key not in _LUT:
-        from .t5_decode import minifloat_table
+        from .quant import minifloat_table
 
         vals = minifloat_table("e2m3")  # 64 float values, code order
         e4 = vals.to(torch.float8_e4m3fn)
@@ -229,7 +229,7 @@ def write_kv(
         return
     _, h_k, d = key.shape
     _, H, N, C = kv_cache.shape
-    assert H == h_k and C == 2 * row_bytes(d) and kv_cache.stride(3) == 1 and kv_cache.dtype == torch.uint8
+    assert h_k == H and 2 * row_bytes(d) == C and kv_cache.stride(3) == 1 and kv_cache.dtype == torch.uint8
     assert key.stride(2) == 1 and value.stride(2) == 1
     _write_kv_kernel[(T, h_k)](
         key,
@@ -271,7 +271,7 @@ def expand_pages_to_fp8(
     P = src_blocks.shape[0]
     _, H, N, C = kv_cache.shape
     d = head_size
-    assert C == 2 * row_bytes(d) and out.shape[1:] == (H, N, 2 * d) and out.stride(3) == 1
+    assert 2 * row_bytes(d) == C and out.shape[1:] == (H, N, 2 * d) and out.stride(3) == 1
     if P == 0:
         return out[:0]
     _expand_kernel[(P, H, N)](
@@ -303,7 +303,7 @@ def expand_pages_to_fp8(
 def quantize_rows_reference(x: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
     """(..., h_k, d) float -> packed bytes (..., h_k, row_bytes(d)) uint8;
     scale per tensor, per KV head or (h_k, d) per channel."""
-    from .t5_decode import minifloat_quantize, pack_subbyte
+    from .quant import minifloat_quantize, pack_subbyte
 
     h_k, d = x.shape[-2], x.shape[-1]
     s = scale.float().cpu()
@@ -318,7 +318,7 @@ def quantize_rows_reference(x: torch.Tensor, scale: torch.Tensor) -> torch.Tenso
 
 def decode_rows_reference(packed: torch.Tensor, d: int) -> torch.Tensor:
     """packed (..., row_bytes) uint8 -> e2m3 code values (..., d) float."""
-    from .t5_decode import minifloat_table
+    from .quant import minifloat_table
 
     b = packed.to(torch.int64).cpu().reshape(*packed.shape[:-1], d // 4, 3)
     w = b[..., 0] | (b[..., 1] << 8) | (b[..., 2] << 16)