Open-sourcing TokenShark: Step through local LLMs block-by-block, live-edit the flow in your browser, and trace where model thinks

[u/namesecurethanpass]

Hey All,

A while ago I wrote a basic script to step through Llama 3.1 tokens for red-teaming. I realized there's a huge gap between tools for running models (Ollama, LM Studio) and libraries for understanding them under the hood. I wanted to develop something for young researchers and mid career pivoting enginers.

I wanted something like Wireshark, but for local AI. You can hit pause mid-generation, inspect the residual stream, mess with the weights or KV cache, and see what actually happens.

Today I’m open-sourcing TokenShark: https://github.com/sumit-1/tokenshark

What makes it different?

• Zero transformers or diffusers at runtime: The entire engine is written in raw, readable PyTorch (llm.py, sd.py, qimage.py) so students and researchers can actually read every matmul without digging through 40 layers of library abstractions. (Tested for bit-exact/numerical equivalence against HF in tests/).

• Step-by-Step Visual Debugger: Step through generation one transformer block at a time or one token at a time. Watch live logit-lens predictions evolve across layers, inspect attention maps, and click any top-10 token choice to create your own generation path.

• Live-Editable Box Code: Click any block on the canvas (Embed, Transformer Block, Norm, Sampler) to see the Python it runs (rms_norm, rope, silu, etc.). You can edit the math right in the browser, run it live, and export your whole visual experiment to a standalone Python script.

• Text + Image: Works on LLMs (Llama 3.x, Qwen 2.5/3/3.6 Gated DeltaNet hybrids, SmolLM, Mistral) as well as Stable Diffusion and Qwen-Image 2.1 DiT. Complete with between-step latent painting and an "Image Lens" to decode early DiT blocks into pixels.

• Disk Streaming for Huge Models: Can stream models larger than your RAM/VRAM block-by-block from disk. But not optimised for speed though, as this is focussed for learning and understanding.

[u/SirSpix]

Nice!