jc-914/prism-quantized

Sarvam-Prism INT4 (≤16 GB GPUs)

An INT4 quantization of sarvamai/prism, a decision model built on Sarvam-105B. The original needs one FP8-capable GPU with ~120 GB of memory. This version runs on any CUDA GPU with about 10 GB of memory or more, for example 16, 12 or 24 GB cards. It does not need FP8 tensor cores.

You give it a context and one or more questions, each with a fixed set of answers. It returns a probability for every answer in one forward pass. The API is the same as the original's: decide(state, questions).

Requirements

• A CUDA GPU with ≥ 10 GB of memory. sm_80+ (Ampere or newer) is recommended. sm_70/75 also work, on a slower path. It was tested with GPU memory hard-capped at 15 GiB and at 11 GiB.
• ≥ 64 GB of host RAM. About 54 GB of expert weights live in pinned host memory.
• Python 3.10+, PyTorch ≥ 2.4, Triton ≥ 3.0, transformers==4.52.1 (see requirements.txt), and a C compiler (Triton builds its kernel on first use).
• About 56 GB of disk.

Quick start

hf download jc-914/prism-quantized --local-dir prism-int4
cd prism-int4
pip install -r requirements.txt
python example.py .            # optional 2nd arg: GPU memory cap in GiB, e.g. `python example.py . 12`

from prism_int4 import Prism

model = Prism.from_pretrained(".")                 # uses the free memory on cuda:0; or gpu_memory_gb=12
answers = model.decide(
    state={"ticket": "My card was charged twice for the same order and I still haven't got a refund."},
    questions={
        "team": {"type": "choice", "instructions": "Which team should handle this ticket?",
                 "criteria": {"billing": "payments, charges and refunds", "tech": "app or website problems",
                              "delivery": "shipping and order status"}},
        "urgent": {"type": "yes_no", "instructions": "Does this need a reply within the hour?"},
        "anger": {"type": "score", "instructions": "How upset is the customer?",
                  "criteria": ["calm", "annoyed", "angry", "furious"]},
    },
)
# Many requests at once (amortises the weight streaming, see Speed):
results = model.decide_batch([(state1, questions1), (state2, questions2), ...])

The request format (choice / yes_no / score) and the output format are the same as the original model card.

What was quantized

part	parameters	format here	vs. original
128 routed experts × 31 MoE layers	~100B	INT4, GPTQ, symmetric, group size 128, fp32 scales	was FP8
attention, dense layer 0, shared experts	~4B	FP8 (e4m3, one scale per output row), widened to bf16 at use	unchanged
embeddings, norms, MoE router, decision head	–	bf16 / fp32	unchanged

• GPTQ calibration. 262,533 tokens of requests in Prism's own input format (state, question, options and answer, with the decision attention mask). The text is multilingual Wikipedia in all 11 supported languages, and the requests cover all three question types. Each layer's MoE input was recorded from the original FP8 model. Every expert's down projection is calibrated on the output of its quantized gate/up projection.
• Rare experts. Experts that were routed fewer than 512 calibration tokens (180 of 3,968) use round-to-nearest with an MSE-optimal clip search instead of GPTQ, so their error feedback is not fitted to a handful of samples.
• Layer error. The Hessian-weighted relative output error, averaged over layers, is 0.0014 for gate/up and 0.0036 for down, against 0.0075 and 0.0109 for plain round-to-nearest. Per-layer numbers are in quantization/quant_stats.json.

Accuracy against the original FP8 model

The suite is 26 requests / 43 questions: English and all 10 Indic languages, all question types, and contexts from a few tokens up to ~4k tokens. It is disjoint from the calibration data. The reference outputs come from the original prism.py on an H200.

metric	INT4 (15 GiB cap)	INT4 (11 GiB cap)
top-1 answer agreement	43 / 43	43 / 43
mean of the largest per-option probability difference	0.012	0.012
worst largest per-option probability difference	0.111	0.111
mean KL(original ‖ int4)	0.0016	0.0016
peak GPU memory	12.3 GiB	9.1 GiB

The largest differences are on borderline questions. For example, a yes/no question at 0.53 / 0.47 in the original comes out at 0.64 / 0.36 here. Clear-cut answers agree closely. This is a small suite: on any task where calibration matters, validate against your own labelled data. Full outputs are in quantization/eval_*.json.

How it runs on a small GPU

• Dense weights (~3.9 GiB) stay on the GPU.
• The INT4 experts (1.59 GiB per layer) stay in pinned host RAM. They are streamed to the GPU one layer at a time on a side CUDA stream, double-buffered, so the next layer's copy overlaps the current layer's compute.
• As many layers as fit in gpu_memory_gb − workspace_gb stay resident. On a GPU with ~60 GB free, everything is resident and nothing streams.
• The experts run in a Triton grouped W4A16 kernel: tensor-core dots with fp32 accumulation, using bf16 on sm_80+ and fp16 on older GPUs. It compiles for sm_70 through sm_120.
• The 2 GB token-embedding table stays in host RAM.

from_pretrained knobs:

knob	default	effect
gpu_memory_gb	all free memory	GPU memory the model may use
workspace_gb	3.5	memory kept for activations. Raise it for long contexts or big batches; lower it to keep more layers resident
batch_size	16	questions per forward pass
max_batch_tokens	16384	padded tokens per forward pass
pin_memory	True	set False if the host can't pin ~54 GB (copies become slower)

Speed

When layers stream, each forward pass moves the streamed expert weights over PCIe: ~50 GB with 2 resident layers, 53 GB with none. Latency is therefore bounded by host-to-GPU bandwidth, not GPU compute.

• Measured on an H200 NVL host (PCIe Gen5) with the GPU capped at 15 GiB and at 11 GiB: median 2.6 s per request.
• Expected at roughly 25 GB/s (PCIe Gen4 x16): about 2–2.5 s per forward pass.
• Expected at roughly 12 GB/s (PCIe Gen4 x8, or Gen3 x16): about 4.5–5 s per forward pass.

A forward pass scores up to batch_size questions together, so put many questions, or many requests via decide_batch, into one call. Throughput then scales almost linearly up to the batch limit. There are no CUDA graphs, because weights move every pass. The first call also compiles the Triton kernel.

Files

file	contents
model-layer-*.safetensors	per layer: FP8 dense tensors (unchanged) plus experts.{gate_up,down}_{qweight,scales} (INT4)
model-embed.safetensors, decision_head.safetensors, tokenizer*, config.json, *_sarvam_moe.py	from the original release; config.json documents the quantization
prism_int4.py, example.py, requirements.txt	inference code
quantization/	the quantization script, calibration generator, per-layer stats and evaluation outputs

INT4 packing: qweight is uint8 [E, N, K/2]. In each 128-wide group along K, byte j holds element j in its low nibble and element j+64 in its high nibble; value = nibble − 8. scales is fp32 [E, N, K/128], and weight = value × scale.

License and credit

Apache 2.0, the same as the original. All model weights derive from sarvamai/prism by Sarvam AI. This is an unofficial quantization: for the original, its evaluation and the hosted API, see Sarvam's release blog.