Ctrlk
Community ForumND4J JavadocDL4J Javadoc
For the complete documentation index, see llms.txt. This page is also available as Markdown.

DSP Execution Engine

Complete guide to the Dynamic Shape Plan execution engine — compiled graph runtime with CUDA graph capture/replay, Triton/NVRTC/PTX JIT, 26-pass optimizer, and multi-backend dispatch

The Dynamic Shape Plan (DSP) engine is the compiled graph execution runtime introduced in DL4J 1.0.0-rewrite. It replaces the previous interpreter-style graph execution with a compile-once, replay-many architecture that delivers substantially lower inference latency — especially on NVIDIA GPUs, where it eliminates kernel launch overhead entirely via CUDA graph capture and replay.

Classifier requirement: DSP's core features — graph compilation, shape freezing, CUDA graph capture/replay, and the 26-pass optimizer — work with the base backend classifier. However, the JIT compilation features described on this page (Triton kernel fusion, NVRTC, PTX, and MLIR CPU JIT) require the -compile classifier variant (e.g., linux-x86_64-compile for CPU, linux-x86_64-cuda-12.9-compile for CUDA). Without -compile, GraphExecutionMode.TRITON, NVRTC, PTX, and MLIR_CPU automatically fall back to CUDA_GRAPHS or SLOT_BY_SLOT. See Hardware Backends — Classifier Variants for the full trade-off analysis.

Overview

Before DSP, executing a SameDiff graph meant traversing the graph node by node on every call, dispatching each operation through the op registry, and paying full kernel launch overhead for every operation. For large transformer models with hundreds of operations per forward pass, this launch tax dominates latency at batch size 1.

DSP changes this model. The first time a SameDiff graph is executed, DSP compiles it into a DynamicShapePlan: a flat, ordered sequence of "slots," one per SameDiff variable. After a warmup phase that learns the concrete shapes of every intermediate tensor, DSP freezes the shapes, captures the entire sequence of GPU operations into a CUDA graph, and replays that single pre-captured graph on every subsequent call. On the replay path, no individual kernel is launched from the Java side — the CUDA driver replays the full captured op sequence in a single API call.

The same compilation pipeline runs on non-CUDA hardware via alternative graph backends (Metal, Vulkan, oneDNN, NNAPI, MLX, and others) that provide analogous captured-graph replay semantics on their respective platforms.

Why Compilation Matters

A compiled plan has several advantages over interpreted execution:

  • Reduced launch overhead. Each kernel launch from Java costs 5–50 µs. A 200-op transformer graph accumulates 1–10 ms of pure overhead per forward pass. CUDA graph replay collapses all 200 launches to a single call costing under 10 µs.

  • Constant memory pointers. Once shapes are frozen, buffer addresses do not change between calls. The CUDA graph can be captured with baked-in pointers and replayed without pointer fixup.

  • Graph-level optimization. The 26-pass optimizer (described below) operates on the full graph and performs cross-op fusions — attention fusion, horizontal QKV projection fusion, algebraic simplifications — that are impossible to apply op-by-op at runtime.

  • JIT specialization. The Triton, NVRTC, and PTX JIT backends generate kernel code specialized to the exact shapes and data layouts of the compiled plan, enabling further micro-optimizations.

Plan Lifecycle

A DynamicShapePlan progresses through four phases. The active phase governs which execution path is taken on each call.

Phase 1: WARMUP

During warmup, DSP executes the graph slot by slot — one operation at a time — using the standard ND4J op dispatch. The purpose of warmup is to observe the concrete output shape of every intermediate tensor for the given input shapes. Each slot records its output shape and stores it in the plan's shape table.

The DynamicShapePlanExecutor maintains one execution stream per thread (tl_dspExecutionStream) plus a secondary stream for "gap" ops (tl_dspGapStream) — operations that fall outside the captured region and must continue running via normal dispatch even after capture.

Warmup ends after the configured number of warmup iterations (default: 1). With dynamic batch sizes, warmup may be repeated for each unique input shape signature.

Phase 2: SHAPES_FROZEN

After warmup, the plan calls sdxFreezeShapes(). At this point, all intermediate buffer allocations are finalized. Buffers are pre-allocated at their frozen sizes and their device pointers are recorded. The plan verifies that subsequent calls supply inputs whose shapes match the frozen shape signature. If an input arrives with a different shape, DSP falls back to slot-by-slot execution for that call.

Variable states also transition: every slot moves from WARMUP to FROZEN state. The PlanIntrospection API can be used to inspect the shape table at this point.

Phase 3: CUDA_GRAPH_CAPTURED

