Runtime Adapters & Integration Layer

Adapter Architecture

A Runtime Adapter (Backend) isolates framework-specific execution details behind a uniform interface. Whether an AI application relies on PyTorch, JAX, TensorFlow, cloud LLM endpoints (OpenAI, Anthropic), open-source vLLM inference engines, or vector database indexes, NeuronScope observes every component using identical platform contracts.

Recorders, probes, evaluation engines, and profiling subroutines never call framework APIs directly. Instead, they interact with the BackendAdapter protocol. This decouples the core platform from framework updates and maintains stable binary serialization in the .nsz telemetry format.

The adapter layer provides three core services to the platform:

• Layer Enumeration: List all named layers, modules, or graph nodes in a model so the platform can attach probes at the right locations without framework-specific knowledge.
• Hook Installation: Install forward-pass intercept hooks that fire on each layer execution and write tensor snapshots into the active Scope telemetry buffer.
• Tensor Snapshots: Convert framework-native tensors (PyTorch Tensor, JAX Array, TF EagerTensor) into portable TensorSnapshot objects that are serializable to the .nsz format without holding a reference to GPU memory.

Supported Runtimes & Adapters

The following adapters ship with NeuronScope and are automatically selected based on the framework of the model object passed to ns.Scope or ns.Model.wrap():

• torch — Hook-based adapter supporting PyTorch, HuggingFace Transformers, PEFT/LoRA, FSDP distributed training, DeepSpeed ZeRO, and torch.compile graph boundaries. Uses register_forward_hook and register_full_backward_hook for non-invasive intercepts.
• jax — Trace-based adapter for Flax and Equinox models. Captures intermediate arrays at JIT compilation boundaries using jax.make_jaxpr tracing. Compatible with pmapand vmap parallelism.
• tf — Graph-based adapter executing inside tf.function tracing loops. Captures layer outputs using Keras callback hooks and eager execution intercepts.
• llm_api — Endpoint adapter for OpenAI GPT-4o/o1, Anthropic Claude, Azure OpenAI, Google Gemini, Mistral, Cohere, and local vLLM/Ollama inference servers. Wraps the API client, capturing prompt tokens, completion tokens, latency, finish reason, and cost per call.
• vector_store — Retrieval adapter for Pinecone, Qdrant, Chroma, Milvus, Weaviate, pgvector, and FAISS. Intercepts search queries, captures retrieved document IDs, scores, and metadata for retrieval precision analysis.
• onnx — Read-only graph analyzer for exported ONNX neural network architectures. Computes architecture fingerprints and layer shape signatures without executing inference.
• hf_transformers — High-level HuggingFace-specific adapter providing tokenizer input/output capture, attention pattern recording, and generation config fingerprinting.

PyTorch Adapter

The torch adapter is the most comprehensive adapter in the platform. It supports every major PyTorch deployment pattern:

import neuronscope as ns
import torch
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2--hf")

# Wrap with NeuronScope — adapter auto-selected as "torch"
ns_model = ns.Model.wrap(model)
print(f"Adapter: {ns_model.adapter_name}")  # "torch"
print(f"Layers:  {len(ns_model.layers)}")   # 320

# Attach layer activation probes at specific layers
with ns.Scope("llama_inference") as scope:
    ns_model.attach_probe(
        layers=["model.layers.0", "model.layers.15", "model.layers.31"],
        kind=ns.probes.LayerActivations,
    )
    output = model.generate(input_ids, max_new_tokens=200)

artifact = scope.save("llama_run.nsz")
print(f"Recorded {len(artifact.telemetry.layer_snapshots)} layer snapshots")

JAX / Flax Adapter

The jax adapter captures intermediate array values at JIT boundary checkpoints and functional transformation boundaries. Unlike PyTorch hooks, JAX intercepts are implemented using custom jax.lax.scan wrappers and pytree traversal:

import neuronscope as ns
import jax
import jax.numpy as jnp
from flax import linen as nn

class TransformerBlock(nn.Module):
    @nn.compact
    def __call__(self, x):
        # NeuronScope intercepts here transparently
        x = nn.SelfAttention(4)(x)
        x = nn.LayerNorm()(x)
        return nn.Dense(256)(x)

model = TransformerBlock()
ns_model = ns.Model.wrap(model, adapter="jax")

# JAX adapter uses traced checkpoints — no Python-level hooks needed
with ns.Scope("jax_inference") as scope:
    params = model.init(jax.random.PRNGKey(0), jnp.ones((1, 16, 64)))
    output = jax.jit(model.apply)(params, inputs)

scope.save("jax_run.nsz")

LLM API Adapters

The llm_api adapter instruments API clients to capture full request/response telemetry without modifying your existing API calls:

import neuronscope as ns
import openai

# Instrument the OpenAI client — all calls are automatically traced
client = ns.hooks.instrument_openai(openai.OpenAI())

