Skip to content

PR12: ND4J Java API & Core Infrastructure#10445

Open
agibsonccc wants to merge 13 commits into
masterfrom
pr/12-nd4j-java-api-core
Open

PR12: ND4J Java API & Core Infrastructure#10445
agibsonccc wants to merge 13 commits into
masterfrom
pr/12-nd4j-java-api-core

Conversation

@agibsonccc

@agibsonccc agibsonccc commented Jun 15, 2026

Copy link
Copy Markdown
Contributor

Summary

PR 12 of 22 PRs in the ag_new_release_updates_2 branch split. Merge after Layer 3 (native platform backends + DSP engine).

  • Multi-device abstraction: DeviceMemoryManager singleton with per-device AtomicLong counters, ReentrantReadWriteLock cap enforcement, selectDeviceForAllocation(bytes) sorted by priority then free memory
  • Device hierarchy: DeviceDescriptor interface → CudaDeviceDescriptor, CpuDeviceDescriptor, StubDeviceDescriptor; StubDeviceContextProvider enables full multi-GPU topology simulation in tests without hardware
  • Framework diagnostics facade: 30+ files in org/nd4j/linalg/framework/; completely lazy tree via Nd4j.framework.execution()/.memory()/.device()/.workspaces()/.profiling()/.lifecycle()/.diagnostics(); costs nothing if unused
  • Memory coherence: HybridDataBuffer tracks CoherenceState (HOST_ONLY/DEVICE_ONLY/SYNCHRONIZED); syncToHost()/syncToDevice() are no-ops when state already valid — eliminates redundant H2D/D2H transfers
  • Graph tracing: Nd4j.graphScope() returns a GraphScope that traces GraphFunction calls into a hidden SameDiff graph; LazyINDArray proxy tracks sdVariableName during tracing while blocking data access; CompiledGraphFunction is the cached result
  • Adam8bit optimizer: INT8 per-block absmax quantization (block size=2048) for m/v state; reduces optimizer state ~4x (~56GB → ~14GB for 7B-parameter model)
  • Dataset curation toolkit: 20 files covering MinHash LSH deduplication, N-gram decontamination, curriculum learning, padding-free batch collation, sequence packing, and weighted domain mixing for LLM pretraining
  • Typed config classes: DspEnvironmentConfig, TritonEnvironmentConfig, and others consolidate scattered System.getProperty("nd4j.dsp.*") calls into typed @Builder classes
  • Backend registry: BackendManager + BackendRegistry replace ad-hoc ServiceLoader calls; plugin-based native ops loading via NativePluginLoader

What Changed

Device Abstraction Layer (10 files)

  • DeviceMemoryManager.java — singleton: per-device AtomicLong counters, memory caps, selectDeviceForAllocation(), configureStubTopology() for test stubs
  • DeviceDescriptor.java — device identity interface (key in all device maps)
  • CudaDeviceDescriptor.java, CpuDeviceDescriptor.java, BaseDeviceDescriptor.java — concrete descriptors with device type and index
  • DeviceContextProvider.java, DeviceContext.java — device-switching with stream handle management
  • StubDeviceDescriptor.java, StubDeviceContextProvider.java — test stubs for multi-GPU simulation without hardware
  • DeviceRoutingConfiguration.java, MultiGpuTracer.java — op routing config and device-switch tracing
  • DeviceType.java — enum: CPU, CUDA, ROCM, TPU, HEXAGON, OPENCL, METAL, VULKAN

Framework Diagnostics Facade (30+ files in org/nd4j/linalg/framework/)

  • Framework.java — lazy-initialized subsystem tree entry point
  • exec/: ExecutionSubsystem.java, DspPlanIntrospection.java, DspStats.java, ExecutionStats.java, KernelStats.java — DSP plan and execution monitoring
  • device/: DeviceSubsystem.java, TransferSubsystem.java, TransferEvent.java, TransferReport.java, TransferStats.java, PointerStabilityGuard.java, ReplicaLeakDetector.java — device transfer tracking and pointer stability
  • memory/: MemorySubsystem.java, MemoryEnvironment.java, MemoryState.java — memory usage snapshots
  • history/: ArrayLifecycleTracker.java, MemorySampler.java, WorkspaceSampler.java — lifecycle and memory sampling
  • leak/: LeakDetector.java, LeakReport.java, PotentialLeak.java — memory leak detection
  • profiling/, constant/, workspace/, stats/: ProfilingSubsystem.java, ConstantCacheSubsystem.java, WorkspaceSubsystem.java, FrameworkSnapshot.java, DiagnosticIssue.java — profiling, constant cache, workspace, and health snapshot subsystems

