DSP Framework Overview

End-to-end overview of the whole Dynamic Shape Plan (DSP) framework — architecture from SameDiff graph to native replay, plan caching, state management, execution backends, LLM decode integration, mul

The Dynamic Shape Plan (DSP) framework is the execution substrate that turns a SameDiff graph into a compiled, cached, natively-executed plan that can be replayed with near-zero per-step overhead. It spans four layers of the codebase:

Java plan layer (org.nd4j.autodiff.samediff.execution) — compiles the graph into a flat plan, manages caches, drives execution.
Native execution engine (libnd4j/include/graph) — slots, segments, shape freezing, CUDA graph capture/replay, JIT backends.
Diagnostics — the DSP_DIAG framework, execution traces, and op timing.
Tooling — benchmark configs, validation suites, and assertion helpers used to keep all of the above correct.

This page is the map of the whole framework and how the pieces interact — especially its state management and DSP_DIAG. The companion page DSP Execution Engine is the deep reference for individual engine features (the 26-pass graph optimizer, buffer coloring, the standalone runtime SDK and language bindings, op timing export). Read this page first.

Why DSP Exists

Classic SameDiff execution (InferenceSession) is an interpreter: every call re-walks the graph, resolves variables through maps keyed by strings, allocates outputs, and dispatches ops one at a time. For a transformer decode step that executes thousands of ops to produce one token, interpreter overhead — Java dispatch, shape re-inference, allocation, and per-op GPU launch latency — dominates the actual compute.

DSP replaces this with a compile-once, replay-many model:

The graph is compiled once into a flat, integer-indexed plan (no string lookups on the hot path).
Shapes are inferred during a short warmup, then frozen.
On GPU, whole regions of the plan are captured into CUDA graphs and/or JIT-compiled into fused Triton/NVRTC/PTX kernels, then replayed with a single launch per region.
Dynamic shapes are handled by caching one plan per shape signature rather than by re-interpreting: a prefill pass at sequence length 512 and a decode pass at sequence length 1 are two cached plans, each individually frozen and replayed.

The result is that the steady-state cost of a step approaches the cost of the GPU work itself.

The Big Picture

sd.output(placeholders, outputNames)
        │
        ▼
InferenceSession  ── DSP disabled or unsupported graph ──► legacy interpreter path
        │ (DSP fast path, default ON)
        ▼
DynamicShapePlanCompiler.compile()          [Java, once per output set]
   graph DAG → DynamicShapePlan
   (slots, wiring, liveness, op traits, loop regions)
        │  serialize (DSP1 binary)  ⇄  plan disk cache (~/.kompile/cache/dsp)
        ▼
NativeOps.dispatchNativePlan()              [every call]
   native NativePlanCache — shape-keyed lookup
        │ hit: reuse NativeDynamicShapePlan     miss: deserialize + build segments
        ▼
NativeDynamicShapePlan.execute()            [native]
   segments → backend dispatch:
     • Triton / NVRTC / PTX JIT kernels
     • CUDA graph capture → replay
     • cuBLAS / batched GEMM gap ops
     • slot-by-slot fallback
        │
        ▼
outputs (optionally zero-copy via sd.outputDirect())

Three facts shape everything else in the framework:

One plan per (requested outputs, input-shape signature, execution mode). Shape changes never mutate a plan; they select or create a different one.
Replay bakes device pointers. Once a CUDA graph is captured, every buffer address it references must stay stable — most of DSP's state management exists to guarantee or detect this.
State advances through explicit phases and only demotes deliberately (error recovery), never silently.

Layer 1: Java — From Graph to Plan

The InferenceSession gate

sd.output(...) routes into DSP when the system property org.nd4j.inference.dynamicShapePlan is true (the default) and the call does not use a secondary placeholder map. Everything else — and any graph the compiler declines — falls back to the legacy interpreter transparently. Runtime toggles:

InferenceSession.setDynamicShapePlanEnabled(boolean) — global kill switch (primarily for tests).
SameDiff.dspAutoCompileEnabled (default true) — when false, DSP is only used if a plan was already compiled; nothing is compiled implicitly.

Plan compilation

DynamicShapePlanCompiler.compile() translates the forward execution DAG into a DynamicShapePlan:

