T5Gemma 2 vLLM plugin

Standalone vLLM plugin for serving google/t5gemma-2-1b-1b and compatible Speculators/DFlash checkpoints trained against that model.

The plugin packages the Python model adapter, text encoder-decoder processor, merged self+cross decoder attention, and the small Triton kernels needed by the adapter. It does not require this repository's vllm-factory checkout at runtime.

Requirements

• Linux or WSL2 with an NVIDIA GPU.
• A vLLM version with v1 speculative decoding and DFlash support. The version used during development was the vLLM nightly wheel built from commit e4b3da3feb20c1854a4b23e431cfb787ee268f72.
• Hugging Face access to google/t5gemma-2-1b-1b.
• For source installs, a compiler toolchain and Python headers are required because Triton/FlashInfer may JIT compile kernels in the target environment.

Install

From this repository:

pip install -e ./t5gemma2-vllm-plugin

If the environment was created with uv and has no pip module:

uv pip install -e ./t5gemma2-vllm-plugin

After installation, vLLM discovers the plugin through the package entry point:

[project.entry-points."vllm.general_plugins"]
t5gemma2_vllm_plugin = "t5gemma2_vllm_plugin:register"

In a clean environment no extra environment variable is required. vLLM will load installed vllm.general_plugins entry points automatically.

If the environment contains multiple vLLM plugins and you want to load only this one, filter plugin loading explicitly:

export VLLM_PLUGINS=t5gemma2_vllm_plugin

This is optional. It is useful for development environments where unrelated plugins may fail to import or register conflicting model architectures.

Serve raw T5Gemma 2

vllm serve google/t5gemma-2-1b-1b \
  --host 127.0.0.1 \
  --port 8000 \
  --served-model-name t5gemma-2-1b-1b \
  --trust-remote-code \
  --no-enable-chunked-prefill \
  --max-model-len 512 \
  --max-num-seqs 1

Serve a T5Gemma 2 DFlash-family checkpoint

The checkpoint must be in Speculators format. The plugin registers T5Gemma 2 plus these DFlash-family draft architectures:

• DFlashDraftModel
• DSparkDraftModel
• DFlareDraftModel

For DFlash, its config.json contains:

{
  "architectures": ["DFlashDraftModel"],
  "speculators_config": {
    "algorithm": "dflash",
    "verifier": {
      "architectures": ["T5Gemma2ForConditionalGeneration"],
      "name_or_path": "google/t5gemma-2-1b-1b"
    }
  }
}

Run:

vllm serve /path/to/t5gemma-2-1b-1b.dflash \
  --host 127.0.0.1 \
  --port 8000 \
  --served-model-name t5gemma-2-1b-1b-dflash \
  --trust-remote-code \
  --no-enable-chunked-prefill \
  --max-model-len 512 \
  --max-num-seqs 1

vLLM reads the Speculators config from the checkpoint, loads the verifier google/t5gemma-2-1b-1b, and enables DFlash speculative decoding.

For example you can try d0rj/t5gemma-2-1b-1b.dflash-dev speculator.

DSpark and DFlare checkpoints are served the same way:

vllm serve /path/to/t5gemma-2-1b-1b.dspark \
  --host 127.0.0.1 \
  --port 8000 \
  --served-model-name t5gemma-2-1b-1b-dspark \
  --trust-remote-code \
  --no-enable-chunked-prefill \
  --max-model-len 512 \
  --max-num-seqs 1

vllm serve /path/to/t5gemma-2-1b-1b.dflare \
  --host 127.0.0.1 \
  --port 8000 \
  --served-model-name t5gemma-2-1b-1b-dflare \
  --trust-remote-code \
  --no-enable-chunked-prefill \
  --max-model-len 512 \
  --max-num-seqs 1

Internally the plugin maps dspark and dflare Speculators configs to vLLM's DFlash scheduler, then dispatches to DSpark/DFlare-specific draft model shims. DSpark uses a sequential Markov-bias sampling loop over the block positions; DFlare uses per-draft-layer target-state fusion before context KV precompute.

CUDA Graph support

Raw T5Gemma 2 and all three DFlash-family runtimes support vLLM CUDA Graph capture for actual batches from 1 through --max-num-seqs. The merged decoder kernel reads self-KV and cross-KV directly from vLLM's paged caches using GPU block tables and sequence lengths, so request reordering and cache-slot reuse do not require host synchronization. Encoder K/V is projected and cached only during prefill. Do not pass --enforce-eager for this setup.

To override the checkpoint's default speculative-token count, serve the verifier and provide the draft checkpoint explicitly:

vllm serve google/t5gemma-2-1b-1b \
  --max-num-seqs 1 \
  --no-enable-chunked-prefill \
  --speculative-config '{"model":"/path/to/checkpoint","method":"dflash","num_speculative_tokens":3}'

Quick request

curl http://127.0.0.1:8000/v1/completions \
  -H 'Content-Type: application/json' \
  -d '{
    "model": "t5gemma-2-1b-1b-dflash",
    "prompt": "Solve the problem step by step.\n\nQuestion: A train travels 60 miles in 2 hours. What is its average speed?\nAnswer:",
    "max_tokens": 64,
    "temperature": 0
  }'

Acceptance metrics

DFlash token acceptance is exposed by vLLM Prometheus metrics:

• vllm:spec_decode_num_drafts_total
• vllm:spec_decode_num_draft_tokens_total
• vllm:spec_decode_num_accepted_tokens_total
• vllm:spec_decode_num_accepted_tokens_per_pos_total

Example:

curl -s http://127.0.0.1:8000/metrics | grep 'spec_decode'

Acceptance rate is:

accepted_tokens / draft_tokens

Mean acceptance length including the bonus token is:

1 + accepted_tokens / drafts

Notes and limitations

• The adapter targets text encoder inputs. The T5Gemma 2 vision path is present in the vendored model code but has not been benchmarked as part of this plugin packaging.
• The merged decoder attention path is designed for FlashAttention/vLLM v1 on a single GPU. Tensor parallel and quantized serving should be validated before production use.
• The package intentionally vendors only the code required to register and serve the T5Gemma 2 generation architecture. It does not include training code.
• DSpark/DFlare serving is implemented for the current single-GPU vLLM setup. Tensor parallel and pipeline parallel serving need separate validation.
• CUDA Graph serving has been validated with batch sizes 1, 2, 3, and 4. Larger values are supported by the same paged-cache path but should be sized to the available KV-cache memory and validated for the target workload.