PR07: C++ Op Implementations#10440
Open
agibsonccc wants to merge 10 commits into
Open
Conversation
Part of the 22-PR split of ag_new_release_updates_2 branch. Merge layer: 2 (native core) Files: 841 See pr-plans/00-master-plan.md for the full split plan and merge order.
This was referenced Jun 15, 2026
…helpers Replace __global__, __device__, __host__, __host__ __device__, and __forceinline__ with SD_KERNEL, SD_DEVICE, SD_HOST, SD_HOST_DEVICE, and SD_INLINE respectively across 28 .cu files under libnd4j/include/ops/declarable/helpers/cuda/. Combined qualifiers (__device__ __forceinline__) are replaced with SD_DEVICE SD_INLINE. __launch_bounds__ attributes are preserved unchanged.
…ops, add TAD warning - dot_product_attention_v2.cpp: replace std::unique_ptr<NDArray> attentionBiasCastOwner with raw NDArray* + explicit delete; remove unused <memory> include - onnx_multi_head_attention.cpp: replace std::unique_ptr<NDArray> attnBiasCastOwner with raw NDArray* + explicit delete - random.cu: replace 3 std::unique_ptr<NDArray> in fillRandomPoisson_ with raw pointers + explicit delete after use; remove unused <memory> include - top_k.cu: replace std::shared_ptr<TadPack> with raw TadPack* (no delete — these are cached singletons owned by ConstantTadHelper) - llm.h: remove two commented-out op blocks (decoder_masked_mha, moe_gate) that referenced /tmp/new_ops_backup/ — dead code with no valid path - Loops.h: replace silent continue workaround with sd_debug warning so out-of-range TAD offsets are visible in debug mode rather than silently swallowed
Fixes type-safety bugs where template functions or ALL_FLOATS ops used hardcoded float instead of the template parameter or input dataType. - dropout.cpp: bufferAsT<float>() -> <T>() in alphaDropOutFunctor_ - selective_scan.cpp: extract templated helper, BUILD_SINGLE_SELECTOR - vlm_ops.cpp: template vlm_patch_embed, vlm_cross_attention, vlm_2d_position_encode with BUILD_SINGLE_SELECTOR dispatch - vision_encode_patches.cpp: template helper with BUILD_SINGLE_SELECTOR - lora_matmul/loha/lokr/multi_lora: create<float> -> create(dataType) - llm_ops.cpp: flash/GQA backward create_<float> -> create_(dataType), apply_alibi e<float>/p<float> -> e<double>/p<double> - paged_attention.cpp: template + BUILD_SINGLE_SELECTOR dispatch - turbo_quant_attention.cpp: template + BUILD_SINGLE_SELECTOR dispatch - decoder_masked_mha.cpp: float accumulators -> T in template function
…_lora - loha_matmul, lokr_matmul, multi_lora_matmul: fix create() argument order to match (char order, vector shape, DataType, LaunchContext*)
Contributor
Author
Architecture OverviewThis PR contains the C++ op implementations — the actual compute kernels for 700+ operations. The Highlights
|
Parallelization (OpenMP pragmas added): - cascade_attention, turbo_quant_attention: COLLAPSE(2) over batch×heads - fused_mrope: COLLAPSE(3) over batch×seq×heads - attention_distillation_loss: reduction for forward, COLLAPSE(2) backward - ema_update, flatten: PARALLEL_FOR_SIMD / PARALLEL_FOR - ggml_dequantize: parallelize all 23 dequant functions - two_way_cross_attention: parallelize softmax rows and batch loop - kl_divergence_per_layer, distillation_kl_loss: reduction pragmas Bug fix: - decoder_masked_mha: fix data race (shared dot product without reduction) - sg_cb: add reduction(+:dot) to SIMD loops std::math → sd::math (SIMD enablement): - fused_mrope, fused_qk_norm_rope, fused_llm_ops: cos/sin/pow - audio: cos/sin/sqrt/pow/log in DFT loops - activation_mul, fusedElementwiseChain, moe_routing, mixture_of_experts Diagnostics: sd_printf → sd_debug in DeclarableOp, PlatformHelper ShapeDescriptor elimination in BroadcastableBoolOp
- Fix sd_tanh/sd_cos/sd_sin template calls to use two params <X, Z> - Fix OMP reduction syntax in distillation_kl_loss.cpp - Fix OpDescriptor.h: hide OpTraits enum from JavaCPP, remove SD_PADDED_NEW_DELETE - Apply THROW_EXCEPTION macro consolidation
…_release_updates_2
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
PR 07 of 22 PRs in the
ag_new_release_updates_2branch split. Merge after Layer 1 (FlatBuffers + memory infra).OpTraitTable.cpp: Single source of truth for op classification; 748 op-name-to-bitmask entries; consumed by NativePlanCompiler, FusionPass, and NativeDynamicShapePlanVIEW_SHAPE_DEPvsVIEWdistinction: VIEW_SHAPE_DEP ops (expand_dims, squeeze, flatten, permute) must NOT carryVALUE_DEPENDENT_SHAPE; VIEW ops (reshape, strided_slice) do — this gates the SHAPES_FROZEN recovery path correctlyautoregressive_decode: Full native decode loop op; plan handle passed as two int32 iArgs (high/low halves of 64-bit C++ pointer)gated_delta_rule:S_t = exp(g_t)*S_{t-1} + beta_t * k_t ⊗ (v_t - exp(g_t)*S_{t-1}^T*k_t); float32 working buffer enforced (no FP16 state per-timestep)paged_attention_forward: Block-pool K/V indirection via page tables; GQA supported vianumKvHeads != numHeadsews()removal, op trait annotations, shape function VIEW vs VALUE_DEPENDENT_SHAPE correctionsWhat Changed
Op infrastructure (10 files)
libnd4j/include/ops/impl/OpTraitTable.cpp— single source of truth; 748 entries; shorthand constants (UNARY_ACT, BINARY_CMP, MATMUL, VIEW_SHAPE_DEP, VIEW, ATTN, CONST_GEN, etc.)libnd4j/include/ops/OpTraitTable.h— OpTraitTable declaration;addTraits()(OR) vssetTraits()(replace) for initialization orderinglibnd4j/include/ops/declarable/MultiPlatformDispatcher.h/impl/MultiPlatformDispatcher.cpp— routes op execution to best platform helper via KernelAutoTunerlibnd4j/include/ops/declarable/KernelManager.h/impl/KernelManager.cpp— per-op kernel lifecycle (warm-up, tune, cache, invalidate)libnd4j/include/ops/declarable/OpExecutionLogger.h— structured execution logging gated on isVerbose/isDebuglibnd4j/include/ops/declarable/headers/llm.h— declares all ~62 LLM opslibnd4j/include/ops/declarable/headers/audio.h— declares all 14 audio opslibnd4j/include/ops/declarable/headers/signal.h— declares 5 signal-processing opslibnd4j/include/ops/declarable/headers/vlm.h— declares 9 VLM opslibnd4j/include/ops/declarable/impl/DeclarableOp.cpp— KernelManager and OpExecutionLogger wired into execution lifecycleLLM ops — core normalization and linear algebra (12 new generic files + helpers)
libnd4j/include/ops/declarable/generic/nn/llm_ops.cpp— rms_norm, skip_rms_norm, rms_norm_linear, rope, quantized_matmul, grouped_query_attention, flash_attention, kv_cache_update, apply_alibi, sliding_window_attentionlibnd4j/include/ops/declarable/generic/nn/fused_llm_ops.cpp— fused_gelu, fused_layer_norm, fused_rope, fused_bias_dropout_residual, fused_rms_norm_swiglu, fused_attention_projection, fused_mropelibnd4j/include/ops/declarable/generic/nn/lora_matmul.cpp—output = input @ weight^T + scaling * (input @ A^T @ B^T)in two passeslibnd4j/include/ops/declarable/generic/nn/loha_matmul.cpp— LoHA (Hadamard Product Adaptation) matmullibnd4j/include/ops/declarable/generic/nn/lokr_matmul.cpp— LoKr (Kronecker Product Adaptation) matmullibnd4j/include/ops/declarable/generic/nn/multi_lora_matmul.cpp— batched multi-adapter LoRAlibnd4j/include/ops/declarable/generic/nn/dora_matmul.cpp— DoRA (magnitude normalization + LoRA delta)libnd4j/include/ops/declarable/generic/nn/loftq_init.cpp— LoftQ initialization (quantized weight + LoRA minimizing quantization error)libnd4j/include/ops/declarable/generic/nn/smooth_quant.cpp— SmoothQuant per-channel scaling for INT8libnd4j/include/ops/declarable/generic/nn/fp8_quantize.cpp/fp8_dequantize.cpp— FP8 (E4M3/E5M2) quantize and dequantizelibnd4j/include/ops/declarable/generic/nn/squared_relu.cpp— squared ReLU (x^2 if x>0 else 0)libnd4j/include/ops/declarable/generic/nn/dual_rope.cpp— dual rotary position embedding (Qwen2-style section-based rope)LLM ops — SSM and recurrent models (4 new files)
libnd4j/include/ops/declarable/generic/nn/selective_scan.cpp— Mamba S6:h_t = A_t*h_{t-1} + B_t*x_t,y_t = C_t*h_t + D*x_t; input-dependent A/B/Clibnd4j/include/ops/declarable/generic/nn/mamba2_ssm.cpp— Mamba-2 SSD; per-head scalar decayA_bar = exp(A[h] * dt[b,t,h]); optional D skip and stateInlibnd4j/include/ops/declarable/generic/nn/gated_delta_rule.cpp— ICLR 2025 (arXiv:2412.06464):S_t = exp(g_t)*S_{t-1} + beta_t*k_t⊗(v_t - exp(g_t)*S_{t-1}^T*k_t)libnd4j/include/ops/declarable/generic/nn/gated_delta_net_block.cpp— full GDN block: QKV projections, gating, delta rule update, output normalizationLLM ops — attention variants (8 new files)
libnd4j/include/ops/declarable/generic/nn/paged_attention.cpp— block-pool K/V indirection via page tables; GQA via numKvHeads != numHeadslibnd4j/include/ops/declarable/generic/nn/attention/mixture_of_experts.cpp— sparse MoE top-K routing (Mixtral/Switch style)libnd4j/include/ops/declarable/generic/nn/attention/windowed_attention.cpp— sliding window attention (Longformer style)libnd4j/include/ops/declarable/generic/nn/attention/relative_position_bias.cpp— T5-style relative position bias added to attention logitslibnd4j/include/ops/declarable/generic/nn/cascade_attention_ops.cpp— cascade attention (local + global patterns)libnd4j/include/ops/declarable/generic/nn/lightning_attention_ops.cpp— Lightning Attention with causal masking and linear complexitylibnd4j/include/ops/declarable/generic/nn/linear_attention_decode.cpp— linear attention decode step for recurrent linear attention inferencelibnd4j/include/ops/declarable/generic/nn/shared_kv_attention.cpp— shared K/V heads across transformer layersLLM ops — generation and sampling (5 new files)
libnd4j/include/ops/declarable/generic/nn/autoregressive_decode.cpp— full native decode loop; plan handle as split int32 iArgs (high/low 32-bit halves of 64-bit pointer)libnd4j/include/ops/declarable/generic/nn/token_sample.cpp— temperature/top-p/top-k sampling from logitslibnd4j/include/ops/declarable/generic/nn/top_k_renorm_ops.cpp— top-K renormalization for constrained samplinglibnd4j/include/ops/declarable/generic/nn/sampling_penalties.cpp— repetition and frequency penalty applicationlibnd4j/include/ops/declarable/generic/nn/turbo_quant_attention.cpp— fused quantized attention with INT8 K/V cacheLLM ops — training utilities (5 new files)
libnd4j/include/ops/declarable/generic/loss/attention_distillation_loss.cpp— attention map distillation losslibnd4j/include/ops/declarable/generic/loss/distillation_kl_loss.cpp— KL divergence distillation (temperature-scaled soft labels)libnd4j/include/ops/declarable/generic/loss/feature_distillation_loss.cpp— intermediate layer feature matching losslibnd4j/include/ops/declarable/generic/loss/contrastive_loss.cpp— contrastive learning loss for embedding similaritylibnd4j/include/ops/declarable/generic/nn/simpo_loss.cpp— SimPO (Simple Preference Optimization) RLHF-style lossLLM ops — KV cache and memory management (4 new files)
libnd4j/include/ops/declarable/generic/nn/kv_scatter.cpp— scatter new KV pairs into paged KV block poolslibnd4j/include/ops/declarable/generic/nn/unpad_input.cpp— unpad_input/pad_input for variable-length sequence batchinglibnd4j/include/ops/declarable/generic/nn/per_layer_embedding.cpp— per-layer learned position embeddingslibnd4j/include/ops/declarable/generic/nn/async_checkpoint_offload.cpp— async CPU offload for activation checkpointing (D2H and H2D prefetch)LLM ops — MoE and mixture models (3 new files)
libnd4j/include/ops/declarable/generic/nn/moe_shared_experts.cpp— shared (always-active) experts in MoE modelslibnd4j/include/ops/declarable/generic/nn/segment_gemm_ops.cpp— segmented GEMM for variable token counts per expertlibnd4j/include/ops/declarable/generic/nn/onnx_multi_head_attention.cpp— ONNX-compatible MHA op matching opset 17 signatureVLM ops (2 new files)
libnd4j/include/ops/declarable/generic/nn/vlm_ops.cpp— vlm_vision_encode, vlm_image_embed, vlm_patch_embed, vlm_cross_attention, vlm_multimodal_fusion, vlm_vision_projection, vlm_image_preprocess, vlm_2d_position_encodelibnd4j/include/ops/declarable/generic/nn/vision_encode_patches.cpp— splits images into fixed-size patches with optional position encodingAudio ops — all new (14 files)
libnd4j/include/ops/declarable/generic/audio/mel_spectrogram.cpp— mel spectrogram from raw waveformlibnd4j/include/ops/declarable/generic/audio/whisper_mel_spectrogram.cpp— 80 bins, 25ms window, 10ms hop, normalized to [-1, 1]libnd4j/include/ops/declarable/generic/audio/mfcc.cpp— 13 or 40 Mel Frequency Cepstral Coefficientslibnd4j/include/ops/declarable/generic/audio/mel_filterbank.cpp— standalone mel filterbank matrix generationlibnd4j/include/ops/declarable/generic/audio/griffin_lim.cpp— iterative spectrogram-to-waveform inversionlibnd4j/include/ops/declarable/generic/audio/spectral_centroid.cpp/spectral_rolloff.cpp— spectral shape descriptorslibnd4j/include/ops/declarable/generic/audio/chroma_features.cpp— 12 pitch class chromagramlibnd4j/include/ops/declarable/generic/audio/zero_crossing_rate.cpp/pre_emphasis.cpp/audio_normalize.cpp/a_weighting.cpp/pitch_detection.cpp/audio_resample.cpp— signal preprocessing and analysislibnd4j/include/ops/declarable/helpers/cpu/audio.cpp/cuda/audio.cu— CPU and CUDA implementations for all 14 audio opsSignal processing ops — all new (5 files)
libnd4j/include/ops/declarable/generic/signal/dft.cpp— DFT with [real, imag] interleaved; supports inverse and one-sidedlibnd4j/include/ops/declarable/generic/signal/stft.cpp— STFT with configurable window, hop, center padding, normalized outputlibnd4j/include/ops/declarable/generic/signal/window_functions.cpp— hann_window, hamming_window, blackman_windowlibnd4j/include/ops/declarable/helpers/cpu/signal.cpp/cuda/signal.cu— CPU and CUDA implementationsNeural network ops — new additions
libnd4j/include/ops/declarable/generic/nn/causal_conv1d.cpp— causal 1D convolution with left-padding (used in Mamba)libnd4j/include/ops/declarable/generic/nn/ctc_greedy_decoder.cpp— CTC greedy decoder (argmax per timestep with blank collapse)libnd4j/include/ops/declarable/generic/nn/center_and_sharpen.cpp— center-and-sharpen image normalization for vision modelslibnd4j/include/ops/declarable/generic/nn/ema_update.cpp— exponential moving average updatelibnd4j/include/ops/declarable/generic/nn/fused_elementwise_chain.cpp— fused chain of elementwise ops specified via iArgslibnd4j/include/ops/declarable/generic/nn/kl_divergence_per_layer.cpp— per-layer KL divergence for layer-wise distillationlibnd4j/include/ops/declarable/generic/nn/pooling/adaptive_pooling.cpp— adaptive average/max pooling to fixed output sizelibnd4j/include/ops/declarable/generic/nn/convo/deformable_conv2d.cpp— deformable convolution with learnable spatial offsets (DCN)Math and broadcastable ops — new additions
libnd4j/include/ops/declarable/generic/broadcastable/xdivy.cpp/xlogy.cpp/xlog1py.cpp— safe x/y, xlog(y), xlog(1+y) with 0 when x==0libnd4j/include/ops/declarable/generic/linalg/einsum.cpp— einsum contractionlibnd4j/include/ops/declarable/generic/images/affine_grid.cpp/grid_sample.cpp— spatial transformer grid generation and bilinear samplinglibnd4j/include/ops/declarable/generic/transforms/cast_and_scale.cpp— fused cast + scale for quantization pipelineExisting op updates (567 files)
All existing op implementations across all categories updated for:
ews()/elementWiseStrideremoval; replaced with stride-based contiguity checksgetOpDescriptor()->addTraits(...)for OpTraitTable classificationKey categories: broadcastable (30 files), blas/matmul (4), boolean ops (11), images (14), linalg (22), list ops (11), loss (10), NLP/skipgram/cbow (2), activations (15), convolutions (12), recurrent ops (10), reductions (16), shape ops (14), transforms/scatters/gathers (40+), updaters (10)
Helper implementations — new (130+ files)
libnd4j/include/ops/declarable/helpers/cpu/— 70+ new CPU helpers (autoregressive_decode, gated_delta_rule, mamba2_ssm, paged_attention, lightning_attention, kv_cache_quantize, speculative_decoding, token_sample, rms_norm, layer_norm, dual_rope, mixture_of_experts, segment_gemm, etc.)libnd4j/include/ops/declarable/helpers/cuda/— 70+ matching CUDA helperslibnd4j/include/ops/declarable/helpers/*.h— 40+ new header files for helper interfacesDeleted files (4 files)
libnd4j/include/ops/declarable/generic/README.md— documentation file removedlibnd4j/include/ops/gemm.h/impl/gemm.cpp— replaced by MmulHelper + CutlassGemmHelperDependencies
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.