Slots. Every op becomes a DynamicShapeSlot with pre-wired integer indices. Inputs are encoded as inputSourceIndices[]: values >= 0 reference a prior slot's output; negative values reference an external input (encoded -(index+1)). Op arguments (iArgs/tArgs/bArgs/dArgs/sArgs) are frozen into the slot.
External inputs are classified by source type: SOURCE_CONSTANT (weights), SOURCE_VARIABLE (mutable model state), SOURCE_PLACEHOLDER (per-call inputs). This classification drives all later staging and freezing behavior.
Op traits. For each op the compiler queries the native trait table via JNI (getOpTraits(opName)). libnd4j/include/ops/impl/OpTraitTable.cpp is the single source of truth (~650 ops) for properties like OP_TRAIT_VALUE_DEPENDENT_SHAPE (shape depends on tensor values, not just input shapes), OP_TRAIT_FULLY_WRITING, OP_TRAIT_ATTENTION, OP_TRAIT_DYNAMIC_OUTPUT_SIZE, and Triton mappability. The compiler additionally promotes or demotes the outputShapeDependsOnInputValues flag for a few special cases (reshape with one input is demoted; concat with axis-in-last-array, expand_dims without iArgs, multi-input argmax/argmin are promoted).
Liveness analysis produces releaseAtStep[][] so intermediate buffers are freed at their last consumer.
Loop regions (Switch/Merge/Enter/Exit/NextIteration) are detected and recorded as LoopRegion metadata so while-loops execute inside the plan.

The compiler returns null (legacy fallback) for graphs it cannot support: external error functions, invoke sub-graphs, no-input tensor/random ops, UTF string inputs, self-referencing inputs, and malformed control flow.

Plan serialization and the disk cache

Plans serialize to a compact binary (DSP1 magic, version 5) consumed directly by C++. DspPlanDiskCache persists serialized plans under ~/.kompile/cache/dsp/dsp_plan_cache, keyed by an FNV-1a hash of the plan bytes and validated against the native build fingerprint — a rebuilt native library invalidates stale plans automatically. An override directory (nd4j.dsp.planCache.overrideDir) can pre-seed plans (highest priority, never overwritten), and nd4j.dsp.planCache.forceRecompile=true bypasses the cache entirely.

Layer 2: Plan Caching and Dispatch

DSP has three caches with different lifetimes:

Cache

Keyed by

Lives

Purpose

Java model-level cache (SameDiff.dynamicShapePlanCache)

set of requested output names

with the SameDiff instance

avoid recompiling the Java plan

Native plan cache (NativePlanCache)

output-set hash + placeholder shape-content hash + placeholder count + execution mode + thread id + plan content hash

native, LRU with pinning

one live native plan per shape signature

Disk cache (DspPlanDiskCache)

FNV-1a of serialized plan + native build fingerprint

across JVM runs

skip Java compilation on restart

Every execution starts with redispatchForCurrentShapes(): the executor hands the cached plan bytes plus the shape-info pointers of all external inputs to dispatchNativePlan(...). The native cache then returns:

the same handle (same shapes — the normal steady-state case),
a different cached handle (a shape signature seen before — e.g. switching between prefill and decode plans), or
a new plan (first time this shape signature is seen).

Cache hits pin the plan against LRU eviction; unpinNativePlan releases it. Eviction respects pins, and cache capacity is bounded (ND4J_DSP_PLAN_CACHE_MAX_PLANS, default 64 on GPU, 4 on CPU; budgeted at a fraction of free GPU memory, default 0.25).

Rule: never clear the plan cache mid-inference. SameDiff.clearDynamicShapePlanCache() destroys compiled plans, frozen shapes, and captured CUDA graphs. Call it only at genuine generation boundaries (new model, new page/prompt with structurally different shapes), and always after resetSession(). Between steps of the same workload, the shape-keyed cache is designed to handle shape changes by itself.

Layer 3: Native Execution

Slots

The native NativeDynamicShapePlan mirrors the Java plan as an array of NativeSlot structures — op identity (hash + resolved op pointer), wiring, frozen args, flags (outputShapeDependsOnInputValues, view/identity capability, fused-chain membership, control-flow role), and a cached shape key. Slot outputs are the plan's working set; a per-slot generation counter increments on every write so readers can detect staleness.

Segments, islands, and gaps

The slot array is partitioned into segments — maximal runs of ops that can execute under one backend. Segment boundaries are forced by non-capturable slots:

ops with dynamic output sizes (where, unique, non_max_suppression, ... — anything requiring a host sync to learn its output shape), and
control-flow ops.

Within a capturable segment, IslandCapturePolicy further partitions slots into islands (captured into a CUDA graph or compiled into a JIT kernel) and gaps (executed eagerly between islands — typically cuBLAS matmuls and ops excluded from capture). Capture profiles are per-platform: CUDA has no island node limit; ROCm caps islands at 128 nodes and excludes attention; Metal caps at 64.

Gap ops run on a dedicated thread-local stream (tl_dspGapStream) unified with the island replay stream, so no per-op synchronization is needed. Contiguous gap matmuls with identical geometry are grouped into batched GEMM calls (cublasGemmBatchedEx), and cuBLAS/cuBLASLt handles use a pre-allocated capture-safe workspace.

Execution backends and `GraphExecutionMode`

GraphExecutionMode expresses the user's backend preference. The important values:

Mode

Meaning

