advanced_features.md

REQUIRED Advanced Features Plan

This document plans the remaining advanced work for the Local LLM Instrumentation, Tracing, and Replay Platform.

The current project already runs a local GGUF model, records real llama.cpp runtime phase telemetry, replays JSONL traces, and displays the data in a keyboard-driven TUI dashboard.

However, the project statement explicitly asks for deeper transformer instrumentation. The features below are therefore marked REQUIRED for full requirement coverage, even though the current MVP is already demoable.

Current Gap

Current real telemetry is based on llama.cpp performance output. It measures broad runtime phases:

• model load
• prompt evaluation
• generated-token evaluation
• sampling
• sampler-chain detail
• total runtime
• unaccounted runtime

This is real and useful, but it is not enough for the full project statement.

The missing required pieces are:

1. REQUIRED: Real attention matrix extraction and visualization
2. REQUIRED: Real internal layer/submodule tracing
3. REQUIRED: Real activation statistics

The dashboard should keep the current runtime metrics, but the placeholder attention panel must eventually become a real heatmap when attention data exists.

Safety Constraints

All advanced work must stay safe for a laptop demo.

• Run CPU-only by default.
• Keep THREADS=2 for normal tests.
• Keep CTX=256 or lower for development.
• Capture a maximum token window of 32 tokens by default.
• Capture only selected layers and heads by default.
• Never dump every layer, every head, and every token unless explicitly requested.
• Store captured data in bounded buffers so memory cannot grow forever.
• If a tensor is not safely readable on CPU, mark it unavailable instead of faking it.
• Keep current MVP commands working while adding advanced mode.

Architecture Direction

The current app treats llama.cpp mostly as an external runtime. Advanced tracing needs deeper access to llama.cpp execution.

The preferred implementation path is:

1. Build against external/llama.cpp as a local source dependency.
2. Use llama_context_params.cb_eval / ggml_backend_sched_eval_callback to observe graph tensor execution.
3. Add our own trace collector around selected tensor names.
4. Add a guarded llama.cpp instrumentation patch only if callbacks cannot expose the required attention/activation tensors.
5. Emit all advanced captures into the same JSONL replay format.
6. Update replay and dashboard to display advanced events when they exist.

The implementation should prefer public callback hooks first. Direct llama.cpp patches should be small, isolated, and clearly named.

REQUIRED 1: Real Attention Matrix Extraction

Goal

Capture a real attention matrix from a local transformer run and display it in the TUI heatmap panel.

Minimum success target:

• one prompt
• one selected layer
• one selected attention head
• a small token window
• a real matrix from model execution, not fake data

Why Required

The project statement lists:

• attention matrix visualization
• extracting metadata and activations from attention submodules
• understanding how data propagates through transformer layers

The current heatmap is only a placeholder. This must become real for full coverage.

Implementation Plan

Start with llama.cpp graph callback instrumentation:

• Set llama_context_params.cb_eval.
• Inspect tensor names and operation types during a tiny prompt run.
• Log candidate tensors related to attention, softmax, KQ, QK, KQV, or attention output.
• Build a tensor-name map for the selected model.
• Disable flash/optimized attention in capture mode if the softmax matrix is fused away.

If callback instrumentation exposes a readable attention-weight tensor:

• Copy only the selected layer/head/window to host memory.
• Convert the values to float.
• Normalize or clamp display values only in the TUI, not in raw telemetry.
• Save the raw bounded matrix in JSONL.

If callback instrumentation does not expose the attention matrix:

• Add a small guarded patch inside external/llama.cpp.
• The patch should activate only when our advanced capture flag is enabled.
• Name/copy the attention softmax tensor immediately after it is produced.
• Avoid changing normal llama.cpp behavior.

Planned CLI Flags

./build/trace run \
  --model models/smollm2-135m-instruct-q4_k_m.gguf \
  --prompt "Explain attention simply" \
  --tokens 16 \
  --threads 2 \
  --ctx-size 256 \
  --advanced-trace \
  --capture-attention \
  --capture-layers 0 \
  --capture-heads 0 \
  --capture-token-window 32

Planned JSONL Fields

Attention events should include:

• capture_source: llama.cpp.cb_eval or llama.cpp.patch
• layer_name: for example layers.0.attention.head0
• layer_index
• submodule: attention
• head_index
• tokens
• attention_shape
• attention_matrix
• capture_truncated
• capture_available
• capture_note

Example shape:

{
  "layer_name": "layers.0.attention.head0",
  "op_type": "attention_matrix",
  "layer_type": "attention",
  "layer_index": 0,
  "head_index": 0,
  "tokens": ["Explain", "attention", "simply"],
  "attention_shape": [3, 3],
  "attention_matrix": [[1.0, 0.0, 0.0], [0.41, 0.59, 0.0], [0.20, 0.35, 0.45]],
  "capture_truncated": false,
  "capture_available": true
}

Dashboard Behavior

• If real attention data exists, panel title should say:
  • Attention Heatmap [real]
• If no real attention data exists, panel title may say:
  • Attention Heatmap [not captured]
• The old placeholder matrix should not appear as if it were real data.
• The selected layer/head should be visible in the panel.
• Rows and columns should use token labels when available.
• Large matrices should show a viewport, not the full matrix.

Tests

• Unit test JSONL parsing for attention matrix fields.
• Add a fixture JSONL file with one known attention matrix.
• Test dashboard/replay reads that fixture without crashing.
• Run a tiny real prompt with capture enabled.
• Verify at least one real attention event exists for selected layer/head.
• Verify matrix dimensions match token window.
• Verify values are finite.
• Verify memory stays bounded.

REQUIRED 2: Real Internal Layer/Submodule Tracing

Goal

Trace transformer internals at the level of embeddings, attention, MLP, norms, and output projection.

Minimum success target:

• emit real events for at least one transformer block
• distinguish attention from MLP
• show layer-level timing in replay and dashboard

Why Required

The project statement asks for:

• layer-by-layer execution latency
• transformer block bottleneck detection
• metadata from embeddings, attention, and MLP
• intercepting the forward pass without modifying model source code

The current app only shows broad llama.cpp runtime phases. That is useful but not true layer/submodule tracing.

Implementation Plan

Use graph callback instrumentation to observe tensor execution:

• Log tensor names, operation types, shapes, dtype, and execution order.
• Build a mapping from tensor names to logical transformer submodules.
• Group tensor-level activity into logical events:
  • embed_tokens
  • layers.N.attention
  • layers.N.mlp
  • layers.N.norm
  • lm_head
• Track start and end timestamps per logical group.
• Compute duration per group.
• Rank layer/submodule latency within the run.

The first implementation does not need every model architecture. It should support the current tiny SmolLM2 GGUF model first, then fall back gracefully for unknown models.

Planned CLI Flags

./build/trace run \
  --model models/smollm2-135m-instruct-q4_k_m.gguf \
  --prompt "Gravity is" \
  --tokens 8 \
  --threads 2 \
  --ctx-size 256 \
  --advanced-trace \
  --capture-layers 0,1

Planned JSONL Fields

Layer/submodule events should include:

• capture_source
• layer_name
• layer_type
• layer_index
• submodule
• op_type
• device
• dtype
• shape
• start_ms
• end_ms
• latency_ms
• duration_ms
• cumulative_ms
• latency_rank
• latency_status
• tensor_count
• capture_note

Example:

{
  "layer_name": "layers.0.attention",
  "layer_type": "attention",
  "layer_index": 0,
  "submodule": "attention",
  "op_type": "forward_submodule",
  "shape": [1, 16, 576],
  "dtype": "float16",
  "device": "CPU",
  "latency_ms": 2.41,
  "tensor_count": 18,
  "latency_status": "normal"
}

Dashboard Behavior

