PR13: Java Op Definitions - #10446

Summary

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

• ~130 new DynamicCustomOp definitions for LLM/VLM inference and training; all registered in ImportClassMapping (ONNX/TF) and DifferentialFunctionClassHolder (SameDiff)
• AutoregressiveDecode: Full decode loop as a single JNI call eliminating per-step Java↔C++ overhead; 17+ iArgs encode plan/context handles, slot indices, stop tokens; outputs: generatedTokenIds, tokenCount, timingInfo[5] (totalMs, avgDecodeMs, tokPerSec, p50Ms, p99Ms)
• 13 attention variants: FlashAttention, GroupedQueryAttention, MLAAttention (DeepSeek-V3), CascadeAttention (chunked prefill), DecoderMaskedMha, LightningAttention, SlidingWindowAttention, OnnxMultiHeadAttention, TwoWayCrossAttention (SmolDocling/SAM), PagedAttentionForward, PagedKvAppend, DotProductAttentionV2 (updated)
• PEFT fused adapters: LoraMatMul, DoraMatMul, LohaMatMul (Hadamard), LokrMatMul (Kronecker); all have differentiable backward ops; scaling as float arg; MultiLoraMatmul batches multiple adapters for serving
• LLM normalization/activation: RmsNorm (LLaMA/Mistral semantics, no mean-centering, epsilon=1e-5), RmsNormLinear, FusedRmsNormSwiGLU, SiluAndMul, GeluAndMul, SquaredReLU, FusedLayerNorm
• Quantization ops: FP8 E4M3/E5M2 (Fp8Quantize/Dequantize/Matmul), AWQ fused matmul, GGML block dequantization (Q4_0/Q4_K/Q8_0), generic INT4/INT8 QuantizedMatmul, SmoothQuant
• Sampling ops: GpuTopKSample, GpuTopPSample, TokenSample, SamplingPenalties, TopKRenorm, TopPRenorm
• Recurrent/SSM ops: SelectiveScan (Mamba), Mamba2SSM, GatedDeltaNetBlock, LinearAttentionDecode, EmaUpdate
• MoE: MixtureOfExperts, MoeGate (top-k routing), MoeSharedExperts (DeepSeek-MoE)
• Audio: 14 ops including WhisperMelSpectrogramOp (80 mel bins, 25ms window, 10ms hop — exact Whisper preprocessing), MFCC, GriffinLim, STFT, windowing functions
• Conv extensions: AdaptiveAvgPooling2D/3D, AdaptiveMaxPooling2D, DeformableConv2D, AffineGrid, GridSample, CausalConv1d, PixelShuffle
• Executioner infrastructure: MultiBackendExecutioner, BackendRoutingStrategy, KernelPluginManager, HelperRouter, routing policy hierarchy (category, data-locality, manual, performance)

What Changed

Fused Attention Ops (12 new files)

• FlashAttention.java / FlashAttentionBp.java — tiled IO-aware attention; Q/K/V + optional causal mask; float args: scale factor
• GroupedQueryAttention.java / GroupedQueryAttentionBp.java — GQA with separate num_heads and num_kv_heads (LLaMA3/Gemma)
• MLAAttention.java — Multi-head Latent Attention (DeepSeek-V3 compressed KV)
• CascadeAttention.java — chunked prefill / long-context decoding
• DecoderMaskedMha.java — decoder masked MHA with KV injection
• LightningAttention.java — linear attention for efficient sequence modeling
• SlidingWindowAttention.java — sliding window attention (Mistral style)
• OnnxMultiHeadAttention.java — ONNX-compatible MHA for import path
• TwoWayCrossAttention.java / TwoWayCrossAttentionBp.java — bidirectional cross-attention for VLMs (SmolDocling/SAM)
• PagedAttentionForward.java / PagedKvAppend.java — paged KV cache forward and slot appending

KV Cache Ops (5 files)

• KVCache.java — in-place KV cache state holder
• KVCacheUpdate.java — updates KV cache at current position (scatter)
• KVCacheQuantize.java / KVCacheDequantize.java — INT8/FP8 KV quantization for memory reduction
• KvScatter.java / SharedKvAttention.java — scatter into paged blocks; multi-query shared KV heads

PEFT Fused Linear Ops (10 files)

• LoraMatMul.java / LoraMatMulBp.java — output = input @ weight + scaling * (input @ A^T @ B^T)
• DoraMatMul.java / DoraMatMulBp.java — weight-decomposed LoRA with per-column magnitude normalization
• LohaMatMul.java / LohaMatMulBp.java — Hadamard product adaptation W + W1 ⊙ W2
• LokrMatMul.java / LokrMatMulBp.java — Kronecker product low-rank adaptation
• MultiLoraMatmul.java — batched multi-adapter LoRA for serving
• ColumnParallelLinear.java / RowParallelLinear.java — tensor-parallel linear layers for multi-GPU

Normalization and Activation Ops (12 files)

• RmsNorm.java / RmsNormBp.java — x * rsqrt(mean(x^2) + eps) * gamma; no mean-centering; epsilon=1e-5 (LLaMA/Mistral)
• RmsNormLinear.java / RmsNormLinearBp.java — fused RMSNorm + linear projection
• SkipRmsNorm.java — RMSNorm with residual skip connection
• FusedRmsNormSwiGLU.java / FusedRmsNormSwiGLUBp.java — fused RMSNorm + SwiGLU gating
• SiLU.java / SiLUBp.java — SiLU: x * sigmoid(x)
• SiluAndMul.java, GeluAndMul.java — fused SiLU/GELU + element-wise multiply (gate computation)
• FusedGELU.java / FusedGELUBp.java, SwishMul.java / SwishMulBp.java, SquaredReLU.java, FusedLayerNorm.java / FusedLayerNormBp.java

