Skip to content

feat!: align non-InfiniLM operator interfaces#799

Merged
voltjia merged 28 commits into
masterfrom
feat/standardize-operator-interfaces
Jul 23, 2026
Merged

feat!: align non-InfiniLM operator interfaces#799
voltjia merged 28 commits into
masterfrom
feat/standardize-operator-interfaces

Conversation

@voltjia

@voltjia voltjia commented Jul 15, 2026

Copy link
Copy Markdown
Collaborator

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 documentation 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-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, together with their existing test coverage. Each old base class has both a Doxygen \deprecated notice 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 private infiniops and ops build targets.

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
85 passed, 15 skipped, 12735 deselected in 10.44s

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 Final merged smoke, main CI, and shadow CI passed 12544 passed, 9315 skipped full suite before final master sync
Iluvatar Yes Main and shadow CI passed on c9fbe36 Common CUDA wrapper and smoke coverage passed
MetaX Yes Main and shadow CI passed on c9fbe36 Common CUDA wrapper and smoke coverage passed
Cambricon Yes Main and shadow CI passed on c9fbe36 Backend build and smoke coverage passed
Moore Yes Main and shadow CI passed on c9fbe36 Common CUDA wrapper and smoke coverage passed
Ascend Yes Main and shadow CI passed on c9fbe36 Backend build and smoke coverage passed
Hygon Yes SDK/hardware unavailable Wrapper codegen only
Additional validation output
# Full build before the final master sync; the final merged head is covered by the smoke and targeted builds below.
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

# Consolidated add/cat/embedding/linear parameter coverage.
python -m pytest tests/test_add.py tests/test_cat.py tests/test_embedding.py tests/test_linear.py -q
570 passed, 144 skipped in 13.63s

# Wrapper extractor, including a class carrying [[deprecated]].
python -m pytest tests/test_generate_wrappers.py -q
13 passed in 0.62s

# The unity-build combination that includes swiglu, silu_and_mul, and embedding builds successfully.
INFINI_OPS_OPS=<smoke operators plus embedding> scripts/dev/build.sh nvidia --smoke --jobs 16
[build] done

# The helper isolation is behavior-neutral for both activation interfaces.
python -m pytest tests/test_swiglu.py tests/test_silu_and_mul.py -q
96 passed in 6.16s

# The conflict-adjacent operators both pass after the latest master sync.
python -m pytest tests/test_moe_sum.py tests/test_silu_and_mul.py -q
79 passed, 6 skipped in 21.89s

# Internal compatibility targets do not flood builds with deprecation warnings.
deprecated_warning_count=0

# An external consumer compiling against the installed headers does receive the warning.
warning: 'Swiglu' is deprecated: Use 'SiluAndMul' instead.

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

# GitHub Actions on c9fbe36.
Documentation build, Ruff, Clang Format, main CI, and CI v2 Shadow passed.
The first full-CI attempt failed only while fetching repositories/actions or communicating with CI agents; rerunning failed jobs on the same head passed every platform.

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; downstream C++ callers get compiler warnings.
  • Additional alpha, broadcast, option, and non-contiguous-layout cases for add, cat, embedding, and linear are consolidated into their existing parametrized test functions.
  • The branch is synchronized with master through feat(nvidia): add moe_sum operator #803 and feat(nvidia): add moe_align_block_size operator #804. The earlier feat(nvidia): add cutlass_scaled_mm operator #801/feat(torch): add flash_attn_varlen_func operator #802 sync exposed a private Sigmoid helper collision; SiluAndMulSigmoid now keeps that implementation local without changing either public API. The latest sync retained both the new moe_sum and existing silu_and_mul smoke matchers.
  • 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 compiler-visible deprecation notices.

voltjia added 27 commits July 15, 2026 10:07
@voltjia
voltjia requested a review from Ziminli July 23, 2026 02:39
@voltjia
voltjia marked this pull request as ready for review July 23, 2026 02:46
@voltjia
voltjia requested a review from a team July 23, 2026 02:46
Ziminli
Ziminli previously approved these changes Jul 23, 2026
@voltjia voltjia changed the title feat!: align non-InfiniLM operator interfaces feat: align non-InfiniLM operator interfaces Jul 23, 2026
@voltjia voltjia changed the title feat: align non-InfiniLM operator interfaces feat!: align non-InfiniLM operator interfaces Jul 23, 2026
@voltjia
voltjia merged commit 799aea8 into master Jul 23, 2026
20 checks passed
@voltjia
voltjia deleted the feat/standardize-operator-interfaces branch July 23, 2026 03:37
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.

2 participants