• Event Stream should show internal layer/submodule events alongside runtime phase events.
• Model/Info should show whether advanced tracing is enabled.
• Runtime Metrics should show shape/dtype/device for selected submodule events.
• Anomaly Log should identify the slowest transformer block/submodule.
• Replay should print internal events clearly, not just llama phase events.

Tests

• Unit test tensor-name to submodule mapping.
• Fixture test for grouped layer events.
• Real tiny run should emit at least:
  • one embedding event or one layer event
  • one attention event
  • one MLP or norm event if available
• Verify submodule latencies are nonnegative.
• Verify selected captured layers respect --capture-layers.
• Verify dashboard still works for old JSONL files without advanced events.

REQUIRED 3: Real Activation Statistics

Goal

Compute real tensor statistics for selected internal activations.

Minimum success target:

• one selected internal tensor
• real shape and dtype
• real mean, min, max, and sparsity
• anomaly status if values look suspicious

Why Required

The project statement asks for:

• tensor shape
• sparsity rate
• runtime metrics inspection
• optional mean/max tracking for numerical anomalies

Since these are directly listed as expected/optional metrics, the implementation should capture real values where possible and clearly label unavailable data.

Implementation Plan

Use the same callback instrumentation as layer tracing:

• Select only safe target tensors.
• Copy tensor data to host memory only when allowed.
• Convert supported numeric dtypes to float for statistics.
• Compute:
  • mean
  • min
  • max
  • sparsity
  • nan_count
  • inf_count
  • sample_count
• Use a sampling limit for large tensors.
• Do not block the model run with unbounded copying.

Sparsity definition:

• Count values as sparse when abs(value) <= epsilon.
• Default epsilon: 1e-6.

If a tensor is unavailable:

• Set activation_stats_available=false.
• Keep numerical stat fields empty or zero.
• Add a note explaining why it was unavailable.

Planned CLI Flags

./build/trace run \
  --model models/smollm2-135m-instruct-q4_k_m.gguf \
  --prompt "Explain gravity simply" \
  --tokens 16 \
  --threads 2 \
  --ctx-size 256 \
  --advanced-trace \
  --capture-activations \
  --capture-layers 0 \
  --activation-sample-limit 4096

Planned JSONL Fields

Activation events should include:

• activation_stats_available
• activation_sample_limit
• activation_sample_count
• mean
• min
• max
• sparsity
• nan_count
• inf_count
• anomaly_flags
• capture_note

Example:

{
  "layer_name": "layers.0.mlp",
  "layer_type": "mlp",
  "layer_index": 0,
  "submodule": "mlp",
  "op_type": "activation_stats",
  "shape": [1, 16, 576],
  "dtype": "float16",
  "activation_stats_available": true,
  "activation_sample_count": 4096,
  "mean": 0.0124,
  "min": -2.125,
  "max": 3.8125,
  "sparsity": 0.184,
  "nan_count": 0,
  "inf_count": 0,
  "anomaly_flags": []
}

Dashboard Behavior

• Runtime Metrics should show real activation stats when available.
• If stats are unavailable, the dashboard should say:
  • activation stats unavailable
• Anomaly Log should show:
  • NaN/Inf present
  • high absolute max
  • high sparsity
  • clipping risk
• Fake/demo values must be visually distinguishable from real captured values.

Tests

• Unit test stats calculations with known arrays.
• Test sparsity epsilon behavior.
• Test NaN/Inf anomaly flags.
• Test JSONL round-trip for activation fields.
• Real tiny run should produce at least one stats-available event or a clear unavailable reason.
• Verify memory does not grow with longer prompts.

Implementation Phases

Phase 1: Trace Discovery

Purpose: learn what llama.cpp exposes for the current model.

Tasks:

• Add an advanced trace discovery mode.
• Register cb_eval.
• Print or record tensor names, shapes, dtypes, op types, and execution order.
• Run one tiny prompt.
• Save discovery output to a local ignored file.