AUTO (default)

resolve the best available backend per segment

SLOT_BY_SLOT

interpret op-by-op inside the plan; no fusion, no capture — the correctness reference

CUDA_GRAPHS

CUDA graph capture/replay only (no JIT kernels)

TRITON, NVRTC_JIT, PTX_JIT

force a specific GPU JIT backend

EMULATED_REPLAY

slot-by-slot execution that still walks the full replay lifecycle (CPU testing of GPU state machine)

SHAPE_INFERENCE_ONLY

propagate shapes without executing kernels

HIP_GRAPHS, LEVEL_ZERO, VULKAN, METAL, MLX, NNAPI, ARM_HYBRID, MLIR_CPU_JIT, ONEDNN, OPENVINO, ACL

alternative hardware backends

AUTO resolves per segment through a fallback chain — on CUDA: Triton → NVRTC → PTX → CUDA graphs → slot-by-slot; on CPU: Triton(-CPU) → MLX → oneDNN → OpenVINO → ACL → NNAPI → ARM hybrid → MLIR CPU JIT → slot-by-slot. If a mode requiring a GPU graph backend is requested on a CPU-only runtime, it is remapped to EMULATED_REPLAY. Set the mode with sd.setGraphExecutionMode(...) or -Dnd4j.dsp.graphExecutionMode=<NAME>.

Each mode carries a mode contract (ModeContract): notably, every mode that participates in replay requires deterministic cuBLAS (CUBLAS_PEDANTIC_MATH), which disables tensor-core accumulation reordering so that SLOT_BY_SLOT reference runs and replayed runs are bit-comparable. TF32 is opt-in via two independent flags: nd4j.cublas.tf32 (gap-op GEMMs) and nd4j.triton.tf32 (Triton kernels).

CUDA graph capture and replay

For a capturable segment the lifecycle is: run slot-by-slot during warmup, then wrap execution in cudaStreamBeginCapture/EndCapture, instantiate the graph, and from then on cudaGraphLaunch replays the whole segment in one call. Key mechanics:

Capture serialization. DeviceCaptureGuard allows only one plan to capture per device at a time (5-second timeout), because capture is process-global per stream.
Capture-safe allocation. Allocations inside capture come from dedicated bump workspaces (default 512 MB device / 32 MB host) and the pooled allocator; a capture that OOMs is deferred (OOM_DEFERRED) and retried up to 3 times, every 4 executions.
Stream authority. All capture-related thread-local state (tl_graphExecutionActive, tl_compositeCaptureStream, tl_dspExecutionStream, ...) is consulted through one authority — DebugHelper::inGraphCapture() / currentCaptureStream() / captureSafeStream() — so any code that must issue async work during capture routes it onto the correct stream.
Replay shapes. A segment can replay as one monolithic graph, or as a composite schedule alternating captured islands and eager gaps; adjacent islands merged at freeze time share a merged replay handle.

Replay validity is enforced by identity keys — see Capture identity and recapture below.

Triton and the JIT backends

The Triton backend (built with -Dlibnd4j.triton=ON; gated by HAVE_TRITON in CMake) JIT-compiles segment sections into fused MLIR→PTX kernels:

Mappability comes from the same OpTraitTable/TritonIRBuilder op table (~200 ops mapped to MLIR dialect ops). Non-mappable slots inside a Triton segment become ordered native ranges executed between kernels.
Section fusion (nd4j.triton.sectionFusion, default on) merges consecutive sections into larger kernels, optionally gated by a fusion-benefit score.
Caching is two-level: an in-memory segment cache keyed by {slot range, shape key, device, internal dtype hash, compileAll, graphCapture, include/exclude-op hashes}, and an on-disk PTX cache (~/.kompile/cache/triton, hash of IR + arch + launch config). Compiled CUmodules are tracked per device under an LRU residency budget (default 512 MB) and reloaded from disk if evicted; all modules are batch-preloaded at plan freeze. A missing execution-time entry reports the requested shape/device/dtype/compileAll/graphCapture identity, the number of live entries for that segment, and a bounded candidate summary (approximately 1 KiB). This makes wrong-device and dtype/shape identity drift much easier to distinguish without producing an unbounded diagnostic.
Argument tables. Kernel argument buffers are refreshed before each replay; options exist to consolidate all sub-kernel arg tables into one pinned H2D copy (nd4j.triton.consolidatedArgTable) and to skip updates for kernels whose arguments are all frozen (nd4j.triton.argDirtyTracking).
Graph composition. By default (nd4j.triton.graphCapture=true) Triton kernel launches and their surrounding gap ops are captured together into a CUDA graph, so steady state is a single graph launch even for mixed Triton+cuBLAS segments.

NVRTC (in-process CUDA C JIT) and PTX (template) backends share a common segment cache (JitSegmentCacheKey) and serve as the AUTO fallbacks when Triton is unavailable.

State Management

DSP's state model is the part most worth understanding. There are three nested state machines — plan, segment, slot — plus the identity keys that guard replay.

Plan phase

SLOT_BY_SLOT  ──setShapesFrozen(true)──►  SHAPES_FROZEN  ──automatic──►  REPLAYING
                                                                             │
                                                              (no segment could capture)
                                                                             ▼
                                                                      REPLAY_BLOCKED

SLOT_BY_SLOT — initial phase. Shapes may change every call; every op is dispatched individually; shape inference runs each time.
SHAPES_FROZEN — entered when the application calls setShapesFrozen(true) (directly, via DspHandle, or implicitly by pipelines such as the LLM GenerationPipeline after prefill). Freezing asserts "all shapes are now constant": shape inference is skipped, and freeze-time transforms run — elementwise-chain fusion, segment merging (nd4j.dsp.freezeMergeSegments), optional recompile (nd4j.dsp.freezeRecompile), and buffer coloring.
REPLAYING — promoted automatically once every segment is pointer-stable and replay-ready for two consecutive executions. This is the steady state where throughput comes from; on the Java side PlanPhase.POINTERS_STABLE survives only as a deprecated alias of REPLAYING.
REPLAY_BLOCKED — terminal Java-visible state when nothing could be captured; execution continues slot-by-slot.

Phases only move forward in normal operation. The native engine can demote a phase deliberately for error recovery (demotePlanPhase(target, reason)) — e.g. when a buffer bound into a frozen plan is invalidated — after which the affected segments re-warm.

nd4j.dsp.noFreeze=true prevents the frozen transition entirely (an isolation/debugging tool, not a production setting).

Segment lifecycle

NEEDS_WARMUP → NEEDS_COMPILE → CAPTURE_PENDING → REPLAYING
                     │                │
                     ▼                ▼
                  FAILED         OOM_DEFERRED (retry)

The typical schedule, indexed by the plan's execution count after freezing:

Execution

What happens

Warmup — segment runs slot-by-slot; shape caches and output buffers materialize

Compile — JIT backends compile asynchronously; detectFrozenConstants() runs

Capture — CUDA graph capture and instantiation

Replay — steady state; the executor may take executeSteadyState(), a fast path that skips lifecycle validation, view refresh, and diagnostics (preconditions: plan is REPLAYING and executions ≥ 3)

A segment that fails compilation permanently is marked FAILED and runs slot-by-slot forever; a capture OOM is retried on the schedule described earlier.

The Java view of this machine is ExecutionPhase (WARMUP, COMPILING, COMPILED, REPLAYING, SLOT_BY_SLOT), observable per segment through sd.dsp().

Slot phase

Individual slots move BUILDING → SEALED. A sealed slot has a stable frozen execution context. Slots proven independent of all external inputs are sealed as frozen constants (SEALED + isConstant): their outputs are computed once and the slot is skipped on every subsequent execution.

detectFrozenConstants() (first frozen execution) computes this by propagating "depends on external input" through the wiring graph from every PLACEHOLDER/VARIABLE external input; ops with the SHAPE_ONLY_OUTPUT trait don't propagate the taint (their output depends only on shapes, which are frozen). The frozen output pointers are snapshotted — if a buffer is later reallocated, the slot falls back to re-execution rather than serving a dangling pointer. Marking an input variable after the fact (markExternalInputVariable) clears the detection so it re-runs.

Capture identity and recapture

A captured CUDA graph is only valid while everything it baked in still holds. Each segment tracks four identity keys computed at capture time:

Key

Hash of

Drift means

cachedShapeKey

segment output shapes

shapes changed → rebuild

capturedInputAddrKey

device addresses of external inputs

an input buffer moved → recapture

capturedCreateValueKey

input values of constant-generation ops

a fill/zeros_like-style op would produce different content

capturedSlotAddrHash

device addresses of all slot outputs

an intermediate buffer moved → recapture

plus a generation counter pair (argTableGeneration vs capturedArgGeneration). markArgsStale() is the single entry point that bumps the generation and resets stability counters whenever any address changes; needsArgRefresh() compares the pair before replay. When a refresh suffices (only argument pointers changed), the graph's baked nodes are updated in place; when an identity key mismatches, the segment goes through resetForWarmup() and re-warms/re-captures from scratch.

This is why "mysterious" recapture loops are almost always address churn: something reallocates a buffer every step (a fresh output array, an unpooled allocation, a view over a moving buffer), each replay check fails, and the plan spends its life re-warming. The GRAPH_REPLAY diagnostics category (below) is built to expose exactly this.

External inputs: constants, variables, placeholders

How an external input is treated in the frozen/replaying phases depends on its source type:

Constants (weights) — assumed immutable; buffers are snapshotted at freeze; replaced weights are detected and rejected.
Placeholders (per-call inputs, host-written) — synced host→device into stable staging buffers before every execution, so the captured graph reads a fixed address regardless of which Java array the caller passes.
Variables (mutable device-resident state, e.g. KV caches) — staged device→device, or registered as device-managed so the plan reads them in place; their raw addresses join capturedInputAddrKey, forcing recapture if they move.

Inputs must be declared mutable for the variable treatment: sd.markDynamicShapePlanInputsMutable(names) or sd.dsp().markVariable(...). A newBorrower flag on dispatch marks views over external inputs for re-validation (the slot is unsealed and the view rebound) when a different caller's buffers back the same plan.

Two runtime guards close the loop: a native execution can return STALE_BUFFER, which makes the Java executor re-dispatch and retry once; and a global frozen-executor count suppresses clearTADCache() while any plan holds baked TAD device pointers (freeing them would leave dangling kernel arguments — CUDA error 700).

Staging-transfer failures

A staging synchronization or transfer failure aborts graph execution and is copied into ND4J's native error reference. Current CUDA builds include the failing stage, DSP status, CUDA error code/string, and—when a specific external input was being copied—the external index/name, source and destination pointers, byte count, source device, and target device:

DSP staging failed: stage=fast_path_transfer ext=3 name='input_ids' src=0x... dst=0x... bytes=... srcDevice=1 targetDevice=0 status=... cudaError=... (...)

Treat this as the primary failure, not as a warning followed by an unrelated decode error. Reproduce with EXECUTE,TRANSFER,STREAM_SYNC,MULTI_DEVICE,MEMORY diagnostics enabled. The stage distinguishes null-buffer validation, direct device transfer, and assignment failures on the fast or slow staging path; the device and pointer fields identify cross-device placement or stale-buffer ownership problems.

For contributor tests, -Dnd4j.dsp.stagingFault=<fault> maps to the native ND4J_DSP_STAGING_FAULT hook. Supported values are cuda_get_device, allocation, cross_stream, d2d_copy, and stream_sync. Each injected failure must throw before graph execution can fall back to raw external arrays; this switch is for fail-closed validation, never production configuration.

Shape changes, unfreezing, and multi-plan switching

A shape change while frozen does not unfreeze anything. The dispatch layer simply returns a different plan handle for the new shape signature; the executor accepts it, marks it frozen too, and each plan maintains its own capture state. This is the prefill/decode pattern: two shapes, two plans, both eventually replaying.

The only legal "unfreeze" is teardown: resetForNextPage() destroys the native handle (releasing CUDA graphs and replay workspaces) and resets the executor to its initial state. For page/prompt turnover where shapes are identical, clearOutputCaches() clears Java-side output caches and trims the memory pool without touching captured graphs — much cheaper.

Thread safety

Plans are compiled once per SameDiff but executed through thread-local executors; the native cache key includes the thread id, and concurrent compilation, freezing, replay, and cache access are exercised by a dedicated regression suite (DspConcurrentPlanSharingTest). Diagnostics recording is lock-protected and safe from any thread.

LLM Decode Integration

DSP is the engine underneath the LLM & VLM stack; autoregressive decoding is its most demanding client and drove several dedicated mechanisms:

In-graph KV cache scatter. configurePlanKvScatter(...) registers pairs of (per-step present KV slot → persistent full-size KV buffer) together with geometry and a device-resident position counter. After every execution the plan scatters the step's K/V into the persistent buffers at the current position (executeKvScatterPostExec), entirely on device. The persistent buffers are registered as device-managed external inputs, so they participate in capture identity without host staging. Position is managed with resetPlanKvCachePosition / getPlanKvCachePosition.
Fixed-size output slots. setPlanOutputSlotMaxSizes(...) (surfaced as nd4j.dsp.maxKvCacheLength for KV shapes) pre-allocates growing outputs — attention masks, position tensors, KV present slots — at their maximum size, keeping addresses stable across decode steps so capture survives step-to-step growth.
Native decode loop. The autoregressive_decode op moves the entire token loop into native code: it takes a plan handle, repeatedly calls executeSteadyState(), samples tokens, applies KV scatter, and exits on EOS — eliminating Java↔native transitions per token. linear_attention_decode covers linear-attention model families.
Pipeline layer. GenerationPipeline (module samediff-llm) orchestrates the pattern: enable DSP, run prefill, freeze shapes, decode through the fixed-buffer path with retained InGraphKvState buffers; GenerationSession resumes decoding from retained KV state without re-prefill.

The plan cache handles the prefill↔decode shape flip automatically — the correct integration is always "keep DSP enabled and let the cache switch plans", never disabling DSP for one phase.

Multi-Device Execution

DSP can shard one plan across multiple GPUs.

Placement

At compile time DynamicShapePlan.assignDevices() distributes slots across devices proportionally to pool-aware free memory (CUDA free memory plus reclaimable pool reserve). Device 0 and P2P-capable peers get their full budget; non-P2P devices get a fraction (nd4j.dsp.nonP2pBudgetFraction) that is zero unless sharding is enabled with -Dnd4j.dsp.multiGpuShard=true. After the proportional split, splitParallelGroups() redistributes sibling ops (Q/K/V projections, gate/up MLP projections) round-robin across devices. DevicePlacementPlanner offers explicit strategies on top: SINGLE_DEVICE, MEMORY_FIT (greedy spill), PIPELINE_PARALLEL (liveness-minimizing cuts), and CUSTOM (variable→device map). Architecture compatibility is enforced via DeviceKey (compute capability), so arch-specific compiled kernels never land on an incompatible GPU. -Dnd4j.dsp.singleGpu=true forces everything onto one device.

Execution across devices

Before each native call the executor determines which device consumes each external input:

Inputs consumed only by a secondary device stay there — their addresses are baked into that device's captured graphs.
Other off-device inputs are migrated with dbAsyncCrossDeviceCopy (a cudaMemcpyPeerAsync on the DSP execution stream; the driver stages through host memory for non-P2P pairs). Views are materialized (dup()) before migration so no view ever aliases a freed source buffer.
Migrated constants are kept in a replica cache so weights are not re-copied every decode step.

The native execution device is chosen once per plan by data locality (the device holding most input bytes) and stays sticky. In shard mode, outputs produced on secondary devices are migrated back to device 0 after execution.

Memory budgets and failover

Device memory allocation follows one failover chain: pooled cudaMallocAsync → proactive soft-limit rerouting (setMemoryPoolSoftLimitPercent) → per-process maxDeviceMemory cap (Nd4j.getEnvironment().setMaxDeviceMemory(bytes)) → pool trim + retry → peer-GPU failover → non-peer-GPU failover (managed memory; non-peer devices are deliberately never excluded) → pinned host memory as last resort. All capture-time allocations must go through the pool so allocation/free pairing keeps baked addresses valid. The executor trims pools at slot-cache teardown points to keep freed memory reusable on-stream.

Testing without hardware: stub topologies

DeviceMemoryManager.configureStubTopology(List<StubDeviceDescriptor>) simulates arbitrary multi-GPU topologies on a single machine — fake devices with configurable memory, peer sets, and modeled bandwidths (300 GB/s peer, 12 GB/s non-peer, 16 GB/s host) plus a StubDeviceContextProvider that records every device switch. Always pair it with clearStubTopology() in test teardown. DspReplayTransferAnalytics records per-step, per-segment cross-device transfers (with TransferReason codes such as CONSTANT_REPLICATION, KV_CACHE_SCATTER) and memory-pressure rerouting decisions; ReplayProfileManager captures and persists replay profiles enriched with those transfer stats. All of it reports through the MULTI_DEVICE and TRANSFER diagnostics categories.

DSP_DIAG: The Diagnostics Framework

DSP_DIAG is the unified, always-compiled-in tracing system for everything above. It is a lock-protected native singleton with a 65,536-event ring buffer, per-category aggregate statistics (which survive ring wrap-around), and up to 64 persistent segment terminal records capturing why each segment ended in its final state.

Header: libnd4j/include/graph/DspDiagnostics.h · Java mirror: org.nd4j.autodiff.samediff.diagnostics.DspDiagnostics.

block the SHAPES_FROZEN transition (debugging)

nd4j.dsp.freezeMergeSegments / nd4j.dsp.freezeRecompile

native: on / off

freeze-time segment merging / recompilation

nd4j.dsp.maxKvCacheLength

0

pre-allocate KV outputs at max length (fixed-buffer decode)

nd4j.dsp.singleGpu / nd4j.dsp.multiGpuShard / nd4j.dsp.nonP2pBudgetFraction

off / off / —

multi-device placement controls

nd4j.dsp.captureWorkspaceMb / nd4j.dsp.captureHostWorkspaceMb

native: 512 / 32

capture workspace sizes

nd4j.dsp.planCache.diskEnabled / .diskDir / .overrideDir / .forceRecompile

on / ~/.kompile/cache/dsp/... / — / off

plan disk cache

nd4j.dsp.diagnostics / .level / .file

off / summary / —

DSP_DIAG (see above)

nd4j.dsp.batchZero / nd4j.dsp.batchedGemm / nd4j.dsp.castElimination / nd4j.dsp.matmulSegmentation

native: on / on / on / off

execution optimizations

nd4j.cublas.tf32 / nd4j.triton.tf32

off / off

TF32 tensor-core math for cuBLAS gaps / Triton kernels

Native environment variables

