All Projects

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.

AI InterpretabilityActiveMITQwen3-1.7BCPU · MPS · CUDAWebSocketnone — live only
2.03B
Parameters
28
Decoder layers
16
Query heads
~0.13s
Per-token latency

Overview

Concept

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.

What it shows

Seven views into one forward pass

Every view reads from the same captured pass — nothing is recomputed per tab.

Architecture flow
embedding → 28 blocks → unembedding, click to expand
Generation timeline
tokens + confidence curve, click to rewind
Logit-lens trajectory
best guess per layer, sharpening with depth
Logit-lens heatmap
layers × candidates, brightness = probability
Next-token candidates
the real distribution behind the chosen token
Attention attribution
rollout or raw per-layer attention
Shareable card
one-click branded PNG of the run
Engine

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.

Why Qwen3-1.7B

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.

Frontend
prompt + device, Next.js / D3
Backend
hook engine, greedy generation
Live views
streamed per token over WebSocket

Built with

FastAPI
backend, WebSocket server
PyTorch
forward hooks, live inference
HuggingFace Transformers
Qwen3-1.7B
orjson
fast per-token streaming
Next.js 16
frontend, App Router
D3.js
attention/logit-lens views
Tailwind v4 + shadcn/ui
UI
Framer Motion
animation

Under the hood

ws /ws/xray protocol
→ {"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

Backend (port 8000)
$ cd backend
$ .venv/bin/pip install -r requirements.txt
$ .venv/bin/uvicorn app.main:app --reload --ws websockets-sansio --ws-ping-timeout 30
Frontend (port 3000)
$ cd frontend
$ npm install
$ npm run dev
quick start
$ curl http://127.0.0.1:8000/api/health
# {"status":"healthy","model_loaded":true}
$ open http://localhost:3000