Add Qwen4-Exp EngramGate and NGramHashMapping support - #32285
Conversation
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: 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>
Tianlei Wu (tianleiwu)
left a comment
There was a problem hiding this comment.
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_idsexists (WebGPU rejects zero-sized storage buffer bindings) and whypast_aliases_presentexists (a buffer cannot be bound read-only and read-write in the same compute pass). - CUDA: why
NGramPresentIdsKerneluses 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_sizesis 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 oneBlockSum3tree 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_idswithreset_on_eos— exactly the case that diverges between CPU and CUDA/WebGPU (see the inline comment ononnxruntime/contrib_ops/cpu/bert/ngram_hash_mapping.cc). - No chunked-vs-full-sequence equivalence test with
reset_on_eosorsegment_ids, even though the schema promises that equivalence.NGramHashMappingChunkedMatchesFullSequence*already provides the harness. - No negative test for requesting
gated_value_normedwithoutconv_norm_scale(the newORT_RETURN_IF_NOTon all three EPs is untested). - No test combines
head_offsetswitheos_token_id/segment_ids.
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>
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>
Tianlei Wu (tianleiwu)
left a comment
There was a problem hiding this comment.
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.
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Description
This PR expands
com.microsoftcontrib-op support needed by the Qwen4-Exp / Qwen3.8-Flash-Next text architecture, focused on Engram gating and n-gram hash mapping. TheGatedRMSNormactivation-mode changes were removed from this PR and are now covered separately by #32512.Changes include:
EngramGate: add optionalconv_norm_scaleinput and optionalgated_value_normedoutput for the normalized gated-value path.NGramHashMapping: add autoregressive decode and packed-sequence support with optionalpast_ids,present_ids,head_offsets,eos_token_id,segment_ids, andreset_on_eosbehavior.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.
GatedRMSNormsigmoid/activation support is intentionally excluded here so that it can be reviewed and merged via #32512.