feat!: align non-InfiniLM operator interfaces

[voltjia]

Summary

• Align canonical add, cat, embedding, linear, and matmul interfaces with pinned PyTorch schemas.
• Add canonical fused_add_rms_norm, silu_and_mul, scaled_dot_product_attention, reshape_and_cache, and rotary_embedding interfaces aligned with vLLM or PyTorch.
• Name handwritten serving-specific softmax and sampling kernels with an explicit Infinilm suffix instead of placing them in the generated-ATen internal namespace.
• Retain the migrated legacy operators and their existing tests with Doxygen \deprecated notices 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-Infinilm operator in scope an upstream-aligned interface, preferring PyTorch, then vLLM. Interfaces were checked against PyTorch commit d84328bd and vLLM commit 56a357ed.

N/A - no linked issue.

API Alignment Matrix

InfiniOps uses an explicit trailing out for tensor-returning operators. It represents the upstream return tensor or .out argument. Upstream in-place APIs preserve their mutation semantics and return no separate tensor.

Canonical APIs

InfiniOps execution API	Alignment target	Upstream interface / semantics	Evidence
add(input, other, out) add(input, other, alpha, out)	PyTorch add.out	add.out(Tensor self, Tensor other, *, Scalar alpha=1, Tensor(a!) out) -> Tensor(a!)	ATen schema
cat(tensors, dim, out)	PyTorch cat.out	cat.out(Tensor[] tensors, int dim=0, *, Tensor(a!) out) -> Tensor(a!)	ATen schema
embedding(weight, indices, out) embedding(weight, indices, padding_idx, scale_grad_by_freq, sparse, out)	PyTorch embedding	embedding(Tensor weight, Tensor indices, SymInt padding_idx=-1, bool scale_grad_by_freq=False, bool sparse=False) -> Tensor	ATen schema
linear(input, weight, bias, out)	PyTorch linear.out	linear.out(Tensor input, Tensor weight, Tensor? bias=None, *, Tensor(a!) out) -> Tensor(a!); computes input @ weight.T + bias	ATen schema
matmul(input, other, out)	PyTorch matmul.out	matmul.out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!)	ATen schema
fused_add_rms_norm(input, residual, weight?, epsilon)	vLLM fused_add_rms_norm	Same order; mutates both input and residual and returns None	vLLM custom op
reshape_and_cache(key, value, key_cache, value_cache, slot_mapping, kv_cache_dtype, k_scale, v_scale)	vLLM reshape_and_cache	Same order; mutates separate key/value caches and observes padding-slot, cache-dtype, and scale semantics	vLLM custom op
rotary_embedding(positions, query, key?, head_size, cos_sin_cache, is_neox, rope_dim_offset=0, inverse=false)	vLLM rotary_embedding	Same order/defaults; mutates query and optional key in place	vLLM custom op
silu_and_mul(input, out)	vLLM SiluAndMul.forward	Packed input: d = input.size(-1) / 2; result is silu(input[..., :d]) * input[..., d:]	vLLM activation
scaled_dot_product_attention(query, key, value, attn_mask, dropout_p, is_causal, scale, enable_gqa, out)	PyTorch SDPA	scaled_dot_product_attention(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, *, scale=None, enable_gqa=False) -> Tensor	ATen schema

InfiniLM-Specific APIs

These remain public custom operators, but their suffix makes clear that they do not claim an exact mainstream public schema.

InfiniOps API	Alignment status	Closest upstream form	Evidence
causal_softmax_infinilm(input, out)	Custom, not exactly aligned	Use SDPA with is_causal=true, or compose a causal mask and softmax	PyTorch SDPA schema
scaled_softmax_infinilm(input, scale, out)	Custom, not exactly aligned	Compose mul and softmax, or use SDPA's scale	PyTorch mul.Tensor, PyTorch _softmax
top_k_top_p_sample_infinilm(logits, k?, p?, seed, offset, out)	Custom, not exactly aligned	vLLM's high-level sampler accepts per-request generators and can return processed logits	vLLM TopKTopPSampler

Deprecated Compatibility APIs

The old operators remain available for migration. Existing tests were restored unchanged where they existed, and each old base class has a \deprecated notice stating that it will be removed in a future release.

Deprecated API	Migration target	Status
add_rms_norm	fused_add_rms_norm	Retained; old out-of-place parameter list is not canonical vLLM
swiglu	silu_and_mul	Retained; old split-input form is not vLLM's packed-input API
flash_attention	scaled_dot_product_attention	Base API retained; use SDPA for standard attention semantics
causal_softmax	causal_softmax_infinilm	Retained under its old name for migration
scaled_softmax	scaled_softmax_infinilm	Retained under its old name for migration
top_k_top_p_sampler	top_k_top_p_sample_infinilm	Retained; migration is not a drop-in parameter substitution