Everything above has an ND4J_DSP_* / ND4J_TRITON_* env-var equivalent read directly by the native library — useful for non-JVM consumers of the runtime SDK. Notable env-only knobs: ND4J_DSP_PLAN_CACHE_MAX_PLANS (64 GPU / 4 CPU), ND4J_DSP_PLAN_CACHE_BUDGET_FRACTION (0.25), ND4J_DSP_CAPTURE_OOM_MAX_RETRIES (3) / ND4J_DSP_CAPTURE_OOM_RETRY_INTERVAL (4), ND4J_DSP_PROACTIVE_EVICT (on), ND4J_REPLAY_CACHE_DIR / ND4J_REPLAY_CACHE_ENABLED (on-disk replay profiles), ND4J_TRITON_CACHE_DIR (PTX cache), ND4J_TRITON_MODULE_RESIDENCY_BUDGET_BYTES (512 MB), ND4J_TRITON_BUILD_THREADS (8), and the debug isolation switches ND4J_DSP_DISABLE_VIEW_FASTPATH / ND4J_DSP_DISABLE_CAST_HWM / ND4J_DSP_DISABLE_WS_SKIP.

Remember that Surefire forks a JVM: in tests, always pass configuration as -D Maven properties (wired to env vars in platform-tests/pom.xml) — exported shell variables do not reach the test JVM.

Observability, Testing, and Benchmarking

`sd.dsp()` — the DspHandle

DspHandle is the runtime introspection facade: plan phase, per-segment execution phase and replay state, capture statistics and telemetry, slot inspection, external-input control (markVariable/markPlaceholder, setExtInput), step snapshots, and nativeJsonReport(). PlanIntrospection formats plans (including Graphviz toDot()), and DspDebugger runs full guided diagnostics (analyzeGraphReplay, trackReplayProgression, validateStreamSync, analyzeMultiDevice).

Assertions for tests

DspPlanAssertions turns lifecycle expectations into hard test assertions with full state dumps on failure: assertPhaseReached(sd, PlanPhase.REPLAYING, ...), assertFullyReplaying, assertPointersStable, assertNoCaptureFailures, assertNoAddressDrift, assertSegmentReplayMode, assertZeroPermCaptureFailures, assertNoHostOnlyOps, assertCompilationSealed, and dozens more. Any test that depends on DSP reaching steady state should assert it rather than assume it.

Named benchmark configurations

BenchmarkConfig (module samediff-llm) bundles coherent property sets; BenchmarkConfigApplier applies them. The canonical ones:

Config

Meaning

OPTIMAL

the production GPU configuration: Triton with section fusion + compile-all + graph capture, consolidated/dirty-tracked arg tables, batched GEMM, freeze-merge, TF32 on

SLOT_BY_SLOT

the correctness reference

CUDA_GRAPHS

capture/replay without JIT kernels

TRITON variants

Triton with specific features toggled (TRITON_NO_GC, TRITON_VERIFY, TRITON_SKIP_KERNELS, ...)

CPU_SLOT_BY_SLOT, CPU_CASCADE, CPU_OPENVINO, ...

CPU-backend configurations

Scripts (all in `platform-tests/`)

run-benchmark.sh — VLM decode benchmark. DSP-relevant flags map straight onto the properties above: --debug (ALL diagnostics), --diag-replay (GRAPH_REPLAY,SEGMENT,EXECUTE at full), --diag-stream (STREAM_SYNC,EXECUTE,TIMING), --diag-device (MULTI_DEVICE,TRANSFER,BACKEND,MEMORY), --diag-all, --diag-json FILE, --no-freeze, --no-triton, --triton-tf32/--cublas-tf32, --op-timing, plus an --audit mode that runs the DSP test suites.
run-llm-benchmarks.sh — multi-model LLM matrix (import/baseline/cuda-graphs/triton/fusion/optimizer/matrix/perplexity/quant tests).
run-validation.sh — TestDspValidation: token-level accuracy of every mode against SLOT_BY_SLOT (outputAccuracy, decodeStep, tf32Isolation, steady-state fast-path isolation), with configurable match rate (default ≥ 90%).
run-dsp-matrix.sh — the 8-config sweep against the golden baseline: SLOT_BY_SLOT_baseline, SLOT_BY_SLOT_batchedGemm, AUTO_defaults, AUTO_frozen, TRITON_sectionFusion, TRITON_compileAll, TRITON_frozen_batchedGemm, CUDA_GRAPHS_frozen.

Benchmark metrics

Generation results report four throughput numbers; know which one you are reading:

Metric

Window

overall tok/s

whole generation including prefill and warmup

decode tok/s

all decode steps (step ≥ 1)

steady tok/s

decode steps 3+ — after the capture schedule settles

lateSteady tok/s

decode steps 20+ — the most stable measurement (reported as 0 for runs under 20 steps)