with ns.Scope("gpt4_rag_query") as scope:
    response = client.chat.completions.create(
        model="gpt-",
        messages=[{"role": "user", "content": query}],
        temperature=0.0,
    )

# Captured telemetry includes:
art = scope.artifact
print(art.llm_spans[0].prompt_tokens)    # e.g. 1024
print(art.llm_spans[0].completion_tokens)# e.g. 187
print(art.llm_spans[0].latency_ms)       # e.g. 842
print(art.llm_spans[0].finish_reason)    # "stop"

Vector Store Adapters

Retrieval quality is one of the biggest drivers of RAG system performance. The vector_storeadapter captures every retrieval step — query embedding, search scores, returned document IDs, and retrieval latency — enabling the comparison engine to isolate retrieval regressions:

import neuronscope as ns
import pinecone

index = pinecone.Index("my-index")
ns_index = ns.Dataset.wrap_vector_store(index)

with ns.Scope("rag_retrieval") as scope:
    query_embedding = embed_model.encode(query)
    results = ns_index.query(vector=query_embedding, top_k=5)
    # Retrieval telemetry is captured automatically:
    # - query embedding fingerprint
    # - top-k document IDs and scores
    # - query latency
    # - retrieved metadata

retrieval_span = scope.artifact.retrieval_spans[0]
print(f"Top score:     {retrieval_span.top_score:.}")
print(f"Query latency: {retrieval_span.latency_ms:.1f}ms")

Custom Adapter Protocol

Enterprise teams with proprietary inference engines, custom model servers, or non-standard frameworks can implement the BackendAdapter protocol to integrate seamlessly with the full NeuronScope platform:

import neuronscope as ns
from typing import Any, Callable, Iterator

class CustomInferenceAdapter(ns.BackendAdapter):
    """Adapter for a proprietary enterprise inference engine."""
    name = "enterprise_vllm"

    def matches(self, model_obj: Any) -> bool:
        """Return True if this adapter handles the given model object."""
        return hasattr(model_obj, "generate_engine")

    def list_layers(self, model_obj: Any) -> list[ns.LayerSpec]:
        """Enumerate all layers the platform can probe."""
        return [
            ns.LayerSpec(id=l.id, name=l.name, shape=l.output_shape)
            for l in model_obj.list_layers()
        ]

    def install_observer_hook(
        self,
        model_obj: Any,
        layer_id: str,
        callback: Callable[[ns.TensorSnapshot], None],
    ) -> None:
        """Attach a non-invasive callback to a named layer."""
        def _hook(raw_tensor):
            snapshot = ns.TensorSnapshot.from_numpy(raw_tensor.to_numpy())
            callback(snapshot)
        model_obj.add_interceptor(layer_id, _hook)

    def remove_hooks(self, model_obj: Any) -> None:
        """Called by the Scope on __exit__ to clean up all hooks."""
        model_obj.clear_interceptors()

    def fingerprint_model(self, model_obj: Any) -> str:
        """Return a stable fingerprint for this model's weights."""
        return ns.fingerprint(model_obj.get_weight_bytes())

# Register adapter — auto-discovered at runtime initialization
ns.registry.register_adapter(CustomInferenceAdapter())

Isolation & Memory Safety

All tensor snapshots are detached and decoupled from framework autograd graphs before being written into the telemetry buffer. This is non-negotiable for memory safety:

• PyTorch: Tensors are captured as .detach().cpu().numpy() copies. No gradient graph references are held.
• JAX: Arrays are materialized via jax.device_get(), which blocks until the asynchronous dispatch is complete and the array is in host memory.
• TensorFlow: EagerTensors are converted via .numpy(), which copies from device memory into host-side NumPy arrays.
• LLM APIs: Only serializable response data (strings, integers, floats) is captured — no framework objects are retained.

Overhead Budget & Safety Valve

NeuronScope continuously measures its own overhead. If the total observability overhead exceeds the configured budget (default: max_overhead_pct=15), the platform automatically:

• Emits a warning with the current overhead percentage and which probe is responsible.
• Switches from full tensor capture to summary statistics (mean, std, norm) automatically.
• If overhead remains above threshold, disables non-critical probes (e.g., attention maps) while retaining core execution traces.
• If overhead exceeds 2x max_overhead_pct, raises ns.OverheadBudgetExceeded and terminates the scope cleanly.

import neuronscope as ns

platform = ns.Platform(
    max_overhead_pct=5,        # strict: abort if overhead > 5%
    overhead_action="warn",    # "warn" | "reduce" | "abort"
)

# Or override per-scope for batch profiling runs where overhead is acceptable
with ns.Scope("profiling_run", max_overhead_pct=50) as scope:
    # Overhead limit is relaxed for this scope only
    output = model.generate(inputs)