With shapes frozen and pointers stable, DSP initiates CUDA graph capture using cudaStreamBeginCapture(). It then replays the slot sequence in capture mode — each op records its GPU work into the capture stream rather than executing immediately. At the end of capture, cudaStreamEndCapture() yields a cudaGraph_t, which is then instantiated into a cudaGraphExec_t for later replay.

Non-CUDA backends perform an equivalent step: HIP records a hipGraph_t, Metal records a MTLCommandBuffer, Vulkan records a VkCommandBuffer, and Level Zero records a ze_command_list_t.

Slots that cannot participate in capture (e.g., operations involving CPU-side conditional logic or ops not yet supported by the backend's capture mechanism) are marked as GAP nodes. Gap slots continue to execute via tl_dspGapStream during replay.

Phase 4: REPLAYING

Once a valid captured graph exists, all subsequent calls enter replay mode. The executor calls cudaGraphLaunch(graphExec, stream) — a single API call that kicks off the entire forward pass — and waits for completion. The per-thread argTableStable fast path skips pointer-stability checks when all input pointers are confirmed stable, shaving a further few microseconds per call.

Slot states in replay: FROZEN slots that are inside the captured region transition to REPLAYING. Gap slots remain in their own state and continue to be dispatched individually.

Slot States

State
Meaning

UNINITIALIZED

Slot has been defined but warmup has not yet run

WARMUP

Actively executing during warmup phase

FROZEN

Shape recorded; buffer allocated; pointer stable

CAPTURING

Inside a CUDA/HIP/Vulkan graph capture

REPLAYING

Being replayed via captured graph

EVICTED

Plan was evicted from cache; slot must be re-warmed

Graph Node Phases

Each node in the compiled plan carries a phase tag that the executor uses to decide how to dispatch it:

Node Phase
Dispatch Method

SLOT_BY_SLOT

Standard ND4J op dispatch (warmup and gap ops)

CAPTURED

Included in CUDA/backend graph capture; replayed atomically

GAP

Op excluded from capture; runs via gap stream each call

TRITON

Dispatched to Triton JIT kernel; participates in captured graph via pre-compiled PTX

Graph Optimizer

Before building the execution plan, DSP runs a 26-pass optimizer over the SameDiff graph. Passes are organized into three outer iterations, allowing earlier simplifications to expose new opportunities for later passes.

The optimizer is configurable:

  • Skip individual passes: -Dnd4j.optimizer.skip=PassName (comma-separated for multiple)

  • Log applied rewrites: -Dnd4j.optimizer.logApplied=true

Pass List

Pass 1 — Dead Code Elimination (UnusedFunctionOptimizations) Removes ops whose outputs are not consumed by any requested output or by any op that transitively feeds a requested output. On large models with optional diagnostic branches, this can remove 10–20% of nodes before any other pass runs.

Pass 2 — Constant Propagation (ConstantFunctionOptimizations) Evaluates ops whose all inputs are constants at compile time. The computed result is folded into a new CONSTANT node and the original op is removed. Shape computations on static inputs are almost always eliminated here.

Pass 3 — Broadcast Elimination and Commutative Canonicalization (BroadcastEliminationOptimizations) Removes unnecessary explicit broadcast ops when the underlying kernel already handles broadcasting natively. Also canonicalizes commutative ops (e.g., swapping operands so the larger tensor is always the first argument) to maximize cache reuse.

Pass 4 — Reassociation and Double Transpose Elimination (ReorderingOptimizations) Reassociates chains of associative operations to minimize intermediate allocations. Eliminates transpose(transpose(x)) → x and transpose(x) immediately followed by an op that accepts a transposed-input flag (e.g., mmul(transpose(A), B) → mmul(A, B, transposeA=true)).

Pass 5 — Algebraic Identities (AlgebraicOptimizations) Rewrites:

  • add(x, 0) → x

  • mul(x, 1) → x

  • mul(x, 0) → zeros_like(x)

  • sqrt(square(x)) → abs(x)

  • pow(x, 0.5) → sqrt(x)

  • log(1) → 0

Pass 6 — Peephole Simplifications (PeepholeOptimizations) Handles composite patterns:

  • relu(relu(x)) → relu(x) (idempotency)

  • log(exp(x)) → x

  • exp(log(x)) → x

  • neg(neg(x)) → x

  • abs(abs(x)) → abs(x)

Pass 7 — Arithmetic Chain Folding (ArithmeticChainOptimizations) Collapses chains of adds or multiplies involving constants into a single op:

  • add(add(x, c1), c2) → add(x, c1+c2)

  • mul(mul(x, c1), c2) → mul(x, c1*c2)

  • add(mul(x, c1), mul(x, c2)) → mul(x, c1+c2)

This is particularly effective after constant propagation exposes scalar constants in weight scaling paths.

Pass 8 — Strength Reduction (StrengthReductionOptimizations) Replaces expensive ops with cheaper equivalents:

  • pow(x, 2) → square(x) (avoids the general pow kernel)

  • div(x, c) → mul(x, 1.0/c) (multiplication is faster than division on most hardware)

  • pow(x, 0.5) → sqrt(x)

  • pow(x, -1) → reciprocal(x)

Pass 9 — Identity Removal (IdentityFunctionOptimizations) Removes identity(x) nodes that are sometimes inserted by the graph-building API or by import pipelines. Replaces all consumer edges with direct references to the identity's input.

Pass 10 — Concat/Split Simplification (ConcatSplitOptimizations)

  • split(concat(a, b), axis=N, numSplits=2) → [a, b] when the split exactly inverts the concat.

  • concat([x]) → x when concatenating a single tensor.

  • split(x, axis=N, numSplits=1) → [x].

Pass 11 — Constant-Condition Select/Where (SelectWhereOptimizations) When the condition tensor of a select or where op is a compile-time constant, replaces the entire op with the statically-chosen branch.

Pass 12 — Redundancy Elimination (RedundancyEliminationOptimizations) Removes duplicate computations that are not caught by CSE (pass 14) because they have distinct node identities in the graph. Specifically handles patterns from model import where the same subexpression may be emitted multiple times for different output consumers.

Pass 13 — Shape Constant Folding (ShapeFunctionOptimizations) Evaluates shape-query ops (shape_of, rank, size_at) whose inputs have statically known shapes, replacing them with constant integers. This often enables pass 2 (constant propagation) to eliminate further downstream computations on shapes.

Pass 14 — Common Subexpression Elimination (CommonSubexpressionElimination) Detects identical subgraphs (same op type, same inputs in the same order, same attributes) and merges them into a single node, redirecting all consumers to the merged output. Effective for attention masks, positional encodings, and layer norm statistics that are computed once but referenced in multiple branches.

Pass 15 — Attention Fusion (AttentionFusionOptimizations) The key pass for transformer acceleration. Detects the attention computation pattern:

and replaces it with a single dot_product_attention_v2 op. On CUDA, this maps to FlashAttention or cuDNN fused attention, reducing memory traffic from O(N²) to O(N) in the sequence length. This pass must run before HorizontalFusion (pass 16) because horizontal fusion may merge the Q, K, V matmuls before attention fusion has a chance to detect the pattern.

Pass 16 — Horizontal Fusion (HorizontalFusionOptimizations) Detects parallel matmul chains that share a common input — most commonly the fused QKV projection in transformer self-attention:

These are fused into a single batched matmul against the stacked weight matrix [Wq | Wk | Wv], followed by a slice to recover Q, K, and V. On Ampere and later GPUs, a single large GEMM achieves much higher FLOP utilization than three separate smaller GEMMs.

Passes 17–26 — Additional Fusions and Specializations The remaining passes include:

  • Normalization fusion (fusing adjacent layer norm components into a single kernel)

  • Gated delta network fusion (for architectures using gated linear units)

  • Quantization optimization (folding quantization scale/zero-point constants and fusing dequant + op + requant)

  • Further platform-specific strength reductions and memory layout canonicalization

GPU JIT Compilation

When the execution mode includes JIT (Triton, NVRTC, or PTX), DSP groups consecutive element-compatible ops into fused kernel segments and compiles them into a single GPU kernel. Fusion eliminates redundant reads and writes of intermediate tensors.

Note: All JIT features in this section require the -compile classifier. With the base classifier, DSP still runs the graph optimizer and captures CUDA graphs, but JIT kernel fusion is unavailable. Use -Djavacpp.platform.extension=-compile with -platform artifacts or add the -compile classifier explicitly.

Kernel Segment Types

The TritonIRBuilder classifies each op into one of the following segment categories, used to determine fusion eligibility:

Category
Examples

ELEMENTWISE

relu, add, mul, exp, log, sigmoid, tanh, cast

MATMUL

mmul, gemm, dot

REDUCTION

sum, mean, max, min, norm2

NORMALIZATION

layer_norm, batch_norm, rms_norm

ATTENTION

dot_product_attention_v2, flash_attention

GATHER

gather, embedding_lookup

SCATTER

scatter_add, scatter_nd

SLICE

slice, strided_slice, gather_nd

TILE

tile, broadcast_to

CONCAT

concat

SPLIT

split

SHAPE_MANIPULATION

reshape, transpose, squeeze, unsqueeze

CONV

conv2d, depthwise_conv2d, conv_transpose

IMAGE2COL

im2col, col2im

Consecutive ELEMENTWISE ops are always fusible. REDUCTION, NORMALIZATION, and ATTENTION ops form their own segment boundaries.

Triton JIT Backend

Triton (MLIR-based GPU compiler) is the highest-fidelity JIT backend. It generates correct, highly optimized kernels for the full range of segment types.

TritonIRBuilder decomposes the fused segment into a sequence of Triton IR operations, then hands the IR to TritonTargetDispatch which selects the appropriate compilation target:

Hardware
Target

NVIDIA GPUs

NVIDIA PTX (via Triton's LLVM NVPTX backend)

AMD GPUs

AMD AMDGCN (via Triton's LLVM AMDGPU backend)

Intel GPUs

Intel SPIR-V (via Triton's LLVM SPIR-V backend)

A background precompilation thread compiles Triton kernels during warmup, so compiled kernels are ready before the CUDA graph capture begins. This avoids any JIT stall on the first post-warmup call.

FusionScoring heuristics evaluate whether fusing a group of ops will yield a net benefit: they estimate the memory-bandwidth savings from eliminating intermediates against the register pressure cost of the larger fused kernel. Groups that score below a threshold are left unfused.

NVRTC JIT Backend

NVRTC (NVIDIA Runtime Compilation) compiles CUDA C++ source code at runtime using the NVRTC library bundled with the CUDA toolkit.

The NVRTC backend:

  1. Generates a CUDA C++ translation unit from the fused segment's slot list.

  2. Calls nvrtcCompileProgram() with architecture flags matching the active GPU's compute capability.

  3. Extracts PTX from the compilation result via nvrtcGetPTX().

  4. Loads the PTX with cuModuleLoadDataEx() and retrieves the kernel entry point.

NVRTC produces code equivalent to Triton for ELEMENTWISE segments but is generally less effective for MATMUL and REDUCTION segments, where Triton's tiling and pipeline scheduling logic produces better code. NVRTC serves as the fallback when Triton is unavailable.

PTX Backend

The PTX backend generates PTX assembly text directly via string templates, bypassing any higher-level compilation step. This gives the fastest compile time (usually under 1 ms for small elementwise kernels) at the cost of lower code quality for complex segments.

PTX assembly is loaded directly via cuModuleLoadDataEx(). The PTX backend is mainly used for simple elementwise fusions where template-generated PTX is already near-optimal, and as the final JIT fallback when neither Triton nor NVRTC are available.

JIT Fallback Order

On CUDA hardware:

Multi-Backend Dispatch

DSP supports 17 execution modes, covering every major hardware platform. The active mode is controlled by the GraphExecutionMode enum.

GraphExecutionMode Enum

Value
ID
Description

AUTO

0

Let DSP choose the best available mode for the current hardware

SLOT_BY_SLOT

1

Interpreted slot-by-slot dispatch; always available

CUDA_GRAPHS

2

CUDA graph capture + replay; no JIT fusion

TRITON

3

Triton JIT + CUDA graph replay

NVRTC

4

NVRTC JIT + CUDA graph replay

PTX

5

PTX string-template JIT + CUDA graph replay

ROCm

6

AMD HIP graph capture + replay

MLX

7

Apple MLX graph backend (Apple Silicon)

VULKAN

8

Vulkan compute graphs

LEVEL_ZERO

9

Intel Level Zero command lists

METAL

10

Apple Metal command buffers

ONNX_RUNTIME

11

Delegate to ONNX Runtime execution provider

OPENVINO

12

Intel OpenVINO inference engine

ACL

13

ARM Compute Library graph backend

NNAPI

14

Android Neural Networks API

ARM_HYBRID

15

Hybrid CPU/GPU dispatch on ARM (Cortex-M + Mali)

MLIR_CPU

16

MLIR CPU dialects for JIT-compiled CPU kernels

AUTO Fallback Chains

When AUTO is selected, DSP probes available hardware and libraries to choose the best mode, falling back to the next option if a backend is unavailable:

On CUDA hardware:

On non-CUDA hardware:

Backend Implementations

Each backend provides a GraphReplayHandle implementation that wraps the platform's captured-graph replay API:

Backend Class
Platform API

CudaGraphReplayHandle

cudaGraphExec_t / cudaGraphLaunch()

HipGraphReplayHandle

hipGraphExec_t / hipGraphLaunch()

MetalReplayHandle

MTLCommandBuffer / [buffer commit]

VulkanReplayHandle

VkCommandBuffer / vkQueueSubmit()

LevelZeroReplayHandle

ze_command_list_t / zeCommandQueueExecuteCommandLists()

TpuReplayHandle

XLA executable replay

HexagonReplayHandle

Qualcomm FastRPC handle

CPU-side graph backends that provide op fusion without captured replay:

Backend Class
Library

OneDnnGraphBackend

Intel oneDNN graph API

AclGraphBackend

ARM Compute Library graph mode

MlxGraphBackend

Apple MLX compute graph

MlirCpuGraphBackend

MLIR CPU dialects (Linalg, Vector, SCF)

ArmHybridGraphBackend

Mixed Cortex + Mali dispatch

NnapiGraphBackend

Android NNAPI model compilation

OpenVinoGraphBackend

OpenVINO CompiledModel

NativePlanCache

The native layer maintains an LRU plan cache (NativePlanCache) keyed on a content-based hash:

outputSetHash encodes which output variables are requested. phShapeContentHash encodes the shapes of all placeholder inputs. Together these two hash components mean that different output sets or different input shapes produce different cache entries, each with their own warmup and captured graph.

Eviction uses a dual policy:

  • Count cap: when the number of cached plans exceeds the configured maximum, the least-recently-used plan is evicted.

  • Memory budget: when total plan memory exceeds a configurable fraction of available device memory, plans are evicted until the memory usage drops below the budget.

On eviction, all slot states for the evicted plan revert to EVICTED, triggering re-warmup on the next call.

Java API Reference

DynamicShapePlanCompiler

DynamicShapePlanCompiler converts a ForwardExecutionDAG into a DynamicShapePlan. You do not call this directly during normal inference — DSP invokes it automatically the first time you call sd.output() with DSP enabled. It is useful to call explicitly when you want to pre-compile the plan at startup.

Compilation pipeline:

  1. Filter ops: remove dead nodes (applies optimizer first if enabled)

  2. Build external input index maps: assign integer indices to placeholder variables

  3. Assign slot indices: give each surviving node a sequential slot index

  4. Build input wiring: record which slot's output each slot reads

  5. Liveness analysis: determine the last step at which each slot's output is needed

  6. Build releaseAtStep table: encode when each intermediate buffer can be freed

  7. Pre-allocate OpContext pool: create a pool of reusable op execution contexts

DynamicShapePlanExecutor

DynamicShapePlanExecutor drives the full lifecycle. It owns the plan-phase state machine and selects the appropriate dispatch path (warmup, slot-by-slot, capture, or replay) for each call.

Key internal fields (useful for diagnostics):

  • tl_dspExecutionStream — per-thread CUDA stream used for main execution

  • tl_dspGapStream — per-thread stream for gap ops

  • argTableStable — boolean fast path: when true, skips pointer-stability verification on replay

PlanIntrospection

PlanIntrospection provides read-only access to the compiled plan's internal structure. Use it to verify compilation results or to diagnose unexpected fallbacks.

DspDiagnostics

DspDiagnostics provides structured performance and correctness diagnostics across 20 categories. It can emit a full JSON report for offline analysis.

Available diagnostic categories: COMPILE, EXECUTE, MEMORY, GRAPH_REPLAY, JIT, OPTIMIZER, CACHE, WARMUP, CAPTURE, LIVENESS, WIRING, SHAPES, GAPS, TENSOR_PARALLEL, PIPELINE_PARALLEL, SLOT_STATE, FALLBACK, STREAM, POINTER, LATENCY.

DspDebugger

DspDebugger enables interactive step-by-step execution for diagnosing numerical issues or unexpected behavior. It inserts breakpoints at specified slot indices and pauses execution so you can inspect intermediate tensor values.

Disk Cache and Diagnostics

DspPlanDiskCache

Compiled plans are serialized to disk so that the warmup and JIT compilation costs are paid only once across JVM restarts. The disk cache uses FNV-1a hashing for cache key construction.

Cache location: ~/.kompile/cache/dsp/ Cache version: DSP_VERSION=5 (cache entries from older versions are automatically invalidated) Write protocol: Atomic temp-file + Files.move(ATOMIC_MOVE) to prevent corrupt entries on crash.

Cache entries are keyed on the same content hash used by NativePlanCache:

When a disk-cached plan is loaded, the executor skips warmup and proceeds directly to the CUDA graph capture phase (if CUDA is available and the shapes match).

NativePlanCache

The in-process native cache is separate from the disk cache. It is an LRU cache with:

  • Maximum entry count (configurable; default varies by available GPU memory)

  • Memory budget: combined size of all cached plans must stay below a fraction of device memory

When device memory is tight (e.g., running multiple models concurrently), the memory-budget eviction policy takes precedence and can evict plans even if the count cap has not been reached.

Parallel Execution

DSP includes native support for both tensor parallelism and pipeline parallelism, allowing large models to span multiple GPUs.

Tensor Parallelism

Tensor parallelism splits individual weight matrices across GPUs. For a linear layer with weight matrix W, ColumnParallelLinear shards W along the column dimension across N GPUs, so each GPU holds W[:,k/N:(k+1)/N]. RowParallelLinear shards along the row dimension.

Each GPU computes a partial result and an all-reduce collective synchronizes the results.

The CollectiveCommunicator interface abstracts the all-reduce communication:

  • LocalCollectiveCommunicator — uses shared CPU memory for multi-GPU machines where all GPUs have NVLink or are on the same PCIe root complex

  • NcclCommunicator — uses NCCL for higher-throughput all-reduce, required for multi-node configurations

Pipeline Parallelism

Pipeline parallelism partitions model layers across GPUs. Each GPU holds a "stage" of the model. Micro-batches are interleaved through the pipeline to maintain GPU utilization.

With interleave=true, PipelineParallelRunner uses a 1F1B (one-forward-one-backward) schedule for training or a fill-and-drain schedule for inference, keeping all stages occupied with different micro-batches simultaneously.

DSP Runtime SDK

The DSP Runtime SDK exposes a stable C ABI defined in dsp_runtime_c.h. The ABI version is SDX_RUNTIME_ABI_VERSION=1.

All structs in the C ABI use a sized-struct pattern for forward compatibility: the first member of every struct is size_t struct_size, which the caller sets to sizeof(the_struct). The runtime checks this field and zero-initializes any fields the caller did not set (because the caller's struct is smaller than the runtime's version). This allows new fields to be added in future ABI versions without breaking existing binaries.

C API Lifecycle

sdx_execution_report_t Fields

Field
Type
Description

struct_size

size_t

Sized-struct discriminator

slotCount

int32_t

Number of slots in the compiled plan

executionMode

int32_t

Actual mode used (may differ from requested AUTO)

replayCount

int64_t

Number of times the captured graph has been replayed

avgLatencyUs

float

Rolling average per-call latency in microseconds

peakMemoryBytes

size_t

Peak device memory consumed by the plan

Shape Control Functions

Language Bindings

Java (JNA)

Kotlin (JNA Wrappers)

Python (ctypes + NumPy)

The Python binding also exposes a gRPC server mode for serving DSP plans over a network:

Rust (FFI + RAII)

The Rust binding uses a thin FFI layer over dsp_runtime_c.h with RAII wrappers (Drop implementations) ensuring correct cleanup order.

Swift (Swift Package Manager)

The Swift binding is distributed as a Swift Package and supports both macOS (Metal backend) and iOS (Metal + CoreML backends).

C# (P/Invoke)

Configuration Reference

All DSP behaviors are controlled via JVM system properties, environment variables, or the DspConfig builder API.

Core Execution

Property
Default
Description

-Dnd4j.dsp.enabled

true

Enable DSP compilation. Set false to fall back to the legacy InferenceSession executor.

-Dnd4j.dsp.executionMode

AUTO

Override the GraphExecutionMode. Accepts mode names (TRITON, CUDA_GRAPHS, SLOT_BY_SLOT, etc.) or integer IDs.

-Dnd4j.dsp.warmupIterations

1

Number of slot-by-slot warmup iterations before shape freeze. Increase if input shapes vary across the first few calls.

-Dnd4j.dsp.argTableStable

true

Enable the pointer-stability fast path on replay. Disable only when input buffer addresses change between calls.

-Dnd4j.dsp.gapStreamEnabled

true

Enable the separate gap stream. Disable to route gap ops to the main execution stream (simpler but may serialize with captured graph).

Graph Optimizer

Property
Default
Description

-Dnd4j.optimizer.skip

(none)

Comma-separated list of optimizer pass class names to skip. Example: AttentionFusionOptimizations,HorizontalFusionOptimizations

-Dnd4j.optimizer.logApplied

false

Log each rewrite rule application (op before → after) to the SLF4J logger at DEBUG level.

-Dnd4j.optimizer.maxIterations

3

Number of outer optimization iterations. Increase to allow more fixpoint convergence at the cost of compile time.

-Dnd4j.optimizer.enabled

true

Enable the graph optimizer entirely. Set false to disable all 26 passes.

JIT Compilation

Property
Default
Description

-Dnd4j.dsp.jit.tritonEnabled

true

Enable the Triton JIT backend.

-Dnd4j.dsp.jit.nvrtcEnabled

true

Enable the NVRTC JIT backend.

-Dnd4j.dsp.jit.ptxEnabled

true

Enable the PTX string-template JIT backend.

-Dnd4j.dsp.jit.backgroundCompile

true

Compile Triton kernels on a background thread during warmup.

-Dnd4j.dsp.jit.fusionScoreThreshold

0.5

Minimum FusionScoring score to apply a fusion. Lower values fuse more aggressively; higher values are more conservative.

Disk Cache

Property
Default
Description

-Dnd4j.dsp.planCache.diskEnabled

true

Enable serializing compiled plans to disk.

-Dnd4j.dsp.planCache.dir

~/.kompile/cache/dsp/

Directory for disk-cached plans.

-Dnd4j.dsp.planCache.maxEntries

64

Maximum number of plans to retain on disk.

-Dnd4j.dsp.planCache.version

5

Cache format version. Do not change; used for automatic invalidation.

Native Plan Cache

Property
Default
Description

-Dnd4j.dsp.nativeCache.maxCount

32

Maximum number of plans in the in-process LRU cache.

-Dnd4j.dsp.nativeCache.memoryBudgetFraction

0.25

Maximum fraction of device memory that in-process cached plans may consume.

Diagnostics

Property
Default
Description

-Dnd4j.dsp.diagnostics.enabled

false

Enable the DspDiagnostics subsystem.

-Dnd4j.dsp.diagnostics.categories

(all)

Comma-separated list of DspDiagnostics.Category names to enable.

-Dnd4j.dsp.diagnostics.level

SUMMARY

Default verbosity level (SUMMARY, DETAILED, or FULL).

-Dnd4j.dsp.diagnostics.jsonOutput

(none)

Path to write the JSON diagnostic report on JVM exit.

Parallel Execution

Property
Default
Description

-Dnd4j.dsp.tensorParallel.numDevices

1

Number of GPUs for tensor parallelism. 1 disables tensor parallelism.

-Dnd4j.dsp.tensorParallel.useNccl

true

Use NCCL for all-reduce collectives. Set false to use LocalCollectiveCommunicator (shared memory).

-Dnd4j.dsp.pipelineParallel.numStages

1

Number of pipeline stages. 1 disables pipeline parallelism.

-Dnd4j.dsp.pipelineParallel.microBatchSize

8

Micro-batch size for pipeline interleaving.

Graph Optimizer — User Controls (ADR 0100)

The graph optimizer described earlier in this page runs automatically during DynamicShapePlanCompiler.compile(). This section covers the user-facing properties for controlling, diagnosing, and tuning it.

Enabling and Disabling

Disabling the optimizer is the fastest way to determine whether an optimizer pass is responsible for a numerical regression or a functional breakage: if the problem disappears with the optimizer disabled, binary-search the pass list to isolate the culprit.

Skipping Individual Passes

Use -Dnd4j.optimizer.skip with a comma-separated list of pass class names:

Complete list of pass names (use these exact strings with -Dnd4j.optimizer.skip):

Category
Pass Name

Simplification

UnusedFunctionOptimizations

Simplification

ConstantFunctionOptimizations

Simplification

BroadcastEliminationOptimizations

Simplification

ReorderingOptimizations

Simplification

AlgebraicOptimizations

Simplification

PeepholeOptimizations

Simplification

ArithmeticChainOptimizations

Simplification

StrengthReductionOptimizations

Simplification

IdentityFunctionOptimizations

Simplification

ConcatSplitOptimizations

Simplification

SelectWhereOptimizations

Simplification

RedundancyEliminationOptimizations

Simplification

ShapeFunctionOptimizations

Simplification

CommonSubexpressionElimination

Fusion

AttentionFusionOptimizations

Fusion

HorizontalFusionOptimizations

Fusion

MatMulChainOptimizations

Fusion

ActivationFusionOptimizations

Fusion

NormalizationFusionOptimizations

Fusion

GatedDeltaNetFusionOptimizations

Fusion

LinearFusionOptimizations

Backend/Quant

RematerializationOptimizations

Backend/Quant

QuantizationOptimizations

Backend/Quant

CuDNNFunctionOptimizations

FP16 Weight Pre-Casting

When enabled, the QuantizationOptimizations pass casts constant weight tensors to FP16 before the plan is compiled. This reduces GPU memory usage for weight storage and allows FP16 GEMM kernels to run without runtime casts. Only enable this if your model was trained with mixed precision or you have validated that FP16 inference accuracy is acceptable.

Iteration Count

The optimizer runs up to maxIterations rounds. It stops early if no pass made a change in the most recent iteration. Increasing this value allows cross-pass opportunities that emerge in later rounds (e.g., a fusion pass exposing a new constant-folding opportunity in the next round) at the cost of longer compile time.

Diagnosing Optimizer Regressions

If you observe a numerical discrepancy or a correctness bug after upgrading and suspect the optimizer:

With -Dnd4j.optimizer.logApplied=true, each pass logs its rewrites at DEBUG level via SLF4J, showing the before and after for every op that was transformed. Pipe this output to a file and search for the op that changed unexpectedly.

Buffer Coloring, Pooling, and Passivation (ADR 0094)

DSP's buffer coloring system is a three-tier memory optimization that reduces peak GPU memory usage — sometimes dramatically — by reusing physical buffers across the slots of a compiled plan and across multiple plans.

Why It Matters

A 1000-slot compiled plan (typical for a 28-layer LLM) naively allocates one dedicated GPU buffer per slot. Most of those buffers have non-overlapping lifetimes: slot 50's output is consumed and dead before slot 200 produces anything. Without coloring, you pay for all 1000 allocations simultaneously even though at most a few dozen are live at any point.

Buffer coloring assigns multiple slots to the same physical buffer when their live ranges do not overlap — analogous to register allocation in a compiler. Combined with a cross-plan buffer pool and a passivation mechanism that releases inactive plan memory under pressure, the three tiers together reduce peak memory far below the naive allocation total.

Tier 1: Buffer Coloring (Within a Plan)

After shape freezing (SHAPES_FROZEN), DSP runs a greedy interval-graph coloring algorithm over all eligible slots:

  • Slots are grouped by (shape, dtype, device) — only same-footprint slots can share a buffer.

  • Within each group, slots are sorted by their producerStep. Two slots get the same "color" only if their live ranges do not overlap.

  • Slots that are view parents, requested outputs, or view children are excluded from coloring to preserve correctness.

The result: instead of N buffers, DSP allocates at most numColors buffers for the colored group. For large LLMs this typically reduces intermediate buffer count by 5–10x.

Tier 2: Buffer Pool (Across Plans)

Released buffers from one plan go into a per-device pool keyed by (numElements, dtype). When another plan needs a buffer of the same shape, it acquires it from the pool at zero allocation cost:

The pool is thread-safe and device-local. You do not interact with it directly; it operates automatically whenever DSP releases or acquires intermediate buffers.

Tier 3: Passivation (Under Memory Pressure)

When the native plan cache is under memory pressure (too many plans or total cached memory exceeds the budget fraction), the cache passivates least-recently-used plans before fully evicting them:

  1. Passivation: GPU intermediate buffers are released back to the pool; the plan's structural metadata (slots, wiring, shapes) stays in CPU heap (~10–50 MB). The plan remains in cache.

  2. Reactivation: On the next cache hit for a passivated plan, buffers are re-acquired from the pool and warmup re-runs if needed.

  3. Full eviction: Only occurs if passivation and pool trimming are insufficient.

This means a large set of plans can coexist in the cache without holding their full GPU allocation simultaneously — only the currently-executing plan needs its buffers.

Inspecting Coloring Results

Use DspHandle to inspect the coloring outcome for a compiled plan:

Diagnostics via DspDiagnostics

Enable the COLORING category in DspDiagnostics to log coloring events:

Op Execution Timing Tracker (ADR 0092)

OpTimingTracker provides low-overhead, always-available timing of individual ND4J operations. Unlike the existing GraphProfile system (which requires a full graph execution context), OpTimingTracker works for any op execution — both inside SameDiff graphs and standalone Nd4j.* calls.

Enabling Timing

Timing Phases

In detailed mode, each op's time is broken down into:

Phase
What It Measures

VALIDATION

Input validation, argument checking, datatype checks

SHAPE_CALC

Output shape inference

MEMORY_ALLOC

Output array allocation

HELPER_CHECK

Platform helper isUsable() probe (oneDNN, cuDNN)

HELPER_EXEC

Platform helper execution

NATIVE_EXEC

Native C++ implementation execution

TOTAL

End-to-end wall-clock time

Viewing Results

Example hotspot report:

Exporting

The Chrome trace shows each op as a horizontal bar (width = duration), organized by thread lane, with helper vs native execution color-coded. Open chrome://tracing in Chrome and load the file.

The CSV includes columns: OpName, Calls, TotalMs, AvgUs, StdDevUs, MinUs, MaxUs, HelperPct, per-phase timing columns, memory columns, and P50Us, P90Us, P99Us percentiles from the logarithmic histogram.

Timing Modes Summary

Mode
Enable Call
Phase Timing
Trace Events
Overhead

Disabled

setOpTimingEnabled(0, 0)

No

No

None

Simple

setOpTimingEnabled(1, 0)

No

No

Low

Detailed

setOpTimingEnabled(1, 1)

Yes

No

Medium

Trace

setOpTimingEnabledWithTrace(1)

Yes

Yes

Higher

Complete Profiling Workflow

Phase Interpretation

Phase High
Consider...

VALIDATION

Simplifying input shapes or reducing input count

MEMORY_ALLOC

Pre-allocating output arrays or using workspaces

HELPER_EXEC

Platform helper is active — check if it is actually faster than native with -Dnd4j.helpers.allowHelpers=false baseline

NATIVE_EXEC

The actual kernel — expected to dominate for large ops

Ring Buffer Considerations

The tracker uses a fixed-size ring buffer (8192 entries). For long-running workloads, call ops.flushOpTiming() periodically — if more than 8192 ops execute between flushes, the oldest entries are overwritten without being aggregated.

See Also

PreviousSerializationNextQuickstart

Last updated

Was this helpful?