drinkmoonshine/parallel-constrained-decoding

1---2language:3- en4license: apache-2.05library_name: mlx6tags:7- structured-generation8- parallel-decoding9- constrained-decoding10- apple-silicon11- mlx12- classification13- json14pipeline_tag: text-generation15base_model: Qwen/Qwen2.5-1.5B-Instruct16---17 18# Qwen2.5-1.5B-Instruct with Parallel Constrained Decoding19 20This repository provides an inference implementation for structured JSON generation and high-cardinality classification using `mlx-community/Qwen2.5-1.5B-Instruct-4bit` on Apple Silicon.21 22Instead of generating structured JSON token-by-token through sequential autoregressive loops, this engine uses **Parallel Constrained Decoding**. It broadcasts the model KV-cache across all schema fields simultaneously, evaluating all decisions in parallel forward passes.23 24## Key Performance Highlights (Apple Silicon M4 Max)25 26- **High-Cardinality Decisions (255 choices)**: 89 ms total latency vs. 500 ms autoregressive baseline (5.6x faster).27- **Enterprise Multi-Field Extraction (28 fields)**: 270 ms total latency vs. 1,900 ms autoregressive baseline (7.0x faster).28- **Guaranteed Schema Validity**: 100% valid JSON syntax without grammar parsers, rejection sampling, or repair loops.29- **Calibrated Field Confidence**: Exact softmax probabilities computed directly over candidate token logits for every field.30- **Unified Memory Footprint**: ~1.1 GB total RAM footprint in 4-bit quantization on Apple Silicon.31 32## How It Works33 34Traditional structured output engines run standard autoregressive decoding. For an N-field JSON schema, the model performs hundreds of sequential forward passes:35 36```37[System + Prompt] -> Token 1 -> Token 2 -> ... -> Token K (O(N) sequential forward passes)38```39 40Parallel Constrained Decoding decomposes the structured generation task into an isolated broadcast pass:41 421. **Prefix Prefill**: The context and semantic schema instructions are prefilled once. The resulting Key-Value (KV) cache is held in Apple Silicon Unified Memory.432. **KV-Cache Broadcasting**: The KV-cache is broadcast across all target fields concurrently.443. **Sub-Vocabulary Projection**: For each field, only valid candidate choices (e.g. enum options or boolean states) are evaluated. Unrelated vocabulary tokens are masked out.454. **Calibrated Softmax**: Probabilities are computed directly via softmax over the candidate logit slice:46   $$P(c_i) = \frac{\exp(z_i / T)}{\sum_j \exp(z_j / T)}$$475. **Collision Disambiguation**: In cases where candidate tokens share prefix strings, the engine follows continuation slices with zero memory reallocation.486. **Programmatic Assembly**: The verified field choices and confidence scores are formatted directly into structured JSON.49 50## Quickstart SDK51 52### Installation53 54```bash55pip install -r requirements.txt56```57 58### Python Usage59 60```python61from core.schema import StructuredSchema62from core.engine import run_parallel_generation63 64# 1. Define schema65schema_definition = {66    "fraud_risk": {67        "type": "enum",68        "choices": ["LOW", "ELEVATED", "SUSPICIOUS", "CRITICAL"],69        "description": "Risk assessment tier for incoming transaction"70    },71    "block_account": {72        "type": "boolean",73        "description": "Whether immediate account restriction is required"74    },75    "recommended_action": {76        "type": "enum",77        "choices": ["ALLOW", "STEP_UP_2FA", "TEMPORARY_HOLD", "TERMINATE_SESSION"],78        "description": "Immediate mitigation action"79    }80}81 82schema = StructuredSchema(schema_definition)83 84# 2. Provide context85context = """86User ID: usr_992187Location: Lagos, Nigeria (usual: Seattle, USA)88Device: Unknown Linux Chromium browser89Action: Wire transfer $49,500 to offshore escrow90Prior velocity: 0 transfers in 90 days91"""92 93# 3. Execute parallel generation94result = run_parallel_generation(context, schema)95 96print(f"Elapsed Time: {result['elapsed_ms']} ms")97print(f"Sequential Passes: {result['sequential_forward_passes']}")98print(f"Parsed JSON: {result['parsed_json']}")99```100 101### Output Example102 103```json104{105  "fraud_risk": { "value": "CRITICAL", "prob": 0.9942 },106  "block_account": { "value": "true", "prob": 0.9881 },107  "recommended_action": { "value": "TEMPORARY_HOLD", "prob": 0.9715 }108}109```110 111## Model Details112 113- **Base Model**: Qwen/Qwen2.5-1.5B-Instruct114- **Quantization**: 4-bit AWQ (mlx-community format)115- **Context Window**: 32,768 tokens116- **Hardware Target**: Apple Silicon (M1, M2, M3, M4 series with unified memory)117- **Supported Field Types**: Categorical Enums (up to 255 choices per field) and Booleans118 119## Benchmark Summary120 121Evaluated on Apple Silicon M4 Max (128GB Unified Memory, MLX 0.22+):122 123| Scenario | Schema Fields | Autoregressive (ms) | Parallel Constrained (ms) | Speedup | Valid Syntax |124| :--- | :--- | :--- | :--- | :--- | :--- |125| Fintech Fraud Routing | 4 fields | 420 ms | 75 ms | **5.6x** | 100% |126| Code Security Audit | 4 fields | 380 ms | 68 ms | **5.6x** | 100% |127| High-Cardinality Tariff | 1 field (255 choices) | 500 ms | 89 ms | **5.6x** | 100% |128| Support Triage Matrix | 28 fields | 1,900 ms | 270 ms | **7.0x** | 100% |129 130