Memory Infrastructure (9 files)

  • HybridDataBuffer.java — DataBuffer mirroring host/device with tracked CoherenceState
  • CoherenceState.javaHOST_ONLY, DEVICE_ONLY, SYNCHRONIZED enum
  • MemoryCoherenceManager.java — manages coherence state across DataBuffers
  • MultiBackendWorkspace.java, MultiBackendWorkspaceManager.java — workspace abstractions spanning multiple backends
  • abstracts/DefaultMultiBackendWorkspace.java, NativeMultiBackendWorkspace.java — default and native workspace implementations
  • deallocation/OpaqueDataBufferDeallocator.java, OpaqueNDArrayDeallocator.java — separate deallocators for opaque native handles
  • conf/DeviceAwareWorkspaceConfiguration.java — per-device workspace configuration

Graph Tracing Primitives (7 files)

  • GraphFunction.java@FunctionalInterface INDArray[] apply(INDArray[] inputs)
  • GraphScope.java — traces GraphFunction calls into a hidden SameDiff graph
  • CompiledGraphFunction.java — cached, compiled graph ready for repeated DSP execution
  • LazyINDArray.java — proxy during tracing: shape()/dataType() work, data access throws IllegalStateException; tracks sdVariableName for graph wiring
  • DeviceAwareNDArrayFactory.java, DeviceAwareNd4j.java, MultiBackendContext.java — device-routed array factory and backend context

Factory and Backend Registration (8 files)

  • BackendManager.java, BackendRegistry.java — centralized backend lifecycle and registry replacing ad-hoc ServiceLoader calls
  • factory/config/: CoreEnvironmentConfig.java, CudaEnvironmentConfig.java, DspEnvironmentConfig.java, LifecycleEnvironmentConfig.java, MemoryConfig.java, TritonEnvironmentConfig.java — typed @Builder configs replacing scattered System.getProperty() calls
  • TritonSectionType.java — Triton kernel section type enum

Adam8bit Optimizer (3 files)

  • config/Adam8bit.java — 8-bit Adam: INT8 per-block absmax quantization, DEFAULT_BLOCK_SIZE=2048, ~4x optimizer state reduction
  • Adam8bitUpdater.java — dequantize → Adam update → requantize per block
  • BlockQuantizationUtils.java — block-wise INT8 absmax quantization utilities

Dataset Curation Toolkit (20 files in org/nd4j/linalg/dataset/curation/)

  • MinHasher.java, TextDeduplicator.java — MinHash LSH deduplication (Jaccard similarity via hashA/hashB arrays, LARGE_PRIME=4294967311L)
  • NGramDecontaminator.java — N-gram decontamination removing test-set leakage
  • LengthBucketingIterator.java, PaddingFreeBatchCollator.java — bucket-sorted batching to minimize padding waste
  • CurriculumIterator.java, DifficultyScorer.java — curriculum learning with pluggable difficulty scoring
  • SequencePacker.java, PackedSequence.java — bin-packing sequences into context windows
  • WeightedDataMixer.java — weighted domain mixing for pretraining blends
  • TextQualityFilter.java, ChatTemplate.java, InstructionDataFormatter.java, StratifiedSplitter.java, AudioDataProcessor.java — quality filtering, chat template formatting, stratified splits, audio preprocessing
  • RawTextDatasetIterator.java, TokenizedTextDataIterator.java — raw and tokenized text iterators

Op Namespace Extensions (updated)

  • NDMath, NDBase, NDBitwise, NDCNN, NDImage, NDLinalg, NDLoss, NDNN, NDRNN, NDRandom — synchronized with new ops from PR13
  • NDAudio.java — audio op namespace (Mel spectrogram, MFCC, pitch detection, Griffin-Lim)
  • NDSignal.java — signal processing (DFT, STFT, windowing functions)
  • NDTraining.java — training ops (gradient accumulation, loss scaling)

Common Infrastructure (8 files)

  • common/config/properties/DspProperties.java, TritonProperties.java — typed property accessors
  • ND4JSystemProperties.java, ND4JEnvironmentVars.java — extended with ~30 new DSP/Triton/multi-device properties
  • schedule/CosineWarmupSchedule.java — cosine annealing + linear warmup LR schedule
  • nativeblas/MultiBackendNativeOpsHolder.java, NativePluginLoader.java — plugin-based native ops loading
  • tools/TritonCacheTool.java — CLI for inspecting/invalidating Triton plan disk cache
  • nd4j-common-tests: TestResultStore, TestResultQuery, TestMilestone for persistent test milestone tracking

