vLLM TML Inkling Support

vLLM TML Inkling Support

vLLM now provides Day-0 support for TML Inkling, a 1T-parameter multimodal model developed by Thinking Machines Lab. This integration enables high-performance inference for both the thinkingmachines/Inkling-NVFP4 and thinkingmachines/Inkling (BF16) versions, supporting text, image, and audio inputs with a native context length of up to 1 million tokens.

Performance and Hardware Benchmarks

On a configuration of 4 NVIDIA GB200 GPUs, vLLM achieves throughput of up to 380 tok/s/user when using Multi-Token Prediction (MTP) and 140 tok/s/user without MTP. These results were measured using prompts of 8K input tokens from SPEED-Bench and 1K output tokens per request.

While current support is optimized for NVIDIA Blackwell and Hopper GPUs, vLLM is actively working on expanding support to other hardware, including AMD GPUs, which currently require a dedicated kernel for the model's relative attention mechanism.

TML Inkling Model Architecture

TML Inkling is a natively multimodal decoder-only Transformer with 1T parameters. Its architecture includes several distinctive components:

  • Multimodal Inputs: The model accepts text, images (via a lightweight hMLP encoder), and audio (via dMel embeddings).
  • Hybrid Attention: The backbone consists of 66 layers: 11 full-attention layers and 55 sliding-window attention layers to maintain efficiency at a 1M context length. All layers utilize Grouped-Query Attention (GQA) with a head size of 128.
  • Relative Attention: Instead of Rotary Positional Embeddings (RoPE), Inkling uses a learned relative-position term added to pre-softmax attention logits.
  • Short Convolution (Sconv): Each layer employs four sconv modules (window size 4) applied to attention keys, attention values, attention output, and MoE output to provide local attention with minimal overhead.
  • Mixture of Experts (MoE): Each layer features 256 routed experts (top-6) and 2 shared experts. Inkling introduces "expert sinks," where the two shared experts participate in routing-score computation to absorb probability mass but are excluded from the top-6 selection.
  • Speculative Decoding (MTP): The model includes 8 chained MTP heads (single-layer Transformers) that allow for the generation of up to 9 tokens per forward step.

In the Inkling-NVFP4 variant, only the routed experts are quantized to NVFP4, while shared experts and qkvr linears remain in BF16.

vLLM Technical Optimizations

vLLM implements several specialized optimizations to handle Inkling's unique architecture:

Sconv Cache and TP Sharding

To manage the short convolution (sconv) cache, vLLM treats it as the KV cache of a virtual sliding-window attention layer, integrating it into the unified KV cache manager.

To avoid duplicating sconv compute and cache across GPUs, vLLM uses Sconv-aware TP sharding. Instead of a standard all-reduce after the output projection, vLLM employs reduce-scatter and all-gather over the channel dimension. This ensures each GPU stores and computes only its own slice of channels, similar to sequence parallelism but applied to the channel dimension.

Low-Latency Collectives and Kernels

  • Fused Collectives: vLLM uses low-latency reduce-scatter and all-gather kernels based on the Lamport protocol (data-value polling), reducing kernel time at batch size 1 from 40 µs to 8 µs.
  • FA4 Kernel: To optimize the relative attention memory access pattern, vLLM integrates a new FA4 kernel with a "sheared-bias" technique. vLLM dynamically selects the num_splits factor based on batch size, TP size, and KV length.
  • MTP KV Cache Management: Since MTP heads depend on previous draft tokens, vLLM caches the base model's hidden states and re-runs MTP heads with accepted tokens following rejection sampling.

Accuracy and Validation

vLLM's implementation matches the reference implementation across multiple benchmarks:

  • Multimodal & Reasoning: Verified via MMAU (Audio), MMMU-Pro (Vision), BFCL (Tool calling), and HLE (Reasoning).
  • Long Context: Validated using NIAH. vLLM matches the reference exactly up to 221K tokens and remains within ~1 percentage point (pp) up to 513K tokens. At extreme lengths (800K+), higher run-to-run variance was observed.

Future Roadmap

vLLM plans the following enhancements for TML Inkling:

  • FP8 Global Attention: Exploring FP8 for global attention via FA4 kernel modifications to reduce compute and KV cache bottlenecks.
  • CUDA Graphs: Applying CUDA graphs to image and audio encoders to eliminate CPU overhead during prefill.
  • AMD Support: Developing the necessary relative attention kernels to enable AMD GPU compatibility.

Sources