TML Inkling on vLLM: Day-0 Support with Optimized Performance

8 min read
vLLM Team

We are thrilled to announce that vLLM officially supports the TML Inkling model on Day 0. Both thinkingmachines/Inkling-NVFP4 and thinkingmachines/Inkling (BF16) models are supported with optimized performance and full feature parity.

TML Inkling is a 1T-parameter multimodal model trained by Thinking Machines Lab. The model natively accepts text, image, and audio inputs and generates text with up to 1M context length. It introduces several novel architecture components—relative attention, short convolution, and shared expert sinks—all of which are now efficiently integrated into vLLM.

With vLLM, the model runs at up to 380 tok/s/user with MTP and 140 tok/s/user without MTP on 4 GB200 GPUs. vLLM also provides full feature parity, including LoRA, TP/DP/EP/PP parallelism, prefix caching, and disaggregated serving. We verified both model accuracy and tool parsing through comprehensive benchmarks.

The integration PR is available here. Run the model as follows:

export VLLM_USE_V2_MODEL_RUNNER=1
export FLASH_ATTENTION_CUTE_DSL_CACHE_ENABLED=1
 
vllm serve thinkingmachines/Inkling-NVFP4 \
      --tokenizer-mode inkling \
      --reasoning-parser inkling \
      --tool-call-parser inkling \
      --enable-auto-tool-choice \
      --tensor-parallel-size 8 \
      --speculative-config '{"method": "mtp", "num_speculative_tokens": 8}' \
      --kernel-config.enable_flashinfer_autotune=False \
      --trust-remote-code

TL;DR

vLLM provides strong Day-0 support for TML Inkling:

  • Models: Both thinkingmachines/Inkling-NVFP4 and thinkingmachines/Inkling (BF16) are supported
  • Hardware: NVIDIA Blackwell and Hopper GPUs
    • Broader hardware support is in progress. Stay tuned!
  • Modality: text/image/audio input → text output
  • Context length: up to 1M tokens natively (Tinker exposes 64K and 256K context windows)
  • Features: LoRA, speculative decoding (MTP), TP/DP/EP/PP, prefix caching, disaggregated serving, and more
  • Optimizations: Sconv-aware TP sharding, low-latency fused collectives, kernel fusion, multi-streaming, PDL, and more
  • Performance: Up to 380 tok/s/user (w/ MTP) and 140 tok/s/user (w/o MTP) on 4 GB200 GPUs
  • Accuracy: Model quality and tool parsing verified with MMAU, MMMU-Pro, BFCL, NIAH-1M, and HLE

Model Architecture

Figure 1. TML Inkling model architecture (some ops such as RMSNorm and residual connections are omitted).
Figure 1. TML Inkling model architecture (some ops such as RMSNorm and residual connections are omitted).

Modality. TML Inkling is a natively multimodal model with 1T parameters. In addition to text and images, it accepts audio input and generates text. The model uses an extremely lightweight image encoder (hMLP) and audio embeddings (dMel), as described in TML's interaction model preview. The resulting embeddings are processed by a decoder-only Transformer backbone.

Attention. The backbone has 66 layers: 11 full-attention layers and 55 sliding-window attention layers. This heavy use of sliding-window attention is what makes the model's 1M context length efficient. All attention layers use grouped-query attention (GQA) with head size 128.

A distinctive design choice in Inkling is relative attention as its positional mechanism. Instead of RoPE, Inkling adds a learned relative-position term to the pre-softmax attention logits. See TML's blog post for details.

Sconv. Inkling makes aggressive use of short convolution (sconv) with window size 4. Each layer contains four sconv modules, applied to the attention keys, attention values, attention output, and MoE output. Sconv acts like a small local attention while incurring minimal compute and memory overhead.

MoE. Each layer has 256 routed experts (top-6) plus 2 shared experts, so every token is processed by 8 experts in total. Unlike existing models, however, Inkling introduces the concept of an expert sink: the two shared experts participate in the routing-score computation (absorbing probability mass) but are excluded as candidates from the top-6 selection.

In thinkingmachines/Inkling-NVFP4, only the routed experts are quantized to NVFP4; all other parameters, including the shared experts and the qkvr linears, remain in BF16. In thinkingmachines/Inkling, the MoE weights are in BF16 as well.

MTP. Inkling ships with 8 MTP heads for speculative decoding, allowing the model to generate up to 9 tokens per forward step. The MTP heads are chained: each head consumes the hidden states and sampled draft token from the previous head. Each MTP head is a single-layer Transformer with full or sliding window attention and a dense MLP. All MTP weights are in BF16.

vLLM Integration & Optimization

vLLM implements the model efficiently through a series of optimizations. Key highlights:

Managing the sconv cache. Short convolution requires keeping the hidden states of the last W-1 tokens. vLLM manages this sconv cache by treating it as the KV cache of a virtual sliding-window attention layer. This lets vLLM handle the sconv cache elegantly through its unified KV cache manager: states that fall outside the window are marked evictable, and prefix caching works seamlessly with the sconv cache.