Serialization / ONNX Export (2 files)

  • OnnxTensorUtils.javaINDArray ↔ ONNX tensor conversion
  • OnnxExporter.java, OnnxExportConfig.java — SameDiff → ONNX export including INT8/FP16 quantized export

Dependencies

  • Depends on: PR01 (root pom), PR02 (nd4j-common), PR03/PR04 (native bindings layer)
  • Required by: PR13 (op class hierarchy uses new base types), PR14 (backends implement DeviceContextProvider/DeviceDescriptor), PR15 (training uses new config classes), PR16 (DSP runtime uses GraphScope, DspEnvironmentConfig, DspPlanDiskCache)

Merge Order

These 22 PRs must merge in layer order. Each layer depends on the layers above it being merged first. PRs within the same layer are independent and can merge in parallel.

This PR: Merge after Layer 3 (native platform backends + DSP engine).

Layer PRs
0 (no deps) PR01, PR02, PR20
1 (build/infra) PR03, PR04
2 (native core) PR05, PR06, PR07
3 (native feat) PR08, PR09, PR10, PR11
4 (java core) PR12, PR13, PR14, PR15
5 (java feat) PR16
6 (import/gen) PR17, PR18, PR19, PR21
7 (validation) PR22

Part of the 22-PR split of ag_new_release_updates_2 branch.
Merge layer: 4 (java core)
Files: 290

See pr-plans/00-master-plan.md for the full split plan and merge order.
…-core

# Conflicts:
#	nd4j/nd4j-backends/nd4j-api-parent/nd4j-api/src/main/java/org/nd4j/linalg/api/shape/LongShapeDescriptor.java

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot wasn't able to review this pull request because it exceeds the maximum number of lines (20,000). Try reducing the number of changed lines and requesting a review from Copilot again.

OnnxExporter and OnnxExportConfig are pure stubs — every method throws
UnsupportedOperationException and nothing in the codebase references them.
Delete both files rather than shipping dead API surface.

In MmulHelper.cpp, upgrade the silent comment acknowledging a known bug in
mmulNxN batch index calculation to a prominent WARNING + TODO so the
workaround path is discoverable and the root cause is not forgotten.
…ion code

Cherry-picked nd4j-profiler, nd4j-common, and nd4j-common-tests changes
from 2fdeaa4 — switches logging/error handling off raw println.
- AttentionFusionOptimizations: remove 17 [ATTN-DIAG*] println calls
  that fired on every optimization pass (log.debug equivalents exist)
- Permute.doDiff(): convert 3 println to log.debug
- Enter: delete bare empty println() in constructor
- UnsortedSegmentSum: delete bare empty println() in constructor
- OpExclusionUtils: convert println/err to slf4j log.debug/warn
Delete 24 commented-out case entries after the default throw in
pairwiseOpClass that were unreachable dead code.
@agibsonccc

Copy link
Copy Markdown
Contributor Author

Architecture Overview

This PR builds the ND4J Java API and core infrastructure that all higher-level Java modules depend on. It introduces the multi-device abstraction layer (DeviceMemoryManager with per-device caps and selection), coherence-aware data buffers (HybridDataBuffer with HOST_ONLY/DEVICE_ONLY/SYNCHRONIZED states), and a completely lazy framework diagnostics facade.

Highlights

  • Multi-device abstraction with stub topology for testingDeviceMemoryManager singleton with per-device AtomicLong counters and selectDeviceForAllocation(bytes) sorted by priority then free memory; StubDeviceDescriptor/StubDeviceContextProvider enable full multi-GPU topology simulation in unit tests without hardware
  • Graph tracing and Adam8bit optimizerNd4j.graphScope() traces GraphFunction calls into a hidden SameDiff graph via LazyINDArray proxy (blocks data access during trace); Adam8bit uses INT8 per-block absmax quantization (block=2048) to reduce optimizer state ~4x (~56 GB → ~14 GB for 7B-parameter models)

…lVersionUID, slf4j, property centralization

- ArrayUtil god class → ArrayTypeConverters + ArrayMathUtils extractions
- AbstractSession god class → SessionExecutionDiagnostics extraction
- ND4JSystemProperties: 11 new centralized property constants
- serialVersionUID additions across evaluation, config, dataset classes
- System.out.println → slf4j logging migration
- Deprecated boxed constructors → valueOf/autobox
Apply Java compilation fixes for ND4J API core files.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants