Skip to content

refactor(pytorch): replace backend operator builders with typed build specs - #4890

Draft
grimoire wants to merge 12 commits into
InternLM:mainfrom
grimoire:refactor/pytorch-op-dispatch
Draft

refactor(pytorch): replace backend operator builders with typed build specs#4890
grimoire wants to merge 12 commits into
InternLM:mainfrom
grimoire:refactor/pytorch-op-dispatch

Conversation

@grimoire

@grimoire grimoire commented Aug 21, 2026

Copy link
Copy Markdown
Collaborator

Motivation

PyTorch operators were constructed through OpType and one-method builder
classes. This separated the operator identity, construction arguments, and
returned implementation type across multiple abstractions, making dispatch
harder to understand and extend.

Note

Since we have change the route of build_op, we better not merge this PR before other PRs that contain new ops.

Changes

  • Introduce typed BuildSpec[ImplT] objects as the operator construction
    contract.
  • Replace get_layer_impl_builder() with OpsBackend.build_op(spec).
  • Migrate all PyTorch operator families, including:
    • attention and linear operators
    • quantized linear operators
    • fused MoE and activations
    • normalization, embedding, and RoPE
    • sampling and cache-copy operators
    • NSA, gated-delta, and DeepSeek-V4-specific operators
  • Remove OpType and legacy one-method builder classes.
  • Preserve Default, CUDA, and DLINFER provider selection, fallback behavior,
    and lazy imports.
  • Preserve implementation ownership, weight loading, state-dict names, and
    public NN APIs.
  • Reserve an internal enable_deterministic construction seam for future
    deterministic kernels. It is currently a no-op and is not exposed through
    the public configuration.
  • Normalize related backend and build-spec naming.

This is a construction-time dispatch refactor. It does not intentionally
change operator forward behavior or kernel selection.

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.

1 participant