valeriow/parallel-constrained-decoding

1---2title: Parallel Constrained Decision Engine3emoji: ⚡4colorFrom: green5colorTo: blue6sdk: gradio7app_file: app.py8pinned: false9license: apache-2.010---11 12# Parallel Constrained Decoding for Apple Silicon13 14A high-throughput inference engine for structured information extraction, decision routing, and categorical classification on Apple Silicon using MLX.15 16Parallel Constrained Decoding evaluates multi-field JSON schemas simultaneously rather than generating tokens sequentially. On an Apple Silicon M4 Max, it delivers **5.6x to 7.0x latency reductions** compared to standard autoregressive decoding with **100% schema validity** and **calibrated field-level confidence scores**.17 18---19 20## Performance Benchmarks (Apple Silicon M4 Max)21 22Evaluated with `mlx-community/Qwen2.5-1.5B-Instruct-4bit` on macOS Sequoia:23 24| Scenario | Fields | Autoregressive Baseline | Parallel Constrained | Latency Speedup | Syntax Validity |25| :--- | :--- | :--- | :--- | :--- | :--- |26| **Fintech Fraud Routing** | 4 fields | 420 ms (120 tok/s) | **75 ms** | **5.6x** | 100% guaranteed |27| **Code Security Audit** | 4 fields | 380 ms (125 tok/s) | **68 ms** | **5.6x** | 100% guaranteed |28| **High-Cardinality Tariff** | 1 field (255 choices) | 500 ms (118 tok/s) | **89 ms** | **5.6x** | 100% guaranteed |29| **Enterprise Support Triage** | 28 fields | 1,900 ms (130 tok/s) | **270 ms** | **7.0x** | 100% guaranteed |30 31---32 33## Why Parallel Constrained Decoding?34 35### The Problem with Autoregressive Structured Generation36 37Standard LLM structured generation (such as JSON mode or grammar-guided sampling) relies on token-by-token autoregressive decoding:38 39```40[Context Prompt] -> "{" -> "\n" -> " " -> "risk" -> ":" -> " " -> "HIGH" -> ...41(Requires 150 to 500 sequential forward passes)42```43 44Each token requires a distinct GPU/NPU forward pass and sequential memory bandwidth roundtrips. As schema size grows, latency scales linearly with output token length:45 46$$T_{\text{autoregressive}} = \sum_{k=1}^{K} t_{\text{step}}(k)$$47 48Additionally, autoregressive decoding is susceptible to syntax degradation, field omission, and hallucinated keys.49 50### The Solution: Parallel Evaluation via KV-Cache Broadcasting51 52In structured extraction and classification, field values belong to bounded candidate sets (booleans or categorical enums). Parallel Constrained Decoding exploits this property:53 54```55                          +---> [Field 1: "risk_level"] -------> Logit Slicing -> Top Choice56                          |57[Context Prefix Prefill] -+---> [Field 2: "requires_review"] ---> Logit Slicing -> Top Choice58(Single KV-Cache State)   |59                          +---> [Field M: "action_tier"] ------> Logit Slicing -> Top Choice60                          61                     (All fields evaluated simultaneously)62```63 641. **Single Broadcast Prefill**: The context document and semantic schema descriptions are prefilled once into an MLX Key-Value (KV) cache.652. **KV-Cache Broadcasting**: The KV-cache is broadcast across all $M$ schema fields in parallel.663. **Sub-Vocabulary Logit Slicing**: For each field, only candidate token IDs belonging to valid schema choices are evaluated. The remaining vocabulary is masked.674. **Calibrated Softmax Probabilities**: Exact normalized probabilities are calculated over the candidate slice:68   $$P(c_i) = \frac{\exp(z_i / T)}{\sum_{j=1}^{C} \exp(z_j / T)}$$695. **Token Tree Disambiguation**: When candidate choices share multi-token prefix roots, the engine executes continuation steps using sliced cache states with zero memory reallocation.706. **Programmatic Assembly**: Output JSON is constructed directly from verified values, guaranteeing 100% valid syntax without JSON parsing errors.71 72---73 74## Installation75 76### Prerequisites77 78- Apple Silicon Mac (M1, M2, M3, M4 series)79- macOS 14.0 or later80- Python 3.10+81 82### Setup83 84Clone the repository and install dependencies:85 86```bash87git clone https://github.com/your-org/parallel-constrained-decoding.git88cd parallel-constrained-decoding89 90python3 -m venv .venv91source .venv/bin/activate92pip install -r requirements.txt93```94 95---96 97## Developer SDK Quickstart98 99### 1. Defining Schemas100 101Schemas are defined using `StructuredSchema`. Each field specifies a `type` (`enum` or `boolean`), a `description` to guide model reasoning, and `choices` (for enum types, supporting up to 255 choices):102 103```python104from core.schema import StructuredSchema, FieldDefinition105 106# Option A: Dictionary-based definition107schema_dict = {108    "priority": {109        "type": "enum",110        "choices": ["P0_CRITICAL", "P1_HIGH", "P2_NORMAL", "P3_LOW"],111        "description": "Urgency tier based on customer business impact"112    },113    "requires_escalation": {114        "type": "boolean",115        "description": "Whether an on-call engineer must be notified immediately"116    },117    "department": {118        "type": "enum",119        "choices": ["BILLING", "INFRASTRUCTURE", "SECURITY", "PRODUCT_SUPPORT"],120        "description": "Target handling department"121    }122}123 124schema = StructuredSchema(schema_dict)125```126 127You can also construct fields explicitly using `FieldDefinition`:128 129```python130fields = {131    "tariff_classification": FieldDefinition(132        name="tariff_classification",133        field_type="enum",134        description="Harmonized System 6-digit tariff category code",135        choices=["0101.21", "0101.29", "8471.30", "8517.12", "8542.31", ...] # Up to 255 choices136    )137}138```139 140### 2. Running Parallel Generation141 142Execute parallel constrained inference on your context string:143 144```python145from core.engine import run_parallel_generation146 147context = """148Incident Report: Production database db-primary-01 CPU at 100%.149Payment gateway failing for 40% of checkout requests.150Tier 1 Enterprise customer affected: Acme Global.151"""152 153result = run_parallel_generation(context, schema)154 155print(f"Latency: {result['elapsed_ms']} ms")156print(f"Prefill Time: {result['prefill_ms']} ms")157print(f"Passes: {result['sequential_forward_passes']}")158print("\nExtracted JSON:")159print(result["parsed_json"])160```161 162### 3. Response Structure163 164The output dictionary provides both the structured JSON and detailed field telemetry:165 166```python167{168    "mode": "parallel_constrained_calibrated",169    "elapsed_ms": 74.5,170    "prefill_ms": 52.1,171    "suffix_eval_ms": 18.2,172    "sequential_forward_passes": 1,173    "is_valid_json": True,174    "schema_match": True,175    "parsed_json": {176        "priority": { "value": "P0_CRITICAL", "prob": 0.9924 },177        "requires_escalation": { "value": "true", "prob": 0.9981 },178        "department": { "value": "INFRASTRUCTURE", "prob": 0.9815 }179    },180    "field_telemetry": {181        "priority": {182            "value": "P0_CRITICAL",183            "confidence": 0.9924,184            "cardinality": 4,185            "top_choices": [186                { "choice": "P0_CRITICAL", "probability": 0.9924 },187                { "choice": "P1_HIGH", "probability": 0.0068 },188                { "choice": "P2_NORMAL", "probability": 0.0006 },189                { "choice": "P3_LOW", "probability": 0.0002 }190            ]191        }192    }193}194```195 196### 4. Streaming Autoregressive Baseline197 198To compare against standard autoregressive generation:199 200```python201from core.engine import stream_naive_generation202 203for event in stream_naive_generation(context, schema):204    if event["type"] == "token":205        print(event["token"], end="", flush=True)206    elif event["type"] == "done":207        print(f"\nCompleted in {event['result']['elapsed_ms']} ms")208```209 210---211 212## Interactive Web Visualizer213 214The repository includes a web interface for side-by-side latency and accuracy comparison.215 216To launch the web server:217 218```bash219bash run.sh220```221 222Or run directly with uvicorn:223 224```bash225python3 -m uvicorn server.app:app --host 0.0.0.0 --port 8000226```227 228Open `http://localhost:8000` in your browser.229 230### Features231 232- **Side-by-Side Comparison**: Parallel Constrained Decoding vs. Autoregressive Streaming.233- **Live Millisecond Timers**: Real-time elapsed latency counters.234- **Synchronized Scrolling**: Matching keys align across both panes.235- **Interactive Row Highlighting**: Hover over any field in either panel to highlight the corresponding key in the other.236- **Hallucination Detection**: Highlights omitted or hallucinated keys in naive autoregressive output.237 238---239 240## Command-Line Benchmark Runner241 242Run the benchmark suite across pre-configured enterprise presets:243 244```bash245python3 -m core.benchmark246```247 248Output example:249 250```text251======================================================================252Parallel Constrained vs. Autoregressive Generation Benchmark253======================================================================254--> Running preset: Fintech Fraud Detection (4 fields)...255    Autoregressive Baseline :    421.3 ms | 148 tokens (122.4 tok/s) | Passes: 148256    Parallel Constrained    :     74.8 ms |   0 tokens (O(1))           | Passes: 1257    >> SPEEDUP: 5.6x faster (Step reduction: 148.0x)258    >> Schema match: Naive=True | Parallel=True (100% guaranteed)259----------------------------------------------------------------------260--> Running preset: Support Triage Matrix (28 fields)...261    Autoregressive Baseline :   1894.2 ms | 312 tokens (131.2 tok/s) | Passes: 312262    Parallel Constrained    :    268.4 ms |   0 tokens (O(1))           | Passes: 1263    >> SPEEDUP: 7.1x faster (Step reduction: 312.0x)264    >> Schema match: Naive=True | Parallel=True (100% guaranteed)265----------------------------------------------------------------------266--> Running preset: High-Cardinality Tariff (1 field, 255 choices)...267    Autoregressive Baseline :    498.7 ms |  42 tokens (116.5 tok/s) | Passes: 42268    Parallel Constrained    :     88.6 ms |   0 tokens (O(1))           | Passes: 1269    >> SPEEDUP: 5.6x faster (Step reduction: 42.0x)270    >> Schema match: Naive=True | Parallel=True (100% guaranteed)271----------------------------------------------------------------------272```273 274---275 276## Repository Structure277 278```text279.280├── core/281│   ├── __init__.py           # SDK package exports282│   ├── engine.py             # Parallel constrained decoding & autoregressive engines283│   ├── schema.py             # Schema definitions, metadata compiler & logit mapping284│   ├── prompt_builder.py     # Prompt templates for prefill catalog and naive baseline285│   └── benchmark.py          # Command-line benchmark runner286├── presets/287│   ├── fintech_fraud.json    # Fraud detection scenario (4 fields)288│   ├── code_security.json    # Vulnerability audit scenario (4 fields)289│   ├── support_triage.json   # Enterprise ticket triage (28 fields)290│   └── high_cardinality_255.json # 255-choice tariff classifier291├── server/292│   ├── app.py                # FastAPI endpoints (/api/run-parallel, /api/stream-naive)293│   └── main.py               # Server launcher294├── web/295│   ├── index.html            # Side-by-side comparison UI296│   ├── app.js                # Frontend streaming & synchronized scrolling297│   └── style.css             # UI styling298├── MODEL_CARD.md             # Hugging Face model card documentation299├── requirements.txt          # Python package requirements300├── run.sh                    # Startup script301└── README.md                 # Project documentation302```303 304---305 306## Supported Models307 308The engine is currently configured for `mlx-community/Qwen2.5-1.5B-Instruct-4bit`.309 310Any decoder LLM supported by `mlx-lm` can be loaded by setting `MODEL_ID` in `core/engine.py`.311 312---313 314## License315 316Apache 2.0317