An OpenAI-compatible chat-completions server backed by Hugging Face Transformers that returns, for every generated token, the complete vocabulary logit vector (raw, pre-sampling) and the exact log probability of the emitted token — nothing approximated from top-k logprobs. Built for computing Jensen–Shannon Divergence between successive Generator iterations in a Generator–Critic refinement loop.

Every non-trivial technical claim in this document carries a bracketed reference [n]; the full list is in References.

hf_logits_server/
├── app/
│   ├── config.py          # env-driven Settings; single source of tunables
│   ├── schemas.py         # Pydantic API contract (OpenAI shape + extensions)
│   ├── model_loader.py    # load model+tokenizer once; GPU placement
│   ├── generation.py      # dynamic batching engine; logit/logprob extraction
│   ├── server.py          # FastAPI routes, lifespan, HTTP error mapping
│   └── utils.py           # logging setup
├── examples/
│   ├── client_python.py   # requests example + JSD computation
│   └── client_curl.sh     # curl examples
├── scripts/
│   └── slurm_server.sbatch
├── requirements.txt
├── run_server.sh
└── README.md

Why this decomposition. Each module owns one responsibility, so each kind of change touches one file: deployment knobs → config.py; API contract → schemas.py; model lifecycle / GPU placement → model_loader.py; inference algorithm → generation.py; HTTP concerns → server.py. The module docstrings repeat and expand these rationales in place.

pip install -r requirements.txt

# 2 GPUs, Qwen2.5-32B-Instruct, 8-hour wall time, port 8000
NUM_GPUS=2 WALL_TIME=8h MODEL_NAME=Qwen/Qwen2.5-32B-Instruct ./run_server.sh

# On SLURM
sbatch --gres=gpu:2 --time=08:00:00 scripts/slurm_server.sbatch

Then point your existing LLMClient at http://<node>:8000 — the endpoint is POST /v1/chat/completions with OpenAI-style messages, and the response carries choices[0].message.content / finish_reason / usage exactly as the OpenAI Chat Completions API defines them [12], which is the same contract vLLM's OpenAI-compatible server implements [13]. Your Generator and Critic need no changes; requesting logits is one extra flag:

{
  "messages": [{"role": "user", "content": "..."}],
  "temperature": 0.0,
  "max_tokens": 256,
  "return_logits": true,
  "return_logprobs": true,
  "logits_encoding": "b64_fp16"
}

generation.py calls:

model.generate(..., return_dict_in_generate=True, output_logits=True)

• output_logits=True makes generate return the unprocessed prediction logit scores [1] — the returned logits field is a tuple with one tensor per generated step, each of shape (batch_size, config.vocab_size) [2]. That is the full vocabulary at every step; no truncation exists anywhere in this path.
• "Unprocessed" matters: at each step the model's raw logits are passed through a LogitsProcessorList (temperature, top-k, top-p, repetition penalty) before token selection [3]. output_scores=True would return those processed scores [2]; to get raw pre-processing logits you use output_logits=True [4]. JSD between generator iterations should be computed on the model's belief, not the sampler's filter, so this server returns the raw tensor and never sets output_scores.
• return_dict_in_generate=True is required for any of the output_* optional outputs to be returned [1].

Payload size is physics, not a design flaw. Qwen2.5's vocabulary is 151,646 tokens (152,064 embedding rows) [14], so one token's logit row is ~600 KB as float32. The server therefore offers logits_encoding: "json" (nested float lists), "b64_fp16", or "b64_fp32" (base64 of little-endian binary per token; decode with numpy.frombuffer, see examples/client_python.py). A MAX_LOGIT_TOKENS cap (HTTP 422 when exceeded) protects you from accidentally requesting multi-GB responses.

For generated position t with raw logit row z_t (full vocab, moved to CPU float32 once), the server computes log_softmax(z_t) over the entire vocabulary and indexes the emitted token:

logprob_t = log_softmax(z_t)[token_t]

log_softmax is used rather than log(softmax(x)) because it is the numerically stabilized formulation of exactly that composition [5]. Because the normalization runs over the full vocabulary, this is the exact model log-probability — not a top-k approximation, and not the probability under the temperature/top-p-filtered sampling distribution.

