Add DeepSeek Engram contrib ops (EngramGate, NGramHashMapping, ShortConv)

[Copilot]

Description

Adds three com.microsoft contrib ops for DeepSeek Engram, each with CPU, CUDA, and WebGPU kernels:

Op	Purpose
NGramHashMapping	Causal n-gram hash ids from compressed tokenizer ids (int32/int64)
EngramGate	Fused key/value projection + dual RMSNorm + gate (float/float16/bfloat16)
ShortConv	Per-hyper-connection RMSNorm + causal depthwise 1D conv with fused SiLU (float/float16/bfloat16)

Layout

• One file per op per EP (engram_gate.*, ngram_hash_mapping.*, short_conv.*), following the existing contrib_ops/<ep>/bert/ convention.
• Shared math/launch helpers live in a new per-EP kernel_helper file rather than being duplicated in each kernel:
  • contrib_ops/cpu/bert/kernel_helper.h — SigmoidFloat, SiluFloat, PositiveMod, WrappedMultiply
  • contrib_ops/cuda/bert/kernel_helper.cuh — same as device functions, plus the shared kThreads / GridSize launch config
  • contrib_ops/webgpu/bert/kernel_helper.h — WGSL snippets (stable_sigmoid, silu, positive_mod)

Type coverage

• CPU: float + MLFloat16 for EngramGate/ShortConv; int32/int64 for NGramHashMapping.
• CUDA: float/MLFloat16/BFloat16.
• WebGPU: WebGpuSupportedFloatTypes(); NGramHashMapping is int32-only.

Notes for review

• WrappedMultiply computes the hash mix through the unsigned counterpart of T so overflow wraps instead of being UB — this is load-bearing for matching reference hash ids, not a style choice.
• New entries are inserted in alphabetical order in each EP's BuildKernelCreateInfo table.
• docs/ContribOperators.md / docs/OperatorKernels.md were updated by hand to match gen_doc output; the Windows CI --gen_doc validate run is the authoritative check here.

Motivation and Context

Engram currently decomposes into long chains of primitive ONNX ops (shifts, hashes, mod, RMSNorm, depthwise conv, sigmoid gating), which is both slow and awkward to export. These three ops collapse those patterns into single kernels so Engram-based models can run efficiently across CPU, CUDA, and WebGPU.

[azure-pipelines]

Azure Pipelines:
There may be pipelines that require an authorized user to comment /azp run to run.

[Copilot]

Pull request overview

Adds three DeepSeek Engram contrib operators across CPU, CUDA, and WebGPU.

Changes:

• Defines and documents EngramGate, NGramHashMapping, and ShortConv schemas.
• Implements and registers kernels for all three execution providers.
• Adds shared helpers and initial operator tests.

Reviewed changes

Copilot reviewed 35 out of 35 changed files in this pull request and generated 9 comments.

Show a summary per file

