LLM X-Ray.
Type a prompt, watch an AI think — live attention, logit-lens predictions, and token probabilities, streamed token by token from a real forward pass.
Overview
An MRI scan, but for a language model
No pre-baked demos — every visual is driven by the user's own prompt, captured from a live PyTorch forward pass and streamed to the browser as it happens. The model answers the question (with an optional <think> reasoning trace first), so what's on screen is the model actually thinking, not a canned autocomplete animation.
Seven views into one forward pass
Every view reads from the same captured pass — nothing is recomputed per tab.
One hook engine, model-agnostic
The core (xray_engine.py) registers forward hooks on every decoder block to capture the residual stream and attention weights, then projects each layer's hidden state to a top-k prediction — a batched logit lens: all layers' last-position vectors through final_norm + lm_head in one matmul, not one pass per layer.
Architecture-specific module paths are resolved by a ModelAdapter, so swapping MODEL_NAME to gpt2 or a different Qwen3 size needs no engine changes. Attention rollout — the cross-layer 'what actually drove this token' attribution — is computed client-side, lazily, only for the selected step.
Small enough to run free, sharp enough to be worth watching
It runs on CPU at roughly 0.13s/token, but is modern and instruction-tuned enough to actually answer — including an optional reasoning trace — rather than just autocomplete. Its logit lens sharpens into coherent words across 28 layers, where GPT-2 small often stayed noisy the whole way through.
How it works
A prompt goes in over WebSocket, a live PyTorch forward pass runs with hooks on every decoder block, and the results stream back token by token to render as D3 views.
Built with
Under the hood
→ {"prompt": "...", "thinking": true, "device": "cpu"}
← {"type": "meta", "data": {"num_layers": 28, "model_label": "Qwen3-1.7B", "device": "cpu"}}
← {"type": "tokens", "data": {"tokens": [...], "token_ids": [...]}}
← {"type": "prompt_attention", "data": {"attention": [...]}}
← {"type": "step", "data": {"step": 0, "token": "...", "prob": 0.94, ...}}
... one "step" message per generated token ...
← {"type": "done", "data": {"generated_text": "...", "num_steps": 42}}Run it yourself
$ cd backend$ .venv/bin/pip install -r requirements.txt$ .venv/bin/uvicorn app.main:app --reload --ws websockets-sansio --ws-ping-timeout 30
$ cd frontend$ npm install$ npm run dev
$ curl http://127.0.0.1:8000/api/health# {"status":"healthy","model_loaded":true}$ open http://localhost:3000