examples/client_python.py shows the full loop: decode base64-fp16 rows → softmax → JSD. Base-2 JSD between two distributions is symmetric and bounded in [0, 1] [15], which makes per-position divergence curves across generator iterations directly comparable.

• The model is loaded with device_map="auto", which lets Accelerate compute an automatic split: layers are assigned to fill GPU 0, then GPU 1, and so on (spilling to CPU/disk only if VRAM runs out) [6], and hooks transfer inter-layer activations between devices during the forward pass [7]. This is layer-sharded model parallelism: devices execute their layers in sequence, which is the simple, correct default for a research server [7].
• GPU count is a launch knob, not code: device_map="auto" shards over every visible device, so run_server.sh controls GPU count via CUDA_VISIBLE_DEVICES (NUM_GPUS=4 or GPU_IDS=0,2,5). CUDA only enumerates the devices listed in that variable [8].
• Inputs are moved to the device of the first sharded module (the embedding layer), read from model.hf_device_map, before generate — with a sharded model, "the model's device" is not a single device, so inputs must target the entry point. Logits come back per step and are moved to CPU in one .to("cpu", torch.float32) call per step, avoiding per-element transfers.
• Sizing: 32B parameters in bf16 ≈ 64 GB of weights alone, so a 32B model needs ≥ 2 of your 80 GB GPUs once KV cache and activations are included; larger future models just need more visible GPUs — no code change.
• Future tensor parallelism: recent Transformers supports tp_plan="auto" for true tensor-parallel inference [9]; Settings.device_map is a free-form string and model_loader.py is the single file to extend.

FastAPI handles requests concurrently on an asyncio event loop, while the GPU work is serialized through one background worker that owns the model:

1. Each request is validated, chat-templated, and enqueued as a job with an asyncio.Future.
2. The worker takes the first job, then waits up to BATCH_TIMEOUT_MS collecting further jobs with an identical sampling configuration (same do_sample, temperature, top-p, top-k, repetition penalty, max-new-tokens), up to MAX_BATCH_SIZE. Why identical: generate applies a single GenerationConfig — and therefore a single logits-processor chain — to the whole batch [1][3]; mixing configs in one batch would silently change some requests' sampling. Grouping by config keeps batching semantically invisible.
3. Prompts are left-padded to a common length with the attention mask passed explicitly. Left padding is required for decoder-only LLMs: they are not trained to continue from pad tokens, and right padding corrupts batched generation [10]. If the checkpoint defines no pad token, the EOS token is used as pad — the standard remedy, safe at inference because padding positions are masked out [10].
4. One model.generate runs for the group inside torch.inference_mode() — inference mode both disables gradient recording and skips autograd's version-counter/view bookkeeping, making it strictly better than no_grad for pure inference [11].
5. The batched output is split back into per-request results (per-row EOS trimming, stop-sequence trimming, logit slicing), and each future is resolved.
6. The blocking generate call runs in a thread via asyncio.to_thread, so the event loop keeps accepting and queueing requests during GPU work.

Seeded requests are never batched. Sampling reproducibility comes from seeding the global torch RNG (torch.manual_seed) [16]; there is no per-request generator stream inside generate, so a seeded request runs solo with the RNG seeded immediately before it — that makes seed actually mean something.

Determinism. temperature: 0.0 (or do_sample: false) selects greedy decoding — do_sample=False means greedy search picks the argmax token at every step [1] — which is deterministic given fixed weights, inputs, and kernels. (Bitwise reproducibility across different GPU models/driver versions is not guaranteed by PyTorch in general [16].)

Model is loaded exactly once. Loading happens in FastAPI's lifespan context, which runs before the application starts accepting requests [17], and run_server.sh pins uvicorn to --workers 1 because multiple uvicorn workers are separate OS processes [18] — each would load its own 64 GB model copy. Concurrency comes from the async queue + batching, not process replication. The tokenizer is likewise initialized once in model_loader.py and reused everywhere.