torch_ops.yaml Rationale

• cat, linear, and matmul are listed because their aligned base overloads can bind to ATen schemas and they have no handwritten Torch backend; codegen therefore emits the Torch backend.
• add stays out of the allowlist because src/torch/ops/add already provides its Torch backend. Generating another one would duplicate the backend.
• The repository's generated-versus-handwritten rule is documented in docs/aten-operators.md; the resulting allowlist is visible in scripts/torch_ops.yaml.

Type of Change

• feat - new feature / new operator / new platform
• fix - bug fix
• perf - performance improvement (no behavioral change)
• refactor - code restructuring without behavior change
• test - adding or fixing tests only
• docs - documentation only
• build / ci - build system or CI configuration
• chore - tooling, formatting, or other non-code changes
• Breaking change (requires a ! in the Conventional Commits prefix or a BREAKING CHANGE: footer)

Platforms Affected

• CPU (WITH_CPU)
• NVIDIA (WITH_NVIDIA)
• Iluvatar (WITH_ILUVATAR)
• MetaX (WITH_METAX)
• Cambricon (WITH_CAMBRICON)
• Moore (WITH_MOORE)
• Ascend (WITH_ASCEND)
• Hygon (WITH_HYGON)
• PyTorch C++ bindings (WITH_TORCH)
• Build system / CMake / CI
• Python bindings / user-facing API

Smoke Test Result

Remote environment: ssh nvidia, accelerator-dev/nvidia:latest (sha256:dd94fce2f83a...), NVIDIA A100-SXM4-80GB, PyTorch 2.10.0a0+b4e4ee8.

INFINI_RT_ROOT=/opt/infinirt scripts/dev/build.sh nvidia --smoke --jobs 16
[build] done

PYTHONPATH=build-nvidia/install python -m pytest tests -m smoke -q
80 passed, 14 skipped, 12577 deselected in 10.36s

The smoke unity build first reproduced the old/new causal-softmax helper collision, then passed after commit 590506e gave the InfiniLM helpers distinct names.

Test Results on Supported Platforms

Platform	Affected	Build / Smoke Result	Full Result / Notes
CPU	Yes	Combined CPU/NVIDIA/Torch build and smoke passed manually	Covered by affected-op and full suites
NVIDIA	Yes	Full/latest smoke builds and main/shadow CI passed	12544 passed, 9315 skipped full suite on A100
Iluvatar	Yes	Main and shadow CI passed on latest head	Common CUDA wrapper and smoke coverage passed
MetaX	Yes	Main and shadow CI passed on latest head	Common CUDA wrapper and smoke coverage passed
Cambricon	Yes	Main and shadow CI passed on latest head	Backend build and smoke coverage passed
Moore	Yes	Main and shadow CI passed on latest head	Common CUDA wrapper and smoke coverage passed
Ascend	Yes	Main and shadow CI passed on latest head	Backend build and smoke coverage passed
Hygon	Yes	SDK/hardware unavailable	Wrapper codegen only

Additional validation output

# Full build at e3ce1c9; latest commits only isolate private helper names and format files.
INFINI_RT_ROOT=/opt/infinirt scripts/dev/build.sh nvidia --jobs 16
[build] done

# All APIs changed by this PR, including old compatibility and new Infinilm names.
python -m pytest <21 affected test files> -q
1381 passed, 255 skipped in 22.34s

# Full suite.
python -m pytest tests -q
12544 passed, 9315 skipped, 106 warnings in 957.47s

# CI formatter used by this repository.
clang-format 21 --dry-run --Werror <changed CUDA headers>
exit 0

# GitHub Actions on 3cf536b.
Documentation build, Ruff, Clang Format, main CI, and CI v2 Shadow passed.

Benchmark / Performance Impact

N/A - this PR standardizes interfaces and correctness behavior; it makes no performance claim.

Notes for Reviewers

• Review canonical signatures and mutation semantics first, especially fused_add_rms_norm, reshape_and_cache, rotary_embedding, and SDPA.
• The Infinilm suffix is intentional for the three handwritten custom APIs. The generated ATen leading-underscore operators remain the only users of the internal namespace.
• Legacy operators and tests are retained for migration and explicitly deprecated rather than deleted.
• torch_ops.yaml contains only cat, linear, and matmul from this remediation because add already has a handwritten Torch backend.

BREAKING CHANGE: canonical signatures or mutation semantics change for add, cat, embedding, linear, matmul, reshape_and_cache, and rotary_embedding. Renamed legacy APIs remain temporarily available with deprecation notices.