Veomni New Op

Before You Start

1. Read .agents/knowledge/constraints.md — especially the "Hardware" section (NPU guards, device-agnostic helpers) and "Module-level OpSlots are shared by every model instance" under "Trainer Extensions".
2. Read docs/design/kernel_selection.md and docs/design/unified_kernel_registry.md — understand the kernel lifecycle, the KERNEL_REGISTRY, and OpSlot dispatch.
3. Familiarize yourself with the ops architecture below.

VeOmni Ops Architecture

Most VeOmni ops in v5 are registry-driven: a kernel registers itself in veomni.ops.kernel_registry.KERNEL_REGISTRY and is dispatched at model-build time through OpSlot instances declared in the patchgen-generated modeling files (see veomni/ops/dispatch.py and _bind_veomni_ops() in veomni/models/auto.py).

veomni/ops/
├── __init__.py          # apply_ops_patch / apply_ops_config entry points
├── kernel_registry.py   # KERNEL_REGISTRY (the single source of truth)
├── dispatch.py          # OpSlot + binding helpers
├── config/              # legacy OpSpec/BackendSpec registry: apply_global_ops()
│                        # + apply_per_model_patches() for device_patch.py models
├── kernels/             # all registry-driven kernels
│   ├── attention/       # FA2/3/4 + sequence-parallel wrappers
│   ├── cross_entropy/   # eager + liger fused CE
│   ├── deepseek_sparse_attention/
│   ├── deepseek_v4/     # TileLang sparse attention / indexer
│   ├── load_balancing_loss/
│   ├── mhc/             # TileKernels DeepSeek V4 adapters
│   ├── moe/             # fused MoE (group_gemm / quack / npu_group_gemm)
│   ├── rms_norm/        # eager / liger / batch-invariant
│   ├── rotary/          # default / triton-deterministic
│   ├── swiglu/          # eager / liger
│   └── gated_delta_rule/
├── batch_invariant_ops/ # ATen-level interception for bitwise determinism
├── liger/               # Liger kernel adapters
└── platform/            # NPU-specific helpers

Three mechanisms coexist. Pick the first one unless you have a concrete reason not to:

1. KERNEL_REGISTRY + OpSlot (preferred for new ops). Each kernel registers itself under a (slot_name, variant) pair (e.g. ("cross_entropy_loss", "causal"), ("moe_experts", "standard")). Patchgen-generated modeling code declares matching OpSlot instances; at model-build time _bind_veomni_ops() walks the generated module, finds each OpSlot, and binds it to the concrete registry entry chosen by OpsImplementationConfig (config/registry.py).
2. Legacy global function pointer shim (kept for a few global ops that are dispatched outside generated modeling). Public-API functions like fused_moe_forward and load_balancing_loss still expose a thin pointer that is rebound by apply_ops_config() so call sites in non-patchgen code (DeepSeek MLA inference paths, NPU custom forwards) can keep importing the public name without going through an OpSlot.
3. Per-model device_patch.py via OpSpec/BackendSpec in ops/config/registry.py. apply_per_model_patches(hf_module, model_name, targets={op: attr}) setattr-replaces attributes on an HF module. Used by the models that have no patchgen-generated file (wan) or that need a runtime device-specific swap after generation (deepseek_v3, deepseek_v4). Those three device_patch.py files are its only callers. Do not extend this for new kernels.

Mechanism 1 covers any kernel living inside a patchgen-generated modeling file. Use 2 only when the kernel must be callable from unpatched (or non-Transformers) Python code, and 3 only when touching a model that already ships a device_patch.py.

Phase 1: Design

1. Determine op category:

   • Registry-driven kernel (the common case, used inside patchgen-generated modeling): register under a (slot_name, variant) in KERNEL_REGISTRY and add a matching OpSlot in the relevant <model>_patch_gen_config.py. No global mutation; selection is driven by OpsImplementationConfig.
   • Global op with public API (e.g. fused_moe_forward, load_balancing_loss): expose a public function in veomni/ops/__init__.py and rebind it from apply_ops_config() based on the active OpsImplementationConfig. Only use this when a non-patchgen call site (NPU MLA forward, manual inference scripts, etc.) needs to import the kernel directly.
   • Library op (no dispatch — called directly by model code): just create the module, no registry entry needed.
   • NPU variant: add alongside the GPU implementation behind an is_torch_npu_available() guard.