Prompts are built with tokenizer.apply_chat_template(messages, add_generation_prompt=True). Chat templates ship inside each model's tokenizer configuration and render system/user/assistant turns in exactly the format the model was trained on; add_generation_prompt=True appends the tokens that begin an assistant turn so the model responds as the assistant [19]. Multi-turn conversations work by passing the full message history — the template formats the entire list [19]. This is why the server supports Qwen2.5, Qwen3, and DeepSeek without any model-specific prompt code: each checkpoint brings its own template.

All of these map onto GenerationConfig fields with the same semantics [1]. Stop sequences are enforced by decode-and-trim after generation (correct and simple for batched research workloads; the returned metadata.stop_sequence_hit tells you which stop fired).

• Pydantic schemas reject malformed requests with automatic, field-precise HTTP 422 responses — FastAPI converts validation failures into structured 422 Unprocessable Entity errors out of the box [20].
• Engine failures map to informative codes in server.py: prompt too long → 413; logit budget exceeded → 422; queue full → 503; CUDA OOM → 507 with concrete remediation hints; anything unhandled → 500 with the exception type, plus a full stack trace in the server log.
• A failed batch fails all of its jobs loudly (futures get the exception); nothing hangs.
• Structured logging throughout (utils.setup_logging), including the resolved hf_device_map at startup so GPU placement is auditable.

• One model load, one tokenizer instance, for the process lifetime (§5).
• torch.inference_mode() around all generation [11].
• Logits leave the GPU once per step, as a whole [batch, vocab] tensor; Python-level lists are materialized only for requests that asked for logits, at serialization time.
• Dynamic batching amortizes the forward pass across compatible concurrent requests (§5).
• attn_implementation="sdpa" by default; set ATTN_IMPLEMENTATION=flash_attention_2 if the cluster has flash-attn installed. Attention implementation changes speed/memory, not the raw logit semantics.
• low_cpu_mem_usage=True at load to avoid materializing a full fp32 copy in RAM during checkpoint loading [6].

run_server.sh:

• NUM_GPUS=n exports CUDA_VISIBLE_DEVICES=0..n-1 (or set GPU_IDS explicitly); CUDA then only sees those devices [8], and device_map="auto" shards over exactly them [6].
• WALL_TIME wraps the server in GNU timeout, which runs the command and sends a signal when the duration elapses [23]; we send SIGTERM (clean FastAPI lifespan shutdown [17]) with a SIGKILL escalation after 30 s. Accepts 3600, 90m, 8h, 2d; 0 disables.
• Under SLURM, use scripts/slurm_server.sbatch: SLURM's --time is the wall clock and SLURM sets CUDA_VISIBLE_DEVICES to the allocated GPUs.

See examples/client_curl.sh (health check, greedy, seeded stochastic with stop sequences, base64 logits) and examples/client_python.py (drop-in requests usage, logit decoding, and end-to-end JSD between two generator iterations).

Each [n] above cites the following (all verified accessible as of 2026-07-10):

