aicogmap

AI Cog Map

Real-time transformer layer activation visualization for SGLang inference servers.

Watch your AI model think. AI Cog Map captures per-layer activation patterns during inference and renders them as a live heatmap — showing which layers, attention heads, and MLP blocks are doing meaningful work on each forward pass.

Features

• Layer Activation Heatmap — Grid where each cell represents a hook point in the transformer. Brightness shows activation intensity, from dark through purple to bright cyan.
• Component-Aware Visualization — Hook points are classified by type (attention, MLP, projections, norms) with distinct color tinting so you can see the model's internal structure.
• Cognitive Depth Slider — Control how much detail you see. Slide from "Layer Outputs" (64 cells for a 64-layer model) through to "Full Detail" (every sub-module, 800+ cells).
• Auto-Interpreter — Computes a cognitive state label from activation patterns: Attending, Reasoning, Structuring, Recalling, Searching, Generating, or Processing. Shows what kind of work the model is doing right now.
• Activation Delta Bar — Shows which layers are changing between forward passes, revealing where the model's attention is shifting as it processes new tokens.
• Live Statistics — Active layer count, mean activation, data freshness, energy breakdown by component type and layer depth.
• Info Tooltip — Built-in reference explaining what each component type does and what high activation means.

Different types of work create different visual signatures. Code generation activates different layers than creative writing. Math reasoning lights up different patterns than summarization.

Architecture

┌─────────────────────┐    shared memory     ┌─────────────────────┐
│                     │    + metadata JSON    │                     │
│   SGLang Server     │  ─────────────────>   │   AI Cog Map UI     │
│                     │  /dev/shm/aicogmap-*  │                     │
│  + forward hooks    │                       │  FastAPI + HTML     │
│  (aicogmap.hook)    │                       │  localhost:7890     │
│                     │                       │                     │
└─────────────────────┘                       └─────────────────────┘

Hook Plugin (aicogmap.hook): SGLang forward hook factory. Attaches to every matched sub-module via SGLang's native --forward-hooks system. Computes the L2 norm of each module's output tensor and writes to a shared memory buffer. Also classifies each hook point by layer index, component type, and category, writing a metadata sidecar JSON. No SGLang source modification required.

Visualization Server (aicogmap.server): Standalone FastAPI server that reads the shared memory buffer and metadata, computes cognitive state, and serves a real-time heatmap UI. Polls at 500ms for smooth animation.

Shared Memory Bridge: Lock-free binary format (4-byte magic ACGM + header + float32 array) at /dev/shm/aicogmap-activations. Metadata sidecar at /dev/shm/aicogmap-metadata.json. Zero disk I/O, zero serialization overhead. The hook writes, the server reads.

Installation

pip install aicogmap

Or from source:

git clone https://github.com/openhelper/ai-cog-map.git
cd ai-cog-map
pip install -e .

Quick Start

1. Start SGLang with hooks enabled

python -m sglang.launch_server \
    --model-path Qwen/Qwen3.5-27B \
    --dtype bfloat16 \
    --port 8000 \
    --forward-hooks '[{
        "name": "aicogmap",
        "target_modules": ["model.layers.*"],
        "hook_factory": "aicogmap.hook:create_activation_hook",
        "config": {"shm_path": "/dev/shm/aicogmap-activations"}
    }]'

2. Start the visualization server

aicogmap --port 7890

3. Open the UI

Navigate to http://localhost:7890 and send requests to your model. Watch the layers light up.

Docker

If SGLang runs in Docker, mount shared memory and the aicogmap package:

volumes:
  - /dev/shm:/dev/shm
  - ./aicogmap:/usr/local/lib/python3.12/dist-packages/aicogmap

See examples/docker-compose.example.yml for a complete setup.

Configuration

Hook Config

Server Config

Target Module Patterns

The target_modules field uses fnmatch patterns. Common patterns:

Component Classification

When using model.layers.*, the hook automatically classifies each sub-module:

The cognitive depth slider controls which categories are visible:

Cognitive States

The auto-interpreter computes a cognitive state label from activation energy patterns:

API

GET /api/activations?depth=N

Returns normalized activation norms, deltas, cognitive state, and hook metadata.

{
  "status": "ok",
  "num_layers": 192,
  "total_hooks": 865,
  "age_s": 0.05,
  "depth": 2,
  "norms": [0.23, 0.45, 0.12, "..."],
  "raw_norms": [145.2, 278.9, 72.1, "..."],
  "deltas": [0.01, 0.15, 0.03, "..."],
  "cognitive_state": {
    "label": "Reasoning",
    "confidence": 0.82,
    "attention_ratio": 0.31,
    "energy": {
      "attention": 1240.5,
      "mlp": 2890.3,
      "early": 980.2,
      "mid": 1650.8,
      "deep": 1499.8,
      "total": 4130.8
    }
  },
  "hooks": [
    {"index": 0, "layer": 0, "type": "layer_output", "category": "layer"},
    {"index": 1, "layer": 0, "type": "attention", "category": "attention"},
    "..."
  ],
  "history": ["..."]
}

GET /api/metadata

Returns the full hook metadata sidecar.

GET /api/health

Returns connection status, hook count, and metadata availability.

Performance

• Hook overhead: ~0.1-0.5% of inference time (one tensor.norm() per hook per forward pass)
• Memory: ~4KB shared memory for 1024 hook points
• Metadata: One-time JSON write per new hook (~4KB for 865 hooks)
• Zero overhead when disabled: Hooks are opt-in via --forward-hooks
• No disk I/O: Shared memory (/dev/shm) is RAM-backed

How It Works

On each forward pass through the model:

1. SGLang's hook manager calls the hook function after each matched sub-module executes
2. On first encounter, the hook resolves the module's fully qualified name, classifies its component type, and writes the metadata sidecar
3. The hook computes output.detach().float().norm().item() — a single scalar capturing the module's activation magnitude
4. The scalar is written to a thread-safe buffer, indexed by hook position
5. A background thread flushes the buffer to shared memory at configurable intervals
6. The visualization server reads shared memory and metadata, normalizes values, computes cognitive state, and serves the UI
7. The frontend renders normalized values as a color-mapped heatmap grid with component tinting

The L2 norm captures overall activation magnitude without being sensitive to individual outlier values, and is extremely cheap to compute on GPU tensors.

Compatibility

• SGLang: Any version with --forward-hooks support (PR #13217)
• Models: Any transformer-based model served by SGLang
• GPU: Any NVIDIA GPU supported by SGLang
• OS: Linux (requires /dev/shm for shared memory)

Author

Sonny Vleisides — CTO at ARRA Networks | VP Development at Alice Research

• LinkedIn
• ARRA Networks
• GitHub

License

Apache 2.0 — see LICENSE and NOTICE.