Metadata-Version: 2.4
Name: sabi-mla-vllm
Version: 0.2.0
Summary: vLLM plugin that serves sabi-noema-mla (shared pre-RoPE rotary MLA, rank 512, on the 16 full-attention layers of a Qwen3.5/3.8 hybrid)
Author: Sabi
License-Expression: Apache-2.0
Project-URL: Model (MLA), https://huggingface.co/sabicap/sabi-noema-mla
Keywords: vllm,mla,multi-head latent attention,kv cache,qwen
Classifier: Programming Language :: Python :: 3
Classifier: Operating System :: POSIX :: Linux
Classifier: Environment :: GPU :: NVIDIA CUDA
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: vllm!=0.25.0,!=0.27.0,!=0.29.*,!=0.30.*,<=0.31.0,>=0.22.1
Requires-Dist: transformers>=5.12
Requires-Dist: safetensors>=0.4
Requires-Dist: huggingface_hub>=0.30
Dynamic: license-file

# sabi-mla-vllm

vLLM plugin that serves [sabi-noema-mla](https://huggingface.co/sabicap/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.

```bash
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").

| GPU | vLLM | TP | KV pool (tokens) | self-check | probe | chat | tool call | forced decode | real decode | long-context answers (3k-103k) | result |
|---|---|---|---|---|---|---|---|---|---|---|---|
| B300 | 0.22.1 | 1 | 9,829,102 | 1/1 | pass | ok | ok | 402/402 | 710/711 | 5/5 | PASS |
| B300 | 0.31.0 | 1 | 9,819,914 | 1/1 | pass | ok | ok | 402/402 | 710/711 | 5/5 | PASS |
| B300 | 0.31.0 | 2 | 11,667,445 | 2/2 | pass | ok | ok | 402/402 | 706/709 | 5/5 | PASS |
| H100 | 0.22.1 | 1 | 838,341 | 1/1 | pass | ok | ok | 401/402 | 706/709 | 5/5 | PASS |
| H100 | 0.22.1 | 2 | 2,336,442 | 2/2 | pass | ok | ok | 402/402 | 710/711 | 5/5 | PASS |
| H100 | 0.23.0 | 1 | 838,341 | 1/1 | pass | ok | ok | 402/402 | 708/710 | 5/5 | PASS |
| H100 | 0.23.0 | 2 | 2,331,065 | 2/2 | pass | ok | ok | 401/402 | 702/707 | 5/5 | PASS |
| H100 | 0.24.0 | 1 | 835,746 | 1/1 | pass | ok | ok | 401/402 | 706/709 | 5/5 | PASS |
| H100 | 0.24.0 | 2 | 2,329,720 | 2/2 | pass | ok | ok | 402/402 | 708/710 | 5/5 | PASS |
| H100 | 0.25.1 | 1 | 835,746 | 1/1 | pass | ok | ok | 402/402 | 704/708 | 5/5 | PASS |
| H100 | 0.25.1 | 2 | 2,308,211 | 2/2 | pass | ok | ok | 402/402 | 704/708 | 5/5 | PASS |
| H100 | 0.26.0 | 1 | 838,341 | 1/1 | pass | ok | ok | 402/402 | 706/709 | 5/5 | PASS |
| H100 | 0.26.0 | 2 | 2,392,904 | 2/2 | pass | ok | ok | 401/402 | 706/709 | 5/5 | PASS |
| H100 | 0.27.1 | 1 | 820,510 | 1/1 | pass | ok | ok | 401/402 | 702/707 | 5/5 | PASS |
| H100 | 0.27.1 | 2 | 2,385,102 | 2/2 | pass | ok | ok | 402/402 | 702/707 | 5/5 | PASS |
| H100 | 0.28.0 | 1 | 851,968 | 1/1 | pass | ok | ok | 402/402 | 704/708 | 5/5 | PASS |
| H100 | 0.28.0 | 2 | 2,379,669 | 2/2 | pass | ok | ok | 402/402 | 704/708 | 5/5 | PASS |

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

| GPU | prefill (latent attention over new tokens + cached context) | decode | notes |
|---|---|---|---|
| H100 / H200 (sm_90) | FlashAttention-3 MLA entry (`q` = rotary 64, `q_v` = absorbed 512) | FlashAttention-3 MLA decode | the eval-server configuration |
| B300 / GB300 (sm_103) | tiled online-softmax in torch (bf16 tensor-core GEMMs, fp32 softmax statistics; tiles 1024x4096) | vLLM's Triton MLA decode, vendored from vLLM 0.31.0 (int64 cache offsets) | FA3 is Hopper-only |
| B200 / GB200 (sm_100) | same as B300 | same as B300 | not tested on hardware: refused unless `SABI_MLA_ALLOW_UNTESTED=1` |

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

| variable | default | meaning |
|---|---|---|
| `SABI_MLA_ALLOW_UNTESTED` | unset | `1` runs an untested vLLM / GPU / TP combination (prints a warning) |
| `SABI_MLA_SKIP_SELFCHECK` | unset | `1` skips the startup kernel self-check |
| `SABI_MLA_SKIP_PROBE` | unset | `1` disables the post-boot probe of `sabi-mla-serve` |
| `SABI_MLA_ARTIFACT` | from config | explicit path to `shared_mla.safetensors` (default: `text_config.shared_mla.artifact` in the model repo) |
| `SABI_MLA_SHA256` | from config | expected artifact sha256 (default: `text_config.shared_mla.artifact_sha256`) |
| `SABI_MLA_ALLOW_COMPILE` | allowed | `0` refuses `torch.compile` (compilation mode must then be 0); mode 3 is the serving config of record |
| `SABI_MLA_WORKSPACE_TOKENS` | 90,112 (32 sequences) | prefill workspace budget; must cover `max_num_seqs x block size` (2,784 at TP1) |
| `SABI_MLA_PREFILL`, `SABI_MLA_KV_DECODE` | per GPU | kernel overrides (see above) |

## 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

| symptom | cause / fix |
|---|---|
| stock vLLM: `model type sabi_mla ... not recognize` or `Model architectures ['SabiMlaForConditionalGeneration'] are not supported` | the plugin is not installed in this environment: `pip install sabi-mla-vllm` |
| `ImportError: sabi-noema-mla needs the sabi-mla-vllm plugin` (with `--trust-remote-code`) | same: install the plugin; `--trust-remote-code` is not needed with it |
| `sabi-mla-vllm: vLLM X is untested ...` / `GPU ... is untested` | use a tested vLLM (`pip install sabi-mla-vllm` picks one) or set `SABI_MLA_ALLOW_UNTESTED=1` at your own risk |
| `kernel 'fa3' is Hopper-only` | remove a `SABI_MLA_PREFILL=fa3` / `SABI_MLA_KV_DECODE=fa3` override on non-Hopper GPUs |
| `self-check FAILED` at startup | the attention kernels give wrong numbers on this GPU / vLLM build; do not override, report it |
| `artifact ... sha256 ... !=` | the factor file is corrupt or from another revision: re-download the model |
| `max_num_seqs x block_size ... exceeds the ... workspace budget` | lower `--max-num-seqs` or raise `SABI_MLA_WORKSPACE_TOKENS` |
| slow prefill at long context on B300 | expected: Blackwell has no FA3 and the tiled torch prefill is compute-bound (see the B300 numbers in the model card) |

## 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.