1. Hugging Face Transformers docs — Generation / GenerationConfig (definitions of output_logits "unprocessed prediction logit scores", output_scores, return_dict_in_generate requirement, do_sample, sampling parameters): https://huggingface.co/docs/transformers/main_classes/text_generation
2. Hugging Face Transformers docs — Utilities for Generation (GenerateDecoderOnlyOutput: logits returned when output_logits=True, tuple of up to max_new_tokens tensors each of shape (batch_size, config.vocab_size); scores = processed prediction scores): https://huggingface.co/docs/transformers/internal/generation_utils
3. Hugging Face Transformers source/docs — Logits processors (raw logits pass through a LogitsProcessorList — temperature, top-k, top-p, repetition penalty — before token selection): https://github.com/huggingface/transformers/blob/main/src/transformers/generation/logits_process.py
4. Hugging Face Forums — Obtain raw logits before decoding scaling is applied (confirms out.scores are post-processed; use output_logits=True for raw pre-temperature/top-k/top-p logits): https://discuss.huggingface.co/t/obtain-raw-logits-before-decoding-scaling-is-applied/171178
5. PyTorch docs — torch.nn.functional.log_softmax (numerically stabilized log-softmax; slower and unstable if computed as separate log + softmax): https://pytorch.org/docs/stable/generated/torch.nn.functional.log_softmax.html
6. Hugging Face Accelerate docs — Big Model Inference / Handling big models (device_map="auto": automatic split filling GPUs in order, then CPU/disk offload; low_cpu_mem_usage): https://huggingface.co/docs/accelerate/concept_guides/big_model_inference
7. Hugging Face Accelerate docs — same guide, dispatch section (hooks move activations between devices; layer-sharded execution is sequential across devices, i.e. model parallelism, not tensor parallelism): https://huggingface.co/docs/accelerate/usage_guides/big_modeling
8. NVIDIA CUDA docs — CUDA Environment Variables, CUDA_VISIBLE_DEVICES (restricts which GPUs CUDA enumerates): https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#env-vars
9. Hugging Face Transformers docs — Multi-GPU inference / Tensor parallelism (tp_plan="auto", launched with torchrun): https://huggingface.co/docs/transformers/perf_infer_gpu_multi
10. Hugging Face Transformers docs — Generation with LLMs, "Wrong padding side" (decoder-only LLMs are not trained to continue from pad tokens; inputs must be left-padded; pass the attention mask; setting pad_token = eos_token when missing): https://huggingface.co/docs/transformers/llm_tutorial
11. PyTorch docs — torch.inference_mode (disables gradient tracking and version-counter/view tracking; recommended over no_grad when the results are never used with autograd): https://pytorch.org/docs/stable/generated/torch.inference_mode.html
12. OpenAI API reference — Chat Completions (request messages with roles; response choices[].message.content, finish_reason, usage): https://platform.openai.com/docs/api-reference/chat
13. vLLM docs — OpenAI-Compatible Server (vLLM implements the same Chat Completions contract; this server mirrors it so only base_url changes in your LLMClient): https://docs.vllm.ai/en/latest/serving/openai_compatible_server.html
14. Qwen2.5 model card / config (vocab_size: 152064 embedding rows; 151,646-token tokenizer vocabulary): https://huggingface.co/Qwen/Qwen2.5-32B-Instruct
15. Lin, J. (1991). Divergence measures based on the Shannon entropy. IEEE Transactions on Information Theory 37(1):145–151 (JSD symmetry and boundedness; ≤ 1 in log base 2): https://ieeexplore.ieee.org/document/61115
16. PyTorch docs — Reproducibility (torch.manual_seed seeds the global RNG; complete determinism across platforms/releases not guaranteed): https://pytorch.org/docs/stable/notes/randomness.html
17. FastAPI docs — Lifespan Events (code before yield runs once, before the app starts receiving requests; after yield on shutdown): https://fastapi.tiangolo.com/advanced/events/
18. Uvicorn docs — Deployment / Server workers (multiple workers are separate processes): https://www.uvicorn.org/deployment/
19. Hugging Face Transformers docs — Chat templates (apply_chat_template, templates stored with the tokenizer, add_generation_prompt appends the assistant-turn header, multi-turn message lists): https://huggingface.co/docs/transformers/chat_templating
20. FastAPI docs — Handling Errors (automatic 422 validation errors from Pydantic models): https://fastapi.tiangolo.com/tutorial/handling-errors/
21. Hugging Face Transformers docs — Streamers (TextIteratorStreamer for token-by-token output): https://huggingface.co/docs/transformers/internal/generation_utils#streamers
22. Hugging Face Transformers docs — Quantization (bitsandbytes 4/8-bit via quantization_config in from_pretrained): https://huggingface.co/docs/transformers/quantization/bitsandbytes
23. GNU coreutils manual — timeout: run a command with a time limit (duration suffixes s/m/h/d; sends a signal on expiry; --kill-after): https://www.gnu.org/software/coreutils/manual/html_node/timeout-invocation.html

Two claims above are engineering arithmetic rather than citable facts, so for honesty they are labeled as such here: (a) "32B params in bf16 ≈ 64 GB" is 32e9 × 2 bytes; (b) "one Qwen logit row ≈ 600 KB fp32" is 151,646 × 4 bytes ≈ 592 KB. Both are direct consequences of [14] and the dtype sizes.