[Performance] OpenAI server logit_bias causes ~2x decode throughput loss at high batch: dense per-request vocab tensor rebuilt and re-uploaded every iteration

[Navjot10]

System Info

• TensorRT-LLM v1.3.0rc22, PyTorch backend (trtllm-serve ... --backend pytorch)
• GPUs: H100 80GB (repro) , TP1/PP1, one GPU per replica
• Model: Qwen3.5-0.8B fine-tune, vocab_size ≈ 249,283, max_seq_len 65,536
• Serve config: cuda_graph_config: {enable_padding: true, max_batch_size: 384},
  kv_cache_config: {dtype: fp8, enable_block_reuse: false}, chunked prefill on,
  max_num_tokens 32768, num_postprocess_workers: 4
• Workload: ~12k-token prompts, ~600–2,700-token completions, sustained
  batch ≈ 100–120 requests per GPU
• Every request carries an OpenAI logit_bias with 1,206 entries at -100
  (banning a contiguous special-token ID range from generation; all requests
  carry the identical dict)

Summary

Sending logit_bias through the OpenAI-compatible server costs us ~2x decode
throughput and +85% time-per-output-token at production batch depth. The
GPU is idle ~55–60% of every iteration (DCGM_FI_PROF_GR_ENGINE_ACTIVE drops
from ~0.78 to ~0.43) while the host rebuilds and re-uploads dense
vocab-sized bias tensors for the entire batch — once per generated token.

The cost is O(batch x vocab_size) per iteration, independent of the number
of bias entries: a 1-entry logit_bias pays the same price as our 1,206-entry
one.

Measurements (A/B/A, 2x H100, batch 224 across 2 replicas)

Same prompts (~12.3k tokens), same concurrency, max_tokens 640; the only
difference between phases is the presence of the logit_bias field:

Phase	logit_bias	ms/output-token (p50)	aggregate tok/s	DCGM GR_ENGINE_ACTIVE
A1	on	34.6	5,503	0.43
B	off	19.3	10,959	0.78
A2	on	35.3	5,652	0.43

• Effect is instant and fully reversible with the flag.
• Implied added host cost: ~17 ms per iteration at batch ~110/GPU.
• Engine /metrics shows iterLatencyMS spiking to 200–300 ms with bias on,
  vs a stable ~18 ms without.
• Penalty grows with batch: +54% ms/token at ~15 requests/GPU, +93% at ~110.

Root cause (v1.3.0rc22 code path)

1. The OpenAI protocol layer densifies the bias per request
(tensorrt_llm/serve/openai_protocol.py, _logit_bias_to_embedding_bias, ~L63):

embedding_bias = torch.zeros(vocab_size, dtype=torch.float32)  # ~1MB @ 249k vocab
for token_str, bias in logit_bias.items():
    embedding_bias[token_id] = bias_value

A sparse 1,206-entry dict becomes a ~1 MB dense fp32 CPU tensor, a fresh
object per request.

2. The sampler deduplicates bias tensors by object identity, so identical
biases are never shared
(tensorrt_llm/_torch/pyexecutor/sampler/sampler.py, _apply_embedding_bias,
~L3934, called from _sample_batched_by_strategy every decode iteration):

# NB: hash(torch.Tensor) is equivalent to id(torch.Tensor), and does not
#     depend on tensor contents ...
bias_to_index: dict[torch.Tensor, int] = defaultdict(provision_bias_index)

The in-code comment ("read-caching is expected to help in typical cases")
anticipates requests sharing one tensor object. But because step 1 creates a
fresh tensor per request, a batch of 110 requests with bit-identical biases
yields 110 "unique" tensors — the dedup never fires for OpenAI-server
traffic.

3. The full dense bias set is then rebuilt on host and re-uploaded to the GPU
every iteration (same function, ~L3963–L4016):

request_steps.tolist()                                        # device sync / iter
biases_tensor = torch.empty((n_unique, vocab), pin_memory=True)  # ~106MB pinned alloc
torch.stack(tuple(bias_to_index.keys()), out=biases_tensor)      # ~106MB host memcpy
biases_tensor_cuda = biases_tensor.to(logits.device, ...)        # ~106MB PCIe H2D
biases_tensor_cuda = torch.index_select(biases_tensor_cuda, ...) # ~106MB on GPU
logits[logits_bias_mask_cuda] += biases_tensor_cuda