Positional Encoding Ops (7 files)

• RoPE.java / RoPEBp.java, FusedRoPE.java / FusedRoPEBp.java — Rotary Position Embedding (LLaMA/Mistral)
• FusedMRoPE.java — multi-modal RoPE for VLMs
• DualRoPE.java — dual RoPE with per-head encodings
• ApplyAlibi.java — ALiBi positional bias
• RelativePositionBias.java — T5-style relative position bias
• PerLayerEmbedding.java — per-layer rotary frequency embedding

Quantization and Sampling Ops (10 files)

• AwqMatmul.java — AWQ fused matmul
• Fp8Quantize.java / Fp8Dequantize.java / Fp8Matmul.java — FP8 E4M3 (forward) / E5M2 (backward) quantization and matmul
• QuantizedMatmul.java — generic INT4/INT8 quantized matmul with scale and zero-point
• GGMLDequantize.java — GGML block-quantized dequantization (Q4_0, Q4_K, Q8_0)
• SmoothQuant.java — per-channel weight/activation scaling for INT8
• GpuTopKSample.java / GpuTopPSample.java, TokenSample.java, SamplingPenalties.java, TopKRenorm.java / TopPRenorm.java — GPU sampling and probability renormalization

Recurrent / SSM Ops (6 files)

• SelectiveScan.java — Mamba SSM selective scan
• Mamba2SSM.java — Mamba-2 state space model recurrence
• GatedDeltaNetBlock.java / GatedDeltaRule.java — Gated Delta Networks recurrent block
• LinearAttentionDecode.java — linear attention decode step
• EmaUpdate.java / EmaUpdateBp.java — exponential moving average update

MoE Ops (3 files)

• MixtureOfExperts.java — sparse routing + expert computation
• MoeGate.java — top-k expert selection gating network
• MoeSharedExperts.java — shared experts for DeepSeek-MoE style models

Fused Fusion Ops (8 files)

• FusedElementwiseChain.java — fused chain of element-wise ops
• FusedBiasDropoutResidual.java — fused bias add + dropout + residual
• FusedGemmSwiglu.java / FusedGemmSwigluBp.java — fused GEMM + SwiGLU (single MLP kernel)
• FusedNormQuantize.java, SegmentGemm.java, CenterAndSharpen.java / CenterAndSharpenBp.java, VisionEmbeddingMerge.java, VisionEncodePatches.java

Autoregressive Decode Op (1 file)

• AutoregressiveDecode.java — single-op decode loop; inputs: prefillEmbeddings, embeddingTable, inputIds, optional mask, posIds, KV buffers, plan external inputs; outputs: generatedTokenIds, tokenCount, timingInfo[5]

Audio and Signal Processing Ops (14 files)

• MelSpectrogram.java, MelFilterbank.java, MFCC.java, WhisperMelSpectrogramOp.java — Mel/MFCC extraction; Whisper: 80 mel bins, 25ms window, 10ms hop
• GriffinLim.java, PitchDetection.java, SpectralCentroid.java, SpectralRolloff.java, ZeroCrossingRate.java — audio feature extraction
• AudioNormalize.java, AudioResample.java, PreEmphasis.java, AWeighting.java, ChromaFeatures.java — audio preprocessing pipeline
• DFT.java, STFT.java, HannWindow.java, HammingWindow.java, BlackmanWindow.java — signal processing

Convolution Extensions (7 files)

• AdaptiveAvgPooling2D.java / AdaptiveAvgPooling2DBp.java, AdaptiveAvgPooling3D.java — adaptive avg pooling (output size, not kernel size)
• AdaptiveMaxPooling2D.java / AdaptiveMaxPooling2DBp.java — adaptive max pooling
• DeformableConv2D.java / config/DeformableConv2DConfig.java — deformable conv with learnable offsets
• AffineGrid.java / GridSample.java — spatial transformer: affine grid + grid sampling
• CausalConv1d.java, PixelShuffle.java — causal 1D conv; pixel shuffle for SR upsampling

Loss Functions for Distillation (8 files)

• DistillationKLLoss.java / DistillationKLLossBp.java — KL divergence for teacher/student softmax
• AttentionDistillationLoss.java / AttentionDistillationLossBp.java — attention map distillation
• FeatureDistillationLoss.java / FeatureDistillationLossBp.java — MSE on intermediate activations
• ContrastiveLoss.java / ContrastiveLossBp.java — contrastive loss for embedding training

Executioner Infrastructure (12 files)

• MultiBackendExecutioner.java / DefaultMultiBackendExecutioner.java — multi-backend op dispatch interface and default
• BackendRoutingStrategy.java / DefaultBackendRoutingStrategy.java — routing strategy for CPU/CUDA/TPU
• DeviceAwareOpExecutioner.java — respects DeviceRoutingConfiguration
• KernelManager.java / KernelPluginManager.java / KernelSelector.java / KernelSelectionConfig.java — kernel plugin system
• OpExecutionDelegator.java / TransferMetrics.java — delegating executioner wrapper
• HelperRouter.java / PlatformHelperDescriptor.java — platform helper routing (cuDNN, MKL, oneDNN)
• RoutingPolicy.java, RoutingDecision.java, CategoryBasedPolicy.java, DataLocalityPolicy.java, ManualPolicy.java, PerformancePolicy.java — routing policy hierarchy

Dependencies

• Depends on: PR12 (base DynamicCustomOp, BaseOp, INDArray, DataType API)
• Required by: PR15 (SameDiff op namespaces wrap these), PR16 (graph optimizer detects and replaces these op patterns), PR17-PR19 (import uses these op classes)

Merge Order

This PR is in Layer 4 (Java op definitions — parallel with PR12, both needed before PR16).