Operator Migration

Operator Migration for cache-dit

Goal

Migrate one operator or kernel family into cache-dit in a way that is:

• semantically correct
• aligned with cache-dit repository conventions
• safe to import when optional native extensions are absent
• validated at multiple layers instead of by one smoke test

This skill is for migration work that touches native code, Python wrappers, operator registration, build packaging, or quantized module integration.

When to Use

Use this skill when you need to:

• migrate a CUDA or Triton operator from another repo into cache-dit
• port a nunchaku operator or kernel family into cache-dit
• decide what native files are actually required for a migration
• design cache-dit public wrappers for a newly migrated operator
• register low-level ops through cache-dit's CUDA registry layer
• add optional-extension build logic, submodule checks, or packaging guards
• design layered validation for a migrated operator
• review whether an operator migration plan is thoughtful or mechanical

Do not use this skill for:

• generic model integration with no operator or kernel work
• pure Python feature work unrelated to kernels or extensions
• blind "copy upstream into csrc" execution

Core Rule

Do not mechanically replay upstream structure.

Treat the source repository as the reference for semantics, not as the required layout.

Before writing code, answer these questions:

1. What behavior is essential to preserve?
2. What is the smallest native and Python closure needed to preserve that behavior?
3. Which names should remain source-compatible, and which should be renamed to match cache-dit conventions?
4. What must be public, and what should remain private implementation detail?
5. Which tests prove the migration works, instead of merely compiling?

If those questions are not answered yet, do not start copying files.

Reference Style Rule

Use portable references only.

• For cache-dit files, use repo-relative paths such as src/cache_dit/kernels/ops.py or tests/kernels/test_svdquant_runtime.py.
• For sibling or external repos, use repository-relative or GitHub-searchable paths such as nunchaku/nunchaku/models/linear.py or deepcompressor/deepcompressor/backend/nunchaku/utils.py.
• Do not write machine-local absolute paths such as /abs/path/to/workspace/... into the skill or its supporting documentation.

Phase 0: Gather Before Coding

Collect the migration inputs first.

Required inputs

1. Source operator and source repo Example: nunchaku/nunchaku/ops/gemm.py plus the native files it depends on.
2. Target cache-dit user-facing surface Example: a low-level op wrapper, a quantized module, or both.
3. Required backends, dtypes, and scope boundaries Example: "INT4 CUDA is required now; FP4 implementation may be retained but not gate current validation."
4. Build and packaging requirements Example: optional extension, submodule dependency, or environment gate.
5. Validation target Example: import safety, low-level parity, module parity, end-to-end inference, or shape rejection.

Gather checklist

• Identify the source operator entry points.
• Identify the native files that implement them.
• Identify helper files that are truly required by those implementations.
• Identify existing cache-dit abstractions that should host the migrated behavior.
• Identify the minimum feature slice that must work first.
• Identify what will explicitly not be validated in the current milestone.

Phase 1: Survey the Existing Design

Inspect both sides before making edits.

Survey the source implementation

Look for:

• the true call chain from public API to kernel launch
• required helper headers, interop layers, dispatch utilities, and packers
• runtime assumptions such as shape, rank, alignment, architecture, or dtype restrictions
• dependency assumptions such as vendored headers, submodules, or environment variables
• test coverage that already encodes behavior worth preserving

Survey cache-dit integration points

Common anchor files include:

• src/cache_dit/kernels/ops.py
• src/cache_dit/kernels/cuda/_ops_registery.py
• src/cache_dit/kernels/cuda/_<feature>.py
• setup.py
• pyproject.toml
• tests/kernels/...

Ask these questions while surveying:

1. Where should the public API live?
2. Where should torch.library registration live?
3. What should remain a private helper module under src/cache_dit/kernels/cuda/?
4. How should optional extension loading fail when the extension is missing?
5. Is there already a naming convention for this operator family?

Phase 2: Decide the Migration Shape

Make the design decisions before editing files.

1. Freeze the public surface first

Define the cache-dit-facing API early.

Examples of questions to settle:

• Which operator names should be exposed publicly?
• Should the public API be low-level only, module-level only, or both?
• Should internal backend toggles be hidden from users?
• Should wrapper functions be explicit rather than partial(...) so editors and type tools can see the real signature?

Default rule: keep backend-selection details private unless there is a strong user-facing reason to expose them.

2. Migrate the minimal viable closure

Do not import an entire subsystem if only one slice is needed.

Usually migrate:

• the kernel implementation files that are actually on the call path
• the minimum helper headers or Python utilities they require
• the registry and wrapper plumbing needed to call them from cache-dit

Usually do not migrate yet:

• unrelated kernels in the same source repo directory
• optimization branches that are not needed for the current milestone
• extra tooling, benchmark harnesses, or framework abstractions with no direct execution path impact

3. Preserve semantics before cleanup

During the first migration pass:

• preserve behavior first
• preserve shape and dtype rules first
• preserve dataflow first

Do not mix the migration with optional cleanups such as naming polish, API reshaping, or algorithmic changes unless they are necessary for repository consistency or import safety.

Phase 3: Implement the Migration

Apply changes from lowest level to highest level.

A. Native code and dependency boundary

When migrating native code:

1. Move only the required native closure into cache-dit's csrc tree.
2. Rename namespaces and top-level identifiers where needed to match cache-dit ownership.
3. Keep dispatch structure if it is functionally necessary; do not rewrite it just because it looks unfamiliar.
4. Decide dependency strategy explicitly:
   • vendored in-tree
   • git submodule
   • preinstalled system dependency
5. Add build-time validation for missing required dependencies.

B. Private CUDA helper layer

Use a private helper module under src/cache_dit/kernels/cuda/ for extension loading and low-level bridging.

Typical responsibilities:

• delayed import of the optional extension
• returning a cached load error
• wrapping direct calls into the extension's ops and utils submodules
• keeping internal details out of the public operator API

If the extension is optional, import cache_dit must remain safe.

C. Registry layer

Put low-level torch.library definitions and implementations in the CUDA registry layer, for example:

• src/cache_dit/kernels/cuda/_ops_registery.py

Typical responsibilities:

• define torch.library schemas
• implement real CUDA behavior
• add fake registrations where compile or tracing paths need them
• keep the public kernel API separate from raw registration details

Registration and fake-implementation conventions:

• name fake registrations explicitly as _fake_<operator_name>; do not use anonymous def _(...) helpers
• apply this naming rule consistently across CUDA, Triton, CuTe DSL, and other operator backends in cache-dit
• when adding or migrating operators, add unit tests in the same change
• tests should cover at least one fake shape or dtype path and one runtime correctness or smoke path

D. Public kernel API layer

Expose user-facing wrappers from src/cache_dit/kernels/ops.py.

Default conventions:

• expose explicit functions instead of partial(...) aliases when signature discoverability matters
• keep public names repository-aligned
• hide internal backend-selection knobs unless users truly need them
• validate backend support centrally instead of scattering checks