Custom models and GPUs

Quail uses model and GPU specifications to estimate query time and set execution limits. The inference backend must also support the model architecture and GPU you select. See Supported models and GPUs for the included configurations.

Model specifications

A model specification records a checkpoint, model dimensions, and numerical formats. Quail represents it with ModelSpec.

Register a Qwen3 checkpoint

Use qwen3_spec to read the model dimensions and weight format from the checkpoint's config.json.

For example, imagine that you have fine tuned Qwen3 4B and saved a complete checkpoint in /models/review-qwen3-4b. The directory contains the weights, model configuration, and tokenizer files.

import quail
from quail.specs import qwen3_spec

spec = qwen3_spec(
    name="review-qwen3-4b",
    hf_name="/models/review-qwen3-4b",
    config="/models/review-qwen3-4b/config.json",
)
registry = quail.ExtensionRegistry.with_built_ins()
registry.register_model(spec)
config = quail.EngineConfig(model=spec.name, device="h100-sxm")
session = quail.Session(config=config, registry=registry)
session.close()

The example registers the checkpoint and creates a session. It does not run a query. The checkpoint directory must be available in the process that loads the model.

For a Hugging Face checkpoint, set hf_name to its repository ID and revision to its commit ID. Omit config to download config.json.

The checkpoint must use BF16 weights or FP8 weights stored in blocks of 128 by 128 values. Models that use sliding window attention are not supported. See weight formats.

You can pass additional ModelSpec fields to qwen3_spec. For example, set w_mem_bytes to a value you have measured to replace the estimate. See qwen3_spec for its arguments.

Other architectures

For another architecture, you can define a ModelSpec and register it with registry.register_model(spec). The specification must match the checkpoint's dimensions, weight format, and prompt format. Registering a specification does not add execution support for the architecture.

See ModelSpec for the fields and Implementing a model architecture for the execution requirements.

GPU specifications

A GPU specification records memory capacity, memory bandwidth, and compute rates. Quail represents it with DeviceSpec.

All values describe one GPU. Use compute rates without structured sparsity. The H100 specification is a complete example.

Register a new specification with registry.register_device(spec). The backend must accept its device name and provide compatible execution code. The Quail backend currently accepts h100-sxm, l40s, b200, and rtx-pro-6000-blackwell-server. Its checks are in backend.py.

Implementing a model architecture

The Quail executor implements the Qwen3 and DiffusionGemma architectures.

Quail loads the checkpoint through vLLM, then selects its forward pass using ModelSpec.arch. Supporting another architecture requires both a compatible model loader and a forward pass that works with Quail's executor.

1. Confirm that vLLM can load the checkpoint. Quail's loader also expects the loaded model to expose an lm_head output layer.
2. Define a ModelSpec with the model dimensions, KV shapes, weight format, and prompt format. The executor uses these values to allocate memory and prepare model inputs.
3. Implement a subclass of ModelPipeline in quail/backends/quail/executor/models/. Its constructor receives (model, arena, *, spec): the loaded model, KV memory allocator, and specification. Its forward_chunk() method returns hidden states after the model's final normalization, for chunk.final_indices.
4. Configure the pipeline's attention support and KV requirements with tree_attention and needs_pages. If gemm_warmup is enabled, implement linears() and provide an engine with the quant() and gemm() methods used to prepare kernels. Models with extra answer positions also need the canvas settings described in the interface.
5. Add the class to PIPELINES in that directory's __init__.py and set ModelSpec.arch to the corresponding key.

The executor allocates and releases KV pages. The forward pass computes outputs using the supplied pages and the selected attention path. Before reusing a kernel, check its operation, input shapes, and number formats against the model. For example, Qwen3's activation kernel computes SiLU(gate) * up. For a model that uses GELU, use that model's GELU implementation in the forward pass.

Check the forward pass on a GPU against the model's reference implementation using identical input tokens and positions. Check both fresh inputs and inputs that reuse saved KV.

See the ModelPipeline interface and the Qwen3 implementation for the execution requirements.