feat!: align non-InfiniLM operator interfaces#799
Merged
Conversation
Ziminli
previously approved these changes
Jul 23, 2026
19 tasks
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
add,cat,embedding,linear, andmatmulinterfaces with pinned PyTorch schemas.fused_add_rms_norm,silu_and_mul,scaled_dot_product_attention,reshape_and_cache, androtary_embeddinginterfaces aligned with vLLM or PyTorch.Infinilmsuffix instead of placing them in the generated-ATeninternalnamespace.\deprecateddocumentation and C++17[[deprecated]]compiler diagnostics for a later removal.Motivation
Several handwritten operators were added quickly for model-serving integration and predated the repository's current interface conventions. Their names, parameter lists, defaults, mutation behavior, or tensor layouts did not exactly match a stable upstream interface.
This PR gives each canonical non-
Infinilmoperator in scope an upstream-aligned interface, preferring PyTorch, then vLLM. Interfaces were checked against PyTorch commitd84328bdand vLLM commit56a357ed.N/A - no linked issue.
API Alignment Matrix
InfiniOps uses an explicit trailing
outfor tensor-returning operators. It represents the upstream return tensor or.outargument. Upstream in-place APIs preserve their mutation semantics and return no separate tensor.Canonical APIs
add(input, other, out)add(input, other, alpha, out)add.outadd.out(Tensor self, Tensor other, *, Scalar alpha=1, Tensor(a!) out) -> Tensor(a!)cat(tensors, dim, out)cat.outcat.out(Tensor[] tensors, int dim=0, *, Tensor(a!) out) -> Tensor(a!)embedding(weight, indices, out)embedding(weight, indices, padding_idx, scale_grad_by_freq, sparse, out)embeddingembedding(Tensor weight, Tensor indices, SymInt padding_idx=-1, bool scale_grad_by_freq=False, bool sparse=False) -> Tensorlinear(input, weight, bias, out)linear.outlinear.out(Tensor input, Tensor weight, Tensor? bias=None, *, Tensor(a!) out) -> Tensor(a!); computesinput @ weight.T + biasmatmul(input, other, out)matmul.outmatmul.out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!)fused_add_rms_norm(input, residual, weight?, epsilon)fused_add_rms_norminputandresidualand returnsNonereshape_and_cache(key, value, key_cache, value_cache, slot_mapping, kv_cache_dtype, k_scale, v_scale)reshape_and_cacherotary_embedding(positions, query, key?, head_size, cos_sin_cache, is_neox, rope_dim_offset=0, inverse=false)rotary_embeddingqueryand optionalkeyin placesilu_and_mul(input, out)SiluAndMul.forwardd = input.size(-1) / 2; result issilu(input[..., :d]) * input[..., d:]scaled_dot_product_attention(query, key, value, attn_mask, dropout_p, is_causal, scale, enable_gqa, out)scaled_dot_product_attention(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, *, scale=None, enable_gqa=False) -> TensorInfiniLM-Specific APIs
These remain public custom operators, but their suffix makes clear that they do not claim an exact mainstream public schema.
causal_softmax_infinilm(input, out)is_causal=true, or compose a causal mask and softmaxscaled_softmax_infinilm(input, scale, out)muland softmax, or use SDPA'sscalemul.Tensor, PyTorch_softmaxtop_k_top_p_sample_infinilm(logits, k?, p?, seed, offset, out)TopKTopPSamplerDeprecated Compatibility APIs
The old operators remain available for migration, together with their existing test coverage. Each old base class has both a Doxygen
\deprecatednotice and a C++17[[deprecated("Use ... instead.")]]attribute, so downstream C++ callers receive a compiler diagnostic. InfiniOps' compatibility implementations suppress those diagnostics only on the privateinfiniopsandopsbuild targets.add_rms_normfused_add_rms_normswiglusilu_and_mulflash_attentionscaled_dot_product_attentioncausal_softmaxcausal_softmax_infinilmscaled_softmaxscaled_softmax_infinilmtop_k_top_p_samplertop_k_top_p_sample_infinilmtorch_ops.yamlRationalecat,linear, andmatmulare listed because their aligned base overloads can bind to ATen schemas and they have no handwritten Torch backend; codegen therefore emits the Torch backend.addstays out of the allowlist becausesrc/torch/ops/addalready provides its Torch backend. Generating another one would duplicate the backend.docs/aten-operators.md; the resulting allowlist is visible inscripts/torch_ops.yaml.Type of Change
feat- new feature / new operator / new platformfix- bug fixperf- performance improvement (no behavioral change)refactor- code restructuring without behavior changetest- adding or fixing tests onlydocs- documentation onlybuild/ci- build system or CI configurationchore- tooling, formatting, or other non-code changes!in the Conventional Commits prefix or aBREAKING CHANGE:footer)Platforms Affected
WITH_CPU)WITH_NVIDIA)WITH_ILUVATAR)WITH_METAX)WITH_CAMBRICON)WITH_MOORE)WITH_ASCEND)WITH_HYGON)WITH_TORCH)Smoke Test Result
Remote environment:
ssh nvidia,accelerator-dev/nvidia:latest(sha256:dd94fce2f83a...), NVIDIA A100-SXM4-80GB, PyTorch2.10.0a0+b4e4ee8.The smoke unity build first reproduced the old/new causal-softmax helper collision, then passed after commit
590506egave the InfiniLM helpers distinct names.Test Results on Supported Platforms
12544 passed,9315 skippedfull suite before final master syncc9fbe36c9fbe36c9fbe36c9fbe36c9fbe36Additional validation output
Benchmark / Performance Impact
N/A - this PR standardizes interfaces and correctness behavior; it makes no performance claim.
Notes for Reviewers
fused_add_rms_norm,reshape_and_cache,rotary_embedding, and SDPA.Infinilmsuffix is intentional for the three handwritten custom APIs. The generated ATen leading-underscore operators remain the only users of theinternalnamespace.add,cat,embedding, andlinearare consolidated into their existing parametrized test functions.masterthrough feat(nvidia): addmoe_sumoperator #803 and feat(nvidia): addmoe_align_block_sizeoperator #804. The earlier feat(nvidia): addcutlass_scaled_mmoperator #801/feat(torch): addflash_attn_varlen_funcoperator #802 sync exposed a privateSigmoidhelper collision;SiluAndMulSigmoidnow keeps that implementation local without changing either public API. The latest sync retained both the newmoe_sumand existingsilu_and_mulsmoke matchers.torch_ops.yamlcontains onlycat,linear, andmatmulfrom this remediation becauseaddalready has a handwritten Torch backend.BREAKING CHANGE: canonical signatures or mutation semantics change for
add,cat,embedding,linear,matmul,reshape_and_cache, androtary_embedding. Renamed legacy APIs remain temporarily available with compiler-visible deprecation notices.