Figure 2. Sconv-aware TP sharding.
Figure 2. Sconv-aware TP sharding.

Sconv-aware TP sharding. A straightforward TP implementation for this model would be: all-reduce (e.g., after o_proj) → sconv → residual connection → RMSNorm. However, this applies sconv to the full hidden states on every GPU, duplicating both the sconv compute and the sconv cache across ranks.

To eliminate this replication, vLLM shards the model differently. Since sconv operates independently along the channel dimension, we shard sconv across channels: instead of an all-reduce, we use a reduce-scatter and all-gather over the channel dimension. Each GPU then stores only a shard of the sconv cache and computes only its own slice of channels. The idea is similar to sequence parallelism, but sharding is applied to the channel dimension rather than the token dimension.

Low-latency fused collectives. vLLM further implements several fused kernels to optimize this new sharding scheme. In particular, we built low-latency reduce-scatter and all-gather kernels (fused with surrounding ops) by extending the Lamport-protocol design of FlashInfer's low-latency all-reduce kernel. The Lamport protocol lets the kernel synchronize via data-value polling instead of explicit barriers, cutting kernel time at batch size 1 from 40 µs to 8 µs (5x).

FA4 with sheared bias. Relative attention complicates the memory access pattern, significantly slowing down the attention kernel's compute pipeline. To overcome this, TML, in collaboration with Colfax Research, released a new FA4 kernel with a sheared-bias technique, which vLLM integrates directly. vLLM additionally selects FA4's num_splits factor per configuration—accounting for batch size, TP size, and KV length—to maximize performance.

Re-computing MTP KV cache. Because each MTP head takes the previous head's draft token as input, its KV cache becomes stale whenever a draft token is rejected. vLLM handles this carefully: it caches the base model's hidden states for the last few tokens and re-runs the MTP heads with the accepted tokens after rejection sampling.

Beyond these, vLLM's model implementation includes additional kernel fusion, PDL, and multi-streaming to achieve speed-of-light performance. For details, please check out our PR.

Performance

Thanks to the extensive optimizations above, vLLM achieves 380 tok/s/user with MTP8 (mean acceptance length 4.5) and 140 tok/s/user without MTP on 4× GB200 GPUs. Results were measured on prompts of 8K input tokens sampled from SPEED-Bench, with 1K output tokens generated per request.

Accuracy Evals

We verified the correctness of vLLM's implementation with comprehensive benchmarks covering every modality and capability:

  • Audio: MMAU
  • Vision: MMMU-Pro
  • Tool calling: BFCL
  • Reasoning: HLE
  • Long context: NIAH

vLLM matches the reference implementation across the board. On long context, vLLM matches the reference exactly through 221K tokens and stays within ~1 pp through 513K. At the most extreme context lengths (800K+), NIAH scores show higher run-to-run variance for this benchmark, and we are working on tightening reproducibility in that regime.

Benchmark / metricvLLM NVFP4Reference NVFP4Delta vs Reference
MMAU overall76.10% (761/1,000)75.50%+0.60 pp
BFCL exact calls78.61% (1,062/1,351)78.16%+0.45 pp
BFCL All-Live macro75.86%73.54%+2.32 pp
MMMU-Pro overall micro71.12% (3,691/5,190)70.52% (3,660/5,190)+0.60 pp
MMMU-Pro Standard 10-option70.23% (1,215/1,730)70.00% (1,211/1,730)+0.23 pp
MMMU-Pro Standard 4-option76.47% (1,323/1,730)76.30% (1,320/1,730)+0.17 pp
MMMU-Pro Vision66.65% (1,153/1,730)65.26% (1,129/1,730)+1.39 pp
HLE29.33% (633/2,158)26.65%+2.68 pp
NIAH (2K-221K)99.09% (436/440)99.09% (436/440)0.00 pp
NIAH (294K-513K)95.68% (421/440)96.82% (426/440)-1.14 pp
NIAH (586K-805K)81.36% (358/440)84.09% (370/440)-2.73 pp
NIAH (878K)70.91% (78/110)80.91% (89/110)-10.00 pp

Roadmap

As described above, vLLM provides strong Day-0 support for TML Inkling. Looking ahead, we see a few potential improvements:

  • FP8 for global attention: Inkling currently uses BF16 for global attention, which can become a bottleneck in both compute and KV cache capacity. We plan to explore FP8 here by modifying the new FA4 kernel.
  • CUDA graphs for the image & audio encoders: The image and audio encoders currently run in eager mode. While this is usually not a major issue since they run during prefill, we plan to apply CUDA graphs to them to eliminate CPU overhead entirely.
  • AMD GPU support: AMD GPUs are not yet supported for this model, since the new relative attention mechanism requires a dedicated kernel. Support is coming soon.

Acknowledgements

We thank the Thinking Machines Lab team for the collaboration. The model support is led by Inferact, a company aiming to grow vLLM into the world's AI inference engine and to accelerate AI progress by making inference cheaper and faster.