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

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.

Start with _logit_bias_to_embedding_bias in tensorrt_llm/serve/openai_protocol.py, then trace _apply_embedding_bias and _sample_batched_by_strategy in tensorrt_llm/_torch/pyexecutor/sampler/sampler.py. Reproduce the high-batch comparison with and without logit_bias, then verify that identical or sparse biases no longer trigger repeated dense host allocation and upload on every decode iteration while preserving bias behavior.