2. Decide selection mechanism: read docs/design/kernel_selection.md and docs/design/unified_kernel_registry.md to determine if you need:

   • Config field in OpsImplementationConfig (veomni/arguments/arguments_types.py)
   • Environment variable
   • Both

3. Determine binding timing:

   • Model build time (default): registry entries are resolved by _bind_veomni_ops() in veomni/models/auto.py when a model is constructed. New kernels just need to register themselves at import time.
   • apply_ops_config() time: legacy global ops (rebound function pointers) are wired in veomni/ops/__init__.py::apply_ops_config(ops_config).

Phase 2: Implement

1. Create the op directory under veomni/ops/kernels/<op_name>/.

2. Implement each kernel variant in its own file (e.g. triton_kernel.py, eager.py, npu_kernel.py). Each variant declares a concrete function with the kernel's canonical signature.

3. Register the kernel in veomni/ops/kernels/<op_name>/__init__.py. One KERNEL_REGISTRY.register(KernelSpec(...)) call per implementation — register() takes a single KernelSpec and returns None, so it is not a decorator:

   from veomni.ops.kernel_registry import KERNEL_REGISTRY, HardwareRequirement, KernelSpec
   
   
   def _my_op_triton_factory():
       from .triton_kernel import my_op_triton  # imported only when selected
   
       return my_op_triton
   
   
   KERNEL_REGISTRY.register(
       KernelSpec(
           name="triton",              # impl name the user selects in the config
           op_name="my_op",            # the logical op — matches the OpSlot
           variant="standard",         # op shape, when one op has several
           factory=_my_op_triton_factory,
           hardware=HardwareRequirement(device_type="gpu"),
           description="Triton my_op",
       )
   )

   factory is a zero-argument callable returning the kernel, not the kernel itself. Keeping it lazy is what stops an optional dependency (Liger, Triton, torch_npu) from being imported just because the module was loaded. hardware is enforced at resolve() time, so an unavailable kernel fails with a clear error instead of at first use.

   Mind the two axes: (op_name, variant) identifies the slot, name identifies the implementation within it. Kernels in different variants never collide.

   Then declare a matching OpSlot in the patchgen config of every model that uses it — the arguments are (op_name, variant), not an implementation:

   from veomni.ops.dispatch import OpSlot
   veomni_my_op = OpSlot("my_op", "standard")

   _bind_veomni_ops() calls slot.bind(impl_name) with the implementation selected by OpsImplementationConfig. See veomni/ops/kernels/rotary/__init__.py for a live example, and veomni/ops/README.md for the op/variant/impl table.

4. Wire the config field (if the user needs to choose an implementation):

   • Add a field to OpsImplementationConfig in veomni/arguments/arguments_types.py.
   • Call register_op(OpSpec(name=..., config_field=..., scope=..., default=..., backends={...})) from the same veomni/ops/kernels/<op_name>/__init__.py — the mapping lives next to the kernel, not inside veomni/ops/config/registry.py, which only defines OpSpec / BackendSpec / register_op. See veomni/ops/kernels/rms_norm/__init__.py, which registers both an OpSpec and its KernelSpecs.

5. For legacy global ops (only when needed): add the public function to veomni/ops/__init__.py and rebind it from apply_ops_config(ops_config).

6. Async Ulysses split wrappers (only for rms_norm and rotary_pos_emb): compound Functions cannot call OpSlot. They use no-autograd (output, saved) / backward pairs in veomni/distributed/sequence_parallel/op_wrappers.py. A new backend or variant must either add a matching wrapper there, or be left off _SUPPORTED_IMPLEMENTATIONS / _SUPPORTED_VARIANTS so get_op_wrapper rejects it. KERNEL_REGISTRY coverage is not enough.

7. NPU support:

   • Always guard NPU imports with is_torch_npu_available().
   • Put NPU implementations in a separate file (e.g., npu_kernel.py).
   • Register the NPU variant under the same slot with a distinct variant name.