File	Description
onnxruntime/test/contrib_ops/engram_ops_test.cc	Adds operator tests.
onnxruntime/core/graph/contrib_ops/ms_opset.h	Registers schemas.
onnxruntime/core/graph/contrib_ops/bert_defs.cc	Defines schemas and inference.
onnxruntime/contrib_ops/webgpu/webgpu_contrib_kernels.cc	Registers WebGPU kernels.
onnxruntime/contrib_ops/webgpu/bert/short_conv.h	Declares WebGPU ShortConv.
onnxruntime/contrib_ops/webgpu/bert/short_conv.cc	Implements WebGPU ShortConv.
onnxruntime/contrib_ops/webgpu/bert/ngram_hash_mapping.h	Declares WebGPU hash mapping.
onnxruntime/contrib_ops/webgpu/bert/ngram_hash_mapping.cc	Implements WebGPU hash mapping.
onnxruntime/contrib_ops/webgpu/bert/kernel_helper.h	Adds shared WGSL helpers.
onnxruntime/contrib_ops/webgpu/bert/engram_gate.h	Declares WebGPU EngramGate.
onnxruntime/contrib_ops/webgpu/bert/engram_gate.cc	Implements WebGPU EngramGate.
onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc	Registers CUDA kernels.
onnxruntime/contrib_ops/cuda/bert/short_conv.h	Declares CUDA ShortConv.
onnxruntime/contrib_ops/cuda/bert/short_conv.cc	Validates and launches ShortConv.
onnxruntime/contrib_ops/cuda/bert/short_conv_impl.h	Declares ShortConv launcher.
onnxruntime/contrib_ops/cuda/bert/short_conv_impl.cu	Implements ShortConv kernel.
onnxruntime/contrib_ops/cuda/bert/ngram_hash_mapping.h	Declares CUDA hash mapping.
onnxruntime/contrib_ops/cuda/bert/ngram_hash_mapping.cc	Validates and launches hashing.
onnxruntime/contrib_ops/cuda/bert/ngram_hash_mapping_impl.h	Declares hash launcher.
onnxruntime/contrib_ops/cuda/bert/ngram_hash_mapping_impl.cu	Implements hash kernel.
onnxruntime/contrib_ops/cuda/bert/kernel_helper.cuh	Adds shared CUDA helpers.
onnxruntime/contrib_ops/cuda/bert/engram_gate.h	Declares CUDA EngramGate.
onnxruntime/contrib_ops/cuda/bert/engram_gate.cc	Validates and launches EngramGate.
onnxruntime/contrib_ops/cuda/bert/engram_gate_impl.h	Declares EngramGate launcher.
onnxruntime/contrib_ops/cuda/bert/engram_gate_impl.cu	Implements EngramGate kernel.
onnxruntime/contrib_ops/cpu/cpu_contrib_kernels.cc	Registers CPU kernels.
onnxruntime/contrib_ops/cpu/bert/short_conv.h	Declares CPU ShortConv.
onnxruntime/contrib_ops/cpu/bert/short_conv.cc	Implements CPU ShortConv.
onnxruntime/contrib_ops/cpu/bert/ngram_hash_mapping.h	Declares CPU hash mapping.
onnxruntime/contrib_ops/cpu/bert/ngram_hash_mapping.cc	Implements CPU hash mapping.
onnxruntime/contrib_ops/cpu/bert/kernel_helper.h	Adds shared CPU helpers.
onnxruntime/contrib_ops/cpu/bert/engram_gate.h	Declares CPU EngramGate.
onnxruntime/contrib_ops/cpu/bert/engram_gate.cc	Implements CPU EngramGate.
docs/OperatorKernels.md	Documents provider type support.
docs/ContribOperators.md	Documents public operator schemas.

---

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

Fixed. std::copysign is gone; the CPU kernel now calls a new shared helper kernel_helper::EngramGateArg() (in onnxruntime/contrib_ops/cpu/bert/kernel_helper.h) which returns exactly 0.0f when dot == 0.0f and otherwise sign(dot) * sqrt(max(abs(dot), 1e-6)). The CUDA kernel got the equivalent helper, and WebGPU now has a matching engram_gate_arg() WGSL function, so all three EPs share one definition.

Added a zero-dot regression test, EngramOpsTest.EngramGateZeroDotProduct, which uses an orthogonal key/query pair so the dot product is exactly zero and asserts the gate is exactly 0.5.

[Copilot]

Fixed. copysignf is replaced by the shared kernel_helper::EngramGateArg() device helper in onnxruntime/contrib_ops/cuda/bert/kernel_helper.cuh, which returns 0.0f for a zero dot product. The CPU helper and the new WGSL engram_gate_arg() use the same definition, so all EPs agree. Covered by the new EngramOpsTest.EngramGateZeroDotProduct test.

[Copilot]

Fixed. The kernel is now one block per (token, g) row instead of one thread per output channel:

