SketchSSM calibration: Nemotron Nano 9B v2 (NVFP4 weights)
calibration.pt is a portable SketchSSM
calibration file for nvidia/NVIDIA-Nemotron-Nano-9B-v2-NVFP4.
It contains the group-shared sketch basis and the Full-Gram allocation scores
from which the per-head rank table and the ordered frames for any mean rank are
derived, and the model layer index of each calibrated recurrent layer. It
contains no model weights.
Calibration weights
Collected with NVFP4 weights: nvidia/NVIDIA-Nemotron-Nano-9B-v2-NVFP4
at revision 8556c9164ddb43fe1f4f4ad730593b3c5e3f7328. The basis and the allocation scores depend on the weights,
so use this file with these weights; for another precision or checkpoint,
calibrate with that checkpoint.
Contents
| Field | Value |
|---|---|
| Base model | Nemotron Nano 9B v2 (Mamba-2) |
| Recurrent layers | 27 (model layers 0, 2, 4, 6, 7, 9, 11, 13, 16, 18, 20, 23, 25, 27, 29, 32, 34, 36, 38, 41, 43, 44, 46, 48, 50, 52, 54) |
| State heads per layer | 128 |
| Key dim K / value dim V | 128 / 80 |
| Basis groups per layer | 8 |
| Window W | 16 |
| Erase factor | no |
| Allocation rank cap | 49 |
Basis omega |
float32, shape (27, 8, 79, 128) |
| File size | 10,952,133 bytes |
The file loads with torch.load(..., weights_only=True). The
calibration guide
documents its keys.
Verified mean ranks
For these mean ranks, the derived tables equal the bundle tables and the exported frames equal those of the bundle export path:
| Mean rank | Dense heads | Sketch heads |
|---|---|---|
| 2 | 0 | 3456 |
| 4 | 12 | 3444 |
| 8 | 106 | 3350 |
| 16 | 536 | 2920 |
Other mean ranks are allocated with the same rule. manifest.json lists the
SHA-256 of calibration.pt and these results.
Usage
With vLLM's SketchSSM support (--sketchssm), the ranks for the requested mean
rank are allocated when the model loads:
vllm serve nvidia/NVIDIA-Nemotron-Nano-9B-v2-NVFP4 \
--sketchssm SketchSSM/Nemotron-Nano-9B-v2-NVFP4 --sketchssm-mean-rank 8 \
--mamba-ssm-cache-dtype float32
With the SketchSSM package, export the frames for a mean rank:
hf download SketchSSM/Nemotron-Nano-9B-v2-NVFP4 calibration.pt --local-dir calibration
python -m sketchssm.calibration export --calibration calibration/calibration.pt \
--mean-rank 8 --out frames.pt
How it was produced
With the public calibration pipeline of the SketchSSM repository (release
0.1.15) and collect.yaml from this repository:
python -m sketchssm.calibration calibrate --config collect.yaml --out out
python -m sketchssm.calibration package --bundle out --out calibration.pt
python -m sketchssm.calibration manifest --bundle out --calibration calibration.pt --out manifest.json
package re-allocates every configured mean rank from the packaged contents
and fails unless each table equals the bundle table.
License
This calibration file is released under the Apache License 2.0, like the SketchSSM repository. It is derived from the base model's weights, so use it under the base model's license as well.