This document is relevant for: Trn2, Trn3

Pooling models on Trainium#

Pooling models are models that return a vector per request instead of generated text. vLLM groups them by task: embed (embeddings), classify and score (sequence classification and scoring), and the token-wise tasks token_embed and token_classify. On Neuron, embed is the task supported today — see Supported tasks for the full matrix.

This topic covers how pooling models execute on AWS Trainium with vLLM Neuron and why their behavior differs from generative models in ways that affect how you configure and size a deployment.

Applies to#

This concept is applicable to:

  • Serving embedding models for semantic search, retrieval, RAG, clustering, and deduplication

  • Understanding why sampling, speculative decoding, and disaggregated inference do not apply to embedding endpoints

  • Sizing and tuning an embedding deployment for throughput or latency

What a pooling model does differently#

A pooling model shares the same transformer backbone as a generative model. The two diverge only at the output stage:

Generative:  forward -> logits        -> sampler -> next token   (prefill + decode)
Pooling:     forward -> hidden states -> pooler  -> embedding    (prefill only)

A generative model projects the final hidden state through its language-model head to vocabulary logits, samples a token, and then repeats — appending to the KV cache once per generated token. A pooling model instead applies a pooler to the backbone’s hidden states to produce one fixed-size vector per request, and stops. There is no token to generate, so there is no decode phase.

vLLM Neuron selects this path from the --runner pooling flag. The endpoints a pooling server exposes depend on the task its pooler reports; because embed is the task supported on Neuron today, the server exposes the OpenAI-compatible /v1/embeddings endpoint in place of the chat and completions endpoints, and the offline LLM.embed() API becomes available.

Everything is shared with the generative path: the same attention kernels, the same weight loading and tensor-parallel sharding, the same compilation pipeline, and the same KV cache implementation. Only the output stage is different, which is why an embedding model inherits the backbone’s performance characteristics on Trainium.

How pooling produces one vector per request#

The pooler decides which of a request’s token positions contribute to its vector:

Pooling method

Reads

Typical model

Last-token

the final token’s hidden state

decoder-based embedders, including Qwen3-Embedding

CLS

the first token’s hidden state

encoder models (BERT family)

Mean

the average over all token states

symmetric embedders

Which method applies is determined by the checkpoint, so you do not normally configure it.

Because the output is L2-normalized, the dot product of two embedding vectors is their cosine similarity. No further normalization is needed on the client.

Some embedding models additionally support returning a truncated vector, which trades a small amount of quality for a much smaller index. Request it per call with the dimensions parameter; the truncated vector is re-normalized, so it remains unit-norm and directly comparable.

What prefill-only means in practice#

Three consequences follow, and they explain most of the configuration differences between an embedding deployment and a generative one.

Startup is shorter. Only prefill graphs are compiled and warmed up. Neuron compiles fixed-shape programs ahead of serving, and an embedding model needs no decode graphs at all, so there are fewer to build.

Features that act on generated tokens do not apply. Sampling parameters, on-device sampling, speculative decoding, and structured outputs all operate on a sampled token, and a pooling model emits none. Disaggregated inference likewise presupposes a decode phase to separate from prefill. Async scheduling, which overlaps host work with the decode loop, has nothing to overlap and is not supported. Requesting an FP8 KV cache is rejected at startup: a prefill-only workload never reads the KV cache back, so there is no bandwidth benefit to gain.

Note: “Prefill-only” means no decode step, not no KV cache. The prompt’s own causal self-attention reads the keys and values of earlier tokens, and that storage is the KV cache. What an embedding model skips is the autoregressive loop in which the cache grows one token at a time.

Converting a generative model to a pooling model#

You do not need a purpose-built embedding checkpoint to produce embeddings on Neuron. Passing --convert embed turns any single-tower text generative model into an embedding model: it drops the language-model head and attaches a pooler, while the backbone, its weights, and its compiled graphs are reused unchanged. This makes embeddings available for architectures that ship no dedicated embedding checkpoint, and for quick prototyping on a model you already have.

--convert embed is the only conversion supported on Neuron today; see Supported tasks for the full matrix.

Supported tasks#

vLLM defines pooling tasks in two families: sequence-wise tasks, which reduce a request to one vector, and token-wise tasks, which return one vector per token. vLLM Neuron officially supports the sequence-wise embed task today.

Task

Family

API

Status on Neuron

embed

sequence-wise

/v1/embeddings, LLM.embed()

✅ Supported

classify

sequence-wise

/classify, LLM.classify()

❌ Not supported

score

sequence-wise

/score, /v1/rerank, LLM.score()

❌ Not supported

token_embed

token-wise

/pooling, LLM.encode()

❌ Not supported

token_classify

token-wise

/pooling, LLM.encode()

❌ Not supported