shreyansh26/Qwen-2.5-1B-RLCD

Parallel Constrained Decoding with Transformers

A local Hugging Face Transformers / PyTorch server comparing sequential JSON generation with batched constrained field decisions. The default model is Qwen/Qwen2.5-1.5B-Instruct, loaded without quantization. The repository contains serving code and presets, not fine-tuned weights.

Setup and serving

Requires Python 3.11+ and uv. The first startup downloads the model into the standard Hugging Face cache.

uv sync --locked
bash run.sh

Open http://localhost:8000 and click Run Comparison. The UI runs tree-based parallel, batched parallel, and streaming naive inference on the same preset, one after another. Each panel shows its output, elapsed time, and forward-pass count; the summary compares all three. Tree results also show shared-prefix and tree-node counts. The first tree run includes FlexAttention compilation, so run again for warm timings.

On a shared CUDA machine, select an available GPU explicitly:

CUDA_VISIBLE_DEVICES=1 bash run.sh

Configuration is read at process startup:

Variable	Default	Purpose
MODEL_ID	Qwen/Qwen2.5-1.5B-Instruct	Hugging Face model ID or local checkpoint directory
DEVICE	cuda when available, otherwise cpu	PyTorch device, such as cuda:0 or cpu
HOST	127.0.0.1	Listening address
PORT	8000	Listening port

CUDA uses bfloat16 when supported, otherwise float16; CPU uses float32. The uv lock uses CUDA 12.8 PyTorch wheels on Linux x86-64. For a CPU-only installation, create the environment and install explicitly:

uv venv --python 3.12
uv pip install torch --index-url https://download.pytorch.org/whl/cpu
uv pip install -r requirements.txt
DEVICE=cpu .venv/bin/python -m server.main

run.sh uses uv run --locked, so it restores the locked environment. Use the direct Python launcher for a custom CPU-only installation. A compatible Qwen2.5 checkpoint can be selected with MODEL_ID; prompt templates are Qwen ChatML, so arbitrary model families are not guaranteed to work.

Inference modes

Sequential: An autoregressive loop uses the Transformers KV cache, emitting one token at a time. /api/run-naive and /api/stream-naive share the same implementation. Temperature 0 is greedy; positive temperatures sample. Generation stops at a complete JSON object, EOS, or max_tokens. This baseline can omit fields or choose invalid values; its response reports JSON validity and schema match.

Parallel: The context is prefilled once. Its KV cache is repeated across fields, and all field suffixes are evaluated in one padded batch with attention masks. Distinct first candidate tokens use a restricted softmax. If candidate tokens collide, complete allowed continuations (including closing delimiters) are scored using summed token log probabilities and normalized across choices. These continuation batches are capped at 16 candidates. Padding is cropped before continuation, and all forward calls are included in sequential_forward_passes (at least two).

Tree parallel: run_parallel_generation_tree prefills once and evaluates a token trie of all field suffixes and candidate continuations in one FlexAttention pass. Both field and candidate prefixes share nodes, and the KV cache remains at batch size one. A compiled BlockMask uses DFS ancestry intervals; no dense token attention mask is stored. Each node sees the prompt and its ancestors, with positions based on depth. Every candidate is scored through its closing delimiter, then its summed log-probability is divided by its token count before the temperature softmax. Shared value-prefix tokens contribute to both the score and length while still sharing trie nodes. This mean token log-probability reduces the automatic preference for shorter labels; it is a heuristic rather than an accuracy or calibration guarantee, and results can differ from the first-token shortcut above. Terminal tokens need no KV entries, and only scoring-parent logits are projected. The result has the same field structure, plus tree_nodes, prefix_tokens, and scoring: "mean_token_log_probability", and always reports two model forwards. This path requires full attention (no sliding-window cache), uses the installed PyTorch/Transformers FlexAttention support, and restores the model's original attention backend afterward. First use includes compilation overhead; warm up before measuring latency.

Values are assembled into typed JSON, so parallel outputs always use allowed enum values or booleans. These normalized model scores are not empirically calibrated probabilities of correctness. Legacy response names such as has_calibrated_probabilities remain for API compatibility. Field decisions are independent and do not enforce cross-field consistency.

Concurrent HTTP clients are supported. Model inference requests queue behind a process-local lock; parallelism refers to fields within one request, not simultaneous requests on the GPU. Use one server worker to avoid duplicate model loads. A disconnected stream releases the inference lock.

API and SDK

All POST routes accept:

{
  "context": "Production checkout is down. Escalate immediately.",
  "schema": {
    "priority": {"type": "enum", "description": "Incident priority", "choices": ["HIGH", "LOW"]},
    "escalate": {"type": "boolean", "description": "Whether escalation is required"}
  },
  "temperature": 0,
  "max_tokens": 700
}

Schemas must have at least one field; enum choices must be unique nonempty strings, with up to 255 choices per field. max_tokens controls sequential generation only (1–4096, default 700). Omitted temperatures default to 0.2 for sequential and 1.0 for parallel inference.

Route	Result
GET /api/presets	Four bundled scenarios
POST /api/run-parallel	Constrained values, field scores, timings
POST /api/run-tree	Tree-based constrained values, scores, timings, node counts
POST /api/run-rlcd	Alias for parallel inference
POST /api/run-naive	Sequential JSON, validation, timings
POST /api/stream-naive	SSE token events followed by done, or error
POST /api/compare	All three results and timing ratios; tree and tree_speedup_multiplier extend the existing response

Parallel parsed_json contains { "field": { "value": ..., "prob": ... } }; sequential parsed_json contains plain field values. Boolean values are JSON booleans.

from core import StructuredSchema, run_parallel_generation, run_parallel_generation_tree, run_naive_generation

schema = StructuredSchema({
    "escalate": {"type": "boolean", "description": "Escalation required"}
})
context = "A critical production outage is affecting every customer."
print(run_parallel_generation(context, schema)["parsed_json"])
print(run_parallel_generation_tree(context, schema)["parsed_json"])
print(run_naive_generation(context, schema, temperature=0)["parsed_json"])

Verification

Fast tests use a tiny randomly initialized Transformers Qwen2 model, requiring no checkpoint download:

uv run pytest -q
node --test tests/test_stream.cjs

They compare cached, padded candidate scores with independent full forwards, exercise cross-thread streaming, schema validation, and server routes. Tree tests also compare complete candidate scores against independent causal forwards, verify shared prefixes and a single KV cache, and cover 255 choices, escaped strings, temperature handling, and backend restoration after errors.

Real-model HTTP smoke checks exercise all four presets, sequential and parallel inference, SSE parity, concurrent requests, disconnect recovery, and static assets. Start the server first:

uv run python tests/smoke_server.py http://127.0.0.1:8000

Run the SDK benchmark with:

CUDA_VISIBLE_DEVICES=1 uv run python -m core.benchmark --presets presets/fintech_fraud.json presets/code_security.json presets/support_triage.json presets/high_cardinality_255.json

Performance depends on hardware, schema size, and candidate collisions. Run the benchmark for measurements on your machine; syntax validity does not establish classification accuracy.

Layout

• core/engine.py: Transformers loading, sequential decoding, shared-prefix parallel inference.
• core/schema.py: Schema validation and candidate metadata.
• core/prompt_builder.py: Qwen prompt templates.
• server/app.py: FastAPI routes, startup model loading, streaming.
• web/: Interactive comparison UI.
• presets/: Example classification scenarios.
• tests/: Offline regression tests and live HTTP smoke checks.

License

Apache 2.0