sabi-mla-vllm

vLLM plugin that serves sabi-noema-mla: the sabi-noema weights with multi-head latent attention (a 512-d latent plus one shared 64-d pre-RoPE rotary key, 576 cached values per token per layer instead of 2,048) on the 16 full-attention layers. The 48 Gated DeltaNet layers, the vision tower and MRoPE stay stock vLLM.

pip install sabi-mla-vllm
vllm serve sabicap/sabi-noema-mla --max-model-len 262144 --reasoning-parser qwen3 --enable-auto-tool-choice --tool-call-parser qwen3_xml

No environment variables are needed: vLLM loads the plugin through its entry point, and the plugin finds the MLA factors through the model's config.json, picks the attention kernels for the GPU and sizes its workspace. Add --tensor-parallel-size 2 for two GPUs. sabi-mla-serve takes the same arguments and additionally checks the server's answers against the reference server after startup (and shuts it down if they differ).

What the plugin does

• Registers the sabi_mla config type and the SabiMlaForConditionalGeneration architecture with vLLM (entry point vllm.general_plugins: sabi_mla). The model repo declares both, so stock vLLM without this plugin refuses the repo instead of silently serving the original attention. Other models in the same environment are untouched.
• Replaces the 16 full-attention layers with MLA on a paged 576-wide bf16 latent cache (no bf16 K/V is ever stored).
• Picks the attention kernels per GPU (table below), runs a startup self-check of those kernels on the real allocated KV cache (first, middle and last block) against an fp32 reference, verifies the factor file's sha256 against config.json, and refuses untested vLLM / GPU / tensor-parallel combinations with an error that names the combination and the override.
• sabi-mla-serve additionally runs a tiny forced-decode probe against the running server (token ids recorded from the H100 eval server, scored only where that server was confident) and shuts the server down if the answers differ.

Tested matrix

Every cell: clean venv pip install of this wheel + vllm==X, boot with examples/serve_vllm.sh (compile mode 3, full CUDA graphs, prefix caching, chunked prefill, 262,144 max length), self-check on every rank, probe, examples/chat.py, a tool call through the qwen3_xml parser, and parity against the H100 eval server (the stack every published number came from): forced decode over 402 positions, real decode over ~710 positions, and greedy answers at 3k, 25k, 52k, 82k and 103k tokens. On B300 the KV cache is first filled with 4.5M tokens so the checks run on cache slots above 3.7M (see "B300 fix").

Rows: development build of 0.2.0 (same kernels and code paths as the release; the release wheel was re-run from a clean room on vLLM 0.31.0, see below). Forced decode = argmax agreement with the eval server over prompt+reference positions (the release gate, >= 0.995); real decode = agreement at generated positions (the eval server agrees with itself 0.9968 across boots; every miss above sits at a near-tie). vLLM 0.29.x and 0.30.x were not tested and are excluded.

Supported = the rows marked PASS. pip install sabi-mla-vllm resolves vLLM within the tested range; anything else raises at startup unless SABI_MLA_ALLOW_UNTESTED=1.

Kernels per GPU

Overrides: SABI_MLA_PREFILL=fa3|fast|reference, SABI_MLA_KV_DECODE=fa3|triton (Hopper-only kernels are refused elsewhere).

B300 fix (0.2.0)

Earlier B300 runs produced a correct first token and then garbage (!!!!) at 19-76k context, followed by CUDA error: an illegal memory access. Cause: vLLM's Triton MLA decode kernel up to the 0.2x releases computed the paged-cache offset block_id x page_size x 576 in int32, which wraps at cache slot 2^31 / 576 = 3,728,270. An H100 pool holds ~0.88M slots and never gets there; a B300 pool holds ~9.8M, so once requests were placed in blocks above ~1,340 decode read the wrong memory. Prefill was unaffected (it gathers the cache with int64 indices), hence the good first token. It is not specific to sm_103: the stock kernel faults on an H100 too when given a 3.9M-slot cache. vLLM 0.31.0 fixed the kernel upstream; this plugin ships that kernel so every supported vLLM version uses it, and the startup self-check exercises the highest cache slot.

Environment variables

Constraints

bf16 weights and bf16 KV cache only (a real FP8 MLA cache is not implemented; do not pass --kv-cache-dtype fp8); tensor parallel 1 or 2; no speculative decoding; text-only by default (--limit-mm-per-prompt '{"image":0,"video":0}'): images pass through the stock vision tower and then the MLA layers but were not evaluated.

Troubleshooting

Contents

src/sabi_mla_vllm/: registration.py (config type + architecture), model.py (MLA layers, tensor-parallel slicing, cache receipt), backend.py (prefill/decode kernels), support.py (tested matrix, guard, per-GPU kernels), selfcheck.py, serve.py (sabi-mla-serve), patches.py (architecture-scoped vLLM config patches), _vendor/ (the vendored Triton MLA decode kernel from vLLM 0.31.0, the artifact loader and reference kernels).

Import package sabi_mla_vllm, vLLM entry point sabi_mla, console script sabi-mla-serve; environment variables SABI_MLA_* (the pre-release Q38_SHARED_MLA_* names are still read). License: Apache-2.0.