Exit criteria:

• We know which tensor names correspond to attention, MLP, and norms.
• We know whether attention weights are visible through callback-only tracing.

Suggested issue:

• Discover llama.cpp graph tensors for advanced tracing

Phase 2: Schema Extension

Purpose: make JSONL capable of storing advanced data.

Tasks:

• Extend telemetry structs for advanced fields.
• Update JSONL writer and parser.
• Keep backward compatibility with existing runtime-only JSONL.
• Add fixtures for attention matrix and activation stats.

Exit criteria:

• Replay can load old and new JSONL files.
• Tests pass for all new fields.

Suggested issue:

• Extend telemetry JSONL schema for advanced trace events

Phase 3: Layer/Submodule Timing

Purpose: produce real internal timing events.

Tasks:

• Map tensor names to logical submodules.
• Aggregate tensor timing into submodule events.
• Add capture filters for layers.
• Display events in replay and dashboard.

Exit criteria:

• Real run shows at least one selected layer with attention/MLP/norm timing.

Suggested issue:

• Capture real layer and submodule timing events

Phase 4: Activation Stats

Purpose: compute real tensor statistics safely.

Tasks:

• Add bounded tensor sampling.
• Compute mean/min/max/sparsity.
• Add anomaly flags.
• Display stats in Runtime Metrics and Anomaly Log.

Exit criteria:

• Real run shows activation stats for at least one selected tensor or a clear unavailable reason.

Suggested issue:

• Capture real activation statistics for selected tensors

Phase 5: Attention Matrix Capture

Purpose: replace placeholder heatmap with real attention data.

Tasks:

• Attempt callback-only attention matrix capture.
• If unavailable, add guarded llama.cpp source instrumentation.
• Store selected layer/head/window in JSONL.
• Render real heatmap in dashboard.

Exit criteria:

• make dashboard-real can show Attention Heatmap [real] for a tiny prompt.

Suggested issue:

• Capture and display real attention heatmap

Phase 6: Final Demo Hardening

Purpose: make the advanced flow reliable for submission.

Tasks:

• Add make advanced-demo.
• Add make test-advanced-telemetry.
• Update README and demo docs.
• Record a short video demo if required by submission form.

Exit criteria:

• One command generates an advanced trace.
• Replay and dashboard both show real internal data.
• Limitations are honest and clearly documented.

Suggested issue:

• Add advanced demo command and documentation

Suggested Small Issues

Create these as separate GitHub issues:

1. REQUIRED: discover llama.cpp tensor names through cb_eval
2. REQUIRED: extend telemetry schema for attention and activation data
3. REQUIRED: add layer/submodule event grouping
4. REQUIRED: add bounded activation stats capture
5. REQUIRED: render real attention matrices in TUI
6. REQUIRED: add advanced telemetry tests and fixtures
7. REQUIRED: update README/demo docs for advanced tracing

Ownership Suggestion

Person A:

• llama.cpp callback integration
• runtime capture flags
• layer/submodule timing
• activation statistics
• validation scripts

Person B:

• TUI rendering for real attention heatmap
• dashboard panel updates
• replay display changes
• fixtures and screenshots
• README/demo updates

Shared:

• final testing
• demo script
• submission explanation

Demo Script After Implementation

Expected future commands:

make build
make model
make advanced-demo PROMPT="Explain attention simply"
make replay-advanced
make dashboard-advanced

In the demo, show:

1. A real prompt running locally.
2. Replay output with internal layer/submodule events.
3. Runtime Metrics with real shape/dtype/activation stats.
4. Anomaly Log with any detected bottlenecks.
5. Attention Heatmap marked [real].

Important Honesty Rule

Do not claim fake or placeholder data is real.

Until attention matrices are actually captured from model execution, the dashboard must keep labeling the heatmap as placeholder or not captured.

Once real capture works, the dashboard should clearly say [real] and show the source layer/head/window.