The effective score used for comparisons prefers lateSteady, then steady, then decode, then overall.

Regression suites

The core DSP correctness batch lives in platform-tests: DspLifecycleValidationTest and DspLifecycleExhaustiveTest (phase progression per mode), DspSlotLifecycleAuditTest (buffer lifetime vs frozen plans), DspConcurrentPlanSharingTest (thread safety), DspFrozenConstantInvariantTest, DspExtInputStalenessTest, DspCompositeReplayTest, DspHandleTest/DspHandleDataModelTest, TestDspShapePrePass, plus the model-level TestDspValidation and TestDspConfigEnumeration (exhaustive mode × optimizer × flag sweep). Changes to DSP internals are expected to pass the whole batch, not just a narrow reproducer.

Opt-in production fixtures

Large model-derived fixtures stay outside source control and their tests skip unless the relevant property is supplied. Run these commands from platform-tests/.

To capture the first non-finite DSP op during a failing generation, set -Dnd4j.dsp.nonFiniteFixtureDir=/absolute/private/path (or ND4J_DSP_NONFINITE_FIXTURE_DIR outside Surefire). GenerationPipeline writes fixture.properties plus binary input/output arrays after the failure is already detected. Model-derived prompts and tensors may be sensitive; inspect the directory and never commit it by default.

Replay a captured gated-delta failure through the independent recurrence, scratch native op, direct-state/actualLen op, one-slot DSP graph, and mixed Triton/native graph:

mvn test -Dbackend.artifactId=nd4j-cuda-12.9 \
  -Dtest=GatedDeltaRuleProductionFixtureTest \
  -Dgdr.production.fixture.dir=/path/to/gdr-fixture

Probe every gated_delta_rule input dtype and shape in an SDZ without executing the model graph:

mvn test -Dbackend.artifactId=nd4j-native \
  -Dtest=GatedDeltaRuleSdzDtypeProbeTest \
  -Dgdr.probe.sdz=/path/to/model.sdz

Validate the retained fixed-buffer transition from constrained Java decoding back to the ordinary native loop:

mvn test -Dbackend.artifactId=nd4j-cuda-12.9 \
  -Dtest=TestFixedBufferDecodeReuse#constrainedThenNativeHandoffUsesValidStagingInputs \
  -Dqwen.sdz.path=/path/to/qwen.sdz

For the larger production envelope and chat-schema-to-native transition, use TestKvMaxAllocationGemmaNames#constrainedSchemaCallsThenNativeGenerateDoNotFailStaging with -Dmodel.sdz.path=/path/to/staged-model.sdz.

For the gated-delta fixture, finiteness alone is insufficient: direct-state and DSP outputs must be bit-identical to the scratch native result for the full active length.

Operational Guidance

Hard-won rules for working with (and on) DSP:

Let the lifecycle run. Never pin SLOT_BY_SLOT (or any mode) to sidestep a capture/replay bug — it hides the bug and costs 5–10× throughput. The only production-legitimate mode override is choosing a backend; the phases themselves are not a tuning knob.
Keep DSP enabled across phases. Don't disable DSP for prefill/warmup and re-enable for decode; the shape-keyed cache exists precisely so both phases coexist as separate frozen plans.
Clear caches only at boundaries. clearDynamicShapePlanCache() belongs at model/page boundaries, after resetSession() — never between steps.
Respect buffer stability. After freezing, anything that reallocates a buffer visible to the plan (closing arrays a plan still references, swapping weights, taking views over moving buffers) either triggers recapture (slow) or is detected and rejected (error). Declare mutable inputs up front with markDynamicShapePlanInputsMutable.
Diagnose with DSP_DIAG, not printf. Targeted categories at detailed level are free enough to leave on while reproducing; full is for watching events live and will distort any timing you collect.
Benchmark in steady state. Quote lateSteady tok/s from ≥ 250-token runs; anything shorter measures warmup and capture, not the framework.

Design History (ADRs)

The framework's evolution is documented in the repository's ADRs/ directory. The most relevant records: 0061 (DynamicShapePlan execution), 0062 (Java-side shape inference), 0071 (Triton graph backend), 0073 (self-contained runtime SDK), 0078 (diagnostic framework extensions — STREAM_SYNC/MULTI_DEVICE/GRAPH_REPLAY), 0079 (native plan structural refactoring), 0080 (Triton fusion replay correctness), 0081 (view shape correctness diagnostics), 0082 (replay pointer stability and frozen steady state), 0084 (execution state simplification), 0089 (capture, replay, and gap-op orchestration), 0093 (plan disk persistence), 0094 (buffer coloring, pooling, passivation), 0096 (LLM generation pipeline), 0097 (decode-path performance), 0103 (execution-mode auction), 0104 (Triton tile-config PGO), 0105 (generation session continuation), 0106 (consolidated decode substrate).