> For the complete documentation index, see [llms.txt](https://deeplearning4j.konduit.ai/llms.txt). Markdown versions of documentation pages are available by appending `.md` to page URLs; this page is available as [Markdown](https://deeplearning4j.konduit.ai/en-1.0.0-rewrite/nd4j/overview-2/dsp-framework.md). # 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](https://github.com/KonduitAI/deeplearning4j-docs/tree/en-1.0.0-rewrite/docs/m2.1/nd4j/samediff/overview/README.md) 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: 1. **Java plan layer** (`org.nd4j.autodiff.samediff.execution`) — compiles the graph into a flat plan, manages caches, drives execution. 2. **Native execution engine** (`libnd4j/include/graph`) — slots, segments, shape freezing, CUDA graph capture/replay, JIT backends. 3. **Diagnostics** — the `DSP_DIAG` framework, execution traces, and op timing. 4. **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](https://github.com/KonduitAI/deeplearning4j-docs/tree/en-1.0.0-rewrite/docs/m2.1/nd4j/samediff/dsp/README.md) 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=`. 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](#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 `CUmodule`s 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 | | --------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | | 0 | **Warmup** — segment runs slot-by-slot; shape caches and output buffers materialize | | 1 | **Compile** — JIT backends compile asynchronously; `detectFrozenConstants()` runs | | 2 | **Capture** — CUDA graph capture and instantiation | | 3+ | **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=` 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](https://github.com/KonduitAI/deeplearning4j-docs/tree/en-1.0.0-rewrite/docs/m2.1/deeplearning4j/llm/overview/README.md); 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)` 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`. ### Categories Enable any comma-separated subset (or `ALL`). All 20 categories: | Category | Traces | | ----------------- | ------------------------------------------------------------------------------------------------------------- | | `COMPILE` | backend compilation, shape-change warmups, compile errors, warmup-complete banners | | `JIT` | kernel generation, PTX/cubin cache hits and misses | | `EXECUTE` | per-step execution flow, dispatch decisions, pointer-stability checks, address snapshots, buffer fingerprints | | `TIMING` | detailed timing breakdowns | | `MEMORY` | allocations, OOM/failover, pool state, slot write/zero events, DataBuffer create/delete | | `BACKEND` | backend selection and device placement decisions | | `SHAPE` | shape analysis, static/dynamic detection, shape-key changes | | `SEGMENT` | segment building and capturability decisions | | `FUSION` | op fusion, identity elimination, section merging | | `VERIFY` | golden comparison, output validation, NaN/Inf detection | | `KV_CACHE` | KV cache configuration, retention, scattering | | `FALLBACK` | fallback events and error recovery (always shown in the end-of-plan report) | | `TRANSFER` | H2D / D2H / D2D transfers | | `EMULATED_REPLAY` | emulated replay lifecycle | | `STREAM_SYNC` | stream synchronization, event ordering, sync-before-read patterns | | `MULTI_DEVICE` | device selection, P2P topology, cross-device migration, pressure rerouting | | `GRAPH_REPLAY` | capture begin/end, instantiation, launch, address validation, capture invalidations | | `SEGMENT_BUCKETS` | classification of non-capturable segments and gap analysis | | `LIFECYCLE` | plan/segment/array lifecycle transitions (plus a per-thread 512-event lifecycle ring) | | `COLORING` | buffer coloring and pool events | ### Levels | Level | Behavior | | ------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | | `summary` (default) | events recorded to the ring only; per-category stats accumulate; report printed at plan teardown | | `detailed` | per-step events into the ring, still no real-time output — zero fast-path distortion | | `full` | **every event echoed to stdout in real time** as `[DSP_DIAG] [CATEGORY] seg[N] slot N (op) Nus message` — required to watch events live, but it flushes stdout per event and forces some synchronous paths. **Never benchmark at `full`.** | When ND4J debug+verbose mode is on (`Nd4j.getEnvironment().setDebug(true)/setVerbose(true)`), everything echoes regardless of category mask or level. ### Enabling it From Maven test runs (mapped to env vars by `platform-tests/pom.xml`): ```bash mvn test -Dtest=MyTest \ -Dnd4j.dsp.diagnostics=GRAPH_REPLAY,SEGMENT,EXECUTE \ -Dnd4j.dsp.diagnostics.level=detailed \ -Dnd4j.dsp.diagnostics.file=/tmp/dsp-report.json ``` Directly via environment variables (any process, no rebuild): `ND4J_DSP_DIAGNOSTICS`, `ND4J_DSP_DIAGNOSTICS_LEVEL`, `ND4J_DSP_DIAGNOSTICS_FILE`. Programmatically: `DspDiagnostics.setCategories(mask)` / `setLevel(...)` / `record(...)` from Java, or the C API (`dspDiagSetCategories`, `dspDiagSetLevel`, ...). Focused helpers: `ND4J_DSP_TRACE_SLOT=` traces one slot's entire lifecycle; `ND4J_DSP_TRACE` / `ND4J_DSP_EXECUTION_TIMING` / `ND4J_DSP_NATIVE_DUMP_OUTPUTS` are legacy switches for the `EXECUTE` / `TIMING` / `VERIFY` categories; `nd4j.triton.verbose` auto-enables `COMPILE|JIT|BACKEND`; `nd4j.triton.verifyKernels` auto-enables `VERIFY` at `full` and cross-checks every Triton kernel against slot-by-slot output. ### Reports At plan teardown the framework prints a human-readable plan report (category stats plus any `FALLBACK` events) and, if a file path is set, writes a JSON report containing `planInfo` (slots, segments, steps, total time), `categoryStats`, `segmentTerminals` (immune to ring overflow), and the last ≤65,536 `events`. Both are also available on demand without teardown: ```java String json = sd.dsp().nativeJsonReport(); // full JSON report DspDebugger.runFullDiagnostics(...); // one-call: enable ALL, run, collect plan report + // replay state + multi-device report + issue list ``` ### Related instrumentation * **`DspExecutionTrace`** — a per-plan, zero-allocation structured event ring (typed events: `SEGMENT_DISPATCH`, `GRAPH_CAPTURED`, `GRAPH_REPLAYED`, `PHASE_TRANSITION`, ...) dumped to stderr on errors; it complements the string-based DSP\_DIAG. * **`OpTimingTracker`** (`nd4j.op.timing`) — independent per-op phase timing (validation/shape-calc/alloc/exec) with histograms, CSV, and Chrome-trace export; see the [DSP Execution Engine](https://github.com/KonduitAI/deeplearning4j-docs/tree/en-1.0.0-rewrite/docs/m2.1/nd4j/samediff/dsp/README.md) page for the full workflow. * **Capture safety.** Diagnostics helpers that read device memory (`dumpSlotBuffer`, `fingerprintArray`, ...) sync to host and must never run inside CUDA graph capture; they are gated behind `EXECUTE` + debug mode. Pure-metadata events (`SLOT_WRITE`, `SLOT_ZERO`) are safe anywhere. *** ## Configuration Reference The authoritative sources are `DspConfig::initFromEnvironment()` and `TritonConfig::initFromEnvironment()` on the native side, and `ND4JSystemProperties` / `DspProperties` on the Java side. The most useful knobs: ### Core (Java system properties) | Property | Default | Effect | | ---------------------------------------------------------------------------------------------------------- | ----------------------------------------- | ----------------------------------------------------------- | | `org.nd4j.inference.dynamicShapePlan` | `true` | master switch for the DSP fast path | | `nd4j.dsp.nativeExecutor.enabled` | `true` | native C++ plan execution (the only supported executor) | | `nd4j.dsp.graphExecutionMode` | `AUTO` | backend preference (see mode table) | | `nd4j.dsp.cudaGraphs.enabled` | `true` | CUDA graph capture/replay within segments | | `nd4j.dsp.jitMode` | `graph` | `graph`, `jit`, or `graph+jit` | | `nd4j.dsp.noFreeze` | off | 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: ```bash 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: ```bash 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: ```bash 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). *** ## See Also * [DSP Execution Engine](https://github.com/KonduitAI/deeplearning4j-docs/tree/en-1.0.0-rewrite/docs/m2.1/nd4j/samediff/dsp/README.md) — deep reference: graph optimizer passes and controls, buffer coloring tiers, op timing tracker, runtime SDK C API and language bindings, full configuration tables * [SameDiff Execution and Inference](https://github.com/KonduitAI/deeplearning4j-docs/tree/en-1.0.0-rewrite/docs/m2.1/nd4j/samediff/execution/README.md) — the `sd.output()` API surface * [SameDiff Overview](https://github.com/KonduitAI/deeplearning4j-docs/tree/en-1.0.0-rewrite/docs/m2.1/nd4j/samediff/overview/README.md) — graph model, variable types, training * [LLM & VLM Stack](https://github.com/KonduitAI/deeplearning4j-docs/tree/en-1.0.0-rewrite/docs/m2.1/deeplearning4j/llm/overview/README.md) — the generation pipeline built on DSP * [CUDA Backend](https://github.com/KonduitAI/deeplearning4j-docs/tree/en-1.0.0-rewrite/docs/m2.1/nd4j/backends/cuda/README.md) — GPU memory management, multi-GPU setup * [Memory and Workspaces](https://github.com/KonduitAI/deeplearning4j-docs/tree/en-1.0.0-rewrite/docs/m2.1/core-concepts/memory-and-workspaces/README.md) — workspace scopes and GPU memory reuse