• Threads in the block stride over d to compute the key projection, accumulating partial key_sum_sq, query_sum_sq and dot_numerator.
• Those three partials are combined with a new shared kernel_helper::BlockSum() block reduction (added to kernel_helper.cuh so ShortConv can reuse it).
• The gate is then computed once and applied in a second strided loop over the output channels, which is where the value projection is computed.

Per-row work drops from hidden_size² × embedding_size to 2 × hidden_size × embedding_size. Launch config is min(rows, 65535) blocks with kThreads threads and kThreads * sizeof(float) bytes of dynamic shared memory.

[Copilot]

Fixed by splitting the shader into two programs:

• EngramGateScalarProgram — one workgroup per (token, g) row. Invocations stride over d, accumulate key_sum_sq / query_sum_sq / dot_numerator into var<workgroup> arrays, do a tree reduction, and invocation 0 writes the single f32 gate for the row into an intermediate {rows} tensor.
• EngramGateProgram — unchanged one-invocation-per-output-element dispatch, but it now only computes the value projection and multiplies by the pre-computed gate read from the intermediate tensor.

The key projection is therefore executed once per row rather than hidden_size times. Input binding order is preserved in both programs (all AddInput() calls, including the optional bias, precede AddOutput() and match the shader order).

[Copilot]

Fixed. ShortConv<T>::Compute now runs a cheap first TryParallelFor over the batch_size * sequence_length * hc_mult rows that fills a std::vector<float> inv_rms with one inverse-RMS value per row. The convolution loop reads inv_rms[source_row] instead of re-reducing the hidden dimension, and norm_scale is hoisted out of the tap loop as well.

Total work drops from O(total × kernel_size × hidden_size) to O(rows × hidden_size) for the normalization plus O(total × kernel_size) for the convolution, and the extra buffer is only rows floats (hidden_size times smaller than the input). The parallel-for cost estimate for the convolution loop was updated to kernel_size accordingly.

[Copilot]

Fixed with a two-kernel launch:

• ShortConvInvRmsKernel — one block per (batch, t, g) row, threads stride over the hidden dimension and the partials are combined with the shared kernel_helper::BlockSum() reduction. It writes one float inverse-RMS per row.
• ShortConvKernel — unchanged grid-stride loop over output elements, but it now reads inv_rms[source_row] per tap instead of re-reducing, and hoists the norm_scale load out of the tap loop.

The scratch buffer is allocated in ShortConv<T>::ComputeInternal via GetScratchBuffer<float>(rows, context->GetComputeStream()), so it is only batch_size * sequence_length * hc_mult floats.

[Copilot]

Fixed. Added a ShortConvInvRmsProgram first pass: one workgroup per (batch, sequence, hc_mult) row, invocations stride over the hidden dimension into a var<workgroup> partials array, tree-reduce, and invocation 0 writes the row's inverse RMS into an intermediate f32 {rows} tensor.

ShortConvProgram now takes that tensor as an extra input and the tap loop simply reads inv_rms[source_row]; norm_scale is also hoisted out of the loop. The epsilon uniform moved to the new program since the conv shader no longer needs it.

[Copilot]

Fixed. The body was extracted into a templated RunNGramHashMappingTest<T>() and there are now two cases: EngramOpsTest.NGramHashMappingInt64 and EngramOpsTest.NGramHashMappingInt32, so the int32 CPU/CUDA paths and the int32-only WebGPU kernel are exercised.

[Copilot]

Fixed. ToTensorType<T> now has a BFloat16 branch calling ToBFloat16(data), matching linear_attention_gates_op_test.cc:37-45.

Because the BFloat16 kernels are CUDA-only, I added a RunOnSupportedProviders<T>() helper: for BFloat16 it runs OpTester against DefaultCudaExecutionProvider() only (and returns false so the caller does GTEST_SKIP() when CUDA is unavailable), and for other types it keeps the default test.Run() path. Added EngramOpsTest.ShortConvBFloat16 and EngramOpsTest.EngramGateBFloat16 with a 2e-2 tolerance.