At batch 110 and vocab 249k that is ~300 MB of allocation/memcpy/PCIe/HBM
traffic per generated token, plus a Python loop over the batch and a device
sync — all on the sampling critical path while the GPU waits. The ~17 ms/iter
we measured matches this arithmetic.

Nothing is cached across iterations even though the bias tensors are constant
for the lifetime of each request.

Expected behavior

logit_bias cost should be negligible: the constraint is a small, per-request-
constant set of (token_id, bias) pairs. For reference, vLLM v1 implements this
as a batch-level logits processor with sparse, GPU-resident state rebuilt only
when batch membership changes; its per-iteration cost is a single indexed-add
kernel over sum(len(bias_i)) cells (microseconds at this scale).

Suggested fixes (in increasing order of effort)

1. Content-based dedup/cache in the protocol layer (~10 lines): key the
   embedding_bias tensor on a hash of the logit_bias dict so identical
   dicts share one tensor object. This alone restores the sampler's existing
   identity-based dedup (n_unique: 110 -> 1) and reduces per-iteration
   traffic ~100x for homogeneous workloads like ours.
2. Persist per-request bias tensors on device: upload each request's
   embedding_bias to GPU once at admission (it is immutable), instead of
   re-stacking and re-uploading the batch's biases from host every iteration.
3. Sparse representation: store (row_indices, token_ids, values) for the
   batch and apply with one logits[rows, cols] += values per iteration,
   updated only on request admission/eviction. Removes the O(vocab) factor
   entirely.

Steps to reproduce

1. trtllm-serve <any ~1B model> --backend pytorch --max_batch_size 384 --max_num_tokens 32768 --max_seq_len 65536 (config as above).
2. Drive sustained concurrency ≈ 100+ per GPU with long prompts (~12k tokens)
   and max_tokens 512–1024, non-streaming /v1/chat/completions.
3. Run once with "logit_bias": {"<id>": -100, ... } (any ~1,000 entries; even
   1 entry shows the effect at reduced magnitude) and once without.
4. Compare per-request latency / completion_tokens, aggregate tokens/s, and
   DCGM_FI_PROF_GR_ENGINE_ACTIVE.

We can share our load-generator script if useful.

Impact

On a 64-GPU production fleet this implementation detail halves effective
serving capacity for any workload that uses logit_bias routinely (e.g.
banning special/control tokens from generation — a common pattern for models
with learned prompt-compression/gist tokens). Because the OpenAI API surface
makes logit_bias look semantically free, the cost is very hard for users to
discover; we found it only via an A/B load test after observing unexplained
GPU-util collapse.

[elmehdi-eljair]

Thanks for the excellent write-up and measurements. This looks very close to the class of GPU memory/transfer amplification problems we’ve been researching in MemForge:

https://github.com/elmehdi-eljair/memforge

Our approach is to measure the ratio between useful data and the actual memory/transfer work performed, then optimize only when the end-to-end cost model proves the transformation is worthwhile.

In this case, the important signal is that a sparse logit_bias is being expanded into a dense vocab-sized representation and then repeatedly rebuilt/re-uploaded on the decode critical path. That creates a very large amplification between the useful bias data and the bytes actually allocated/copied/transferred.

I see two complementary paths:

Minimal/low-risk: content-based interning/cache so identical logit_bias values share one immutable representation instead of one tensor per request.
Structural fix: keep the bias sparse and GPU-resident, updating the (request/token/value) state only when batch membership changes, then apply it with an indexed GPU operation during decode.

We’d also include the preprocessing/cache/update cost in the benchmark rather than measuring only the final kernel. MemForge has caught several cases where an inner operation looked dramatically faster but the full transformation was slower end-to-end, so that total-cost check is important.

I’d be interested in reproducing this through MemForge and preparing a benchmarked patch against TensorRT-LLM. If useful, I can share before/after measurements for host time, H2D traffic, GPU utilization, and decode throughput, and contribute the resulting fix upstream