Skip to content

Add Qwen4-Exp EngramGate and NGramHashMapping support - #32285

Merged
kunal-vaishnavi merged 39 commits into
mainfrom
copilot/add-op-support-for-qwen-3-8-flash-next
Sep 15, 2026
Merged

kunal-vaishnavi merged 39 commits into
mainfrom
copilot/add-op-support-for-qwen-3-8-flash-next

Conversation

Copilot AI commented Aug 26, 2026 •

Copy link
Copy Markdown
Contributor

Description

This PR expands com.microsoft contrib-op support needed by the Qwen4-Exp / Qwen3.8-Flash-Next text architecture, focused on Engram gating and n-gram hash mapping. The GatedRMSNorm activation-mode changes were removed from this PR and are now covered separately by #32512.

Changes include:

  • EngramGate: add optional conv_norm_scale input and optional gated_value_normed output for the normalized gated-value path.
  • NGramHashMapping: add autoregressive decode and packed-sequence support with optional past_ids, present_ids, head_offsets, eos_token_id, segment_ids, and reset_on_eos behavior.
  • CPU, CUDA, and WebGPU implementations for the retained EngramGate and NGramHashMapping behavior.
  • Generated contrib operator/kernel documentation for the retained schema and kernel changes.
  • Extended contrib-op tests covering CPU/CUDA/WebGPU behavior, in-place state handling, packed sequences, EOS reset behavior, and EngramGate normalized output.

Motivation and Context

Qwen4-Exp requires n-gram embeddings and Engram-style gating variants that were not fully covered by the existing contrib operators. This PR extends the relevant operators incrementally while preserving prior behavior by default.

GatedRMSNorm sigmoid/activation support is intentionally excluded here so that it can be reviewed and merged via #32512.

Copilot AI and others added 13 commits August 25, 2026 20:18
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
…eek-engram' into copilot/add-op-support-for-qwen-3-8-flash-next

Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
@azure-pipelines

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

…ids, head_offsets

Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
…t-for-qwen-3-8-flash-next

Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
@kunal-vaishnavi kunal-vaishnavi changed the title Add Qwen4-Exp op support for gated RMSNorm, n-gram hashing, and EngramGate Add Qwen4-Exp EngramGate and NGramHashMapping support Sep 11, 2026

@tianleiwu Tianlei Wu (tianleiwu) 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.

Review summary

The combined = past_ids ++ input_ids reformulation is a clean way to express both the EOS and segment resets, conv_norm_scale/gated_value_normed are validated consistently on all three EPs, and the generated docs match the schema. Two things I would like to see addressed before merge: a cross-EP correctness divergence in the reset-boundary scan, and a CPU throughput regression for batch_size == 1. Details are in the inline comments; the cross-cutting notes are below.

Schema documentation

The NGramHashMapping doc block dropped the paragraph documenting that past_ids and present_ids may share an allocation, together with the transaction-safety caveat, even though CUDA (NGramPresentIdsKernel's chunked barrier) and WebGPU (past_aliases_present) still implement that contract. That paragraph was the only place telling an integrator the aliasing is supported and under what conditions — please keep it.

Comment removal unrelated to the feature

The PR deletes roughly 40 lines of rationale comments across all six implementation files that have nothing to do with the new inputs/outputs. The highest-risk removals are the ones that explain non-obvious constraints a future edit would otherwise violate:

  • WebGPU: why has_input_ids exists (WebGPU rejects zero-sized storage buffer bindings) and why past_aliases_present exists (a buffer cannot be bound read-only and read-write in the same compute pass).
  • CUDA: why NGramPresentIdsKernel uses one block per row with a chunked barrier (past/present aliasing), and why the two lookup tables are staged into shared memory with a 16 KB cap.
  • CPU: why vocab_sizes is validated on the host instead of guarded on device (a device guard turns the mistake into a silent constant hash id).
  • CUDA engram_gate_impl.cu: why the three partials are fused into one BlockSum3 tree reduction.

These are invariants, not narration, so they fall on the keep side of the repo's comment guidance. Separating cosmetic comment churn from the functional change would also make this diff much easier to review.

Test coverage gaps

The new tests pin exact expected values and cover both int32 and int64, which is good. Missing:

  • No test combines past_ids with reset_on_eos — exactly the case that diverges between CPU and CUDA/WebGPU (see the inline comment on onnxruntime/contrib_ops/cpu/bert/ngram_hash_mapping.cc).
  • No chunked-vs-full-sequence equivalence test with reset_on_eos or segment_ids, even though the schema promises that equivalence. NGramHashMappingChunkedMatchesFullSequence* already provides the harness.
  • No negative test for requesting gated_value_normed without conv_norm_scale (the new ORT_RETURN_IF_NOT on all three EPs is untested).
  • No test combines head_offsets with eos_token_id / segment_ids.

Comment thread onnxruntime/contrib_ops/cpu/bert/ngram_hash_mapping.cc Outdated
Comment thread onnxruntime/contrib_ops/cpu/bert/ngram_hash_mapping.cc Outdated
Comment thread onnxruntime/contrib_ops/cpu/bert/ngram_hash_mapping.cc Outdated
Comment thread onnxruntime/contrib_ops/cpu/bert/ngram_hash_mapping.cc Outdated
Comment thread onnxruntime/contrib_ops/cpu/bert/ngram_hash_mapping.cc
Comment thread onnxruntime/core/graph/contrib_ops/bert_defs.cc Outdated
Comment thread onnxruntime/contrib_ops/webgpu/bert/ngram_hash_mapping.cc
Comment thread onnxruntime/contrib_ops/cuda/bert/engram_gate_impl.cu Outdated
Comment thread onnxruntime/contrib_ops/cpu/bert/ngram_hash_mapping.cc Outdated
Comment thread onnxruntime/contrib_ops/cuda/bert/ngram_hash_mapping_impl.cu Outdated
Comment thread onnxruntime/contrib_ops/cpu/bert/engram_gate.cc Outdated
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Copilot AI and others added 3 commits September 14, 2026 22:06
…t-for-qwen-3-8-flash-next

Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>

@tianleiwu Tianlei Wu (tianleiwu) 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.

The cross-provider fixes at the current head address the earlier EOS-reset, aliasing, invalid-vocabulary, offset, and normalization concerns, and all prior threads from my earlier round are resolved. I found two remaining non-blocking graph-time validation issues; details are inline.

Comment thread onnxruntime/core/graph/contrib_ops/bert_defs.cc Outdated
Comment thread onnxruntime/core/graph/contrib_ops/bert_defs.cc
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
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.

5 participants