From c572251c0898248b235282c17dc214815ec67cff Mon Sep 17 00:00:00 2001 From: root Date: Thu, 30 Jul 2026 16:18:19 +0800 Subject: [PATCH 1/6] refactor(kernel): register-dispatch architecture for kernelize (registry/config/ops split) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Rework kernelize into a register-dispatch model, replacing the hard-coded builtin/liger module pair: - registry: ops REGISTER themselves with per-backend impls (lazy load + availability check); KernelChoice DISPATCHES to the first available backend in the priority chain - config: DEFAULT_KERNEL_CONFIG as pure data (dotted-path targets only) — the mapping declares WHAT to replace and WHICH op; never imports transformers/torch_npu/liger_kernel - core: kernelize() walks the mapping and dispatches each entry to an installer — HOW to install is decided by priority KernelChoice.installer -> OpDefinition.installer -> default_installer (class swap for nn.Module targets, setattr for functions/methods); transformers.* family-missing skips are DEBUG on the default path and WARNING (typo hint) for explicit mappings - ops: swiglu / geglu / rms_norm / rotary / moe / sdpa_attention / fla moved from npu_impls & liger_impls into per-op packages; new ep op (NPU grouped-matmul) replaces the inlined EP expert loop in expert_parallel.py - moe_experts default chain is ('npu',) — LigerExperts is opt-in only, absorbing the old _prefer_cann_on_npu drop semantics - hub(): HubRef values resolve lazily via the optional kernels package Cookbooks: kernelize is gated only by the torch-baseline escape hatch; --enable-liger is narrowed to gating the fused-linear-CE loss. EP utils: advanced indexing instead of index_select (aclnnIndexAdd broken on some torch_npu/CANN), sync D2H copies for split sizes. --- .../transformers/ep_fsdp2_lora_qwen3_5_moe.py | 16 +- cookbook/transformers/fsdp2.py | 20 +- cookbook/transformers/sp_fsdp_dense.py | 16 +- docs/source_en/Components/Kernel/Kernel.md | 267 ++++++++------ .../\345\206\205\346\240\270/Kernel.md" | 192 +++++------ src/twinkle/cli/cli.py | 4 +- src/twinkle/kernel/__init__.py | 15 +- src/twinkle/kernel/builtin.py | 216 ------------ src/twinkle/kernel/config.py | 117 +++++++ src/twinkle/kernel/core.py | 199 +++++++---- src/twinkle/kernel/liger.py | 325 ------------------ src/twinkle/kernel/liger_impls/__init__.py | 28 -- src/twinkle/kernel/liger_impls/swiglu.py | 32 -- src/twinkle/kernel/npu_impls/__init__.py | 28 -- src/twinkle/kernel/ops/__init__.py | 22 ++ src/twinkle/kernel/ops/ep/__init__.py | 112 ++++++ src/twinkle/kernel/ops/ep/loop.py | 66 ++++ src/twinkle/kernel/ops/ep/npu.py | 103 ++++++ src/twinkle/kernel/ops/fla/__init__.py | 36 ++ .../{npu_impls/fla.py => ops/fla/npu.py} | 0 src/twinkle/kernel/ops/geglu/__init__.py | 15 + src/twinkle/kernel/ops/geglu/liger.py | 20 ++ src/twinkle/kernel/ops/moe/__init__.py | 37 ++ src/twinkle/kernel/ops/moe/liger.py | 12 + .../{npu_impls/moe.py => ops/moe/npu.py} | 0 src/twinkle/kernel/ops/rms_norm/__init__.py | 53 +++ .../rms_norm.py => ops/rms_norm/liger.py} | 0 .../rms_norm.py => ops/rms_norm/npu.py} | 0 src/twinkle/kernel/ops/rotary/__init__.py | 36 ++ src/twinkle/kernel/ops/rotary/liger.py | 14 + .../rotary.py => ops/rotary/npu.py} | 0 .../kernel/ops/sdpa_attention/__init__.py | 46 +++ .../sdpa_attention/npu.py} | 0 src/twinkle/kernel/ops/swiglu/__init__.py | 19 + src/twinkle/kernel/ops/swiglu/liger.py | 19 + .../swiglu.py => ops/swiglu/npu.py} | 0 src/twinkle/kernel/registry.py | 178 ++++++++++ .../loss/liger_fused_linear_cross_entropy.py | 2 +- .../model/transformers/moe/ep_utils.py | 22 +- .../model/transformers/moe/expert_parallel.py | 49 +-- .../sequence_parallel/linear_attention_sp.py | 2 +- src/twinkle/patch/gdn_padding_free.py | 2 +- src/twinkle/patch/transformers_fused_ce.py | 3 +- 43 files changed, 1359 insertions(+), 984 deletions(-) delete mode 100644 src/twinkle/kernel/builtin.py create mode 100644 src/twinkle/kernel/config.py delete mode 100644 src/twinkle/kernel/liger.py delete mode 100644 src/twinkle/kernel/liger_impls/__init__.py delete mode 100644 src/twinkle/kernel/liger_impls/swiglu.py delete mode 100644 src/twinkle/kernel/npu_impls/__init__.py create mode 100644 src/twinkle/kernel/ops/__init__.py create mode 100644 src/twinkle/kernel/ops/ep/__init__.py create mode 100644 src/twinkle/kernel/ops/ep/loop.py create mode 100644 src/twinkle/kernel/ops/ep/npu.py create mode 100644 src/twinkle/kernel/ops/fla/__init__.py rename src/twinkle/kernel/{npu_impls/fla.py => ops/fla/npu.py} (100%) create mode 100644 src/twinkle/kernel/ops/geglu/__init__.py create mode 100644 src/twinkle/kernel/ops/geglu/liger.py create mode 100644 src/twinkle/kernel/ops/moe/__init__.py create mode 100644 src/twinkle/kernel/ops/moe/liger.py rename src/twinkle/kernel/{npu_impls/moe.py => ops/moe/npu.py} (100%) create mode 100644 src/twinkle/kernel/ops/rms_norm/__init__.py rename src/twinkle/kernel/{liger_impls/rms_norm.py => ops/rms_norm/liger.py} (100%) rename src/twinkle/kernel/{npu_impls/rms_norm.py => ops/rms_norm/npu.py} (100%) create mode 100644 src/twinkle/kernel/ops/rotary/__init__.py create mode 100644 src/twinkle/kernel/ops/rotary/liger.py rename src/twinkle/kernel/{npu_impls/rotary.py => ops/rotary/npu.py} (100%) create mode 100644 src/twinkle/kernel/ops/sdpa_attention/__init__.py rename src/twinkle/kernel/{npu_impls/attention.py => ops/sdpa_attention/npu.py} (100%) create mode 100644 src/twinkle/kernel/ops/swiglu/__init__.py create mode 100644 src/twinkle/kernel/ops/swiglu/liger.py rename src/twinkle/kernel/{npu_impls/swiglu.py => ops/swiglu/npu.py} (100%) create mode 100644 src/twinkle/kernel/registry.py diff --git a/cookbook/transformers/ep_fsdp2_lora_qwen3_5_moe.py b/cookbook/transformers/ep_fsdp2_lora_qwen3_5_moe.py index 71302b3fb..71a6658d7 100644 --- a/cookbook/transformers/ep_fsdp2_lora_qwen3_5_moe.py +++ b/cookbook/transformers/ep_fsdp2_lora_qwen3_5_moe.py @@ -18,7 +18,7 @@ from twinkle.model import TransformersModel from twinkle.preprocessor import SelfCognitionProcessor from twinkle.utils.framework import Torch -from twinkle.kernel import kernelize, liger_builtin, npu_builtin +from twinkle.kernel import kernelize logger = get_logger() args = CLI.from_args() @@ -93,16 +93,14 @@ def train(): }, ) # Kernel mode: torch (TWINKLE_TORCH_BASELINE=1, no fusion) | npu (default, - # CANN + FLA) | npu+liger (--enable-liger, Liger per-layer on top of CANN). + # CANN + FLA via DEFAULT_KERNEL_CONFIG). `--enable-liger` gates the fused-CE + # loss below; per-layer kernels stay on CANN on NPU (the default config's + # npu-first chains). To force Liger per-layer kernels, pass a custom + # mapping with liger-first KernelChoice chains — see Kernel.md. # Sharding/batch are unchanged across modes. _torch_baseline = os.environ.get('TWINKLE_TORCH_BASELINE', '').lower() in ('1', 'true', 'yes') - kernel_mapping = {} - if Torch.is_npu_available() and not _torch_baseline: - kernel_mapping.update(npu_builtin(model)) - if args.model.enable_liger and not _torch_baseline: - kernel_mapping.update(liger_builtin(model)) - if kernel_mapping: - model = kernelize(model, kernel_mapping) + if not _torch_baseline: + model = kernelize(model) _use_fused_ce = args.model.enable_liger and not _torch_baseline and args.model.enable_fused_ce _task = 'fused_lm_ce' if _use_fused_ce else 'causal_lm' lora_cfg = _build_lora_config() diff --git a/cookbook/transformers/fsdp2.py b/cookbook/transformers/fsdp2.py index 397df79a5..c501f6132 100644 --- a/cookbook/transformers/fsdp2.py +++ b/cookbook/transformers/fsdp2.py @@ -13,7 +13,7 @@ from twinkle.model import TransformersModel from twinkle.preprocessor import SelfCognitionProcessor from twinkle.utils.framework import Torch -from twinkle.kernel import kernelize, liger_builtin, npu_builtin +from twinkle.kernel import kernelize logger = get_logger() args = CLI.from_args() @@ -66,19 +66,15 @@ def train(): discovered = {type(m).__name__ for m in model.model.modules() if type(m).__name__.endswith('DecoderLayer')} model.model._no_split_modules = list(discovered) or [model.model.config.model_type.title() + 'DecoderLayer'] - # Compose the kernel mapping: NPU built-ins first, then Liger on top so - # `--enable-liger` opts into Liger's cross-device Triton/Ascend kernels - # (later keys win on overlap — see twinkle.kernel Kernel.md). - kernel_mapping = {} # Kernel mode: torch (TWINKLE_TORCH_BASELINE=1, no fusion) | npu (default, - # CANN + FLA) | npu+liger (--enable-liger, Liger per-layer + fused-CE). + # CANN + FLA via DEFAULT_KERNEL_CONFIG). `--enable-liger` gates the fused-CE + # loss below; per-layer kernels stay on CANN on NPU (the default config's + # npu-first chains), matching the old npu+liger mode after CANN preference. + # To force Liger per-layer kernels, pass a custom mapping with liger-first + # KernelChoice chains — see Kernel.md. _torch_baseline = os.environ.get('TWINKLE_TORCH_BASELINE', '').lower() in ('1', 'true', 'yes') - if Torch.is_npu_available() and not _torch_baseline: - kernel_mapping.update(npu_builtin(model)) - if args.model.enable_liger and not _torch_baseline: - kernel_mapping.update(liger_builtin(model)) - if kernel_mapping: - model = kernelize(model, kernel_mapping) + if not _torch_baseline: + model = kernelize(model) # `--enable-liger` turns on BOTH the per-layer Liger/CANN kernels (above) # AND, by default (`enable_fused_ce=True`), the LigerFusedLinearCrossEntropyLoss diff --git a/cookbook/transformers/sp_fsdp_dense.py b/cookbook/transformers/sp_fsdp_dense.py index dc7f6bc63..d63c2cd5c 100644 --- a/cookbook/transformers/sp_fsdp_dense.py +++ b/cookbook/transformers/sp_fsdp_dense.py @@ -10,7 +10,7 @@ from twinkle.model import TransformersModel from twinkle.preprocessor import SelfCognitionProcessor from twinkle.utils.framework import Torch -from twinkle.kernel import kernelize, liger_builtin, npu_builtin +from twinkle.kernel import kernelize logger = get_logger() args = CLI.from_args() @@ -70,16 +70,14 @@ def train(): strategy=args.model.strategy, ) # Kernel mode: torch (TWINKLE_TORCH_BASELINE=1, no fusion) | npu (default, - # CANN + FLA) | npu+liger (--enable-liger, Liger per-layer on top of CANN). + # CANN + FLA via DEFAULT_KERNEL_CONFIG). `--enable-liger` gates the fused-CE + # loss below; per-layer kernels stay on CANN on NPU (the default config's + # npu-first chains). To force Liger per-layer kernels, pass a custom + # mapping with liger-first KernelChoice chains — see Kernel.md. # Sharding/batch are unchanged across modes. _torch_baseline = os.environ.get('TWINKLE_TORCH_BASELINE', '').lower() in ('1', 'true', 'yes') - kernel_mapping = {} - if Torch.is_npu_available() and not _torch_baseline: - kernel_mapping.update(npu_builtin(model)) - if args.model.enable_liger and not _torch_baseline: - kernel_mapping.update(liger_builtin(model)) - if kernel_mapping: - model = kernelize(model, kernel_mapping) + if not _torch_baseline: + model = kernelize(model) _use_fused_ce = args.model.enable_liger and not _torch_baseline and args.model.enable_fused_ce _task = 'fused_lm_ce' if _use_fused_ce else 'causal_lm' lora_config = LoraConfig(**args.get_lora_args()) diff --git a/docs/source_en/Components/Kernel/Kernel.md b/docs/source_en/Components/Kernel/Kernel.md index 7080f45a4..18a1d76c1 100644 --- a/docs/source_en/Components/Kernel/Kernel.md +++ b/docs/source_en/Components/Kernel/Kernel.md @@ -1,95 +1,105 @@ # Twinkle Kernel -`twinkle.kernel` exposes a mapping-driven kernel replacement API. Replacing one -implementation with another collapses to a single `kernelize(model, mapping)` -call. +`twinkle.kernel` exposes a mapping-driven kernel replacement API: swapping one +implementation for another in a model collapses to a single +`kernelize(model, mapping)` call. Three concerns are fully separated: **what to +replace** (mapping key), **which implementation to pick** (`KernelChoice`), and +**how to install it** (installer). The public surface is exactly four symbols: | Symbol | Purpose | | --- | --- | -| `kernelize(model, mapping=None)` | Apply ``mapping`` to ``model`` (in place) and return it. If ``mapping`` is omitted, it is auto-detected from the current platform (see below) | -| `npu_builtin(model=None)` | Return the Ascend NPU built-in mapping (composes with user mappings) | -| `liger_builtin(model=None)` | Return the Liger Kernel built-in mapping — cross-device (CUDA Triton + Ascend NPU). Bare impls, no device gating | -| `hub(ref, *, revision=None, version=None, backend=None, trust_remote_code=False)` | Build a ``HubRef`` for use as a mapping value; the actual Hub download is deferred to ``kernelize`` | +| `kernelize(model, mapping=None)` | Apply `mapping` to `model` in place and return it. If `mapping` is omitted, the built-in `DEFAULT_KERNEL_CONFIG` is applied | +| `KernelChoice(op, backends, installer=None)` | Mapping value: which op to use for this replacement, and in what priority order to pick a backend | +| `DEFAULT_KERNEL_CONFIG` | Built-in default mapping (target → `KernelChoice`); copy/merge/override it to customize | +| `hub(ref, *, revision=None, version=None, backend=None)` | Build a `HubRef` for use as a mapping value; the actual Hub download is deferred to `kernelize` | ## Mapping semantics -`mapping` keys describe the target to replace: +**Keys (targets) come in two forms**, freely combinable with any value form: -- `type[nn.Module]` subclass — replace **every** instance whose exact type matches (`m.__class__ = impl`; subclasses are **not** touched) -- `str` of the form `'pkg.sub.attr'` or `'pkg.sub.ClassName.attr'` — `setattr(target, attr, impl)` +- `type[nn.Module]` subclass: replace **every** instance whose exact type + matches (`m.__class__ = impl_class`; subclasses are **not** touched) +- Dotted `str` path: `'pkg.mod.ClassName'` (resolves to an `nn.Module` + subclass → same as a class key), `'pkg.mod.ClassName.forward'` (class method + → `setattr`), `'pkg.mod.attr'` (module function → `setattr`). When a + `transformers.*` family is not installed, the entry is silently skipped at + DEBUG level (a missing family is the normal case) -`mapping` values describe the replacement: +**Values come in three forms**: -- `type[nn.Module]` subclass — used as the impl class. The class' `__init__` is **never** invoked; its forward must work against the attributes the original instance already has -- `Callable` — assigned with `setattr` -- `dict[str, V]` — device → impl dispatch. Device is inferred from the model; entries without a matching key are **silently skipped** -- `HubRef` — built via `hub(...)`; resolved lazily +- A direct impl (class or function): no backend selection is performed; the + generic installer installs it as-is. An impl class is **never `__init__`ed** — + its forward must work against the attributes the original instance already has +- `KernelChoice(op='rms_norm', backends=('npu', 'liger'))`: pick the first + available implementation in `backends` order — an unregistered backend, a + failed `available()` check (with reason), or a load exception all fall through + to the next backend; if none is available the original implementation is kept +- `HubRef`: a Hub reference built via `hub(...)`; loaded lazily -Device is inferred from `next(model.parameters()).device.type` (falling back to buffers, then `'cpu'`). - -## Auto-detection (mapping omitted) - -When `mapping` is `None`, `kernelize` auto-detects the current platform via `Platform.device_prefix()` and applies the matching built-in bundle. Platforms without a built-in bundle are a safe no-op (the model is returned unchanged). - -## Examples - -### Enable the built-in bundle for the current platform +## Default config and customization ```python -from twinkle.kernel import kernelize - -model = kernelize(model) # auto-detects the platform and applies its built-in bundle +model = kernelize(model) +# equivalent to +model = kernelize(model, DEFAULT_KERNEL_CONFIG) ``` -The explicit form is still supported: +`DEFAULT_KERNEL_CONFIG` covers: RMSNorm / RoPE / SwiGLU / MoE for Qwen2 / Qwen3 +/ Qwen3-MoE / Qwen2.5-VL / Qwen3.5 / Qwen3.5-MoE (chain `('npu', 'liger')`, +CANN preferred on NPU); the 8 llama families and 4 gemma families (chain +`('liger',)`); plus the logical targets `'sdpa'` (global SDPA installation) and +`'fla'` (per-instance Qwen3.5 Flash Linear Attention patching). -```python -import torch -from twinkle.kernel import kernelize, npu_builtin - -if torch.npu.is_available(): - model = kernelize(model, npu_builtin(model)) -``` - -### Custom class replacement +**A user mapping fully replaces the default config — it is not merged.** Passing +only your own entries means no default entry is applied. To tweak on top of the +defaults, copy and merge: ```python -from transformers.models.qwen2.modeling_qwen2 import Qwen2RMSNorm -from twinkle.kernel import kernelize +from twinkle.kernel import DEFAULT_KERNEL_CONFIG, KernelChoice, kernelize -model = kernelize(model, {Qwen2RMSNorm: MyRMSNorm}) +model = kernelize(model, {**DEFAULT_KERNEL_CONFIG, + # make rms_norm liger-first (falls back to npu if liger is unavailable) + 'transformers.models.qwen3.modeling_qwen3.Qwen3RMSNorm': + KernelChoice(op='rms_norm', backends=('liger', 'npu')), +}) ``` -### Built-in + custom override +`DEFAULT_KERNEL_CONFIG` itself must not be mutated in place at runtime. -```python -from twinkle.kernel import kernelize, npu_builtin +**Log levels**: with `mapping=None` (the default-config path), all +fallback/failure logs are DEBUG (a CUDA machine or a missing family is the +normal case); with an explicitly passed mapping (even a copy-merge of the +defaults) they are raised to WARNING — you stated an intent, so a no-op must be +reported. Every successful installation emits one INFO line: -model = kernelize(model, {**npu_builtin(model), Qwen2RMSNorm: MyRMSNorm}) +```text +[kernelize] target=...Qwen3MLP.forward op=swiglu backend=npu installer=default +[kernelize] target=sdpa op=sdpa_attention backend=npu installer=install_sdpa ``` -Plain dict merge — later keys override earlier ones. +**CUDA behavior**: the default config is platform-uniform. On CUDA, the npu +backend's `available()` in a `('npu', 'liger')` chain returns False, so +selection falls through to liger; if liger is not installed the whole chain +fails (DEBUG skip under the default config ≈ no-op). -### Hub kernel (HF Hub format) +## Scenarios -```python -from twinkle.kernel import kernelize, hub -from my_pkg import SiluAndMul +### Pick an implementation by op name (with fallback order) -model = kernelize(model, { - SiluAndMul: hub('kernels-community/activation:SiluAndMul', version=1), +```python +model = kernelize(model, {**DEFAULT_KERNEL_CONFIG, + 'transformers.models.qwen3.modeling_qwen3.Qwen3MLP.forward': + KernelChoice(op='swiglu', backends=('liger', 'npu')), # liger first, npu as fallback }) ``` -Exactly one of `revision` / `version` must be passed. The `kernels` package is imported lazily; absence raises a clear "install kernels" error. - -### Function-level replacement +### String key + direct impl (bypass selection) ```python from twinkle.kernel import kernelize -from twinkle.kernel.npu_impls.rotary import npu_apply_rotary_pos_emb +from twinkle.kernel.ops.rotary.npu import npu_apply_rotary_pos_emb model = kernelize(model, { 'transformers.models.qwen2.modeling_qwen2.apply_rotary_pos_emb': @@ -97,89 +107,134 @@ model = kernelize(model, { }) ``` -### Cross-device mapping (NPU enabled, CUDA skipped) +### Custom class replacement ```python -from twinkle.kernel import kernelize +from transformers.models.qwen2.modeling_qwen2 import Qwen2RMSNorm -model = kernelize(model, { - Qwen2RMSNorm: {'npu': NpuRMSNorm, 'cuda': CudaRMSNorm}, -}) +model = kernelize(model, {Qwen2RMSNorm: MyRMSNorm}) ``` -Safe to run on CUDA — entries whose dict misses the current device just skip. +### Hub kernel (HF Hub format) -## NPU built-in coverage +```python +from twinkle.kernel import hub, kernelize -`npu_builtin(model)` returns a dict that (as available transformers modules permit) covers: +model = kernelize(model, { + SiluAndMul: hub('kernels-community/activation:SiluAndMul', version=1), +}) +``` + +Exactly one of `revision` or `version` is required. `hub(...)` triggers a lazy +import of the `kernels` package; if it is missing you will be told to +`pip install kernels`. -- RMSNorm class replacement for Qwen2 / Qwen3 / Qwen3-MoE / Qwen2.5-VL / Qwen3.5 / Qwen3.5-MoE families -- `apply_rotary_pos_emb` function replacement (fused RoPE) for the same families -- SwiGLU fused replacement for the MLP variants -- `Experts.forward` and `SparseMoeBlock.forward` for Qwen3-MoE / Qwen3.5-MoE -- GatedRMSNorm forward for Qwen3.5 / Qwen3.5-MoE -- `apply_multimodal_rotary_pos_emb` for Qwen2.5-VL -- Global SDPA replacement (one-shot side effect on `ALL_ATTENTION_FUNCTIONS['sdpa']`) -- Qwen3.5 Flash Linear Attention enablement (one-shot side effect + per-instance traversal, triggered inside `npu_builtin(model)`). Delegates to fla's native operators (`fla.modules.convolution.causal_conv1d` and `fla.ops.gated_delta_rule.chunk_gated_delta_rule`); on NPU, fla's `triton_ascend` backend dispatch handles the Ascend-specific kernels. Requires `flash-linear-attention` >= 0.5.2 +### GMM (grouped matmul) opt-in -**Not included by default:** the NPU replacement for `transformers.integrations.moe._grouped_mm`. Without Expert Parallelism the contiguous-copy overhead is ~8x. Opt in explicitly when EP is enabled: +The default config deliberately does **not** include the NPU replacement for +`transformers.integrations.moe._grouped_mm` (≈8x overhead without Expert +Parallelism). Add it explicitly when needed (note: a direct impl has no platform +gating — only do this on a confirmed NPU environment): ```python -from twinkle.kernel import kernelize, npu_builtin -from twinkle.kernel.npu_impls.moe import npu_grouped_mm - -mapping = { - **npu_builtin(model), - 'transformers.integrations.moe._grouped_mm': {'npu': npu_grouped_mm}, -} -model = kernelize(model, mapping) -``` +from twinkle.kernel.ops.moe.npu import npu_grouped_mm -## Liger Kernel built-in +model = kernelize(model, {**DEFAULT_KERNEL_CONFIG, + 'transformers.integrations.moe._grouped_mm': npu_grouped_mm, +}) +``` -`liger_builtin(model)` returns a dict of **bare** Liger impls (no device gating) covering RoPE, RMSNorm, SwiGLU/GeGLU, and MoE experts across the Qwen / Llama / Mistral / Mixtral / Phi3 / Gemma / Olmo2 / GLM4 / Granite / InternVL families. Values are bare because Liger self-dispatches across CUDA (Triton) and Ascend NPU (the auto-applied `backends/_ascend` backend) via `liger_kernel.utils.infer_device` — wrapping in `{'cuda': impl}` would wrongly skip Liger on NPU, where it is fully supported. +## The op registry (how to add an op / backend) -The fused-linear-CE `forward` replacement and the global `nn.functional.cross_entropy` swap are **excluded** — they belong to the loss layer (`twinkle.loss`), not the kernel layer. +Built-in implementations are organized per operator under +`twinkle/kernel/ops//`: `__init__.py` performs registration (lazy references +and lightweight availability checks only — **it must not import optional +dependencies directly**), and `.py` holds the implementation itself. +Registering means declaring "backend X can provide an implementation for op Y": ```python -from twinkle.kernel import kernelize, liger_builtin - -model = kernelize(model, liger_builtin(model)) +# twinkle/kernel/ops/swiglu/__init__.py +from ...registry import KernelImpl, is_liger_available, is_npu_available, lazy_import, register_op + +register_op( + 'swiglu', + implementations={ + 'npu': KernelImpl( + load=lazy_import('twinkle.kernel.ops.swiglu.npu:npu_swiglu_forward'), + available=is_npu_available, + ), + 'liger': KernelImpl( + load=lazy_import('twinkle.kernel.ops.swiglu.liger:liger_swiglu_forward'), + available=is_liger_available, + ), + }, +) ``` -### Composing Liger with the NPU bundle +- `KernelImpl.load(target)`: lazy loading factory, called only when the + implementation is selected; receives the mapping target and may specialize per + model family (e.g. liger's RMSNorm dispatches between gemma / qwen3_5 variants) +- `KernelImpl.available()`: returns `(True, None)` when usable, or + `(False, reason)` to fall through; all platform and dependency checks live here +- Registering the same op name twice, or with empty implementations → `ValueError` +- **Vendor rule**: when an impl incorporates an external kernel, the file + docstring must state the source repository@commit, the list of local changes, + and a re-sync reminder + +### Custom installers -On NPU both `npu_builtin` and `liger_builtin` are NPU implementations of the same operators. Plain dict merge lets you pick precedence (later keys win): +Installations beyond plain class/attr replacement are provided by the op's own +installer (signature `(model, target, impl)`). Priority: +`KernelChoice.installer` → `OpDefinition.installer` → generic installer. +Example — SDPA writes into transformers' global attention registry: ```python -from twinkle.kernel import kernelize, liger_builtin, npu_builtin +def install_sdpa(model, target, impl) -> None: + from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, AttentionInterface + AttentionInterface._global_mapping['sdpa'] = impl + ALL_ATTENTION_FUNCTIONS['sdpa'] = impl -# Liger wins on overlap -model = kernelize(model, {**npu_builtin(model), **liger_builtin(model)}) -# Twinkle-NPU wins on overlap -model = kernelize(model, {**liger_builtin(model), **npu_builtin(model)}) +register_op('sdpa_attention', implementations={...}, installer=install_sdpa) ``` -### NPU precedence: CANN wins over Liger-Triton for per-layer ops - -On Ascend NPU, `liger_builtin()` post-processes itself via `_prefer_cann_on_npu`: Liger's Triton-on-Ascend kernels for the bandwidth-bound per-layer ops (RMSNorm, SwiGLU, RoPE) are swapped for the faster CANN vendor ops from `npu_impls`, and the `LigerExperts` / `LigerQwen3MoeSwiGLUMLP` class replacements are dropped so `npu_builtin`'s forward-level CANN grouped-matmul MoE expert path takes effect. Non-Qwen families without a CANN equivalent keep their Liger impls. So on NPU `liger_builtin` composes *additively* with `npu_builtin` (it only contributes ops CANN lacks), and `npu + liger` is not slower than `npu` alone for these ops. The fused-linear-CE loss (`twinkle.loss`) is unaffected — Liger's fused kernel is still used there because it has no CANN equivalent. On CUDA the bundle is unchanged. +The default config references such ops via logical targets +(`'sdpa': KernelChoice(op='sdpa_attention', backends=('npu',))`); a logical +target is only an identifier passed through to the installer and is not resolved +by the generic replacer. If an installer raises, the exception propagates — the +model must never be left in a half-installed state the user doesn't know about. -### RMSNorm attribute migration +## Migration guide (old API → new) -Liger's `LigerRMSNorm.forward` reads instance attributes (`offset` / `casting_mode` / `in_place` / `row_mode`) that HuggingFace RMSNorm variants do not define. Liger's monkey-patch sets these eagerly via `_patch_rms_norm_module`; the `liger_impls.rms_norm` adapters do the same **lazily** inside `forward` (with per-family defaults: llama-style `offset=0.0, casting_mode="llama"`, gemma-style `offset=1.0, casting_mode="gemma"`, gemma4 `offset=0.0`). No global state is mutated — contrast with `npu_builtin`'s SDPA install. +| Old usage | New equivalent | +| --- | --- | +| `kernelize(model, npu_builtin(model))` | `kernelize(model)` (default config is CANN-first on NPU) | +| `kernelize(model, liger_builtin(model))` | Build an all-liger mapping, or `{**DEFAULT_KERNEL_CONFIG, ...}` with the targets you want rewritten to `('liger', ...)` chains | +| Manually merging `{**npu_builtin(m), **liger_builtin(m)}` | `{**DEFAULT_KERNEL_CONFIG, ...}` with per-target overrides | +| `{Qwen3RMSNorm: {'npu': NpuRMSNorm}}` (device-conditional dict) | `{Qwen3RMSNorm: NpuRMSNorm}` or `KernelChoice(op='rms_norm', backends=('npu',))` | +| `from twinkle.kernel.npu_impls.x import ...` | `from twinkle.kernel.ops..npu import ...` | +| `from twinkle.kernel.liger_impls.x import ...` | `from twinkle.kernel.ops..liger import ...` | + +`npu_builtin()` / `liger_builtin()` / device-conditional dicts (`{'npu': impl}`) +have been removed with no compatibility shim; platform checks moved into +`KernelImpl.available()`. On NPU, `kernelize(model)`'s default replacement +result matches the old version; on CUDA it changes from a no-op to applying the +default config (effective wherever liger is installed, ≈ no-op otherwise). ## Environment variables Only two remain: -- `TWINKLE_NPU_FLA` — Qwen3.5 FLA switch (default on; `0`/`false` to disable) -- `TWINKLE_NPU_GATED_RMSNorm_FP32` — force FP32 in Gated RMSNorm forward (default off) - -The legacy `TWINKLE_NPU_PATCH` / `TWINKLE_NPU_FUSED_OPS` / `TWINKLE_NPU_GMM_PATCH` / `TWINKLE_USE_KERNELS` are gone — they're now "include the entry in the mapping or don't" decisions. +- `TWINKLE_NPU_FLA`: Qwen3.5 FLA switch (on by default; set to `0`/`false` to disable) +- `TWINKLE_NPU_GATED_RMSNorm_FP32`: force Gated RMSNorm computation up to FP32 (off by default) ## Caveats -- `m.__class__ = impl_cls` is Python class-replacement magic. The impl class **must** override only `forward` (and helpers); defining `__init__` is incompatible with the contract -- Exact match: `type(m) is target_cls`. Subclasses of `target_cls` are not replaced — add them to the mapping yourself -- `kernelize` is idempotent under repeated calls +- `m.__class__ = impl_cls` is Python class-swap magic. An impl class **must** + only override `forward` (plus helper methods) and must not define `__init__`, + otherwise the original instance's attributes will mismatch the impl's + expectations +- Matching is exact: `type(m) is target_cls`. Subclasses of `target_cls` are not + replaced; if you need them replaced, list them in the mapping too +- Calling `kernelize` multiple times is idempotent (setting `__class__` to the + impl again is harmless) - There is no `unkernelize` — replacement is one-way diff --git "a/docs/source_zh/\347\273\204\344\273\266/\345\206\205\346\240\270/Kernel.md" "b/docs/source_zh/\347\273\204\344\273\266/\345\206\205\346\240\270/Kernel.md" index c32943248..7933db5cc 100644 --- "a/docs/source_zh/\347\273\204\344\273\266/\345\206\205\346\240\270/Kernel.md" +++ "b/docs/source_zh/\347\273\204\344\273\266/\345\206\205\346\240\270/Kernel.md" @@ -1,93 +1,78 @@ # Twinkle Kernel 模块 -`twinkle.kernel` 提供一个 mapping 驱动的内核替换接口，把“用一种实现替换模型里的另一种实现”压缩为一次 `kernelize(model, mapping)` 调用。 +`twinkle.kernel` 提供一个 mapping 驱动的内核替换接口，把"用一种实现替换模型里的另一种实现"压缩为一次 `kernelize(model, mapping)` 调用。三件事被彻底拆开：**替换什么**（mapping 的 key）、**选什么实现**（`KernelChoice`）、**怎么安装**（installer）。 公开符号只有四个： | 符号 | 作用 | | --- | --- | -| `kernelize(model, mapping=None)` | 在 `model` 上应用 `mapping`，原地修改后返回。省略 `mapping` 时按当前平台自动检测（见下文） | -| `npu_builtin(model=None)` | 返回 Ascend NPU 内置替换的 mapping dict（可与用户 mapping 自由组合） | -| `liger_builtin(model=None)` | 返回 Liger Kernel 内置替换的 mapping dict —— 跨设备（CUDA Triton + Ascend NPU）。值为裸 impl，不做设备门控 | -| `hub(ref, *, revision=None, version=None, backend=None, trust_remote_code=False)` | 构造一个 `HubRef`，用作 mapping value；真实下载推迟到 `kernelize` 执行 | +| `kernelize(model, mapping=None)` | 在 `model` 上应用 `mapping`，原地修改后返回。省略 `mapping` 时应用内置的 `DEFAULT_KERNEL_CONFIG` | +| `KernelChoice(op, backends, installer=None)` | mapping value：本次替换用哪个 op、按什么优先级选 backend | +| `DEFAULT_KERNEL_CONFIG` | 内置默认 mapping（target → `KernelChoice`），复制/合并/覆盖即可自定义 | +| `hub(ref, *, revision=None, version=None, backend=None)` | 构造一个 `HubRef`，用作 mapping value；真实下载推迟到 `kernelize` 执行 | ## Mapping 语义 -`mapping` 的 **key** 表示要替换的目标： +**key（target）两种形态**，与 value 自由组合： - `type[nn.Module]` 子类：替换模型里**所有**该精确类型的实例（`m.__class__ = impl_class`，**不包含**子类） -- `str` 形如 `'pkg.sub.attr'` 或 `'pkg.sub.ClassName.attr'`：`setattr(target, attr, impl)` +- `str` 点路径：`'pkg.mod.ClassName'`（解析为 `nn.Module` 子类 → 等价类替换）、`'pkg.mod.ClassName.forward'`（类方法 → `setattr`）、`'pkg.mod.attr'`（模块函数 → `setattr`）。`transformers.*` 家族未安装时 DEBUG 静默跳过（家族缺失 = 常态） -**value** 表示用什么替换： +**value 三种形式**： -- `type[nn.Module]` 子类：直接作为 impl 类。该类**不会被 `__init__` 调用**，必须只依赖原 instance 已经有的 attribute（weight / eps / ...）正确工作 -- `Callable`：直接 `setattr` 上去 -- `dict[str, V]`：device → impl 嵌套分派。从 `model` 推断当前 device，未匹配则**静默跳过** +- 直接 impl（类或函数）：不执行 backend 选择，直接用通用 installer 安装。impl 类**不会被 `__init__` 调用**，必须只依赖原 instance 已有的 attribute 正确工作 +- `KernelChoice(op='rms_norm', backends=('npu', 'liger'))`：按 `backends` 顺序挑第一个可用实现——backend 未注册 / `available()` 不通过（含原因）/ 加载异常都自动顺延；全部不可用则保留原始实现 - `HubRef`：通过 `hub(...)` 构造的 Hub 引用，延迟加载 -device 从 `next(model.parameters()).device.type` 推断（无参数则用 buffers，再无则为 `'cpu'`）。 - -## 自动检测（省略 mapping） - -当 `mapping` 为 `None` 时，`kernelize` 通过 `Platform.device_prefix()` 自动检测当前平台，并应用对应平台的内置 bundle。没有内置 bundle 的平台为安全空操作，原样返回 model。 - -## 场景示例 - -### 启用当前平台的内置优化 +## 默认配置与自定义 ```python -from twinkle.kernel import kernelize - -model = kernelize(model) # 自动检测当前平台并应用其内置 bundle +model = kernelize(model) +# 等价于 +model = kernelize(model, DEFAULT_KERNEL_CONFIG) ``` -显式写法依然支持： - -```python -import torch -from twinkle.kernel import kernelize, npu_builtin - -if torch.npu.is_available(): - model = kernelize(model, npu_builtin(model)) -``` +`DEFAULT_KERNEL_CONFIG` 覆盖：Qwen2 / Qwen3 / Qwen3-MoE / Qwen2.5-VL / Qwen3.5 / Qwen3.5-MoE 的 RMSNorm / RoPE / SwiGLU / MoE（链 `('npu', 'liger')`，NPU 上 CANN 优先）；llama 系 8 族、gemma 系 4 族（链 `('liger',)`）；以及逻辑 target `'sdpa'`（全局 SDPA 安装）与 `'fla'`（Qwen3.5 Flash Linear Attention 逐实例 patch）。 -### 自定义类替换 +**mapping 是全量替换默认配置而非增量**——只传自己那几条就意味着默认条目全部不应用。在默认基础上微调用复制合并： ```python -from transformers.models.qwen2.modeling_qwen2 import Qwen2RMSNorm -from twinkle.kernel import kernelize +from twinkle.kernel import DEFAULT_KERNEL_CONFIG, KernelChoice, kernelize -model = kernelize(model, {Qwen2RMSNorm: MyRMSNorm}) +model = kernelize(model, {**DEFAULT_KERNEL_CONFIG, + # rms_norm 改成 liger 优先（liger 不可用自动退 npu） + 'transformers.models.qwen3.modeling_qwen3.Qwen3RMSNorm': + KernelChoice(op='rms_norm', backends=('liger', 'npu')), +}) ``` -### 内置 + 自定义混合 +`DEFAULT_KERNEL_CONFIG` 本身不应在运行时原地修改。 -```python -from twinkle.kernel import kernelize, npu_builtin +**日志分级**：`mapping=None`（默认配置路径）时，回退/失败日志全部是 DEBUG（CUDA 机器、家族缺失是常态）；显式传入 mapping（哪怕是默认配置的复制合并）则升为 WARNING——明确表达了意图，没生效必须告知。每次成功安装有一条 INFO： -model = kernelize(model, {**npu_builtin(model), Qwen2RMSNorm: MyRMSNorm}) +```text +[kernelize] target=...Qwen3MLP.forward op=swiglu backend=npu installer=default +[kernelize] target=sdpa op=sdpa_attention backend=npu installer=install_sdpa ``` -后写入的 key 会覆盖前面的，普通 dict 合并语义。 +**CUDA 行为**：默认配置跨平台统一。CUDA 上 `('npu', 'liger')` 链中 npu 的 `available()` 为 False 自动落到 liger；liger 未装则整链失败（默认配置下 DEBUG 跳过 ≈ no-op）。 -### Hub Kernel（HF Hub 格式） +## 场景示例 -```python -from twinkle.kernel import kernelize, hub -from my_pkg import SiluAndMul +### 按算子点名选实现（带替补顺序） -model = kernelize(model, { - SiluAndMul: hub('kernels-community/activation:SiluAndMul', version=1), +```python +model = kernelize(model, {**DEFAULT_KERNEL_CONFIG, + 'transformers.models.qwen3.modeling_qwen3.Qwen3MLP.forward': + KernelChoice(op='swiglu', backends=('liger', 'npu')), # 首选 liger，不行退 npu }) ``` -`revision` 与 `version` 二选一必传。`hub(...)` 触发 `kernels` 包的延迟 import，未安装时会提示 `pip install kernels`。 - -### 函数级替换 +### 字符串 key + 直接 impl（绕过选择） ```python from twinkle.kernel import kernelize -from twinkle.kernel.npu_impls.rotary import npu_apply_rotary_pos_emb +from twinkle.kernel.ops.rotary.npu import npu_apply_rotary_pos_emb model = kernelize(model, { 'transformers.models.qwen2.modeling_qwen2.apply_rotary_pos_emb': @@ -95,76 +80,93 @@ model = kernelize(model, { }) ``` -### 跨设备 mapping（NPU 启用、CUDA 跳过） +### 自定义类替换 ```python -from twinkle.kernel import kernelize +from transformers.models.qwen2.modeling_qwen2 import Qwen2RMSNorm -model = kernelize(model, { - Qwen2RMSNorm: {'npu': NpuRMSNorm, 'cuda': CudaRMSNorm}, -}) +model = kernelize(model, {Qwen2RMSNorm: MyRMSNorm}) ``` -在 CUDA 模型上跑也安全：未匹配 device 的 entry 不会替换、不会报错。 +### Hub Kernel（HF Hub 格式） -## 内置 NPU 优化 +```python +from twinkle.kernel import hub, kernelize -`npu_builtin(model)` 返回的 dict 至少包含以下覆盖（实际条目随 transformers 已安装的 modeling 模块动态收集）： +model = kernelize(model, { + SiluAndMul: hub('kernels-community/activation:SiluAndMul', version=1), +}) +``` + +`revision` 与 `version` 二选一必传。`hub(...)` 触发 `kernels` 包的延迟 import，未安装时会提示 `pip install kernels`。 -- Qwen2 / Qwen3 / Qwen3-MoE / Qwen2.5-VL / Qwen3.5 / Qwen3.5-MoE 系列的 RMSNorm 类替换 -- 同上系列的 `apply_rotary_pos_emb` 函数替换（融合 RoPE） -- 同上系列 MLP 的 SwiGLU 融合替换 -- Qwen3-MoE / Qwen3.5-MoE 的 `Experts.forward` 与 `SparseMoeBlock.forward` 替换 -- Qwen3.5 / Qwen3.5-MoE 的 GatedRMSNorm forward 替换 -- Qwen2.5-VL 的 `apply_multimodal_rotary_pos_emb` 替换 -- 全局 SDPA 替换（一次性副作用，写入 `ALL_ATTENTION_FUNCTIONS['sdpa']`） -- Qwen3.5 Flash Linear Attention 启用（一次性副作用 + 实例遍历，由 `npu_builtin(model)` 内部触发）。直接委托给 fla 原生算子（`fla.modules.convolution.causal_conv1d` 与 `fla.ops.gated_delta_rule.chunk_gated_delta_rule`）；在 NPU 上由 fla 的 `triton_ascend` 后端 dispatch 处理 Ascend 专用 kernel。需要 `flash-linear-attention` >= 0.5.2 +### GMM（grouped matmul）opt-in -**未默认包含** `transformers.integrations.moe._grouped_mm` 的 NPU 替换（在没有 Expert Parallelism 时会带来约 8x 开销）。需要时手动加入： +默认配置**不包含** `transformers.integrations.moe._grouped_mm` 的 NPU 替换（没有 Expert Parallelism 时约 8x 开销）。需要时显式加入（注意：直接 impl 无平台门控，仅在确认 NPU 环境时使用）： ```python -from twinkle.kernel import kernelize, npu_builtin -from twinkle.kernel.npu_impls.moe import npu_grouped_mm - -mapping = { - **npu_builtin(model), - 'transformers.integrations.moe._grouped_mm': {'npu': npu_grouped_mm}, -} -model = kernelize(model, mapping) -``` +from twinkle.kernel.ops.moe.npu import npu_grouped_mm -## Liger Kernel 内置 +model = kernelize(model, {**DEFAULT_KERNEL_CONFIG, + 'transformers.integrations.moe._grouped_mm': npu_grouped_mm, +}) +``` -`liger_builtin(model)` 返回**裸** Liger impl（不做设备门控）的 mapping，覆盖 Qwen / Llama / Mistral / Mixtral / Phi3 / Gemma / Olmo2 / GLM4 / Granite / InternVL 各族的 RoPE、RMSNorm、SwiGLU/GeGLU 及 MoE experts。值为裸 impl，是因为 Liger 自身就跨设备分派：CUDA 走 Triton、Ascend NPU 走自动应用的 `backends/_ascend` 后端（经 `liger_kernel.utils.infer_device`）。若包成 `{'cuda': impl}` 会在 NPU 上错误跳过——而 Liger 在 NPU 上是完全支持的。 +## op 注册机制（如何新增 op / backend） -融合线性 CE 的 `forward` 替换与全局 `nn.functional.cross_entropy` 替换**不在此 bundle 内**——它们属于 loss 层（`twinkle.loss`），不属于 kernel 层。 +内置实现按算子组织在 `twinkle/kernel/ops//` 下：`__init__.py` 负责注册（只含惰性引用与轻量可用性检查，**不得直接 import 可选依赖**），`.py` 是实现本体。注册 = 登记"某 backend 能给某 op 提供实现"： ```python -from twinkle.kernel import kernelize, liger_builtin - -model = kernelize(model, liger_builtin(model)) +# twinkle/kernel/ops/swiglu/__init__.py +from ...registry import KernelImpl, is_liger_available, is_npu_available, lazy_import, register_op + +register_op( + 'swiglu', + implementations={ + 'npu': KernelImpl( + load=lazy_import('twinkle.kernel.ops.swiglu.npu:npu_swiglu_forward'), + available=is_npu_available, + ), + 'liger': KernelImpl( + load=lazy_import('twinkle.kernel.ops.swiglu.liger:liger_swiglu_forward'), + available=is_liger_available, + ), + }, +) ``` -### Liger 与 NPU bundle 组合 +- `KernelImpl.load(target)`：惰性加载工厂，仅在实现被选中时调用；接收 mapping target，可按目标家族特化（如 liger 的 RMSNorm 按 gemma / qwen3_5 变体分派） +- `KernelImpl.available()`：返回 `(True, None)` 可用 / `(False, reason)` 不可用并顺延；平台与依赖判断都收在这里 +- 同名 op 重复注册、空 implementations → `ValueError` +- **vendor 规范**：impl 引入外部 kernel 时，文件 docstring 必须注明来源仓库@commit、本地改动清单，以及 re-sync 提醒 + +### 自定义 installer -在 NPU 上 `npu_builtin` 与 `liger_builtin` 都是同一批算子的 NPU 实现。普通 dict 合并即可选择优先级（后写的 key 胜出）： +普通 class / attr 替换之外的安装方式，由 op 自带 installer（签名 `(model, target, impl)`）。优先级：`KernelChoice.installer` → `OpDefinition.installer` → 通用 installer。例子——SDPA 写入 transformers 全局 attention 注册表： ```python -from twinkle.kernel import kernelize, liger_builtin, npu_builtin +def install_sdpa(model, target, impl) -> None: + from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, AttentionInterface + AttentionInterface._global_mapping['sdpa'] = impl + ALL_ATTENTION_FUNCTIONS['sdpa'] = impl -# 重叠算子由 Liger 胜出 -model = kernelize(model, {**npu_builtin(model), **liger_builtin(model)}) -# 重叠算子由 Twinkle-NPU 胜出 -model = kernelize(model, {**liger_builtin(model), **npu_builtin(model)}) +register_op('sdpa_attention', implementations={...}, installer=install_sdpa) ``` -### NPU 优先级：逐层算子上 CANN 胜过 Liger-Triton +默认配置用逻辑 target 引用这类 op（`'sdpa': KernelChoice(op='sdpa_attention', backends=('npu',))`）；逻辑 target 仅作标识传给 installer，不由通用替换器解析。installer 执行失败会上抛异常——不让模型处于用户不知情的半安装状态。 -在 Ascend NPU 上，`liger_builtin()` 会通过 `_prefer_cann_on_npu` 做后处理：对带宽敏感的逐层算子（RMSNorm、SwiGLU、RoPE），Liger 的 Triton-on-Ascend 内核会被替换为 `npu_impls` 里更快的 CANN 厂商算子；`LigerExperts` / `LigerQwen3MoeSwiGLUMLP` 的类替换也会被丢弃，从而让 `npu_builtin` 的 forward 级 CANN 分组矩阵乘 MoE expert 路径生效。没有 CANN 对应实现的非 Qwen 族仍保留 Liger impl。因此在 NPU 上 `liger_builtin` 与 `npu_builtin` 是**叠加**关系（只贡献 CANN 缺少的算子），`npu + liger` 在这些算子上不会比单独 `npu` 更慢。融合线性 CE loss（`twinkle.loss`）不受影响——该处没有 CANN 对应实现，仍用 Liger 的融合内核。CUDA 上 bundle 不变。 +## 迁移指南（旧接口 → 新写法） -### RMSNorm 属性迁移 +| 旧用法 | 新写法 | +| --- | --- | +| `kernelize(model, npu_builtin(model))` | `kernelize(model)`（默认配置 NPU 上 CANN 优先） | +| `kernelize(model, liger_builtin(model))` | 自建全 liger mapping，或 `{**DEFAULT_KERNEL_CONFIG, ...}` 把想改的 target 链写成 `('liger', ...)` | +| `{**npu_builtin(m), **liger_builtin(m)}` 手动合并 | `{**DEFAULT_KERNEL_CONFIG, ...}` 按 target 覆盖 | +| `{Qwen3RMSNorm: {'npu': NpuRMSNorm}}`（设备条件 dict） | `{Qwen3RMSNorm: NpuRMSNorm}` 或 `KernelChoice(op='rms_norm', backends=('npu',))` | +| `from twinkle.kernel.npu_impls.x import ...` | `from twinkle.kernel.ops..npu import ...` | +| `from twinkle.kernel.liger_impls.x import ...` | `from twinkle.kernel.ops..liger import ...` | -Liger 的 `LigerRMSNorm.forward` 读取的实例属性（`offset` / `casting_mode` / `in_place` / `row_mode`）在 HuggingFace 的 RMSNorm 变体上并不存在。Liger 的 monkey-patch 通过 `_patch_rms_norm_module` 急切地设置这些属性；`liger_impls.rms_norm` 适配器改为在 `forward` 内**懒设置**（按族默认值：llama 风格 `offset=0.0, casting_mode="llama"`，gemma 风格 `offset=1.0, casting_mode="gemma"`，gemma4 `offset=0.0`）。不污染任何全局状态——与 `npu_builtin` 的 SDPA 全局 install 形成对比。 +`npu_builtin()` / `liger_builtin()` / 设备条件 dict（`{'npu': impl}`）已删除，无兼容层；平台判断收进 `KernelImpl.available()`。NPU 上 `kernelize(model)` 的默认替换结果与旧版一致；CUDA 上从 no-op 变为应用默认配置（liger 可用即生效，未装 ≈ no-op）。 ## 环境变量 @@ -173,8 +175,6 @@ Liger 的 `LigerRMSNorm.forward` 读取的实例属性（`offset` / `casting_mod - `TWINKLE_NPU_FLA`：Qwen3.5 FLA 开关（默认开，设为 `0`/`false` 关闭） - `TWINKLE_NPU_GATED_RMSNorm_FP32`：将 Gated RMSNorm 强制升到 FP32 计算（默认关） -旧的 `TWINKLE_NPU_PATCH` / `TWINKLE_NPU_FUSED_OPS` / `TWINKLE_NPU_GMM_PATCH` / `TWINKLE_USE_KERNELS` 已移除——这些都改成"是否把对应 entry 写进 mapping"的显式选择。 - ## 注意事项 - `m.__class__ = impl_cls` 是 Python class 替换魔法。impl 类**必须**只覆盖 `forward`（以及辅助方法），不能定义 `__init__`，否则原 instance 的 attribute 会与 impl 的预期错位 diff --git a/src/twinkle/cli/cli.py b/src/twinkle/cli/cli.py index a12594d00..c09264658 100644 --- a/src/twinkle/cli/cli.py +++ b/src/twinkle/cli/cli.py @@ -34,8 +34,8 @@ class ModelArgs: ddp_config: dict[str, Any] | None = None fsdp_config: dict[str, Any] | None = None grad_scaler_config: dict[str, Any] | None = None - # Liger Kernel toggle: when True, the cookbook applies `liger_builtin()` via - # `kernelize`. Off by default — opt in with --enable-liger / TWINKLE_ENABLE_LIGER. + # Liger Kernel toggle: gates the fused-linear-CE loss in the cookbooks. + # Off by default — opt in with --enable-liger / TWINKLE_ENABLE_LIGER. enable_liger: bool = False # Fused-linear-CE loss toggle. Only meaningful when `enable_liger` is True. # Defaults True so `--enable-liger` turns on BOTH the per-layer Liger/CANN diff --git a/src/twinkle/kernel/__init__.py b/src/twinkle/kernel/__init__.py index eb738edcd..a2c63806f 100644 --- a/src/twinkle/kernel/__init__.py +++ b/src/twinkle/kernel/__init__.py @@ -3,13 +3,14 @@ Public symbols: -- :func:`kernelize` apply ``mapping`` to a model -- :func:`hub` build a Hub kernel reference -- :func:`npu_builtin` the Ascend NPU built-in bundle -- :func:`liger_builtin` the Liger Kernel built-in bundle (cross-device) +- :func:`kernelize` apply ``mapping`` to a model +- :func:`hub` build a Hub kernel reference +- :class:`KernelChoice` per-target op + backend-priority selection +- :data:`DEFAULT_KERNEL_CONFIG` the built-in default mapping (copy/merge to customize) """ -from .builtin import npu_builtin +from . import ops # noqa: F401 triggers built-in op registration (must happen before the first kernelize() call) +from .config import DEFAULT_KERNEL_CONFIG from .core import hub, kernelize -from .liger import liger_builtin +from .registry import KernelChoice -__all__ = ['kernelize', 'hub', 'npu_builtin', 'liger_builtin'] +__all__ = ['kernelize', 'hub', 'KernelChoice', 'DEFAULT_KERNEL_CONFIG'] diff --git a/src/twinkle/kernel/builtin.py b/src/twinkle/kernel/builtin.py deleted file mode 100644 index ee031df12..000000000 --- a/src/twinkle/kernel/builtin.py +++ /dev/null @@ -1,216 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""``npu_builtin()`` returns the bundle of Ascend NPU replacements. - -All values are wrapped in ``{'npu': impl}`` so the bundle composes safely on -CUDA/CPU systems — non-NPU devices silently skip every entry. - -GMM is **not** included by default (without EP it causes ~8x slowdown). Opt -in by merging: - - {**npu_builtin(model), 'transformers.integrations.moe._grouped_mm': - {'npu': npu_grouped_mm}} -""" -from __future__ import annotations - -import importlib -import torch.nn as nn -from typing import Any - -from twinkle import get_logger -from twinkle.utils.device_mesh import Platform - -logger = get_logger() - - -def _import_optional(name: str): - try: - return importlib.import_module(name) - except ImportError: - return None - - -def npu_builtin(model: nn.Module | None = None) -> dict[Any, dict[str, Any]]: - """Return the NPU builtin mapping; optionally apply per-instance FLA.""" - from .npu_impls.attention import npu_sdpa_attention_forward - from .npu_impls.fla import apply_qwen3_5_fla - from .npu_impls.moe import npu_packed_moe_experts_forward, npu_qwen3_5_moe_sparse_block_forward - from .npu_impls.rms_norm import NpuRMSNorm, npu_gated_rms_norm_forward - from .npu_impls.rotary import npu_apply_multimodal_rotary_pos_emb, npu_apply_rotary_pos_emb - from .npu_impls.swiglu import npu_swiglu_forward - - bundle: dict[Any, dict[str, Any]] = {} - - is_npu_platform = Platform.device_prefix() == 'npu' - - # Apply SDPA install eagerly (one-shot module-level mutation) on NPU - # platforms. The NPU impl inverts boolean masks, which is wrong for - # CUDA/CPU execution, so non-NPU platforms must not mutate the global HF - # registry even if ``torch_npu`` is importable in the environment. - if is_npu_platform: - _install_sdpa(npu_sdpa_attention_forward) - - # === per-family class + function entries === - _add_qwen2_entries(bundle, NpuRMSNorm, npu_apply_rotary_pos_emb, npu_swiglu_forward) - _add_qwen3_entries(bundle, NpuRMSNorm, npu_apply_rotary_pos_emb, npu_swiglu_forward) - _add_qwen3_moe_entries( - bundle, - NpuRMSNorm, - npu_apply_rotary_pos_emb, - npu_swiglu_forward, - npu_packed_moe_experts_forward, - npu_qwen3_5_moe_sparse_block_forward, - ) - _add_qwen2_5_vl_entries( - bundle, - NpuRMSNorm, - npu_apply_rotary_pos_emb, - npu_swiglu_forward, - npu_apply_multimodal_rotary_pos_emb, - ) - _add_qwen3_5_entries( - bundle, - NpuRMSNorm, - npu_gated_rms_norm_forward, - npu_apply_rotary_pos_emb, - npu_swiglu_forward, - ) - _add_qwen3_5_moe_entries( - bundle, - NpuRMSNorm, - npu_gated_rms_norm_forward, - npu_apply_rotary_pos_emb, - npu_swiglu_forward, - npu_packed_moe_experts_forward, - npu_qwen3_5_moe_sparse_block_forward, - ) - - # === FLA (side-effect; mapping-incompatible) === - if is_npu_platform: - from twinkle.utils.import_utils import exists - if exists('flash-linear-attention'): - apply_qwen3_5_fla(model) - else: - logger.warning('[NPU] [FLA] flash-linear-attention is not installed; ' - 'FLA patch for Qwen3.5 requires the fla package (which ships ' - 'its own Ascend-backend dispatch). ' - 'Install it with: pip install flash-linear-attention. ' - 'Other NPU patches will still apply.') - return bundle - - -def _install_sdpa(impl) -> None: - """One-shot install of SDPA attention forward (global modeling_utils dict). - - ``AttentionInterface._global_mapping`` is a private transformers attribute; - guard against its removal so an upstream change can't take down the rest - of ``npu_builtin()``. - """ - try: - from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, AttentionInterface - except ImportError: - return - try: - AttentionInterface._global_mapping['sdpa'] = impl - except AttributeError: - logger.warning('[NPU] [SDPA] AttentionInterface._global_mapping unavailable; skipping') - ALL_ATTENTION_FUNCTIONS['sdpa'] = impl - - -# ---- helpers that conditionally add entries based on module availability ---- - - -def _add_class_if_present(bundle, module_path, class_name, impl_cls): - mod = _import_optional(module_path) - if mod is None: - return - cls = getattr(mod, class_name, None) - if isinstance(cls, type): - bundle[cls] = {'npu': impl_cls} - - -def _add_swiglu_if_present(bundle, module_path, class_name, fn): - mod = _import_optional(module_path) - if mod is None: - return - cls = getattr(mod, class_name, None) - if isinstance(cls, type): - # Function-level: wrap as string-keyed forward replacement. - # We override on the *class object*, not the module attribute, by - # using a class-key with a synthetic impl wrapping the forward. - # The simplest way is to subclass and reassign __class__, but here - # we follow the legacy approach of overwriting the class's forward: - bundle[f'{module_path}.{class_name}.forward'] = {'npu': fn} - - -def _add_attr_if_present(bundle, module_path, attr_name, impl): - mod = _import_optional(module_path) - if mod is None: - return - if '.' in attr_name: - # Dotted attr like 'Qwen3MoeExperts.forward': resolve the class on - # the module, then check the trailing member on the class. - head, _, tail = attr_name.partition('.') - owner = getattr(mod, head, None) - if owner is None or not hasattr(owner, tail): - return - else: - if not hasattr(mod, attr_name): - return - bundle[f'{module_path}.{attr_name}'] = {'npu': impl} - - -def _add_qwen2_entries(bundle, rms_cls, rope_fn, swiglu_fn): - # Qwen2 (used by Qwen2.5-VL etc. via inheritance) - _add_class_if_present(bundle, 'transformers.models.qwen2.modeling_qwen2', 'Qwen2RMSNorm', rms_cls) - _add_attr_if_present(bundle, 'transformers.models.qwen2.modeling_qwen2', 'apply_rotary_pos_emb', rope_fn) - _add_swiglu_if_present(bundle, 'transformers.models.qwen2.modeling_qwen2', 'Qwen2MLP', swiglu_fn) - - -def _add_qwen3_entries(bundle, rms_cls, rope_fn, swiglu_fn): - base = 'transformers.models.qwen3.modeling_qwen3' - _add_class_if_present(bundle, base, 'Qwen3RMSNorm', rms_cls) - _add_attr_if_present(bundle, base, 'apply_rotary_pos_emb', rope_fn) - _add_swiglu_if_present(bundle, base, 'Qwen3MLP', swiglu_fn) - - -def _add_qwen3_moe_entries(bundle, rms_cls, rope_fn, swiglu_fn, experts_fn, sparse_fn): - base = 'transformers.models.qwen3_moe.modeling_qwen3_moe' - _add_class_if_present(bundle, base, 'Qwen3MoeRMSNorm', rms_cls) - _add_attr_if_present(bundle, base, 'apply_rotary_pos_emb', rope_fn) - _add_swiglu_if_present(bundle, base, 'Qwen3MoeMLP', swiglu_fn) - _add_attr_if_present(bundle, base, 'Qwen3MoeExperts.forward', experts_fn) - _add_attr_if_present(bundle, base, 'Qwen3MoeSparseMoeBlock.forward', sparse_fn) - - -def _add_qwen2_5_vl_entries(bundle, rms_cls, rope_fn, swiglu_fn, multimodal_rope_fn): - base = 'transformers.models.qwen2_5_vl.modeling_qwen2_5_vl' - _add_class_if_present(bundle, base, 'Qwen2_5_VLRMSNorm', rms_cls) - _add_attr_if_present(bundle, base, 'apply_rotary_pos_emb', rope_fn) - _add_attr_if_present(bundle, base, 'apply_multimodal_rotary_pos_emb', multimodal_rope_fn) - _add_swiglu_if_present(bundle, base, 'Qwen2MLP', swiglu_fn) - _add_swiglu_if_present(bundle, base, 'Qwen2_5_VLMLP', swiglu_fn) - - -def _add_qwen3_5_entries(bundle, rms_cls, gated_rms_fn, rope_fn, swiglu_fn): - base = 'transformers.models.qwen3_5.modeling_qwen3_5' - if _import_optional(base) is None: - return - _add_class_if_present(bundle, base, 'Qwen3_5RMSNorm', rms_cls) - _add_class_if_present(bundle, base, 'Qwen3_5VisionRMSNorm', rms_cls) - _add_attr_if_present(bundle, base, 'apply_rotary_pos_emb', rope_fn) - _add_swiglu_if_present(bundle, base, 'Qwen3_5MLP', swiglu_fn) - _add_swiglu_if_present(bundle, base, 'Qwen3_5VisionMLP', swiglu_fn) - # Qwen3_5GatedRMSNorm: forward-level replacement - _add_attr_if_present(bundle, base, 'Qwen3_5GatedRMSNorm.forward', gated_rms_fn) - - -def _add_qwen3_5_moe_entries(bundle, rms_cls, gated_rms_fn, rope_fn, swiglu_fn, experts_fn, sparse_fn): - base = 'transformers.models.qwen3_5_moe.modeling_qwen3_5_moe' - if _import_optional(base) is None: - return - _add_class_if_present(bundle, base, 'Qwen3_5MoeRMSNorm', rms_cls) - _add_attr_if_present(bundle, base, 'apply_rotary_pos_emb', rope_fn) - _add_swiglu_if_present(bundle, base, 'Qwen3_5MoeMLP', swiglu_fn) - _add_attr_if_present(bundle, base, 'Qwen3_5MoeExperts.forward', experts_fn) - _add_attr_if_present(bundle, base, 'Qwen3_5MoeSparseMoeBlock.forward', sparse_fn) - _add_attr_if_present(bundle, base, 'Qwen3_5MoeGatedRMSNorm.forward', gated_rms_fn) diff --git a/src/twinkle/kernel/config.py b/src/twinkle/kernel/config.py new file mode 100644 index 000000000..8e79a70b2 --- /dev/null +++ b/src/twinkle/kernel/config.py @@ -0,0 +1,117 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Built-in default mapping: target -> replacement (KernelChoice). + +Pure data file: never imports transformers / torch_npu / liger_kernel, no side effects; +targets are uniformly dotted-path strings (class targets are resolved by default_installer into __class__ replacement, +targets whose family is not installed are skipped at DEBUG level during install). + +``kernelize(model)`` ≡ ``kernelize(model, DEFAULT_KERNEL_CONFIG)``。 +User customization = copy + override: ``{**DEFAULT_KERNEL_CONFIG, : }`` -- +the mapping fully replaces the default config, it is not incremental. This dict must not be mutated in place at runtime. +""" +from __future__ import annotations + +from typing import Any + +from .registry import KernelChoice + +# ── Qwen families (npu first, liger fallback) ────────────────────────── +# (module, rms_cls, mlp_classes, extra_rms_classes) +_QWEN_DENSE = [ + ('transformers.models.qwen2.modeling_qwen2', 'Qwen2RMSNorm', ['Qwen2MLP'], []), + ('transformers.models.qwen3.modeling_qwen3', 'Qwen3RMSNorm', ['Qwen3MLP'], []), + ('transformers.models.qwen2_5_vl.modeling_qwen2_5_vl', 'Qwen2_5_VLRMSNorm', ['Qwen2MLP', 'Qwen2_5_VLMLP'], []), + ('transformers.models.qwen3_5.modeling_qwen3_5', 'Qwen3_5RMSNorm', ['Qwen3_5MLP', 'Qwen3_5VisionMLP'], + ['Qwen3_5VisionRMSNorm']), +] +# (module, rms_cls, mlp_cls, experts_cls, block_cls) +_QWEN_MOE = [ + ('transformers.models.qwen3_moe.modeling_qwen3_moe', 'Qwen3MoeRMSNorm', 'Qwen3MoeMLP', 'Qwen3MoeExperts', + 'Qwen3MoeSparseMoeBlock'), + ('transformers.models.qwen3_5_moe.modeling_qwen3_5_moe', 'Qwen3_5MoeRMSNorm', 'Qwen3_5MoeMLP', + 'Qwen3_5MoeExperts', 'Qwen3_5MoeSparseMoeBlock'), +] +# qwen3_5 Gated RMSNorm (forward-level replacement, no liger impl) +_QWEN_GATED_RMS = { + 'transformers.models.qwen3_5.modeling_qwen3_5': 'Qwen3_5GatedRMSNorm', + 'transformers.models.qwen3_5_moe.modeling_qwen3_5_moe': 'Qwen3_5MoeGatedRMSNorm', +} +# Qwen3.5 uses partial_rotary_factor=0.25 + mrope_interleaved; liger's full-rotation +# rotate_half impl is incompatible -> rotary chains for these two families are ('npu',) only +_QWEN35_ROPE_NPU_ONLY = ( + 'transformers.models.qwen3_5.modeling_qwen3_5', + 'transformers.models.qwen3_5_moe.modeling_qwen3_5_moe', +) + +# ── llama families x8 (liger only) ────────────────────────────────────── +_LLAMA_STYLE = [ + ('transformers.models.llama.modeling_llama', 'LlamaRMSNorm', 'LlamaMLP'), + ('transformers.models.mistral.modeling_mistral', 'MistralRMSNorm', 'MistralMLP'), + ('transformers.models.mixtral.modeling_mixtral', 'MixtralRMSNorm', 'MixtralBlockSparseTop2MLP'), + ('transformers.models.phi3.modeling_phi3', 'Phi3RMSNorm', 'Phi3MLP'), + ('transformers.models.glm4.modeling_glm4', 'Glm4RMSNorm', 'Glm4MLP'), + ('transformers.models.olmo2.modeling_olmo2', 'Olmo2RMSNorm', 'Olmo2MLP'), + ('transformers.models.granite.modeling_granite', 'GraniteRMSNorm', 'GraniteMLP'), + ('transformers.models.internvl.modeling_internvl', 'InternVLRMSNorm', 'InternVLMLP'), +] + +# ── gemma families x4 (liger only: rms_norm + geglu) ──────────────────── +_GEMMA_STYLE = [ + ('transformers.models.gemma.modeling_gemma', 'GemmaRMSNorm', 'GemmaMLP'), + ('transformers.models.gemma2.modeling_gemma2', 'Gemma2RMSNorm', 'Gemma2MLP'), + ('transformers.models.gemma3.modeling_gemma3', 'Gemma3RMSNorm', 'Gemma3MLP'), + ('transformers.models.gemma4.modeling_gemma4', 'Gemma4RMSNorm', 'Gemma4TextMLP'), +] + + +def _rope_backends(module: str) -> tuple[str, ...]: + return ('npu', ) if module in _QWEN35_ROPE_NPU_ONLY else ('npu', 'liger') + + +def _build() -> dict[Any, Any]: + cfg: dict[Any, Any] = {} + + # Qwen dense：rms_norm / rotary / swiglu + for mod, rms, mlps, extra_rms in _QWEN_DENSE: + cfg[f'{mod}.{rms}'] = KernelChoice(op='rms_norm', backends=('npu', 'liger')) + for extra in extra_rms: + cfg[f'{mod}.{extra}'] = KernelChoice(op='rms_norm', backends=('npu', 'liger')) + cfg[f'{mod}.apply_rotary_pos_emb'] = KernelChoice(op='rotary', backends=_rope_backends(mod)) + for mlp in mlps: + cfg[f'{mod}.{mlp}.forward'] = KernelChoice(op='swiglu', backends=('npu', 'liger')) + + # Qwen MoE: dense trio + moe_experts / moe_block + # moe_experts default chain ('npu',): liger's LigerExperts class replacement only + # takes effect when explicitly chosen by the user; it never blocks the npu CANN + # grouped-matmul fast path (absorbs the drop semantics of the old _prefer_cann_on_npu) + for mod, rms, mlp, experts, block in _QWEN_MOE: + cfg[f'{mod}.{rms}'] = KernelChoice(op='rms_norm', backends=('npu', 'liger')) + cfg[f'{mod}.apply_rotary_pos_emb'] = KernelChoice(op='rotary', backends=_rope_backends(mod)) + cfg[f'{mod}.{mlp}.forward'] = KernelChoice(op='swiglu', backends=('npu', 'liger')) + cfg[f'{mod}.{experts}.forward'] = KernelChoice(op='moe_experts', backends=('npu', )) + cfg[f'{mod}.{block}.forward'] = KernelChoice(op='moe_block', backends=('npu', )) + + # Qwen gated rms_norm (npu only) + for mod, gated in _QWEN_GATED_RMS.items(): + cfg[f'{mod}.{gated}.forward'] = KernelChoice(op='gated_rms_norm', backends=('npu', )) + + # Qwen2.5-VL multimodal rope (npu only) + cfg['transformers.models.qwen2_5_vl.modeling_qwen2_5_vl.apply_multimodal_rotary_pos_emb'] = KernelChoice( + op='multimodal_rotary', backends=('npu', )) + + # llama families x3 ops, gemma families x2 ops (liger only) + for mod, rms, mlp in _LLAMA_STYLE: + cfg[f'{mod}.{rms}'] = KernelChoice(op='rms_norm', backends=('liger', )) + cfg[f'{mod}.apply_rotary_pos_emb'] = KernelChoice(op='rotary', backends=('liger', )) + cfg[f'{mod}.{mlp}.forward'] = KernelChoice(op='swiglu', backends=('liger', )) + for mod, rms, mlp in _GEMMA_STYLE: + cfg[f'{mod}.{rms}'] = KernelChoice(op='rms_norm', backends=('liger', )) + cfg[f'{mod}.{mlp}.forward'] = KernelChoice(op='geglu', backends=('liger', )) + + # logical target: handled by a custom installer (never resolved by the generic replacer) + cfg['sdpa'] = KernelChoice(op='sdpa_attention', backends=('npu', )) + cfg['fla'] = KernelChoice(op='fla', backends=('npu', )) + return cfg + + +DEFAULT_KERNEL_CONFIG = _build() diff --git a/src/twinkle/kernel/core.py b/src/twinkle/kernel/core.py index 38df38293..803e35efc 100644 --- a/src/twinkle/kernel/core.py +++ b/src/twinkle/kernel/core.py @@ -1,17 +1,26 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -"""Minimal mapping-driven kernel replacement. +"""Mapping-driven kernel replacement. + +``kernelize(model, mapping)`` installs the mapping entry by entry onto the +model: keys are replacement targets (HF class objects or dotted-path +strings), values are direct impls / ``HubRef`` / ``KernelChoice`` (picks an +impl by backend priority, see ``twinkle.kernel.registry``). How each entry +is installed is decided by the installer: +``KernelChoice.installer`` → ``OpDefinition.installer`` → ``default_installer``. Public API: ``kernelize``, ``hub`` (re-exported from ``twinkle.kernel``). """ from __future__ import annotations import importlib +import logging import torch.nn as nn from dataclasses import dataclass from typing import Any from twinkle import get_logger -from twinkle.utils.device_mesh import Platform + +from .registry import KernelChoice, get_op, resolve_impl logger = get_logger() @@ -49,19 +58,6 @@ def hub( return HubRef(repo_id, layer_name, revision, version, backend) -def _resolve_value(value: Any, device: str) -> Any | None: - """Resolve a mapping value against the selected device. - - - ``dict``: device-conditional; recurse into ``value[device]`` or return None. - - anything else (including ``HubRef``): pass through. - """ - if isinstance(value, dict): - if device not in value: - return None - return _resolve_value(value[device], device) - return value - - def _replace_class(model: nn.Module, target_cls: type, impl_cls: type) -> None: """Rewrite ``__class__`` of every module whose exact type is ``target_cls``. @@ -73,15 +69,12 @@ def _replace_class(model: nn.Module, target_cls: type, impl_cls: type) -> None: m.__class__ = impl_cls -def _replace_attr(dotted_path: str, impl) -> None: - """``setattr`` ``impl`` onto the attribute identified by the dotted path. - - Supports two forms: - - ``pkg.mod.attr`` (set module attribute) - - ``pkg.mod.ClassName.attr`` (set class attribute / method) +def _resolve_dotted(dotted_path: str) -> tuple[Any, str]: + """Resolve ``pkg.mod[.Class].attr`` to ``(owner, final_attr)``. The split is found by walking the prefix from the longest importable - module backwards until ``importlib.import_module`` succeeds. + module backwards until ``importlib.import_module`` succeeds; the remaining + attributes (except the final one) are walked with ``getattr``. """ parts = dotted_path.split('.') if len(parts) < 2: @@ -103,11 +96,10 @@ def _replace_attr(dotted_path: str, impl) -> None: if module is None: raise ImportError(f'Could not import any prefix of {dotted_path!r}') from last_err - # Walk remaining attributes; the last one is the target. obj = module for attr in parts[module_depth:-1]: obj = getattr(obj, attr) - setattr(obj, parts[-1], impl) + return obj, parts[-1] def _load_hub_ref(ref: HubRef): @@ -133,54 +125,145 @@ def _load_hub_ref(ref: HubRef): return impl +# ── installer ───────────────────────────────────────────────────────────── + + +def resolve_direct_value(replacement: Any) -> Any: + """Non-KernelChoice mapping values: HubRef -> lazy download; anything else passes through unchanged.""" + if isinstance(replacement, HubRef): + return _load_hub_ref(replacement) + return replacement + + +def _install_dotted(model: nn.Module, target: str, impl, *, warn: bool = False) -> bool: + """Dispatch a dotted-path target. Returns True = installed, False = family missing, skipped. + + - resolves to an ``nn.Module`` subclass -> ``_replace_class(model, cls, impl)`` + (exactly equivalent to a class-object key, exact type match) + - otherwise -> ``setattr`` (module function / class method) + - unresolvable ``transformers.*`` family path (missing module/attr) -> + skip, return False (a missing family is normal). ``warn=False`` + (default config path) -> DEBUG; ``warn=True`` (explicit mapping) -> + WARNING with a typo hint — explicit entries are the user's + responsibility and a silent skip would hide spelling mistakes + - unresolvable non-family string (e.g. a logical target mistakenly + routed to the default installer) -> explicit error + """ + try: + owner, final_attr = _resolve_dotted(target) + resolved = getattr(owner, final_attr) + except (ImportError, AttributeError, ValueError) as e: + if target.startswith('transformers.'): + level = logging.WARNING if warn else logging.DEBUG + hint = ' (explicit mapping: check for typos)' if warn else '' + logger.log(level, "[kernelize] target %r unresolvable (%r); family not installed, skipping%s", + target, e, hint) + return False + raise ValueError(f"Cannot resolve mapping target {target!r} with the default installer " + f'(logical targets require a custom installer): {e!r}') from e + if isinstance(resolved, type) and issubclass(resolved, nn.Module): + _replace_class(model, resolved, impl) + else: + setattr(owner, final_attr, impl) + return True + + +def default_installer(model: nn.Module, target: Any, impl, *, warn: bool = False) -> bool: + """Generic installer: class objects -> ``_replace_class``; dotted-path strings -> ``_install_dotted``. + + ``warn`` is forwarded to ``_install_dotted`` for the family-skip log + level; custom installers keep the ``(model, target, impl)`` three-arg + signature, with kernelize dispatching between the two call shapes. + """ + if isinstance(target, type) and issubclass(target, nn.Module): + _replace_class(model, target, impl) + return True + if isinstance(target, str): + return _install_dotted(model, target, impl, warn=warn) + raise TypeError(f'Unsupported mapping target: {target!r}') + + +# ── kernelize ───────────────────────────────────────────────────────────── + + +def _target_name(target: Any) -> str: + if isinstance(target, type): + return f'{target.__module__}.{target.__qualname__}' + return str(target) + + +def _installer_name(installer) -> str: + if installer is default_installer: + return 'default' + return getattr(installer, '__name__', repr(installer)) + + +def _log_all_unavailable(target: Any, choice: KernelChoice, *, warn: bool) -> None: + level = logging.WARNING if warn else logging.DEBUG + logger.log(level, "[kernelize] target %s: no available backend for op '%s' (tried: %s); " + 'keeping the original implementation', _target_name(target), choice.op, + ', '.join(choice.backends)) + + def kernelize(model: nn.Module, mapping: dict | None = None) -> nn.Module: """Apply ``mapping`` to ``model`` and return it (modified in place). - Keys: + Keys (targets): - ``type[nn.Module]``: replace ``m.__class__`` for every module of the exact type (no subclass walking). - - ``str`` (dotted path ``pkg.mod.attr``): ``setattr`` the impl onto the - identified module attribute. + - ``str`` dotted path: resolved by the default installer — an + ``nn.Module`` subclass resolves to class replacement, anything else + (module function / class method) to ``setattr``. Unresolvable + ``transformers.*`` family paths are skipped: DEBUG on the default + path, WARNING (with typo hint) for explicit mappings. Values: - - ``dict[str, V]``: device-conditional dispatch using the current - Twinkle platform device prefix; non-matching devices skip. + - ``KernelChoice``: pick the first available backend impl for the op + (see ``twinkle.kernel.registry``); installer priority is + ``KernelChoice.installer`` → ``OpDefinition.installer`` → default. - ``HubRef``: lazy-resolved via the optional ``kernels`` package. - - anything else: used directly as the impl. + - anything else: used directly as the impl (default installer). - If ``mapping`` is ``None`` it is auto-detected from the current platform - via ``Platform.device_prefix()``: on NPU the built-in ``npu_builtin(model)`` - bundle is applied (including its side effects); on any other platform this - is a no-op and the model is returned unchanged. + ``mapping=None`` applies the built-in ``DEFAULT_KERNEL_CONFIG``; on this + default path all fallback/skip logs are DEBUG. Passing any explicit + mapping (even a copy of the default) raises them to WARNING. The mapping + fully *replaces* the default config — merge with + ``{**DEFAULT_KERNEL_CONFIG, ...}`` to customize on top of it. """ - device = Platform.device_prefix() - - if mapping is None: - if device == 'npu': - from .builtin import npu_builtin - mapping = npu_builtin(model) - else: - logger.debug(f'[kernelize] No mapping provided and device {device!r} ' - f'has no built-in bundle; returning model unchanged') - return model + default_mapping = mapping is None + if default_mapping: + from .config import DEFAULT_KERNEL_CONFIG # in-function import, avoids the core<->config cycle + mapping = DEFAULT_KERNEL_CONFIG if not mapping: return model - for key, value in mapping.items(): - impl = _resolve_value(value, device) - if impl is None: + warn = not default_mapping + for target, replacement in mapping.items(): + if isinstance(replacement, KernelChoice): + op = get_op(replacement.op) # unregistered -> ValueError + impl, backend = resolve_impl(op, replacement.backends, warn=warn, target=target) + if impl is None: + _log_all_unavailable(target, replacement, warn=warn) + continue # keep the original implementation + installer = replacement.installer or op.installer or default_installer + else: + impl = resolve_direct_value(replacement) + backend = None + installer = default_installer + + # custom installers keep the three-arg signature; warn is only forwarded to default_installer + if installer is default_installer: + installed = installer(model, target, impl, warn=warn) + else: + installed = installer(model, target, impl) # failures propagate, never swallowed (half-installed state stays visible) + if installed is False: continue - if isinstance(impl, HubRef): - impl = _load_hub_ref(impl) - if isinstance(key, type) and issubclass(key, nn.Module): - _replace_class(model, key, impl) - logger.info(f'[kernelize] Replaced {key.__module__}.{key.__qualname__} ' - f'with {impl.__module__}.{impl.__qualname__}') - elif isinstance(key, str): - _replace_attr(key, impl) - impl_repr = getattr(impl, '__qualname__', repr(impl)) - logger.info(f'[kernelize] Replaced {key!r} with {impl_repr}') + if isinstance(replacement, KernelChoice): + logger.info('[kernelize] target=%s op=%s backend=%s installer=%s', _target_name(target), op.name, + backend, _installer_name(installer)) else: - raise TypeError(f'Unsupported mapping key: {key!r}') + impl_repr = getattr(impl, '__qualname__', repr(impl)) + logger.info('[kernelize] target=%s impl=%s installer=%s', _target_name(target), impl_repr, + _installer_name(installer)) return model diff --git a/src/twinkle/kernel/liger.py b/src/twinkle/kernel/liger.py deleted file mode 100644 index c8366a71d..000000000 --- a/src/twinkle/kernel/liger.py +++ /dev/null @@ -1,325 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""``liger_builtin()`` returns the Liger Kernel replacement bundle. - -This mirrors ``npu_builtin()`` but for Liger Kernel. Two key differences from -the NPU bundle, both consequences of Liger being cross-device: - - 1. **No device gating.** Values are *bare* impls (not ``{'cuda': impl}``). - Liger's modules self-dispatch across CUDA (Triton) and Ascend NPU (the - auto-applied ``backends/_ascend`` backend) via ``infer_device`` / - ``select_impl`` in ``liger_kernel``. Device-conditional wrapping would - wrongly skip Liger on NPU, where it is fully supported. - - 2. **No side effects.** Unlike ``npu_builtin`` (which installs global SDPA - and per-instance FLA), this bundle contains only class/function - replacements consumable by ``kernelize``. The process-level - ``nn.functional.cross_entropy`` swap and fused-linear-CE ``forward`` - replacement are deliberately *excluded* — they belong to the loss layer, - not the kernel layer (see ROADMAP / ``twinkle.loss``). - -Composing with ``npu_builtin`` on NPU: both bundles are NPU implementations of -the same operators. Plain dict merge lets the user pick precedence (later keys -win) — see ``Kernel.md``: - - model = kernelize(model, {**npu_builtin(model), **liger_builtin(model)}) # Liger wins - model = kernelize(model, {**liger_builtin(model), **npu_builtin(model)}) # Twinkle-NPU wins - model = kernelize(model, liger_builtin(model)) # Liger only - -RMSNorm uses Liger-aware adapter classes (``liger_impls.rms_norm``) that set the -Liger-specific instance attributes lazily in ``forward``; SwiGLU/GeGLU use -function-level forward replacements; RoPE and the Qwen3-MoE expert/MLP classes -are wired as bare Liger classes (their forwards read only attributes the HF -instances already provide). -""" -from __future__ import annotations - -import importlib -import torch.nn as nn -from typing import Any - -from twinkle import get_logger -from twinkle.utils.device_mesh import Platform - -logger = get_logger() - - -def _import_optional(name: str): - try: - return importlib.import_module(name) - except ImportError: - return None - - -def _has(maybe_none, attr_path: str) -> bool: - """Return True if ``attr_path`` (may contain one dot, e.g. ``Cls.member``) - exists on ``maybe_none`` (a module or None).""" - if maybe_none is None: - return False - if '.' in attr_path: - head, _, tail = attr_path.partition('.') - owner = getattr(maybe_none, head, None) - return owner is not None and hasattr(owner, tail) - return hasattr(maybe_none, attr_path) - - -def _add_class_if_present(bundle, module_path, class_name, impl_cls): - mod = _import_optional(module_path) - if mod is None: - return - cls = getattr(mod, class_name, None) - if isinstance(cls, type): - bundle[cls] = impl_cls - - -def _add_fwd_if_present(bundle, module_path, class_name, fn): - """Function-level forward replacement: ``bundle['..forward'] = fn``. - - Only added when the HF module imports AND the target class exists, so the - bundle stays safe when a model family is not installed. - """ - mod = _import_optional(module_path) - if mod is None: - return - cls = getattr(mod, class_name, None) - if isinstance(cls, type): - bundle[f'{module_path}.{class_name}.forward'] = fn - - -def _add_attr_if_present(bundle, module_path, attr_name, impl): - mod = _import_optional(module_path) - if mod is None: - return - if _has(mod, attr_name): - bundle[f'{module_path}.{attr_name}'] = impl - - -# ── family registrations ────────────────────────────────────────────────────── -# Each helper mirrors the subset of liger_kernel's apply_liger_kernel_to_ -# that is expressible as a pure class/function replacement (RoPE, RMSNorm, -# SwiGLU/GeGLU). Fused-linear-CE and the global cross_entropy swap are excluded. - - -def _add_llama_style(bundle, module_path, rms_name, mlp_name, with_rope=True): - """llama / qwen / mistral / mixtral / phi3 / olmo2 / glm4 / granite / internvl: - llama-cast RMSNorm + SwiGLU MLP + (optional) RoPE.""" - from liger_kernel.transformers import liger_rotary_pos_emb - - from .liger_impls import LigerRMSNormReplacement, liger_swiglu_forward - - _add_class_if_present(bundle, module_path, rms_name, LigerRMSNormReplacement) - _add_fwd_if_present(bundle, module_path, mlp_name, liger_swiglu_forward) - if with_rope: - _add_attr_if_present(bundle, module_path, 'apply_rotary_pos_emb', liger_rotary_pos_emb) - - -def _add_gemma_style(bundle, module_path, rms_name, mlp_name): - """gemma / gemma2 / gemma3_text: gemma-cast RMSNorm (offset=1.0) + GeGLU MLP. - RoPE differs per gemma variant and is left to the model-specific entries.""" - from .liger_impls import LigerRMSNormGemmaReplacement, liger_geglu_forward - - _add_class_if_present(bundle, module_path, rms_name, LigerRMSNormGemmaReplacement) - _add_fwd_if_present(bundle, module_path, mlp_name, liger_geglu_forward) - - -def _add_qwen2(bundle): - _add_llama_style(bundle, 'transformers.models.qwen2.modeling_qwen2', 'Qwen2RMSNorm', 'Qwen2MLP') - - -def _add_qwen3(bundle): - _add_llama_style(bundle, 'transformers.models.qwen3.modeling_qwen3', 'Qwen3RMSNorm', 'Qwen3MLP') - - -def _add_qwen3_moe(bundle): - base = 'transformers.models.qwen3_moe.modeling_qwen3_moe' - from liger_kernel.transformers import LigerExperts, LigerQwen3MoeSwiGLUMLP, liger_rotary_pos_emb - - from .liger_impls import LigerRMSNormReplacement, liger_swiglu_forward - - _add_class_if_present(bundle, base, 'Qwen3MoeRMSNorm', LigerRMSNormReplacement) - _add_attr_if_present(bundle, base, 'apply_rotary_pos_emb', liger_rotary_pos_emb) - # transformers v5+: experts are a separate fused class; v4 uses Qwen3MoeMLP. - _add_class_if_present(bundle, base, 'Qwen3MoeExperts', LigerExperts) - _add_fwd_if_present(bundle, base, 'Qwen3MoeMLP', liger_swiglu_forward) - # Qwen3MoeMLP (v4) is structurally identical to LigerQwen3MoeSwiGLUMLP; offer - # class replacement too so v4 users get the fused SiLU kernel without a - # forward-level swap shadowing it. - mod = _import_optional(base) - if mod is not None and isinstance(getattr(mod, 'Qwen3MoeMLP', None), type): - if not any(k == getattr(mod, 'Qwen3MoeMLP') for k in bundle): - bundle[getattr(mod, 'Qwen3MoeMLP')] = LigerQwen3MoeSwiGLUMLP - - -def _add_qwen3_5(bundle): - from .liger_impls import LigerRMSNormQwen35Replacement, liger_swiglu_forward - - base = 'transformers.models.qwen3_5.modeling_qwen3_5' - if _import_optional(base) is None: - return - _add_class_if_present(bundle, base, 'Qwen3_5RMSNorm', LigerRMSNormQwen35Replacement) - _add_class_if_present(bundle, base, 'Qwen3_5VisionRMSNorm', LigerRMSNormQwen35Replacement) - # RoPE intentionally NOT replaced: Qwen3.5 uses partial_rotary_factor=0.25 + - # mrope_interleaved=True; Liger's liger_rotary_pos_emb assumes full-rotation - # with the rotate_half convention, which is incompatible. Let npu_builtin's - # npu_apply_rotary_pos_emb (which handles Partial-RoPE) take precedence. - _add_fwd_if_present(bundle, base, 'Qwen3_5MLP', liger_swiglu_forward) - _add_fwd_if_present(bundle, base, 'Qwen3_5VisionMLP', liger_swiglu_forward) - - -def _add_qwen3_5_moe(bundle): - from liger_kernel.transformers import LigerExperts - - from .liger_impls import LigerRMSNormQwen35Replacement, liger_swiglu_forward - - base = 'transformers.models.qwen3_5_moe.modeling_qwen3_5_moe' - if _import_optional(base) is None: - return - _add_class_if_present(bundle, base, 'Qwen3_5MoeRMSNorm', LigerRMSNormQwen35Replacement) - # RoPE intentionally NOT replaced (same reason as _add_qwen3_5). - _add_class_if_present(bundle, base, 'Qwen3_5MoeExperts', LigerExperts) - _add_fwd_if_present(bundle, base, 'Qwen3_5MoeMLP', liger_swiglu_forward) - - -def _add_qwen2_5_vl(bundle): - from liger_kernel.transformers import liger_rotary_pos_emb - - from .liger_impls import LigerRMSNormReplacement, liger_swiglu_forward - - base = 'transformers.models.qwen2_5_vl.modeling_qwen2_5_vl' - if _import_optional(base) is None: - return - _add_class_if_present(bundle, base, 'Qwen2_5_VLRMSNorm', LigerRMSNormReplacement) - _add_attr_if_present(bundle, base, 'apply_rotary_pos_emb', liger_rotary_pos_emb) - _add_fwd_if_present(bundle, base, 'Qwen2MLP', liger_swiglu_forward) - _add_fwd_if_present(bundle, base, 'Qwen2_5_VLMLP', liger_swiglu_forward) - - -def _add_gemma4(bundle): - base = 'transformers.models.gemma4.modeling_gemma4' - if _import_optional(base) is None: - return - from .liger_impls import LigerRMSNormGemma4Replacement, liger_geglu_forward - - _add_class_if_present(bundle, base, 'Gemma4RMSNorm', LigerRMSNormGemma4Replacement) - _add_fwd_if_present(bundle, base, 'Gemma4TextMLP', liger_geglu_forward) - - -def liger_builtin(model: nn.Module | None = None) -> dict[Any, Any]: - """Return the Liger Kernel built-in mapping; composes with ``kernelize``. - - Args: - model: Accepted for API symmetry with ``npu_builtin(model)``; Liger's - pure class/function replacements have no per-instance side effects, - so the model is not traversed here. - - Returns: - A ``dict`` whose keys are HF ``nn.Module`` subclasses (class replacement) - or dotted paths (function/forward replacement) and whose values are - *bare* Liger impls (no device gating). Missing model families are - silently skipped — the bundle only contains entries for installed - transformers model modules. - - On NPU, Liger's Triton-on-Ascend kernels are slower than the CANN vendor - ops in ``npu_impls`` for the bandwidth-bound per-layer ops (RMSNorm, - SwiGLU, RoPE). ``_prefer_cann_on_npu`` post-processes the bundle to swap - those Liger impls for their CANN equivalents and drops the LigerExperts - class replacement (which shadows ``npu_builtin``'s faster forward-level - MoE expert replacement). Non-Qwen families without a CANN equivalent keep - their Liger impls. On CUDA, the bundle is unchanged. - - Raises: - ImportError: if ``liger_kernel`` is not importable (caught at the first - ``from liger_kernel ...`` statement inside the family helpers). - """ - bundle: dict[Any, Any] = {} - - # ── Qwen family (primary on Twinkle) ────────────────────────────────────── - _add_qwen2(bundle) - _add_qwen3(bundle) - _add_qwen3_moe(bundle) - _add_qwen3_5(bundle) - _add_qwen3_5_moe(bundle) - _add_qwen2_5_vl(bundle) - - # ── Llama-style dense / MoE families ───────────────────────────────────── - _add_llama_style(bundle, 'transformers.models.llama.modeling_llama', 'LlamaRMSNorm', 'LlamaMLP') - _add_llama_style(bundle, 'transformers.models.mistral.modeling_mistral', 'MistralRMSNorm', 'MistralMLP') - _add_llama_style(bundle, 'transformers.models.mixtral.modeling_mixtral', 'MixtralRMSNorm', - 'MixtralBlockSparseTop2MLP') - _add_llama_style(bundle, 'transformers.models.phi3.modeling_phi3', 'Phi3RMSNorm', 'Phi3MLP') - _add_llama_style(bundle, 'transformers.models.glm4.modeling_glm4', 'Glm4RMSNorm', 'Glm4MLP') - _add_llama_style(bundle, 'transformers.models.olmo2.modeling_olmo2', 'Olmo2RMSNorm', 'Olmo2MLP') - _add_llama_style(bundle, 'transformers.models.granite.modeling_granite', 'GraniteRMSNorm', 'GraniteMLP') - _add_llama_style(bundle, 'transformers.models.internvl.modeling_internvl', 'InternVLRMSNorm', 'InternVLMLP') - - # ── Gemma family (gemma-cast RMSNorm + GeGLU) ──────────────────────────── - _add_gemma_style(bundle, 'transformers.models.gemma.modeling_gemma', 'GemmaRMSNorm', 'GemmaMLP') - _add_gemma_style(bundle, 'transformers.models.gemma2.modeling_gemma2', 'Gemma2RMSNorm', 'Gemma2MLP') - _add_gemma_style(bundle, 'transformers.models.gemma3.modeling_gemma3', 'Gemma3RMSNorm', 'Gemma3MLP') - _add_gemma4(bundle) - - if not bundle: - logger.warning('[liger_builtin] No Liger entries were registered — ' - 'is liger_kernel installed and are transformers model modules importable?') - - # On NPU, prefer CANN vendor ops over Liger's Triton-on-Ascend kernels for - # the per-layer ops where CANN is significantly faster (RMSNorm: single-pass - # vs Liger's 2-pass tiled; SwiGLU/RoPE: one CANN op vs Triton launch + extra - # allocations). Also drop LigerExperts class replacements that shadow - # npu_builtin's faster forward-level MoE expert replacement. - if Platform.device_prefix() == 'npu': - _prefer_cann_on_npu(bundle) - return bundle - - -def _prefer_cann_on_npu(bundle: dict[Any, Any]) -> None: - """In-place: swap Liger Triton impls for CANN vendor ops on NPU. - - Replaces: - - LigerRMSNorm* class values -> NpuRMSNorm (single-pass CANN aclnn) - - liger_swiglu_forward values -> npu_swiglu_forward (one CANN fused op) - - liger_rotary_pos_emb values -> npu_apply_rotary_pos_emb (no Q/K copies) - - Removes: - - LigerExperts / LigerQwen3MoeSwiGLUMLP class keys — they replace - ``m.__class__`` which shadows npu_builtin's forward-level - ``npu_packed_moe_experts_forward`` (a different dict key), leaving - the fast CANN grouped-matmul MoE path unused. - """ - from .npu_impls import NpuRMSNorm, npu_apply_rotary_pos_emb, npu_swiglu_forward - - # Stringify once for fast membership checks. - def _val_name(v) -> str: - return getattr(v, '__name__', getattr(v, '__qualname__', '')) - - keys_to_drop: list[Any] = [] - for key, val in list(bundle.items()): - name = _val_name(val) - - # Drop LigerExperts / LigerQwen3MoeSwiGLUMLP class replacements. - if isinstance(val, type) and any( - s in name for s in ('LigerExperts', 'LigerQwen3MoeSwiGLUMLP', 'LigerQwen3_5MoeSwiGLUMLP')): - keys_to_drop.append(key) - continue - - # Swap LigerRMSNorm* class replacements -> NpuRMSNorm. - if isinstance(val, type) and name.startswith('LigerRMSNorm'): - bundle[key] = NpuRMSNorm - continue - - # Swap liger_swiglu_forward -> npu_swiglu_forward. - if name == 'liger_swiglu_forward': - bundle[key] = npu_swiglu_forward - continue - - # Swap liger_rotary_pos_emb -> npu_apply_rotary_pos_emb. - if name == 'liger_rotary_pos_emb': - bundle[key] = npu_apply_rotary_pos_emb - continue - - for key in keys_to_drop: - del bundle[key] - - if keys_to_drop: - logger.info( - '[liger_builtin] NPU: dropped %d LigerExperts/LigerMoeMLP class replacement(s) ' - 'so npu_builtin\'s CANN grouped-matmul MoE forward takes effect', len(keys_to_drop)) diff --git a/src/twinkle/kernel/liger_impls/__init__.py b/src/twinkle/kernel/liger_impls/__init__.py deleted file mode 100644 index 0ea8f34d1..000000000 --- a/src/twinkle/kernel/liger_impls/__init__.py +++ /dev/null @@ -1,28 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Class-replacement adapters that bridge Liger Kernel modules onto -HuggingFace instances when applied via ``kernelize``. - -``kernelize`` swaps ``m.__class__ = impl_cls`` *without* calling ``__init__`` -(see ``Kernel.md`` caveats). Liger's own ``LigerRMSNorm.forward`` reads -instance attributes (``offset`` / ``casting_mode`` / ``in_place`` / ``row_mode``) -that HuggingFace RMSNorm variants do not define. Liger's monkey-patch path -sets those attributes eagerly via ``_patch_rms_norm_module``; the adapters here -do the same lazily inside ``forward`` so the class-replacement contract is -honoured and no global state is mutated. - -SwiGLU / Experts / RoPE need no adapter — Liger's classes/functions read only -attributes (``gate_proj`` / ``up_proj`` / ``down_proj`` / ``weight``) that the -HuggingFace instances already provide, so they are re-exported verbatim. -""" -from .rms_norm import (LigerRMSNormGemma4Replacement, LigerRMSNormGemmaReplacement, LigerRMSNormQwen35Replacement, - LigerRMSNormReplacement) -from .swiglu import liger_geglu_forward, liger_swiglu_forward - -__all__ = [ - 'LigerRMSNormReplacement', - 'LigerRMSNormGemmaReplacement', - 'LigerRMSNormGemma4Replacement', - 'LigerRMSNormQwen35Replacement', - 'liger_swiglu_forward', - 'liger_geglu_forward', -] diff --git a/src/twinkle/kernel/liger_impls/swiglu.py b/src/twinkle/kernel/liger_impls/swiglu.py deleted file mode 100644 index 97b134b42..000000000 --- a/src/twinkle/kernel/liger_impls/swiglu.py +++ /dev/null @@ -1,32 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""SwiGLU forward-replacement for Liger Kernel. - -Used as a *function-level* mapping value (string key ``'..forward'``) -so it composes with the existing SwiGLU forward-replacement pattern in -``builtin.py`` (``_add_swiglu_if_present``). Reads only ``gate_proj`` / -``up_proj`` / ``down_proj``, which every HuggingFace SwiGLU MLP variant -(Qwen2MLP, Qwen3MLP, LlamaMLP, MistralMLP, ...) already defines, so no -``__init__`` and no per-instance attribute setup is required. - -For Qwen3-MoE the expert/MLP classes differ; class replacement with Liger's -own ``LigerQwen3MoeSwiGLUMLP`` / ``LigerExperts`` is wired directly in -``liger_builtin`` since those read matching attributes. -""" -from __future__ import annotations - -from liger_kernel.ops import LigerGELUMulFunction, LigerSiLUMulFunction - - -def liger_swiglu_forward(self, x): - return self.down_proj(LigerSiLUMulFunction.apply(self.gate_proj(x), self.up_proj(x))) - - -def liger_geglu_forward(self, x): - """GeGLU forward replacement for the gemma family (gemma / gemma2 / gemma3 / gemma4). - - Reads only ``gate_proj`` / ``up_proj`` / ``down_proj`` — the same attributes - HuggingFace GeGLU MLP variants define — so it is a safe function-level - mapping value. Uses the tanh GELU approximation, matching Liger's own - ``LigerGEGLUMLP`` and HF's gemma activation choice. - """ - return self.down_proj(LigerGELUMulFunction.apply(self.gate_proj(x), self.up_proj(x))) diff --git a/src/twinkle/kernel/npu_impls/__init__.py b/src/twinkle/kernel/npu_impls/__init__.py deleted file mode 100644 index ebf34c6cc..000000000 --- a/src/twinkle/kernel/npu_impls/__init__.py +++ /dev/null @@ -1,28 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Per-layer NPU implementations consumed by ``npu_builtin()``. - -Each impl is contracted to be applied via ``m.__class__ = ImplCls`` (class -replacement) or ``setattr(module, attr, fn)`` (function replacement). No impl -here is meant to be instantiated directly. -""" -from .attention import npu_sdpa_attention_forward -from .fla import apply_qwen3_5_fla, npu_causal_conv1d_fn -from .moe import GmmFunction, npu_grouped_mm, npu_packed_moe_experts_forward, npu_qwen3_5_moe_sparse_block_forward -from .rms_norm import NpuRMSNorm, npu_gated_rms_norm_forward -from .rotary import npu_apply_multimodal_rotary_pos_emb, npu_apply_rotary_pos_emb -from .swiglu import npu_swiglu_forward - -__all__ = [ - 'NpuRMSNorm', - 'npu_gated_rms_norm_forward', - 'npu_apply_rotary_pos_emb', - 'npu_apply_multimodal_rotary_pos_emb', - 'npu_swiglu_forward', - 'npu_sdpa_attention_forward', - 'GmmFunction', - 'npu_grouped_mm', - 'npu_packed_moe_experts_forward', - 'npu_qwen3_5_moe_sparse_block_forward', - 'apply_qwen3_5_fla', - 'npu_causal_conv1d_fn', -] diff --git a/src/twinkle/kernel/ops/__init__.py b/src/twinkle/kernel/ops/__init__.py new file mode 100644 index 000000000..0d9d39925 --- /dev/null +++ b/src/twinkle/kernel/ops/__init__.py @@ -0,0 +1,22 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Backend-agnostic operator interfaces with per-platform implementations. + +Each op family defines an abstract base class (e.g. ``EpExpertsGmm``) plus a +dispatcher that picks the first eligible backend implementation from a +lazily-built registry. Callers invoke the dispatcher only; backend details +(NPU, GPU, ...) stay hidden behind the interface. + +Importing this package also registers all built-in kernel ops (swiglu, +rms_norm, rotary, geglu, moe, sdpa_attention, fla) into +``twinkle.kernel.registry`` — registration modules are lightweight (lazy +references + availability checks only, no optional-dependency imports). +""" +from .ep import EpExpertsGmm, ep_forward + +# Trigger built-in op registration (must happen before the first kernelize() call) +from . import fla, geglu, moe, rms_norm, rotary, sdpa_attention, swiglu # noqa: F401,E402 + +__all__ = [ + 'EpExpertsGmm', + 'ep_forward', +] diff --git a/src/twinkle/kernel/ops/ep/__init__.py b/src/twinkle/kernel/ops/ep/__init__.py new file mode 100644 index 000000000..c011d3ced --- /dev/null +++ b/src/twinkle/kernel/ops/ep/__init__.py @@ -0,0 +1,112 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Grouped-matmul interface for expert-parallel (EP) experts compute. + +``EpExpertsGmm`` is the abstract backend interface; each platform provides a +subclass (e.g. ``NpuEpExpertsGmm`` in ``ep_experts_gmm_npu.py``) that is +registered lazily. The generic per-expert loop (``LoopEpExpertsGmm`` in +``ep_experts_gmm_loop.py``) is registered last as the always-eligible +fallback. ``ep_forward`` is the single entry point used by the EP +forward: it dispatches to the first eligible backend, so it always returns a +result. +""" +from __future__ import annotations + +import torch +from abc import ABC, abstractmethod +from torch import nn + +from twinkle import get_logger + +logger = get_logger() + + +class EpExpertsGmm(ABC): + """Backend interface for batched EP experts compute via grouped matmul. + + A backend computes all local experts in one shot from the permuted token + buffer (tokens sorted by expert, contiguous chunks per local expert), + replacing the per-expert Python loop. + """ + + name: str + # Fallback backends (e.g. the per-expert loop) are always eligible and are + # registered last; hitting one means every accelerated backend declined. + fallback: bool = False + + @abstractmethod + def ineligible_reason(self, experts_mod: nn.Module) -> str | None: + """Return why this backend cannot handle ``experts_mod``, or None if it can.""" + + @abstractmethod + def forward( + self, + experts_mod: nn.Module, + permuted_tokens: torch.Tensor, + num_global_sum_tokens_per_local_expert: torch.Tensor, + experts_per_rank: int, + ) -> torch.Tensor: + """Compute local experts on the permuted token buffer.""" + + +_IMPLS: list[EpExpertsGmm] | None = None +_PATH_LOGGED = False +_WARN_LOGGED = False + + +def _get_impls() -> list[EpExpertsGmm]: + """Build the backend registry on first use. + + Each backend module is imported defensively: platforms lacking its + dependencies (e.g. no torch_npu) simply skip that backend. + """ + global _IMPLS + if _IMPLS is None: + _IMPLS = [] + try: + from .npu import NpuEpExpertsGmm + _IMPLS.append(NpuEpExpertsGmm()) + except ImportError: + pass + # The loop fallback has no optional dependencies and is always last: + # it guarantees the dispatcher never runs out of backends. + from .loop import LoopEpExpertsGmm + _IMPLS.append(LoopEpExpertsGmm()) + return _IMPLS + + +def ep_forward( + experts_mod: nn.Module, + permuted_tokens: torch.Tensor, + num_global_sum_tokens_per_local_expert: torch.Tensor, + experts_per_rank: int, +) -> torch.Tensor: + """Dispatch to the first eligible backend (the loop fallback guarantees a hit). + + One-time logging: INFO on the first accelerated dispatch (process-wide); + WARNING with the ineligibility reasons when falling back to the loop + (once per experts instance, as reasons may differ per MoE block). + """ + if permuted_tokens.numel() == 0: + # Preserve the autograd edge to token_pre_all2all. Returning a new + # empty tensor can make this rank skip the matching backward + # all-to-all, causing EP collective order divergence. + return permuted_tokens + + global _PATH_LOGGED, _WARN_LOGGED + ineligible: list[str] = [] + for impl in _get_impls(): + reason = impl.ineligible_reason(experts_mod) + if reason is not None: + ineligible.append(f'{impl.name}: {reason}') + continue + if impl.fallback: + if ineligible and not _WARN_LOGGED: + detail = '; '.join(ineligible) + logger.warning(f'EP experts compute: grouped matmul disabled ({detail}); ' + f'falling back to {impl.name}.') + _WARN_LOGGED = True + elif not _PATH_LOGGED: + logger.info(f'EP experts compute: using {impl.name} grouped matmul (per-expert loop replaced).') + _PATH_LOGGED = True + return impl.forward(experts_mod, permuted_tokens, num_global_sum_tokens_per_local_expert, experts_per_rank) + raise RuntimeError('no EP experts backend registered (loop fallback missing)') diff --git a/src/twinkle/kernel/ops/ep/loop.py b/src/twinkle/kernel/ops/ep/loop.py new file mode 100644 index 000000000..c08d03040 --- /dev/null +++ b/src/twinkle/kernel/ops/ep/loop.py @@ -0,0 +1,66 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Generic per-expert Python loop backend for the EP experts interface. + +This is the always-eligible fallback: it runs on any platform (pure PyTorch) +and is registered last, so the dispatcher can never run out of backends. +""" +from __future__ import annotations + +import torch +import torch.nn.functional as F +from torch import nn + +from . import EpExpertsGmm + + +class LoopEpExpertsGmm(EpExpertsGmm): + """Per-expert Python loop over the permuted token buffer (generic fallback).""" + + name = 'per-expert loop' + fallback = True + + def ineligible_reason(self, experts_mod: nn.Module) -> str | None: + # The loop is the reference implementation: always eligible. + return None + + def forward( + self, + experts_mod: nn.Module, + permuted_tokens: torch.Tensor, + num_global_sum_tokens_per_local_expert: torch.Tensor, + experts_per_rank: int, + ) -> torch.Tensor: + input_dtype = permuted_tokens.dtype + + cumsum = torch.zeros(experts_per_rank + 1, dtype=torch.long) + for i in range(experts_per_rank): + cumsum[i + 1] = cumsum[i] + int(num_global_sum_tokens_per_local_expert[i].item()) + + output_chunks = [] + for i in range(experts_per_rank): + start = int(cumsum[i].item()) + end = int(cumsum[i + 1].item()) + expert_in = permuted_tokens[start:end] + if expert_in.numel() == 0: + output_chunks.append(expert_in) + continue + + gate_up = experts_mod.gate_up_proj[i] + down = experts_mod.down_proj[i] + compute_dtype = gate_up.dtype + if expert_in.dtype != compute_dtype: + expert_in = expert_in.to(compute_dtype) + gate_up_out = F.linear(expert_in, gate_up) + if hasattr(experts_mod, '_apply_gate'): + out = experts_mod._apply_gate(gate_up_out) + else: + gate, up = gate_up_out.chunk(2, dim=-1) + out = experts_mod.act_fn(gate) * up + out = F.linear(out, down) + + if out.dtype != input_dtype: + out = out.to(input_dtype) + output_chunks.append(out) + + return torch.cat( + output_chunks, dim=0) if output_chunks else permuted_tokens.new_empty(0, permuted_tokens.size(-1)) diff --git a/src/twinkle/kernel/ops/ep/npu.py b/src/twinkle/kernel/ops/ep/npu.py new file mode 100644 index 000000000..36ab249fa --- /dev/null +++ b/src/twinkle/kernel/ops/ep/npu.py @@ -0,0 +1,103 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""NPU backend for the EP experts grouped-matmul interface.""" +from __future__ import annotations + +import torch +from torch import nn + +from . import EpExpertsGmm + + +class NpuEpExpertsGmm(EpExpertsGmm): + """Batched EP experts compute via ``torch_npu.npu_grouped_matmul``.""" + + name = 'NPU' + + def ineligible_reason(self, experts_mod: nn.Module) -> str | None: + """Check whether the tensor-experts EP compute can use NPU grouped matmul. + + Falls back to the per-expert Python loop when: + - torch_npu / NPU is unavailable; + - there is no way to compute the gated activation: no ``_apply_gate`` + hook and a non-SiLU ``act_fn`` (``npu_swiglu`` == silu(gate) * up); + - the packed weights are not plain 3D floating tensors (e.g. DTensor + that was not unsharded, quantized weights, etc.). + + Note on ``_apply_gate``: transformers injects a default implementation + (``_default_apply_gate``: chunk + act_fn(gate) * up) onto many MoE + expert classes. With SiLU activation that is exactly ``npu_swiglu``, so + it does not block the fused path. A *custom* gate hook is still + supported — it is called eagerly on the grouped-matmul output. + """ + apply_gate = getattr(experts_mod, '_apply_gate', None) + is_default_gate = getattr(apply_gate, '__name__', None) == '_default_apply_gate' + if apply_gate is None or is_default_gate: + # Activation will be computed by npu_swiglu == silu(gate) * up. + act_fn = getattr(experts_mod, 'act_fn', None) + if act_fn is None: + return 'no act_fn attribute' + # nn.SiLU and transformers' SiLUActivation both compute F.silu + # exactly; either is equivalent to the silu inside npu_swiglu. + if act_fn.__class__.__name__ not in ('SiLU', 'SiLUActivation'): + return f'act_fn is {act_fn.__class__.__name__}, not SiLU' + # else: custom gate hook is called eagerly in the GMM path — allowed. + try: + import torch_npu # noqa: F401 + except ImportError: + return 'torch_npu not importable' + if not torch.npu.is_available(): + return 'NPU not available' + for name in ('gate_up_proj', 'down_proj'): + weight = getattr(experts_mod, name, None) + # Note: accept plain tensors and nn.Parameter. Under FSDP2 the + # pre-forward hook has already unsharded params by the time + # ep_forward runs; under PEFT target_parameters the parametrized + # property returns a merged plain tensor. Anything else (DTensor, + # quantized types) falls back to the loop. + if not isinstance(weight, torch.Tensor) or weight.__class__.__name__ == 'DTensor': + return f'{name} is {type(weight).__name__}, not a plain tensor' + if weight.ndim != 3 or not weight.dtype.is_floating_point: + return f'{name} has shape {tuple(weight.shape)} dtype {weight.dtype}' + return None + + def forward( + self, + experts_mod: nn.Module, + permuted_tokens: torch.Tensor, + num_global_sum_tokens_per_local_expert: torch.Tensor, + experts_per_rank: int, + ) -> torch.Tensor: + """Compute local experts with one grouped matmul instead of a Python loop. + + Mathematically equivalent to the per-expert loop: grouped matmul + applies each expert's weight to its contiguous token chunk, and + ``npu_swiglu(x) == silu(gate) * up`` for gate/up = x.chunk(2, -1). + The group list stays on device, eliminating the per-expert ``.item()`` + host synchronizations of the loop implementation. + """ + import torch_npu + + from twinkle.kernel.ops.moe.npu import GmmFunction, _get_cached_expert_weights + + input_dtype = permuted_tokens.dtype + hidden_dim = permuted_tokens.size(-1) + # counts arrives as a CPU tensor from preprocess(); grouped matmul needs + # the group list on device. + group_list = num_global_sum_tokens_per_local_expert.to(device=permuted_tokens.device, dtype=torch.int64) + gate_up_weight, down_weight = _get_cached_expert_weights(experts_mod, input_dtype, hidden_dim) + + expert_in = permuted_tokens + if expert_in.dtype != gate_up_weight.dtype: + expert_in = expert_in.to(gate_up_weight.dtype) + intermediate = GmmFunction.apply(expert_in, group_list, gate_up_weight) + apply_gate = getattr(experts_mod, '_apply_gate', None) + if apply_gate is not None and getattr(apply_gate, '__name__', '') != '_default_apply_gate': + # Custom gate hook: honor it eagerly instead of the fused swiglu. + activated = apply_gate(intermediate) + else: + # No hook, or transformers' default gate: silu(gate) * up == npu_swiglu. + activated = torch_npu.npu_swiglu(intermediate, dim=-1) + out = GmmFunction.apply(activated, group_list, down_weight) + if out.dtype != input_dtype: + out = out.to(input_dtype) + return out diff --git a/src/twinkle/kernel/ops/fla/__init__.py b/src/twinkle/kernel/ops/fla/__init__.py new file mode 100644 index 000000000..ecbc45e76 --- /dev/null +++ b/src/twinkle/kernel/ops/fla/__init__.py @@ -0,0 +1,36 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""fla op registration: Qwen3.5 Flash Linear Attention per-instance patching +(custom installer, needs the model passed to kernelize). + +Availability = NPU platform + flash-linear-attention installed; when fla is +missing the whole chain fails and falls back per the warn/debug rules, +preserving the "missing fla only warns, other patches proceed" semantics. +""" +from __future__ import annotations + +from ...registry import KernelImpl, exists, is_npu_available, lazy_import, register_op + + +def install_fla(model, target, impl) -> None: + impl(model) # impl = apply_qwen3_5_fla + + +def _fla_available() -> tuple[bool, str | None]: + ok, reason = is_npu_available() + if not ok: + return ok, reason + if not exists('flash-linear-attention'): + return False, 'flash-linear-attention not installed' + return True, None + + +register_op( + 'fla', + implementations={ + 'npu': KernelImpl( + load=lazy_import('twinkle.kernel.ops.fla.npu:apply_qwen3_5_fla'), + available=_fla_available, + ), + }, + installer=install_fla, +) diff --git a/src/twinkle/kernel/npu_impls/fla.py b/src/twinkle/kernel/ops/fla/npu.py similarity index 100% rename from src/twinkle/kernel/npu_impls/fla.py rename to src/twinkle/kernel/ops/fla/npu.py diff --git a/src/twinkle/kernel/ops/geglu/__init__.py b/src/twinkle/kernel/ops/geglu/__init__.py new file mode 100644 index 000000000..bb1e8c47c --- /dev/null +++ b/src/twinkle/kernel/ops/geglu/__init__.py @@ -0,0 +1,15 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""geglu op registration: forward-level replacement for gemma-family MLPs, liger only.""" +from __future__ import annotations + +from ...registry import KernelImpl, is_liger_available, lazy_import, register_op + +register_op( + 'geglu', + implementations={ + 'liger': KernelImpl( + load=lazy_import('twinkle.kernel.ops.geglu.liger:liger_geglu_forward'), + available=is_liger_available, + ), + }, +) diff --git a/src/twinkle/kernel/ops/geglu/liger.py b/src/twinkle/kernel/ops/geglu/liger.py new file mode 100644 index 000000000..a482fcc1a --- /dev/null +++ b/src/twinkle/kernel/ops/geglu/liger.py @@ -0,0 +1,20 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""GeGLU forward-replacement for Liger Kernel. + +Used as a *function-level* mapping value (string key ``'..forward'``). +Reads only ``gate_proj`` / ``up_proj`` / ``down_proj`` — the same attributes +HuggingFace GeGLU MLP variants define — so it is a safe function-level +mapping value. +""" +from __future__ import annotations + +from liger_kernel.ops import LigerGELUMulFunction + + +def liger_geglu_forward(self, x): + """GeGLU forward replacement for the gemma family (gemma / gemma2 / gemma3 / gemma4). + + Uses the tanh GELU approximation, matching Liger's own ``LigerGEGLUMLP`` + and HF's gemma activation choice. + """ + return self.down_proj(LigerGELUMulFunction.apply(self.gate_proj(x), self.up_proj(x))) diff --git a/src/twinkle/kernel/ops/moe/__init__.py b/src/twinkle/kernel/ops/moe/__init__.py new file mode 100644 index 000000000..1839c48e1 --- /dev/null +++ b/src/twinkle/kernel/ops/moe/__init__.py @@ -0,0 +1,37 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""moe_experts / moe_block op registration. + +moe_experts: the npu backend is a forward-level replacement (CANN +grouped-matmul fast path); the liger backend is a LigerExperts class +replacement that only takes effect when the user explicitly puts 'liger' +in the backend chain — the default chain is ('npu',), so the liger class +replacement never shadows the npu fast path (absorbs the drop semantics +of the old _prefer_cann_on_npu). +""" +from __future__ import annotations + +from ...registry import KernelImpl, is_liger_available, is_npu_available, lazy_import, register_op + +register_op( + 'moe_experts', + implementations={ + 'npu': KernelImpl( + load=lazy_import('twinkle.kernel.ops.moe.npu:npu_packed_moe_experts_forward'), + available=is_npu_available, + ), + 'liger': KernelImpl( + load=lazy_import('twinkle.kernel.ops.moe.liger:LigerExperts'), + available=is_liger_available, + ), + }, +) + +register_op( + 'moe_block', + implementations={ + 'npu': KernelImpl( + load=lazy_import('twinkle.kernel.ops.moe.npu:npu_qwen3_5_moe_sparse_block_forward'), + available=is_npu_available, + ), + }, +) diff --git a/src/twinkle/kernel/ops/moe/liger.py b/src/twinkle/kernel/ops/moe/liger.py new file mode 100644 index 000000000..b55f8ac00 --- /dev/null +++ b/src/twinkle/kernel/ops/moe/liger.py @@ -0,0 +1,12 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Liger MoE experts adapter: class replacement via ``LigerExperts``. + +Only effective when the user explicitly puts ``'liger'`` in the +``moe_experts`` backend chain — the default chain is ``('npu',)`` so the +CANN grouped-matmul forward-level replacement is never shadowed. +""" +from __future__ import annotations + +from liger_kernel.transformers import LigerExperts + +__all__ = ['LigerExperts'] diff --git a/src/twinkle/kernel/npu_impls/moe.py b/src/twinkle/kernel/ops/moe/npu.py similarity index 100% rename from src/twinkle/kernel/npu_impls/moe.py rename to src/twinkle/kernel/ops/moe/npu.py diff --git a/src/twinkle/kernel/ops/rms_norm/__init__.py b/src/twinkle/kernel/ops/rms_norm/__init__.py new file mode 100644 index 000000000..19c058d3b --- /dev/null +++ b/src/twinkle/kernel/ops/rms_norm/__init__.py @@ -0,0 +1,53 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""rms_norm / gated_rms_norm op registration. + +Liger's RMSNorm is parameterized per family (casting_mode / offset, see the +adapter classes in liger.py), so the liger load factory dispatches variants +by mapping target: gemma4 -> Gemma4Replacement; gemma* -> GemmaReplacement; +qwen3_5* -> Qwen35Replacement (residual parameterization — llama-cast would +produce NaN); everything else -> the default llama-cast Replacement. +""" +from __future__ import annotations + +from typing import Any + +from ...registry import KernelImpl, is_liger_available, is_npu_available, lazy_import, register_op + + +def _liger_rms_norm_load(target: Any): + name = target if isinstance(target, str) else f'{target.__module__}.{target.__qualname__}' + if 'gemma4' in name: + from .liger import LigerRMSNormGemma4Replacement as cls + elif 'gemma' in name: + from .liger import LigerRMSNormGemmaReplacement as cls + elif 'qwen3_5' in name: + from .liger import LigerRMSNormQwen35Replacement as cls + else: + from .liger import LigerRMSNormReplacement as cls + return cls + + +register_op( + 'rms_norm', + implementations={ + 'npu': KernelImpl( + load=lazy_import('twinkle.kernel.ops.rms_norm.npu:NpuRMSNorm'), + available=is_npu_available, + ), + 'liger': KernelImpl( + load=_liger_rms_norm_load, + available=is_liger_available, + ), + }, +) + +# Gated RMSNorm (qwen3_5 family, forward-level replacement); no liger impl +register_op( + 'gated_rms_norm', + implementations={ + 'npu': KernelImpl( + load=lazy_import('twinkle.kernel.ops.rms_norm.npu:npu_gated_rms_norm_forward'), + available=is_npu_available, + ), + }, +) diff --git a/src/twinkle/kernel/liger_impls/rms_norm.py b/src/twinkle/kernel/ops/rms_norm/liger.py similarity index 100% rename from src/twinkle/kernel/liger_impls/rms_norm.py rename to src/twinkle/kernel/ops/rms_norm/liger.py diff --git a/src/twinkle/kernel/npu_impls/rms_norm.py b/src/twinkle/kernel/ops/rms_norm/npu.py similarity index 100% rename from src/twinkle/kernel/npu_impls/rms_norm.py rename to src/twinkle/kernel/ops/rms_norm/npu.py diff --git a/src/twinkle/kernel/ops/rotary/__init__.py b/src/twinkle/kernel/ops/rotary/__init__.py new file mode 100644 index 000000000..25399e14e --- /dev/null +++ b/src/twinkle/kernel/ops/rotary/__init__.py @@ -0,0 +1,36 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""rotary / multimodal_rotary op registration. + +The liger exclusion for the qwen3_5 family is NOT expressed at the registry +layer — it is decided in config.py, where the rotary targets of those two +families use a ('npu',)-only chain (liger's full-rotation implementation is +incompatible with partial_rotary). +""" +from __future__ import annotations + +from ...registry import KernelImpl, is_liger_available, is_npu_available, lazy_import, register_op + +register_op( + 'rotary', + implementations={ + 'npu': KernelImpl( + load=lazy_import('twinkle.kernel.ops.rotary.npu:npu_apply_rotary_pos_emb'), + available=is_npu_available, + ), + 'liger': KernelImpl( + load=lazy_import('twinkle.kernel.ops.rotary.liger:liger_rotary_pos_emb'), + available=is_liger_available, + ), + }, +) + +# Qwen2.5-VL multimodal rope; no liger impl +register_op( + 'multimodal_rotary', + implementations={ + 'npu': KernelImpl( + load=lazy_import('twinkle.kernel.ops.rotary.npu:npu_apply_multimodal_rotary_pos_emb'), + available=is_npu_available, + ), + }, +) diff --git a/src/twinkle/kernel/ops/rotary/liger.py b/src/twinkle/kernel/ops/rotary/liger.py new file mode 100644 index 000000000..8e8d55961 --- /dev/null +++ b/src/twinkle/kernel/ops/rotary/liger.py @@ -0,0 +1,14 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Liger RoPE adapter: re-export liger_kernel's rotary implementation. + +Loaded lazily via ``registry.lazy_import`` — importing this module pulls in +``liger_kernel``, so it must only be imported after ``is_liger_available()`` +has passed. Compatible with full-rotation families only; partial-RoPE +families (qwen3_5) are excluded by the per-target backend chains in +``config.py``. +""" +from __future__ import annotations + +from liger_kernel.transformers import liger_rotary_pos_emb + +__all__ = ['liger_rotary_pos_emb'] diff --git a/src/twinkle/kernel/npu_impls/rotary.py b/src/twinkle/kernel/ops/rotary/npu.py similarity index 100% rename from src/twinkle/kernel/npu_impls/rotary.py rename to src/twinkle/kernel/ops/rotary/npu.py diff --git a/src/twinkle/kernel/ops/sdpa_attention/__init__.py b/src/twinkle/kernel/ops/sdpa_attention/__init__.py new file mode 100644 index 000000000..589b86065 --- /dev/null +++ b/src/twinkle/kernel/ops/sdpa_attention/__init__.py @@ -0,0 +1,46 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""sdpa_attention op registration: installs the NPU SDPA forward into the +transformers global attention registry (not a plain class/attr replacement, +hence the custom installer). + +The default config references this op via the logical target 'sdpa'; the +logical target is only a mapping label passed to install_sdpa and is never +resolved by the generic replacer. +""" +from __future__ import annotations + +from twinkle import get_logger + +from ...registry import KernelImpl, is_npu_available, lazy_import, register_op + +logger = get_logger() + + +def install_sdpa(model, target, impl) -> None: + """One-shot install of SDPA attention forward (global modeling_utils dict). + + ``AttentionInterface._global_mapping`` is a private transformers attribute; + guard against its removal so an upstream change can't take down the rest + of kernelize. + """ + try: + from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, AttentionInterface + except ImportError: + return + try: + AttentionInterface._global_mapping['sdpa'] = impl + except AttributeError: + logger.warning('[SDPA] AttentionInterface._global_mapping unavailable; skipping') + ALL_ATTENTION_FUNCTIONS['sdpa'] = impl + + +register_op( + 'sdpa_attention', + implementations={ + 'npu': KernelImpl( + load=lazy_import('twinkle.kernel.ops.sdpa_attention.npu:npu_sdpa_attention_forward'), + available=is_npu_available, + ), + }, + installer=install_sdpa, +) diff --git a/src/twinkle/kernel/npu_impls/attention.py b/src/twinkle/kernel/ops/sdpa_attention/npu.py similarity index 100% rename from src/twinkle/kernel/npu_impls/attention.py rename to src/twinkle/kernel/ops/sdpa_attention/npu.py diff --git a/src/twinkle/kernel/ops/swiglu/__init__.py b/src/twinkle/kernel/ops/swiglu/__init__.py new file mode 100644 index 000000000..694d71a2f --- /dev/null +++ b/src/twinkle/kernel/ops/swiglu/__init__.py @@ -0,0 +1,19 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""swiglu op registration: npu / liger implementations, class-forward replacement.""" +from __future__ import annotations + +from ...registry import KernelImpl, is_liger_available, is_npu_available, lazy_import, register_op + +register_op( + 'swiglu', + implementations={ + 'npu': KernelImpl( + load=lazy_import('twinkle.kernel.ops.swiglu.npu:npu_swiglu_forward'), + available=is_npu_available, + ), + 'liger': KernelImpl( + load=lazy_import('twinkle.kernel.ops.swiglu.liger:liger_swiglu_forward'), + available=is_liger_available, + ), + }, +) diff --git a/src/twinkle/kernel/ops/swiglu/liger.py b/src/twinkle/kernel/ops/swiglu/liger.py new file mode 100644 index 000000000..466c77aea --- /dev/null +++ b/src/twinkle/kernel/ops/swiglu/liger.py @@ -0,0 +1,19 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""SwiGLU forward-replacement for Liger Kernel. + +Used as a *function-level* mapping value (string key ``'..forward'``) +so it composes with the existing SwiGLU forward-replacement pattern. Reads only +``gate_proj`` / ``up_proj`` / ``down_proj``, which every HuggingFace SwiGLU MLP +variant (Qwen2MLP, Qwen3MLP, LlamaMLP, MistralMLP, ...) already defines, so no +``__init__`` and no per-instance attribute setup is required. + +For Qwen3-MoE the expert/MLP classes differ; class replacement with Liger's +own ``LigerExperts`` is wired in ``ops/moe/liger.py``. +""" +from __future__ import annotations + +from liger_kernel.ops import LigerSiLUMulFunction + + +def liger_swiglu_forward(self, x): + return self.down_proj(LigerSiLUMulFunction.apply(self.gate_proj(x), self.up_proj(x))) diff --git a/src/twinkle/kernel/npu_impls/swiglu.py b/src/twinkle/kernel/ops/swiglu/npu.py similarity index 100% rename from src/twinkle/kernel/npu_impls/swiglu.py rename to src/twinkle/kernel/ops/swiglu/npu.py diff --git a/src/twinkle/kernel/registry.py b/src/twinkle/kernel/registry.py new file mode 100644 index 000000000..d556f48f4 --- /dev/null +++ b/src/twinkle/kernel/registry.py @@ -0,0 +1,178 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Operator implementation registry: lazy loading and availability of each +backend's impl per op, plus KernelChoice resolution and fallback. + +Only "implementations" are registered here, never "targets" — targets live in +config.py's default mapping and in user mappings. +""" +from __future__ import annotations + +import importlib +import importlib.util +import logging +from dataclasses import dataclass +from typing import Any, Callable + +import torch.nn as nn + +from twinkle import get_logger +from twinkle.utils.device_mesh import Platform +from twinkle.utils.import_utils import exists + +logger = get_logger() + +__all__ = [ + 'KernelImpl', + 'OpDefinition', + 'KernelChoice', + 'Installer', + 'register_op', + 'get_op', + 'resolve_impl', + 'lazy_import', + 'is_npu_available', + 'is_liger_available', + 'exists', +] + + +# ── Data structures ─────────────────────────────────────────────────────── + + +@dataclass(frozen=True) +class KernelImpl: + """One backend's implementation of one op. + + ``load``: lazily loads and returns the final class or function; impls + must not import optional deps (``torch_npu`` / ``liger_kernel`` ...) at + registration time. Receives the mapping target (dotted-path string or + class object) so a backend can specialize per target family (e.g. liger's + RMSNorm dispatches gemma/qwen3_5 variants); factories that don't care + simply ignore it (see ``lazy_import``). + ``available``: whether the current platform/deps/hardware permit this + impl; ``(True, None)`` = usable, ``(False, reason)`` = not usable, fall + through to the next backend. + """ + load: Callable[[Any], Any] + available: Callable[[], tuple[bool, str | None]] + + +Installer = Callable[[nn.Module | None, Any, Any], None] # (model, target, impl) -> None + + +@dataclass(frozen=True) +class OpDefinition: + """All backend implementations of one op, plus its default installer.""" + name: str # 'swiglu' | 'rms_norm' | 'sdpa_attention' | ... + implementations: dict[str, KernelImpl] # backend name -> impl + installer: Installer | None = None # None = use the generic installer + + +@dataclass(frozen=True) +class KernelChoice: + """Selection descriptor in a mapping: which op to use and the backend priority.""" + op: str # references a registered OpDefinition + backends: tuple[str, ...] # priority-ordered, at least one element + installer: Installer | None = None # advanced override; usually None + + +# installer priority: KernelChoice.installer -> OpDefinition.installer -> default_installer + +# ── Registry ────────────────────────────────────────────────────────────── + +_OPS: dict[str, OpDefinition] = {} + + +def register_op( + name: str, + *, + implementations: dict[str, KernelImpl], + installer: Installer | None = None, +) -> None: + """Register an op. Duplicate name / empty implementations -> ValueError.""" + if name in _OPS: + raise ValueError(f'op {name!r} is already registered') + if not implementations: + raise ValueError(f'op {name!r} has no implementations') + _OPS[name] = OpDefinition(name=name, implementations=dict(implementations), installer=installer) + + +def get_op(name: str) -> OpDefinition: + """Fetch an op definition; unregistered -> ValueError (includes the op name).""" + try: + return _OPS[name] + except KeyError: + raise ValueError(f'op {name!r} is not registered') from None + + +def resolve_impl( + op: OpDefinition, + backends: tuple[str, ...], + *, + warn: bool, + target: Any = None, +) -> tuple[Any, str | None]: + """Pick the first available impl in ``backends`` order. + + In turn: backend not registered -> log and skip; available() False -> + log (with reason) and skip; load() raises -> log (with the exception) + and skip; first success -> (impl, backend_name). All failed -> + (None, None): no installer call, original implementation kept. + + warn=True (explicit mapping) -> fallback/failure logs at WARNING; + warn=False (default config path) -> all DEBUG. + ``target`` is passed verbatim to KernelImpl.load() for family-specialized + dispatch. + """ + level = logging.WARNING if warn else logging.DEBUG + for backend in backends: + impl_entry = op.implementations.get(backend) + if impl_entry is None: + logger.log(level, "[kernelize] op '%s': backend '%s' not registered, skipping", op.name, backend) + continue + ok, reason = impl_entry.available() + if not ok: + logger.log(level, "[kernelize] op '%s': backend '%s' unavailable (%s), skipping", op.name, backend, + reason) + continue + try: + return impl_entry.load(target), backend + except Exception as e: + logger.log(level, "[kernelize] op '%s': backend '%s' failed to load (%r), skipping", op.name, backend, e) + return None, None + + +# ── Lazy references and availability helpers (shared by ops/*/__init__.py) ─ + + +def lazy_import(spec: str) -> Callable[[Any], Any]: + """'pkg.mod:attr' -> lazy-load factory; imports the module and getattr + only when called. + + Missing module / attribute -> raises ImportError/AttributeError, caught + by resolve_impl's load() exception branch (falls through). + The factory accepts (and ignores) the mapping target argument, matching + the KernelImpl.load signature. + """ + module_path, _, attr = spec.partition(':') + + def _load(_target: Any = None) -> Any: + return getattr(importlib.import_module(module_path), attr) + + return _load + + +def is_npu_available() -> tuple[bool, str | None]: + """Platform is npu and torch_npu is importable.""" + if Platform.device_prefix() != 'npu': + return False, f"platform is '{Platform.device_prefix()}', not 'npu'" + if importlib.util.find_spec('torch_npu') is None: + return False, 'torch_npu not installed' + return True, None + + +def is_liger_available() -> tuple[bool, str | None]: + """liger_kernel is importable (find_spec only, no real import).""" + if importlib.util.find_spec('liger_kernel') is None: + return False, 'liger_kernel not installed' + return True, None diff --git a/src/twinkle/loss/liger_fused_linear_cross_entropy.py b/src/twinkle/loss/liger_fused_linear_cross_entropy.py index 408d73041..d72fb8369 100644 --- a/src/twinkle/loss/liger_fused_linear_cross_entropy.py +++ b/src/twinkle/loss/liger_fused_linear_cross_entropy.py @@ -40,7 +40,7 @@ dispatch: this file contains no ``torch.cuda`` / ``Platform.is_npu`` probes. The same loss class runs on CUDA (Triton kernel) and Ascend NPU (the ``backends/_ascend`` fused-linear-CE backend), mirroring the bare-impl -philosophy of ``liger_builtin``. +philosophy of the Liger per-layer impls under ``twinkle.kernel.ops``. FSDP2 collective ordering ------------------------ diff --git a/src/twinkle/model/transformers/moe/ep_utils.py b/src/twinkle/model/transformers/moe/ep_utils.py index f64c17964..bb2bed884 100644 --- a/src/twinkle/model/transformers/moe/ep_utils.py +++ b/src/twinkle/model/transformers/moe/ep_utils.py @@ -8,6 +8,8 @@ import torch.distributed as dist from typing import Optional +from twinkle import torch_util + # ========================== comm ========================== class _AllToAll(torch.autograd.Function): @@ -113,7 +115,11 @@ def permute(tokens: torch.Tensor, expert_mask: torch.Tensor): sorted_indices = token_indices.masked_select(expert_mask) # use the mapping to permute the tokens - permuted_input = tokens.index_select(0, sorted_indices) + # NOTE: use advanced indexing instead of index_select — index_select's + # backward (index_add) is broken on some torch_npu/CANN versions + # (aclnnIndexAdd fails for any dtype), while x[idx] backward (index_put) + # works. Mathematically identical. + permuted_input = tokens[sorted_indices] return permuted_input, sorted_indices @@ -203,6 +209,12 @@ def preprocess( ) dist.all_gather_into_tensor(num_global_tokens_per_expert, num_local_tokens_per_expert, group=ep_group) + # The collective may run on a separate stream (e.g. HCCL on NPU) whose + # completion is not always ordered with the host reads (.tolist()/.item()) + # below — observed in practice as garbage split sizes. Force a device sync + # before reading gathered results back on the host. + torch_util.synchronize() + # [ep_size, num_local_experts] start_idx, end_idx = rank * num_local_experts, (rank + 1) * num_local_experts num_global_tokens_per_local_expert = num_global_tokens_per_expert[:, start_idx:end_idx].contiguous() @@ -211,11 +223,13 @@ def preprocess( output_splits = num_global_tokens_per_local_expert.sum(dim=1).tolist() # [num_local_expert] - num_global_sum_tokens_per_local_expert = num_global_tokens_per_local_expert.sum(dim=0).to( - torch.device('cpu'), non_blocking=True) + # NOTE: keep these copies synchronous. With non_blocking=True the CPU + # tensors are read (tolist()/item()) before the async D2H copy lands, + # producing garbage split sizes (observed on NPU). + num_global_sum_tokens_per_local_expert = num_global_tokens_per_local_expert.sum(dim=0).to(torch.device('cpu')) num_global_tokens_per_local_expert = num_global_tokens_per_local_expert.view(-1, num_local_experts).to( - torch.device('cpu'), non_blocking=True) + torch.device('cpu')) return input_splits, output_splits, num_global_tokens_per_local_expert, num_global_sum_tokens_per_local_expert diff --git a/src/twinkle/model/transformers/moe/expert_parallel.py b/src/twinkle/model/transformers/moe/expert_parallel.py index f043e547c..218e7b337 100644 --- a/src/twinkle/model/transformers/moe/expert_parallel.py +++ b/src/twinkle/model/transformers/moe/expert_parallel.py @@ -4,11 +4,11 @@ import inspect import torch import torch.distributed as dist -import torch.nn.functional as F from dataclasses import dataclass from torch import nn from typing import Any, Dict, Iterable, List, Optional, Tuple +from twinkle.kernel.ops import ep_forward from twinkle.model.transformers.moe.ep_utils import preprocess, token_pre_all2all, tokens_post_all2all from twinkle.utils import DeviceMesh @@ -330,53 +330,6 @@ def _install_ep_forward(experts_mod: nn.Module, experts_per_rank: int) -> None: if getattr(experts_mod, '_ep_forward_installed', False): return - def ep_forward( - self, - permuted_tokens: torch.Tensor, - num_global_sum_tokens_per_local_expert: torch.Tensor, - experts_per_rank: int, - ) -> torch.Tensor: - if permuted_tokens.numel() == 0: - # Preserve the autograd edge to token_pre_all2all. Returning a new - # empty tensor can make this rank skip the matching backward - # all-to-all, causing EP collective order divergence. - return permuted_tokens - - input_dtype = permuted_tokens.dtype - - cumsum = torch.zeros(experts_per_rank + 1, dtype=torch.long) - for i in range(experts_per_rank): - cumsum[i + 1] = cumsum[i] + int(num_global_sum_tokens_per_local_expert[i].item()) - - output_chunks = [] - for i in range(experts_per_rank): - start = int(cumsum[i].item()) - end = int(cumsum[i + 1].item()) - expert_in = permuted_tokens[start:end] - if expert_in.numel() == 0: - output_chunks.append(expert_in) - continue - - gate_up = self.gate_up_proj[i] - down = self.down_proj[i] - compute_dtype = gate_up.dtype - if expert_in.dtype != compute_dtype: - expert_in = expert_in.to(compute_dtype) - gate_up_out = F.linear(expert_in, gate_up) - if hasattr(self, '_apply_gate'): - out = self._apply_gate(gate_up_out) - else: - gate, up = gate_up_out.chunk(2, dim=-1) - out = self.act_fn(gate) * up - out = F.linear(out, down) - - if out.dtype != input_dtype: - out = out.to(input_dtype) - output_chunks.append(out) - - return torch.cat( - output_chunks, dim=0) if output_chunks else permuted_tokens.new_empty(0, permuted_tokens.size(-1)) - import types experts_mod.forward = types.MethodType(ep_forward, experts_mod) experts_mod._ep_forward_installed = True diff --git a/src/twinkle/model/transformers/strategy/sequence_parallel/linear_attention_sp.py b/src/twinkle/model/transformers/strategy/sequence_parallel/linear_attention_sp.py index 754849558..ab8db95ed 100644 --- a/src/twinkle/model/transformers/strategy/sequence_parallel/linear_attention_sp.py +++ b/src/twinkle/model/transformers/strategy/sequence_parallel/linear_attention_sp.py @@ -107,7 +107,7 @@ def _torch_causal_conv1d_fn( return out.transpose(1, 2).contiguous() # NPU: fla native causal_conv1d and chunk_gated_delta_rule - # are both patched by twinkle.kernel.npu_impls.fla at model initialization. + # are both patched by twinkle.kernel.ops.fla.npu at model initialization. # No need to set them here - they are already bound on the module. if getattr(mod, '_twinkle_npu_patched', False): return False diff --git a/src/twinkle/patch/gdn_padding_free.py b/src/twinkle/patch/gdn_padding_free.py index 6409990eb..b4e4d1955 100644 --- a/src/twinkle/patch/gdn_padding_free.py +++ b/src/twinkle/patch/gdn_padding_free.py @@ -71,7 +71,7 @@ def _patch_gdn_kernels_for_cu_seqlens( ) -> torch.Tensor: is_npu = getattr(mod, '_twinkle_npu_patched', False) if is_npu: - from twinkle.kernel.npu_impls.fla import npu_causal_conv1d_fn + from twinkle.kernel.ops.fla.npu import npu_causal_conv1d_fn else: causal_conv1d, chunk_gated_delta_rule = _get_flash_linear_attention_kernels() diff --git a/src/twinkle/patch/transformers_fused_ce.py b/src/twinkle/patch/transformers_fused_ce.py index a1eeb47ab..d1e682043 100644 --- a/src/twinkle/patch/transformers_fused_ce.py +++ b/src/twinkle/patch/transformers_fused_ce.py @@ -19,7 +19,8 @@ (``torch.cuda`` / ``Platform.is_npu`` / ``infer_device``). Device selection of the fused-CE kernel is handled by Liger's own ``infer_device`` / ``select_impl`` dispatch (CUDA-Triton vs Ascend backend), exactly like the -``liger_builtin`` bundle emits bare impls. See ``liger.py`` docstring. +Liger per-layer impls under ``twinkle.kernel.ops`` are emitted bare (no +device-conditional wrapping). The patch is applied per-forward via ``apply_context`` (see ``_resolve_task_context`` with ``task='fused_lm_ce'``), so it composes with From a598f67a8314bb2eb87552b0affa6557e726b964 Mon Sep 17 00:00:00 2001 From: root Date: Thu, 30 Jul 2026 17:03:51 +0800 Subject: [PATCH 2/6] style(kernel): apply pre-commit lint (yapf/isort/quote-fixer) Formatting only, no behavior change: yapf + isort + double-quote-string -fixer per .pre-commit-config.yaml; fixes E501 in core.py. --- src/twinkle/kernel/config.py | 8 ++++---- src/twinkle/kernel/core.py | 20 +++++++++---------- src/twinkle/kernel/ops/__init__.py | 3 +-- src/twinkle/kernel/ops/geglu/__init__.py | 3 ++- src/twinkle/kernel/ops/moe/__init__.py | 9 ++++++--- src/twinkle/kernel/ops/rms_norm/__init__.py | 3 ++- src/twinkle/kernel/ops/rotary/__init__.py | 9 ++++++--- .../kernel/ops/sdpa_attention/__init__.py | 4 ++-- src/twinkle/kernel/ops/swiglu/__init__.py | 6 ++++-- src/twinkle/kernel/registry.py | 7 ++----- 10 files changed, 39 insertions(+), 33 deletions(-) diff --git a/src/twinkle/kernel/config.py b/src/twinkle/kernel/config.py index 8e79a70b2..67d623e7a 100644 --- a/src/twinkle/kernel/config.py +++ b/src/twinkle/kernel/config.py @@ -21,15 +21,15 @@ ('transformers.models.qwen2.modeling_qwen2', 'Qwen2RMSNorm', ['Qwen2MLP'], []), ('transformers.models.qwen3.modeling_qwen3', 'Qwen3RMSNorm', ['Qwen3MLP'], []), ('transformers.models.qwen2_5_vl.modeling_qwen2_5_vl', 'Qwen2_5_VLRMSNorm', ['Qwen2MLP', 'Qwen2_5_VLMLP'], []), - ('transformers.models.qwen3_5.modeling_qwen3_5', 'Qwen3_5RMSNorm', ['Qwen3_5MLP', 'Qwen3_5VisionMLP'], - ['Qwen3_5VisionRMSNorm']), + ('transformers.models.qwen3_5.modeling_qwen3_5', 'Qwen3_5RMSNorm', ['Qwen3_5MLP', + 'Qwen3_5VisionMLP'], ['Qwen3_5VisionRMSNorm']), ] # (module, rms_cls, mlp_cls, experts_cls, block_cls) _QWEN_MOE = [ ('transformers.models.qwen3_moe.modeling_qwen3_moe', 'Qwen3MoeRMSNorm', 'Qwen3MoeMLP', 'Qwen3MoeExperts', 'Qwen3MoeSparseMoeBlock'), - ('transformers.models.qwen3_5_moe.modeling_qwen3_5_moe', 'Qwen3_5MoeRMSNorm', 'Qwen3_5MoeMLP', - 'Qwen3_5MoeExperts', 'Qwen3_5MoeSparseMoeBlock'), + ('transformers.models.qwen3_5_moe.modeling_qwen3_5_moe', 'Qwen3_5MoeRMSNorm', 'Qwen3_5MoeMLP', 'Qwen3_5MoeExperts', + 'Qwen3_5MoeSparseMoeBlock'), ] # qwen3_5 Gated RMSNorm (forward-level replacement, no liger impl) _QWEN_GATED_RMS = { diff --git a/src/twinkle/kernel/core.py b/src/twinkle/kernel/core.py index 803e35efc..fea01fe15 100644 --- a/src/twinkle/kernel/core.py +++ b/src/twinkle/kernel/core.py @@ -19,7 +19,6 @@ from typing import Any from twinkle import get_logger - from .registry import KernelChoice, get_op, resolve_impl logger = get_logger() @@ -156,10 +155,10 @@ def _install_dotted(model: nn.Module, target: str, impl, *, warn: bool = False) if target.startswith('transformers.'): level = logging.WARNING if warn else logging.DEBUG hint = ' (explicit mapping: check for typos)' if warn else '' - logger.log(level, "[kernelize] target %r unresolvable (%r); family not installed, skipping%s", - target, e, hint) + logger.log(level, '[kernelize] target %r unresolvable (%r); family not installed, skipping%s', target, e, + hint) return False - raise ValueError(f"Cannot resolve mapping target {target!r} with the default installer " + raise ValueError(f'Cannot resolve mapping target {target!r} with the default installer ' f'(logical targets require a custom installer): {e!r}') from e if isinstance(resolved, type) and issubclass(resolved, nn.Module): _replace_class(model, resolved, impl) @@ -200,9 +199,9 @@ def _installer_name(installer) -> str: def _log_all_unavailable(target: Any, choice: KernelChoice, *, warn: bool) -> None: level = logging.WARNING if warn else logging.DEBUG - logger.log(level, "[kernelize] target %s: no available backend for op '%s' (tried: %s); " - 'keeping the original implementation', _target_name(target), choice.op, - ', '.join(choice.backends)) + logger.log( + level, "[kernelize] target %s: no available backend for op '%s' (tried: %s); " + 'keeping the original implementation', _target_name(target), choice.op, ', '.join(choice.backends)) def kernelize(model: nn.Module, mapping: dict | None = None) -> nn.Module: @@ -256,12 +255,13 @@ def kernelize(model: nn.Module, mapping: dict | None = None) -> nn.Module: if installer is default_installer: installed = installer(model, target, impl, warn=warn) else: - installed = installer(model, target, impl) # failures propagate, never swallowed (half-installed state stays visible) + installed = installer(model, target, + impl) # failures propagate, never swallowed (half-installed state stays visible) if installed is False: continue if isinstance(replacement, KernelChoice): - logger.info('[kernelize] target=%s op=%s backend=%s installer=%s', _target_name(target), op.name, - backend, _installer_name(installer)) + logger.info('[kernelize] target=%s op=%s backend=%s installer=%s', _target_name(target), op.name, backend, + _installer_name(installer)) else: impl_repr = getattr(impl, '__qualname__', repr(impl)) logger.info('[kernelize] target=%s impl=%s installer=%s', _target_name(target), impl_repr, diff --git a/src/twinkle/kernel/ops/__init__.py b/src/twinkle/kernel/ops/__init__.py index 0d9d39925..14ce9cd1e 100644 --- a/src/twinkle/kernel/ops/__init__.py +++ b/src/twinkle/kernel/ops/__init__.py @@ -11,10 +11,9 @@ ``twinkle.kernel.registry`` — registration modules are lightweight (lazy references + availability checks only, no optional-dependency imports). """ -from .ep import EpExpertsGmm, ep_forward - # Trigger built-in op registration (must happen before the first kernelize() call) from . import fla, geglu, moe, rms_norm, rotary, sdpa_attention, swiglu # noqa: F401,E402 +from .ep import EpExpertsGmm, ep_forward __all__ = [ 'EpExpertsGmm', diff --git a/src/twinkle/kernel/ops/geglu/__init__.py b/src/twinkle/kernel/ops/geglu/__init__.py index bb1e8c47c..5a8c7c0a8 100644 --- a/src/twinkle/kernel/ops/geglu/__init__.py +++ b/src/twinkle/kernel/ops/geglu/__init__.py @@ -7,7 +7,8 @@ register_op( 'geglu', implementations={ - 'liger': KernelImpl( + 'liger': + KernelImpl( load=lazy_import('twinkle.kernel.ops.geglu.liger:liger_geglu_forward'), available=is_liger_available, ), diff --git a/src/twinkle/kernel/ops/moe/__init__.py b/src/twinkle/kernel/ops/moe/__init__.py index 1839c48e1..0233849fe 100644 --- a/src/twinkle/kernel/ops/moe/__init__.py +++ b/src/twinkle/kernel/ops/moe/__init__.py @@ -15,11 +15,13 @@ register_op( 'moe_experts', implementations={ - 'npu': KernelImpl( + 'npu': + KernelImpl( load=lazy_import('twinkle.kernel.ops.moe.npu:npu_packed_moe_experts_forward'), available=is_npu_available, ), - 'liger': KernelImpl( + 'liger': + KernelImpl( load=lazy_import('twinkle.kernel.ops.moe.liger:LigerExperts'), available=is_liger_available, ), @@ -29,7 +31,8 @@ register_op( 'moe_block', implementations={ - 'npu': KernelImpl( + 'npu': + KernelImpl( load=lazy_import('twinkle.kernel.ops.moe.npu:npu_qwen3_5_moe_sparse_block_forward'), available=is_npu_available, ), diff --git a/src/twinkle/kernel/ops/rms_norm/__init__.py b/src/twinkle/kernel/ops/rms_norm/__init__.py index 19c058d3b..97b9d263d 100644 --- a/src/twinkle/kernel/ops/rms_norm/__init__.py +++ b/src/twinkle/kernel/ops/rms_norm/__init__.py @@ -45,7 +45,8 @@ def _liger_rms_norm_load(target: Any): register_op( 'gated_rms_norm', implementations={ - 'npu': KernelImpl( + 'npu': + KernelImpl( load=lazy_import('twinkle.kernel.ops.rms_norm.npu:npu_gated_rms_norm_forward'), available=is_npu_available, ), diff --git a/src/twinkle/kernel/ops/rotary/__init__.py b/src/twinkle/kernel/ops/rotary/__init__.py index 25399e14e..f47a31910 100644 --- a/src/twinkle/kernel/ops/rotary/__init__.py +++ b/src/twinkle/kernel/ops/rotary/__init__.py @@ -13,11 +13,13 @@ register_op( 'rotary', implementations={ - 'npu': KernelImpl( + 'npu': + KernelImpl( load=lazy_import('twinkle.kernel.ops.rotary.npu:npu_apply_rotary_pos_emb'), available=is_npu_available, ), - 'liger': KernelImpl( + 'liger': + KernelImpl( load=lazy_import('twinkle.kernel.ops.rotary.liger:liger_rotary_pos_emb'), available=is_liger_available, ), @@ -28,7 +30,8 @@ register_op( 'multimodal_rotary', implementations={ - 'npu': KernelImpl( + 'npu': + KernelImpl( load=lazy_import('twinkle.kernel.ops.rotary.npu:npu_apply_multimodal_rotary_pos_emb'), available=is_npu_available, ), diff --git a/src/twinkle/kernel/ops/sdpa_attention/__init__.py b/src/twinkle/kernel/ops/sdpa_attention/__init__.py index 589b86065..386d6b463 100644 --- a/src/twinkle/kernel/ops/sdpa_attention/__init__.py +++ b/src/twinkle/kernel/ops/sdpa_attention/__init__.py @@ -10,7 +10,6 @@ from __future__ import annotations from twinkle import get_logger - from ...registry import KernelImpl, is_npu_available, lazy_import, register_op logger = get_logger() @@ -37,7 +36,8 @@ def install_sdpa(model, target, impl) -> None: register_op( 'sdpa_attention', implementations={ - 'npu': KernelImpl( + 'npu': + KernelImpl( load=lazy_import('twinkle.kernel.ops.sdpa_attention.npu:npu_sdpa_attention_forward'), available=is_npu_available, ), diff --git a/src/twinkle/kernel/ops/swiglu/__init__.py b/src/twinkle/kernel/ops/swiglu/__init__.py index 694d71a2f..de9ca3af5 100644 --- a/src/twinkle/kernel/ops/swiglu/__init__.py +++ b/src/twinkle/kernel/ops/swiglu/__init__.py @@ -7,11 +7,13 @@ register_op( 'swiglu', implementations={ - 'npu': KernelImpl( + 'npu': + KernelImpl( load=lazy_import('twinkle.kernel.ops.swiglu.npu:npu_swiglu_forward'), available=is_npu_available, ), - 'liger': KernelImpl( + 'liger': + KernelImpl( load=lazy_import('twinkle.kernel.ops.swiglu.liger:liger_swiglu_forward'), available=is_liger_available, ), diff --git a/src/twinkle/kernel/registry.py b/src/twinkle/kernel/registry.py index d556f48f4..adb52a4de 100644 --- a/src/twinkle/kernel/registry.py +++ b/src/twinkle/kernel/registry.py @@ -10,11 +10,10 @@ import importlib import importlib.util import logging +import torch.nn as nn from dataclasses import dataclass from typing import Any, Callable -import torch.nn as nn - from twinkle import get_logger from twinkle.utils.device_mesh import Platform from twinkle.utils.import_utils import exists @@ -35,7 +34,6 @@ 'exists', ] - # ── Data structures ─────────────────────────────────────────────────────── @@ -132,8 +130,7 @@ def resolve_impl( continue ok, reason = impl_entry.available() if not ok: - logger.log(level, "[kernelize] op '%s': backend '%s' unavailable (%s), skipping", op.name, backend, - reason) + logger.log(level, "[kernelize] op '%s': backend '%s' unavailable (%s), skipping", op.name, backend, reason) continue try: return impl_entry.load(target), backend From f7db8442ff2e56d326ee233fce4f544e9ad453c7 Mon Sep 17 00:00:00 2001 From: clc Date: Mon, 3 Aug 2026 16:35:39 +0800 Subject: [PATCH 3/6] test(kernel): adapt tests to register-dispatch refactor - npu_impls/ -> ops/: update imports to twinkle.kernel.ops.*.npu (impl functions moved, not removed; all CPU-importable via lazy torch_npu), plus NPU-skipif numerical parity cases - test_builtin.py / test_liger_builtin.py: dropped, builtin module retired - test_replace.py: drop _replace_attr cases (helper removed), keep _replace_class coverage - test_resolve_value.py: retarget to resolve_direct_value (passthrough + HubRef delegation); device-dict mechanism superseded by KernelChoice chains - test_kernelize.py: drop device-dict cases, npu_builtin auto-apply case now covers DEFAULT_KERNEL_CONFIG path - test_public_api.py: expect new __all__ (DEFAULT_KERNEL_CONFIG, KernelChoice, hub, kernelize) - test_registry.py: new coverage for registry/config resolution, installer priority, dotted targets, KernelChoice fallback chains --- tests/kernel/npu_impls/test_swiglu.py | 12 - tests/kernel/{npu_impls => ops}/__init__.py | 0 .../{npu_impls => ops}/test_attention.py | 4 +- tests/kernel/{npu_impls => ops}/test_fla.py | 10 +- tests/kernel/{npu_impls => ops}/test_moe.py | 2 +- .../{npu_impls => ops}/test_rms_norm.py | 26 +- .../kernel/{npu_impls => ops}/test_rotary.py | 4 +- tests/kernel/ops/test_swiglu.py | 43 ++ tests/kernel/test_builtin.py | 90 ---- tests/kernel/test_kernelize.py | 68 +-- tests/kernel/test_liger_builtin.py | 116 ----- tests/kernel/test_public_api.py | 9 +- tests/kernel/test_registry.py | 419 ++++++++++++++++++ tests/kernel/test_replace.py | 40 +- tests/kernel/test_resolve_value.py | 42 +- 15 files changed, 523 insertions(+), 362 deletions(-) delete mode 100644 tests/kernel/npu_impls/test_swiglu.py rename tests/kernel/{npu_impls => ops}/__init__.py (100%) rename tests/kernel/{npu_impls => ops}/test_attention.py (73%) rename tests/kernel/{npu_impls => ops}/test_fla.py (89%) rename tests/kernel/{npu_impls => ops}/test_moe.py (89%) rename tests/kernel/{npu_impls => ops}/test_rms_norm.py (51%) rename tests/kernel/{npu_impls => ops}/test_rotary.py (84%) create mode 100644 tests/kernel/ops/test_swiglu.py delete mode 100644 tests/kernel/test_builtin.py delete mode 100644 tests/kernel/test_liger_builtin.py create mode 100644 tests/kernel/test_registry.py diff --git a/tests/kernel/npu_impls/test_swiglu.py b/tests/kernel/npu_impls/test_swiglu.py deleted file mode 100644 index d4ec2da9a..000000000 --- a/tests/kernel/npu_impls/test_swiglu.py +++ /dev/null @@ -1,12 +0,0 @@ -def test_swiglu_imports(): - from twinkle.kernel.npu_impls.swiglu import npu_swiglu_forward - assert callable(npu_swiglu_forward) - - -def test_swiglu_signature(): - import inspect - - from twinkle.kernel.npu_impls.swiglu import npu_swiglu_forward - - params = list(inspect.signature(npu_swiglu_forward).parameters) - assert params == ['self', 'hidden_state'] \ No newline at end of file diff --git a/tests/kernel/npu_impls/__init__.py b/tests/kernel/ops/__init__.py similarity index 100% rename from tests/kernel/npu_impls/__init__.py rename to tests/kernel/ops/__init__.py diff --git a/tests/kernel/npu_impls/test_attention.py b/tests/kernel/ops/test_attention.py similarity index 73% rename from tests/kernel/npu_impls/test_attention.py rename to tests/kernel/ops/test_attention.py index ed916dba1..668cda508 100644 --- a/tests/kernel/npu_impls/test_attention.py +++ b/tests/kernel/ops/test_attention.py @@ -1,12 +1,12 @@ def test_attention_imports(): - from twinkle.kernel.npu_impls.attention import npu_sdpa_attention_forward + from twinkle.kernel.ops.sdpa_attention.npu import npu_sdpa_attention_forward assert callable(npu_sdpa_attention_forward) def test_attention_signature(): import inspect - from twinkle.kernel.npu_impls.attention import npu_sdpa_attention_forward + from twinkle.kernel.ops.sdpa_attention.npu import npu_sdpa_attention_forward sig = inspect.signature(npu_sdpa_attention_forward) params = list(sig.parameters) diff --git a/tests/kernel/npu_impls/test_fla.py b/tests/kernel/ops/test_fla.py similarity index 89% rename from tests/kernel/npu_impls/test_fla.py rename to tests/kernel/ops/test_fla.py index a4fb9e4ef..f13340a84 100644 --- a/tests/kernel/npu_impls/test_fla.py +++ b/tests/kernel/ops/test_fla.py @@ -1,11 +1,11 @@ def test_fla_imports(): - from twinkle.kernel.npu_impls.fla import apply_qwen3_5_fla + from twinkle.kernel.ops.fla.npu import apply_qwen3_5_fla assert callable(apply_qwen3_5_fla) def test_fla_disabled_by_env(monkeypatch): monkeypatch.setenv('TWINKLE_NPU_FLA', '0') - from twinkle.kernel.npu_impls.fla import apply_qwen3_5_fla + from twinkle.kernel.ops.fla.npu import apply_qwen3_5_fla # With env=0, function returns 0 (no-op) without raising assert apply_qwen3_5_fla(None) == 0 @@ -14,7 +14,7 @@ def test_fla_skips_when_no_torch_npu(monkeypatch): import sys monkeypatch.setenv('TWINKLE_NPU_FLA', '1') monkeypatch.setitem(sys.modules, 'torch_npu', None) # forces ImportError on import - from twinkle.kernel.npu_impls import fla as fla_mod + from twinkle.kernel.ops.fla import npu as fla_mod # Reload-tolerant: should return 0 when torch_npu is missing. assert fla_mod.apply_qwen3_5_fla(None) == 0 @@ -34,7 +34,7 @@ def test_fla_does_not_flip_flag_when_fla_missing(monkeypatch): spec = importlib.util.spec_from_loader('torch_npu', loader=None) fake_npu = importlib.util.module_from_spec(spec) monkeypatch.setitem(sys.modules, 'torch_npu', fake_npu) - # Force the fla-backed operator import to fail. twinkle.kernel.npu_impls.fla + # Force the fla-backed operator import to fail. twinkle.kernel.ops.fla.npu # imports ``fla.modules.convolution`` and ``fla.ops.gated_delta_rule`` # lazily inside ``apply_qwen3_5_fla``; stubbing the top-level ``fla`` package # as None makes both imports raise ImportError. @@ -42,7 +42,7 @@ def test_fla_does_not_flip_flag_when_fla_missing(monkeypatch): original_flag = tui.is_flash_linear_attention_available try: - from twinkle.kernel.npu_impls.fla import apply_qwen3_5_fla + from twinkle.kernel.ops.fla.npu import apply_qwen3_5_fla assert apply_qwen3_5_fla(None) == 0 assert tui.is_flash_linear_attention_available is original_flag, ( 'is_flash_linear_attention_available was flipped to True while the ' diff --git a/tests/kernel/npu_impls/test_moe.py b/tests/kernel/ops/test_moe.py similarity index 89% rename from tests/kernel/npu_impls/test_moe.py rename to tests/kernel/ops/test_moe.py index 34452b61c..4b210029c 100644 --- a/tests/kernel/npu_impls/test_moe.py +++ b/tests/kernel/ops/test_moe.py @@ -1,5 +1,5 @@ def test_moe_imports(): - from twinkle.kernel.npu_impls.moe import ( + from twinkle.kernel.ops.moe.npu import ( GmmFunction, npu_grouped_mm, npu_packed_moe_experts_forward, diff --git a/tests/kernel/npu_impls/test_rms_norm.py b/tests/kernel/ops/test_rms_norm.py similarity index 51% rename from tests/kernel/npu_impls/test_rms_norm.py rename to tests/kernel/ops/test_rms_norm.py index 184d7ef70..78800417b 100644 --- a/tests/kernel/npu_impls/test_rms_norm.py +++ b/tests/kernel/ops/test_rms_norm.py @@ -11,21 +11,21 @@ def test_imports(): """NpuRMSNorm and npu_gated_rms_norm_forward import without torch_npu.""" - from twinkle.kernel.npu_impls.rms_norm import NpuRMSNorm, npu_gated_rms_norm_forward + from twinkle.kernel.ops.rms_norm.npu import NpuRMSNorm, npu_gated_rms_norm_forward assert NpuRMSNorm is not None assert callable(npu_gated_rms_norm_forward) def test_npu_rmsnorm_has_no_init(): """Class-replacement contract: NpuRMSNorm must not define its own __init__.""" - from twinkle.kernel.npu_impls.rms_norm import NpuRMSNorm + from twinkle.kernel.ops.rms_norm.npu import NpuRMSNorm # If NpuRMSNorm defines __init__, it'd appear in NpuRMSNorm.__dict__ assert '__init__' not in NpuRMSNorm.__dict__ @pytest.mark.skipif(not _NPU_OK, reason='torch_npu unavailable') def test_npu_rmsnorm_forward_runs_on_npu(): - from twinkle.kernel.npu_impls.rms_norm import NpuRMSNorm + from twinkle.kernel.ops.rms_norm.npu import NpuRMSNorm class _Orig(nn.Module): def __init__(self): @@ -37,4 +37,22 @@ def __init__(self): m.__class__ = NpuRMSNorm x = torch.randn(2, 8, device='npu') y = m(x) - assert y.shape == (2, 8) \ No newline at end of file + assert y.shape == (2, 8) + + +@pytest.mark.skipif(not _NPU_OK, reason='torch_npu unavailable') +def test_npu_rmsnorm_matches_torch_reference(): + """Numerical parity: npu_rms_norm output ~= x * rsqrt(mean(x^2)+eps) * weight.""" + from twinkle.kernel.ops.rms_norm.npu import NpuRMSNorm + + class _Orig(nn.Module): + def __init__(self): + super().__init__() + self.weight = nn.Parameter(torch.randn(64)) + self.variance_epsilon = 1e-6 + + m = _Orig().to('npu') + m.__class__ = NpuRMSNorm + x = torch.randn(4, 64, device='npu') + ref = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + 1e-6) * m.weight + torch.testing.assert_close(m(x), ref, rtol=1e-4, atol=1e-5) \ No newline at end of file diff --git a/tests/kernel/npu_impls/test_rotary.py b/tests/kernel/ops/test_rotary.py similarity index 84% rename from tests/kernel/npu_impls/test_rotary.py rename to tests/kernel/ops/test_rotary.py index 460d0fc33..b74acdc64 100644 --- a/tests/kernel/npu_impls/test_rotary.py +++ b/tests/kernel/ops/test_rotary.py @@ -1,5 +1,5 @@ def test_rotary_imports(): - from twinkle.kernel.npu_impls.rotary import ( + from twinkle.kernel.ops.rotary.npu import ( npu_apply_multimodal_rotary_pos_emb, npu_apply_rotary_pos_emb, ) @@ -11,7 +11,7 @@ def test_rotary_signature_compat(): """Signature must match HF apply_rotary_pos_emb so setattr swap is safe.""" import inspect - from twinkle.kernel.npu_impls.rotary import npu_apply_rotary_pos_emb + from twinkle.kernel.ops.rotary.npu import npu_apply_rotary_pos_emb sig = inspect.signature(npu_apply_rotary_pos_emb) params = list(sig.parameters) diff --git a/tests/kernel/ops/test_swiglu.py b/tests/kernel/ops/test_swiglu.py new file mode 100644 index 000000000..06009c084 --- /dev/null +++ b/tests/kernel/ops/test_swiglu.py @@ -0,0 +1,43 @@ +import pytest +import torch +import torch.nn as nn + +try: + import torch_npu # noqa: F401 + _NPU_OK = True +except ImportError: + _NPU_OK = False + + +def test_swiglu_imports(): + from twinkle.kernel.ops.swiglu.npu import npu_swiglu_forward + assert callable(npu_swiglu_forward) + + +def test_swiglu_signature(): + import inspect + + from twinkle.kernel.ops.swiglu.npu import npu_swiglu_forward + + params = list(inspect.signature(npu_swiglu_forward).parameters) + assert params == ['self', 'hidden_state'] + + +@pytest.mark.skipif(not _NPU_OK, reason='torch_npu unavailable') +def test_npu_swiglu_matches_torch_reference(): + """Numerical parity: npu_swiglu(cat(gate, up)) ~= silu(gate(x)) * up(x), then down_proj.""" + import torch.nn.functional as F + + from twinkle.kernel.ops.swiglu.npu import npu_swiglu_forward + + class _Mlp(nn.Module): + def __init__(self): + super().__init__() + self.gate_proj = nn.Linear(32, 64, bias=False) + self.up_proj = nn.Linear(32, 64, bias=False) + self.down_proj = nn.Linear(64, 32, bias=False) + + m = _Mlp().to('npu') + x = torch.randn(2, 32, device='npu') + ref = m.down_proj(F.silu(m.gate_proj(x)) * m.up_proj(x)) + torch.testing.assert_close(npu_swiglu_forward(m, x), ref, rtol=1e-4, atol=1e-5) \ No newline at end of file diff --git a/tests/kernel/test_builtin.py b/tests/kernel/test_builtin.py deleted file mode 100644 index 38d83915d..000000000 --- a/tests/kernel/test_builtin.py +++ /dev/null @@ -1,90 +0,0 @@ -import importlib.machinery -import sys -import types - -import torch -import torch.nn as nn - -import pytest - - -def _fake_module(name: str): - module = types.ModuleType(name) - module.__spec__ = importlib.machinery.ModuleSpec(name, loader=None) - return module - - -def test_npu_builtin_returns_dict(): - from twinkle.kernel.builtin import npu_builtin - bundle = npu_builtin() - assert isinstance(bundle, dict) - assert len(bundle) > 0 - - -def test_npu_builtin_values_are_npu_gated(): - """Every value in npu_builtin() must be wrapped in {'npu': ...} so it's - safely no-op on CUDA/CPU.""" - from twinkle.kernel.builtin import npu_builtin - for key, value in npu_builtin().items(): - assert isinstance(value, dict), f'value for {key!r} is not a device-dict' - assert 'npu' in value, f'value for {key!r} is missing npu entry' - - -def test_npu_builtin_compose_with_user_override(): - """User-supplied keys override the builtin (via plain dict merge).""" - from twinkle.kernel.builtin import npu_builtin - sentinel = object() - merged = {**npu_builtin(), 'fake.module.path.fn': sentinel} - assert merged['fake.module.path.fn'] is sentinel - - -def test_npu_builtin_safe_on_cpu_model(): - """kernelize(cpu_model, npu_builtin()) must not raise and not modify.""" - from twinkle.kernel import kernelize - from twinkle.kernel.builtin import npu_builtin - - m = nn.Sequential(nn.Linear(2, 2)) - pre_type = type(m[0]) - out = kernelize(m, npu_builtin()) - assert out is m - assert type(m[0]) is pre_type # no replacement happened (cpu device) - - -def test_npu_builtin_skips_missing_modeling_modules(): - """If transformers.models.qwen3_5 is not installed, the bundle must - still produce a dict (with whatever subset is available).""" - from twinkle.kernel.builtin import npu_builtin - bundle = npu_builtin() # must not raise - assert isinstance(bundle, dict) - - -def test_npu_builtin_does_not_overwrite_global_sdpa_on_non_npu_host(monkeypatch): - """Calling npu_builtin() on a CUDA/CPU host must not contaminate the - global HF SDPA registry. The NPU impl inverts boolean masks, which is - wrong for non-NPU execution.""" - from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS - from twinkle.kernel.builtin import npu_builtin - from twinkle.utils.device_mesh import Platform - - monkeypatch.setattr(Platform, 'device_prefix', staticmethod(lambda platform=None: 'cuda')) - original = ALL_ATTENTION_FUNCTIONS.get('sdpa') - npu_builtin() - assert ALL_ATTENTION_FUNCTIONS.get('sdpa') is original - - -def test_npu_builtin_skips_side_effects_on_non_npu_platform(monkeypatch): - from twinkle.kernel import builtin - from twinkle.kernel.npu_impls import fla - from twinkle.utils.device_mesh import Platform - - installs = [] - fla_calls = [] - monkeypatch.setattr(Platform, 'device_prefix', staticmethod(lambda platform=None: 'cuda')) - monkeypatch.setitem(sys.modules, 'torch_npu', _fake_module('torch_npu')) - monkeypatch.setattr(builtin, '_install_sdpa', lambda impl: installs.append(impl)) - monkeypatch.setattr(fla, 'apply_qwen3_5_fla', lambda model: fla_calls.append(model)) - - builtin.npu_builtin(nn.Linear(1, 1)) - - assert installs == [] - assert fla_calls == [] diff --git a/tests/kernel/test_kernelize.py b/tests/kernel/test_kernelize.py index 323d2ee7b..e52bc67ca 100644 --- a/tests/kernel/test_kernelize.py +++ b/tests/kernel/test_kernelize.py @@ -49,74 +49,30 @@ def test_kernelize_string_key_calls_setattr(): sys.modules.pop(mod_name, None) -def test_kernelize_device_dict_match(monkeypatch): - from twinkle.utils.device_mesh import Platform - - parent = nn.Sequential(_SrcLayer()) - monkeypatch.setattr(Platform, 'device_prefix', staticmethod(lambda platform=None: 'cpu')) - - kernelize(parent, {_SrcLayer: {'cpu': _DstLayer, 'npu': nn.Identity}}) - - assert type(parent[0]) is _DstLayer - - -def test_kernelize_uses_platform_device_prefix(monkeypatch): - from twinkle.utils.device_mesh import Platform - - parent = nn.Sequential(_SrcLayer()) # params may still be CPU before FSDP placement - monkeypatch.setattr(Platform, 'device_prefix', staticmethod(lambda platform=None: 'npu')) - - kernelize(parent, {_SrcLayer: {'npu': _DstLayer}}) - - assert type(parent[0]) is _DstLayer - - -def test_kernelize_device_dict_miss_skips_silently(monkeypatch): - from twinkle.utils.device_mesh import Platform - - parent = nn.Sequential(_SrcLayer()) - monkeypatch.setattr(Platform, 'device_prefix', staticmethod(lambda platform=None: 'cpu')) - - kernelize(parent, {_SrcLayer: {'npu': _DstLayer}}) - - assert type(parent[0]) is _SrcLayer - - def test_kernelize_rejects_unknown_key_type(): - with pytest.raises(TypeError, match='Unsupported mapping key'): + with pytest.raises(TypeError, match='Unsupported mapping target'): kernelize(nn.Linear(1, 1), {42: _DstLayer}) -def test_kernelize_no_mapping_on_npu_uses_npu_builtin(monkeypatch): - """kernelize(model) with no mapping auto-detects NPU and applies npu_builtin.""" - from twinkle.utils.device_mesh import Platform - import twinkle.kernel.builtin as builtin - - monkeypatch.setattr(Platform, 'device_prefix', staticmethod(lambda platform=None: 'npu')) - monkeypatch.setattr(builtin, 'npu_builtin', lambda model=None: {_SrcLayer: _DstLayer}) - - parent = nn.Sequential(_SrcLayer()) - out = kernelize(parent) - assert out is parent - assert type(parent[0]) is _DstLayer +def test_kernelize_no_mapping_applies_default_config(monkeypatch, caplog): + """kernelize(model) with no mapping applies DEFAULT_KERNEL_CONFIG. + On a CPU platform with no liger_kernel installed, every default entry is + unavailable or its family is missing -> model unchanged, and (default + config path) no WARNING-level noise. + """ + import logging -def test_kernelize_no_mapping_on_non_npu_is_noop(monkeypatch): - """kernelize(model) with no mapping on a non-NPU device must not touch the - model and must not invoke npu_builtin (avoiding its side effects).""" from twinkle.utils.device_mesh import Platform - import twinkle.kernel.builtin as builtin - called = [] - monkeypatch.setattr(Platform, 'device_prefix', staticmethod(lambda platform=None: 'cuda')) - monkeypatch.setattr(builtin, 'npu_builtin', - lambda model=None: called.append(model) or {_SrcLayer: _DstLayer}) + monkeypatch.setattr(Platform, 'device_prefix', staticmethod(lambda platform=None: 'cpu')) parent = nn.Sequential(_SrcLayer()) - out = kernelize(parent) + with caplog.at_level(logging.WARNING): + out = kernelize(parent) assert out is parent assert type(parent[0]) is _SrcLayer - assert called == [] + assert not [r for r in caplog.records if r.levelno >= logging.WARNING] def test_kernelize_loads_hub_ref(monkeypatch): diff --git a/tests/kernel/test_liger_builtin.py b/tests/kernel/test_liger_builtin.py deleted file mode 100644 index 893df2059..000000000 --- a/tests/kernel/test_liger_builtin.py +++ /dev/null @@ -1,116 +0,0 @@ -import os - -import torch -import torch.nn as nn -import pytest - -import twinkle.kernel as k - -try: - import liger_kernel # noqa: F401 - _HAS_LIGER = True -except ImportError: - _HAS_LIGER = False - -requires_liger = pytest.mark.skipif(not _HAS_LIGER, reason='liger_kernel not installed') - - -@requires_liger -def test_liger_builtin_returns_dict(): - bundle = k.liger_builtin() - assert isinstance(bundle, dict) - assert len(bundle) > 0 - - -@requires_liger -def test_liger_builtin_values_are_bare_impls(): - """unlike npu_builtin (which wraps in {'npu': impl}), liger_builtin emits - *bare* impls because Liger self-dispatches across CUDA/NPU internally.""" - for key, value in k.liger_builtin().items(): - assert not isinstance(value, dict), f'value for {key!r} must be bare, got a device-dict' - - -@requires_liger -def test_liger_builtin_compose_with_user_override(): - sentinel = object() - merged = {**k.liger_builtin(), 'fake.module.path.fn': sentinel} - assert merged['fake.module.path.fn'] is sentinel - - -@requires_liger -def test_liger_builtin_skips_missing_modeling_modules(monkeypatch): - """If every transformers modeling module is unimportable, the bundle must - still return a dict (empty) rather than raising.""" - import twinkle.kernel.liger as liger_mod - - monkeypatch.setattr(liger_mod, '_import_optional', lambda name: None) - bundle = liger_mod.liger_builtin() - assert isinstance(bundle, dict) - assert bundle == {} - - -@requires_liger -def test_liger_builtin_no_global_side_effects(): - """Liger's bundle must not mutate process-global state — contrast with - npu_builtin which installs a global SDPA override. Concretely, the - ``torch.nn.functional.cross_entropy`` callable (the thing a fused-linear-CE - swap would touch) must be byte-identical before and after.""" - import torch.nn.functional as F - - before = F.cross_entropy - k.liger_builtin() - assert F.cross_entropy is before - - -def test_kernelize_class_replacement_applies(): - """kernelize must swap __class__ on a module whose type matches a bundle - class key (independent of which families are installed).""" - import torch.nn as nn - - class _FakeHFNorm(nn.Module): - def __init__(self): - super().__init__() - self.weight = nn.Parameter(torch.zeros(4)) - - def forward(self, x): - return x - - class _FakeLigerImpl(nn.Module): - def forward(self, x): - return x + 1 - - parent = nn.Sequential(_FakeHFNorm()) - out = k.kernelize(parent, {_FakeHFNorm: _FakeLigerImpl}) - assert out is parent - assert type(parent[0]) is _FakeLigerImpl - - -# ── CLI toggle (liger-independent: part of twinkle core) ────────────────────── - - -def test_enable_liger_defaults_false(): - from twinkle.cli import ModelArgs - - assert ModelArgs().enable_liger is False - - -def test_enable_liger_toggle_via_cli(): - from twinkle.cli import CLI - - on = CLI.from_args(['--enable-liger']) - off = CLI.from_args(['--no-enable-liger']) - assert on.model.enable_liger is True - assert off.model.enable_liger is False - - -def test_enable_liger_env_var(): - from twinkle.cli import CLI - - old = os.environ.copy() - try: - os.environ.clear() - os.environ['TWINKLE_ENABLE_LIGER'] = 'true' - assert CLI.from_args(argv=[]).model.enable_liger is True - finally: - os.environ.clear() - os.environ.update(old) diff --git a/tests/kernel/test_public_api.py b/tests/kernel/test_public_api.py index 2d01f116c..955e7dfeb 100644 --- a/tests/kernel/test_public_api.py +++ b/tests/kernel/test_public_api.py @@ -1,14 +1,14 @@ def test_public_exports_exactly_four_symbols(): import twinkle.kernel as k - assert sorted(k.__all__) == ['hub', 'kernelize', 'liger_builtin', 'npu_builtin'] + assert sorted(k.__all__) == ['DEFAULT_KERNEL_CONFIG', 'KernelChoice', 'hub', 'kernelize'] assert callable(k.kernelize) - assert callable(k.npu_builtin) - assert callable(k.liger_builtin) assert callable(k.hub) + assert isinstance(k.DEFAULT_KERNEL_CONFIG, dict) + assert k.DEFAULT_KERNEL_CONFIG # non-empty built-in default mapping def test_no_legacy_symbols(): - """Legacy registrar / patch helpers must be gone.""" + """Legacy registrar / patch helpers and the retired builtin bundles must be gone.""" import twinkle.kernel as k legacy = [ 'kernelize_model', 'register_layer_kernel', 'register_function_kernel', @@ -18,6 +18,7 @@ def test_no_legacy_symbols(): 'get_global_layer_registry', 'get_global_function_registry', 'get_global_external_layer_registry', 'LayerRegistry', 'ExternalLayerRegistry', 'FunctionRegistry', + 'npu_builtin', 'liger_builtin', ] for name in legacy: assert not hasattr(k, name), f'unexpected legacy symbol: {name}' \ No newline at end of file diff --git a/tests/kernel/test_registry.py b/tests/kernel/test_registry.py new file mode 100644 index 000000000..897394d3d --- /dev/null +++ b/tests/kernel/test_registry.py @@ -0,0 +1,419 @@ +"""tests/kernel/test_registry.py -- registry/KernelChoice/installer behavior matrix. + +Everything runs on CPU: fake ops / fake backends verify selection, fallback, log +levels and installer priority; no real NPU needed (VeOmni B2). +""" +import logging +import sys +import types + +import pytest +import torch.nn as nn + +from twinkle.kernel import core +from twinkle.kernel.core import default_installer, kernelize +from twinkle.kernel.registry import _OPS, KernelChoice, KernelImpl, get_op, register_op, resolve_impl + + +class _SrcLayer(nn.Module): + def forward(self, x): + return x + + +class _DstLayer(nn.Module): + def forward(self, x): + return x + 100 + + +@pytest.fixture +def twinkle_log(): + """Capture records of the 'twinkle' logger (propagate=False, caplog cannot see it).""" + logger = logging.getLogger('twinkle') + old_level = logger.level + logger.setLevel(logging.DEBUG) + records = [] + + class _H(logging.Handler): + def emit(self, record): + records.append(record) + + handler = _H(logging.DEBUG) + logger.addHandler(handler) + try: + yield records + finally: + logger.removeHandler(handler) + logger.setLevel(old_level) + + +@pytest.fixture +def fake_ops(): + """Snapshot and clear _OPS, register fake ops, restore after the test.""" + saved = dict(_OPS) + _OPS.clear() + yield + _OPS.clear() + _OPS.update(saved) + + +def _fake_impl(name='fake_impl'): + fn = types.FunctionType(compile(f'def {name}(): pass', '', 'exec'), {}) + return fn + + +def _register_fake_op(name='fake_op', backends=('fake_a', 'fake_b'), installer=None, available_a=(True, None), + available_b=(True, None), calls=None): + """Register one fake op; return (impl_a, impl_b, calls). calls records available/load invocations.""" + if calls is None: + calls = [] + impl_a, impl_b = _fake_impl('impl_a'), _fake_impl('impl_b') + + def avail_a(): + calls.append('available_a') + return available_a + + def avail_b(): + calls.append('available_b') + return available_b + + def load_a(_target=None): + calls.append('load_a') + return impl_a + + def load_b(_target=None): + calls.append('load_b') + return impl_b + + register_op( + name, + implementations={ + 'fake_a': KernelImpl(load=load_a, available=avail_a), + 'fake_b': KernelImpl(load=load_b, available=avail_b), + }, + installer=installer, + ) + return impl_a, impl_b, calls + + +# ── 1. direct impl pass-through ──────────────────────────────────────────── + + +def test_direct_impl_passthrough(fake_ops): + parent = nn.Sequential(_SrcLayer()) + kernelize(parent, {_SrcLayer: _DstLayer}) + assert type(parent[0]) is _DstLayer + + +# ── 2. HubRef pass-through ───────────────────────────────────────────────── + + +def test_hub_ref_passthrough(fake_ops, monkeypatch): + monkeypatch.setattr(core, '_load_hub_ref', lambda ref: _DstLayer) + parent = nn.Sequential(_SrcLayer()) + kernelize(parent, {_SrcLayer: core.HubRef('org/repo', 'X', revision='main')}) + assert type(parent[0]) is _DstLayer + + +# ── 3. KernelChoice first backend hits ───────────────────────────────────── + + +def test_choice_first_backend_wins(fake_ops): + """First backend hits: default_installer receives impl_a; fake_b.available is never called.""" + impl_a, _, calls = _register_fake_op() + mod_name = 'tests.kernel._tmp_choice_first' + mod = types.ModuleType(mod_name) + mod.slot = None + sys.modules[mod_name] = mod + try: + kernelize(nn.Linear(1, 1), {f'{mod_name}.slot': KernelChoice(op='fake_op', backends=('fake_a', 'fake_b'))}) + assert mod.slot is impl_a + assert 'available_b' not in calls + finally: + sys.modules.pop(mod_name, None) + + +# ── 4. available failure falls through ───────────────────────────────────── + + +def test_available_false_falls_back_with_reason(fake_ops, twinkle_log): + impl_a, impl_b, _ = _register_fake_op(available_a=(False, 'fake_a needs CUDA 12.9')) + mod_name = 'tests.kernel._tmp_avail_fb' + mod = types.ModuleType(mod_name) + mod.slot = None + sys.modules[mod_name] = mod + try: + kernelize(nn.Linear(1, 1), { + f'{mod_name}.slot': KernelChoice(op='fake_op', backends=('fake_a', 'fake_b'))}) + assert mod.slot is impl_b + assert any('fake_a needs CUDA 12.9' in r.getMessage() and r.levelno >= logging.WARNING + for r in twinkle_log) + finally: + sys.modules.pop(mod_name, None) + + +# ── 5. load exception falls through ──────────────────────────────────────── + + +def test_load_exception_falls_back(fake_ops): + impl_b = _fake_impl('impl_b') + register_op( + 'fake_op', + implementations={ + 'fake_a': KernelImpl( + load=lambda _t=None: (_ for _ in ()).throw(ImportError('boom')), + available=lambda: (True, None)), + 'fake_b': KernelImpl(load=lambda _t=None: impl_b, available=lambda: (True, None)), + }) + mod_name = 'tests.kernel._tmp_load_fb' + mod = types.ModuleType(mod_name) + mod.slot = None + sys.modules[mod_name] = mod + try: + kernelize(nn.Linear(1, 1), {f'{mod_name}.slot': KernelChoice(op='fake_op', backends=('fake_a', 'fake_b'))}) + assert mod.slot is impl_b + finally: + sys.modules.pop(mod_name, None) + + +# ── 6. all fail -> keep original + explicit mapping WARNING ──────────────── + + +def test_all_backends_fail_keeps_original_and_warns(fake_ops, twinkle_log): + _register_fake_op(available_a=(False, 'no a'), available_b=(False, 'no b')) + called = [] + parent = nn.Sequential(_SrcLayer()) + kernelize( + parent, + { + _SrcLayer: + KernelChoice( + op='fake_op', + backends=('fake_a', 'fake_b'), + installer=lambda m, t, i: called.append(i)), + }) + assert called == [] # installer never called + assert type(parent[0]) is _SrcLayer # original kept + assert any('no available backend' in r.getMessage() and r.levelno >= logging.WARNING for r in twinkle_log) + + +# ── 7. default config path log levels (P1: fallbacks/failures all DEBUG) ─── + + +def test_default_config_path_logs_debug_only(fake_ops, monkeypatch, twinkle_log): + _register_fake_op(available_a=(False, 'no a'), available_b=(False, 'no b')) + monkeypatch.setattr( + 'twinkle.kernel.config.DEFAULT_KERNEL_CONFIG', + {_SrcLayer: KernelChoice(op='fake_op', backends=('fake_a', 'fake_b'))}) + parent = nn.Sequential(_SrcLayer()) + kernelize(parent) # mapping=None -> default config path + assert type(parent[0]) is _SrcLayer + assert not [r for r in twinkle_log if r.levelno >= logging.WARNING] + assert any('no available backend' in r.getMessage() for r in twinkle_log if r.levelno == logging.DEBUG) + + +# ── 8. unregistered backend -> warning, falls through ────────────────────── + + +def test_unregistered_backend_warns_and_falls_back(fake_ops, twinkle_log): + _, impl_b, _ = _register_fake_op() + mod_name = 'tests.kernel._tmp_unreg_be' + mod = types.ModuleType(mod_name) + mod.slot = None + sys.modules[mod_name] = mod + try: + kernelize(nn.Linear(1, 1), { + f'{mod_name}.slot': KernelChoice(op='fake_op', backends=('ghost', 'fake_b'))}) + assert mod.slot is impl_b + assert any("'ghost' not registered" in r.getMessage() and r.levelno >= logging.WARNING + for r in twinkle_log) + finally: + sys.modules.pop(mod_name, None) + + +# ── 9. unregistered op -> ValueError ─────────────────────────────────────── + + +def test_unregistered_op_raises(fake_ops): + with pytest.raises(ValueError, match="'nope' is not registered"): + kernelize(nn.Linear(1, 1), {_SrcLayer: KernelChoice(op='nope', backends=('fake_a', ))}) + + +# ── 10. installer priority: choice > op > default ────────────────────────── + + +def test_installer_priority(fake_ops): + calls = [] + choice_installer = lambda m, t, i: calls.append('choice') # noqa: E731 + op_installer = lambda m, t, i: calls.append('op') # noqa: E731 + + _register_fake_op(installer=op_installer) + parent = nn.Sequential(_SrcLayer()) + + # choice.installer overrides op.installer + kernelize(parent, {_SrcLayer: KernelChoice(op='fake_op', backends=('fake_a', ), installer=choice_installer)}) + assert calls == ['choice'] + + # choice.installer empty -> op.installer + calls.clear() + kernelize(parent, {_SrcLayer: KernelChoice(op='fake_op', backends=('fake_a', ))}) + assert calls == ['op'] + + # op.installer empty -> default_installer (class replacement really happens) + register_op('fake_op2', implementations={ + 'fake_a': KernelImpl(load=lambda _t=None: _DstLayer, available=lambda: (True, None)), + }) + parent2 = nn.Sequential(_SrcLayer()) + kernelize(parent2, {_SrcLayer: KernelChoice(op='fake_op2', backends=('fake_a', ))}) + assert type(parent2[0]) is _DstLayer + + +# ── 11. custom installer signature & logical target pass-through ─────────── + + +def test_custom_installer_receives_model_target_impl(fake_ops): + received = [] + impl_a = _fake_impl('impl_a') + register_op( + 'sdpa_like', + implementations={'fake_a': KernelImpl(load=lambda _t=None: impl_a, available=lambda: (True, None))}, + installer=lambda m, t, i: received.append((m, t, i))) + model = nn.Linear(1, 1) + kernelize(model, {'sdpa': KernelChoice(op='sdpa_like', backends=('fake_a', ))}) + assert received == [(model, 'sdpa', impl_a)] # logical target passed through unresolved + + +# ── 12. string target dispatch ───────────────────────────────────────────── + + +def test_dotted_target_dispatch(fake_ops, twinkle_log): + mod_name = 'tests.kernel._tmp_dispatch' + mod = types.ModuleType(mod_name) + + class Foo(nn.Module): + def forward(self, x): + return x + + mod.Foo = Foo + mod.fn = lambda x: x + sys.modules[mod_name] = mod + try: + # resolves to an nn.Module subclass -> __class__ swap + parent = nn.Sequential(Foo()) + kernelize(parent, {f'{mod_name}.Foo': _DstLayer}) + assert type(parent[0]) is _DstLayer + + # resolves to a function -> setattr + new_fn = lambda x: x * 3 # noqa: E731 + kernelize(nn.Linear(1, 1), {f'{mod_name}.fn': new_fn}) + assert mod.fn is new_fn + + # transformers family missing + explicit mapping -> WARNING skip (typo hint), no raise + kernelize(nn.Linear(1, 1), + {'transformers.models.no_such_family.modeling_x.NoSuchRMSNorm': _DstLayer}) + assert any('family not installed' in r.getMessage() and 'typos' in r.getMessage() + for r in twinkle_log if r.levelno == logging.WARNING) + finally: + sys.modules.pop(mod_name, None) + + +def test_family_skip_default_path_stays_debug(twinkle_log): + """Default config path (warn=False): a missing family stays DEBUG, not escalated to WARNING.""" + default_installer(nn.Linear(1, 1), + 'transformers.models.no_such_family.modeling_x.NoSuchRMSNorm', + _DstLayer, + warn=False) + matches = [r for r in twinkle_log if 'family not installed' in r.getMessage()] + assert matches and all(r.levelno == logging.DEBUG for r in matches) + + +# ── 13. unresolvable non-family string + default installer -> explicit error ── + + +def test_unresolvable_non_family_string_raises(fake_ops): + with pytest.raises(ValueError, match='Cannot resolve mapping target'): + kernelize(nn.Linear(1, 1), {'no_such_pkg_zzz.mod.attr': _DstLayer}) + + +# ── 14. installer exception propagates ───────────────────────────────────── + + +def test_installer_exception_propagates(fake_ops): + def bad_installer(m, t, i): + raise RuntimeError('half-installed state must be visible') + + register_op( + 'fake_op', + implementations={'fake_a': KernelImpl(load=lambda _t=None: _DstLayer, available=lambda: (True, None))}, + installer=bad_installer) + with pytest.raises(RuntimeError, match='half-installed'): + kernelize(nn.Linear(1, 1), {_SrcLayer: KernelChoice(op='fake_op', backends=('fake_a', ))}) + + +# ── 15. registration guards ──────────────────────────────────────────────── + + +def test_register_op_defense(fake_ops): + register_op('dup', implementations={'a': KernelImpl(load=lambda _t=None: 1, available=lambda: (True, None))}) + with pytest.raises(ValueError, match='already registered'): + register_op('dup', + implementations={'a': KernelImpl(load=lambda _t=None: 1, available=lambda: (True, None))}) + with pytest.raises(ValueError, match='no implementations'): + register_op('empty', implementations={}) + with pytest.raises(ValueError, match="'ghost' is not registered"): + get_op('ghost') + + +# ── 16. logical target misused with default installer (single-segment string) -> explicit error ── + + +def test_logical_target_single_segment_raises(fake_ops): + """A logical name like 'sdpa' paired with an op that has no dedicated installer -> + default_installer fails to resolve it and raises ValueError mentioning a custom installer.""" + register_op( + 'fake_op', + implementations={'fake_a': KernelImpl(load=lambda _t=None: _DstLayer, available=lambda: (True, None))}) + with pytest.raises(ValueError, match='logical targets require a custom installer'): + kernelize(nn.Linear(1, 1), {'sdpa': KernelChoice(op='fake_op', backends=('fake_a', ))}) + + +# ── 17. mixed mapping: KernelChoice + direct impl in one install ─────────── + + +class _SrcLayer2(nn.Module): + def forward(self, x): + return x + + +def test_mixed_mapping_choice_and_direct(fake_ops): + """One kernelize call: KernelChoice entries (via registry) and bare impl entries (pass-through) both take effect.""" + register_op( + 'fake_op', + implementations={'fake_a': KernelImpl(load=lambda _t=None: _DstLayer, available=lambda: (True, None))}) + parent = nn.Sequential(_SrcLayer(), _SrcLayer2()) + kernelize(parent, { + _SrcLayer: KernelChoice(op='fake_op', backends=('fake_a', )), + _SrcLayer2: _DstLayer, + }) + assert type(parent[0]) is _DstLayer + assert type(parent[1]) is _DstLayer + + +# ── 18. family skip produces no success log ──────────────────────────────── + + +def test_family_skip_produces_no_success_info(fake_ops, twinkle_log): + """After a family-missing skip (installed=False), the main loop must not emit an INFO success log for that target.""" + kernelize(nn.Linear(1, 1), {'transformers.models.no_such_family.modeling_x.Foo': _DstLayer}) + assert not [r for r in twinkle_log if r.levelno == logging.INFO] + + +# ── 19. kernelize idempotent: safe to call repeatedly ────────────────────── + + +def test_kernelize_idempotent(fake_ops): + """Second call: instances of the target class no longer exist (already swapped); zero replacements, no exception.""" + parent = nn.Sequential(_SrcLayer()) + kernelize(parent, {_SrcLayer: _DstLayer}) + kernelize(parent, {_SrcLayer: _DstLayer}) + kernelize(parent, {_SrcLayer: _DstLayer}) + assert type(parent[0]) is _DstLayer diff --git a/tests/kernel/test_replace.py b/tests/kernel/test_replace.py index e649b2e3c..916d5ef15 100644 --- a/tests/kernel/test_replace.py +++ b/tests/kernel/test_replace.py @@ -1,9 +1,6 @@ -import sys -import types - import torch.nn as nn -from twinkle.kernel.core import _replace_attr, _replace_class +from twinkle.kernel.core import _replace_class class _Target(nn.Module): @@ -38,37 +35,4 @@ def test_replace_class_idempotent(): m = nn.Sequential(_Target()) _replace_class(m, _Target, _Impl) _replace_class(m, _Target, _Impl) # second call must be safe - assert type(m[0]) is _Impl - - -def test_replace_attr_sets_module_attribute(): - mod_name = 'tests.kernel._tmp_replace_attr' - mod = types.ModuleType(mod_name) - mod.target_fn = lambda x: x - sys.modules[mod_name] = mod - try: - new_fn = lambda x: x * 2 # noqa: E731 - _replace_attr(f'{mod_name}.target_fn', new_fn) - assert mod.target_fn is new_fn - finally: - sys.modules.pop(mod_name, None) - - -def test_replace_attr_supports_class_attribute(): - import sys - import types - - mod_name = 'tests.kernel._tmp_class_attr' - mod = types.ModuleType(mod_name) - - class Foo: - def forward(self, x): - return x - mod.Foo = Foo - sys.modules[mod_name] = mod - try: - new_forward = lambda self, x: x + 7 # noqa: E731 - _replace_attr(f'{mod_name}.Foo.forward', new_forward) - assert Foo.forward is new_forward - finally: - sys.modules.pop(mod_name, None) \ No newline at end of file + assert type(m[0]) is _Impl \ No newline at end of file diff --git a/tests/kernel/test_resolve_value.py b/tests/kernel/test_resolve_value.py index 652783f5c..db08d1259 100644 --- a/tests/kernel/test_resolve_value.py +++ b/tests/kernel/test_resolve_value.py @@ -1,48 +1,26 @@ import torch.nn as nn -from twinkle.kernel.core import HubRef, _resolve_value +from twinkle.kernel.core import HubRef, resolve_direct_value class _ImplA(nn.Module): pass -class _ImplB(nn.Module): - pass - - def test_passthrough_class_value(): - assert _resolve_value(_ImplA, 'cuda') is _ImplA + assert resolve_direct_value(_ImplA) is _ImplA def test_passthrough_callable_value(): f = lambda x: x # noqa: E731 - assert _resolve_value(f, 'npu') is f - - -def test_passthrough_hubref(): - ref = HubRef('org/repo', 'Layer', revision='main') - assert _resolve_value(ref, 'cuda') is ref + assert resolve_direct_value(f) is f -def test_device_dict_match(): - val = {'npu': _ImplA, 'cuda': _ImplB} - assert _resolve_value(val, 'npu') is _ImplA - assert _resolve_value(val, 'cuda') is _ImplB +def test_hubref_delegates_to_load(monkeypatch): + """HubRef values are resolved via _load_hub_ref (lazy hub download).""" + import twinkle.kernel.core as core - -def test_device_dict_miss_returns_none(): - val = {'npu': _ImplA} - assert _resolve_value(val, 'cuda') is None - - -def test_device_dict_nested(): - # nested dict -> recursive resolve - val = {'npu': {'npu': _ImplA}} - assert _resolve_value(val, 'npu') is _ImplA - - -def test_device_dict_miss_then_passthrough(): - # nested dict whose inner is also a dict that misses -> None - val = {'npu': {'cuda': _ImplA}} - assert _resolve_value(val, 'npu') is None \ No newline at end of file + sentinel = object() + ref = HubRef('org/repo', 'Layer', revision='main') + monkeypatch.setattr(core, '_load_hub_ref', lambda r: sentinel) + assert resolve_direct_value(ref) is sentinel From 54621dba922018df3246eaf1d90d8a00bae96e16 Mon Sep 17 00:00:00 2001 From: clc Date: Tue, 4 Aug 2026 09:32:21 +0800 Subject: [PATCH 4/6] fix(kernel): replace owner class when a Module-class impl targets a function path _install_dotted dispatched any non-class target (e.g. ...Experts.forward) to setattr, so a Module-class impl like liger's LigerExperts was assigned onto Experts.forward and failed on the first call. The dispatcher now replaces the owner class in that case (function impls keep the setattr path); family-missing DEBUG/WARN semantics are unchanged. Also fixes a stale liger_builtin reference in the rms_norm liger adapter docstring, and adds a regression test covering the class-impl dispatch. --- src/twinkle/kernel/core.py | 14 ++++++++++-- src/twinkle/kernel/ops/rms_norm/liger.py | 2 +- tests/kernel/test_registry.py | 28 ++++++++++++++++++++++++ 3 files changed, 41 insertions(+), 3 deletions(-) diff --git a/src/twinkle/kernel/core.py b/src/twinkle/kernel/core.py index fea01fe15..f7f728b7b 100644 --- a/src/twinkle/kernel/core.py +++ b/src/twinkle/kernel/core.py @@ -139,6 +139,10 @@ def _install_dotted(model: nn.Module, target: str, impl, *, warn: bool = False) - resolves to an ``nn.Module`` subclass -> ``_replace_class(model, cls, impl)`` (exactly equivalent to a class-object key, exact type match) + - resolves to a function / class method but the impl is an ``nn.Module`` + subclass (e.g. liger's ``LigerExperts`` on a ``...Experts.forward`` + target) -> ``_replace_class(model, owner_cls, impl)``; setattr'ing a + class onto ``forward`` would fail on the first call - otherwise -> ``setattr`` (module function / class method) - unresolvable ``transformers.*`` family path (missing module/attr) -> skip, return False (a missing family is normal). ``warn=False`` @@ -162,6 +166,11 @@ def _install_dotted(model: nn.Module, target: str, impl, *, warn: bool = False) f'(logical targets require a custom installer): {e!r}') from e if isinstance(resolved, type) and issubclass(resolved, nn.Module): _replace_class(model, resolved, impl) + elif isinstance(impl, type) and issubclass(impl, nn.Module) and isinstance(owner, type): + # function target (e.g. ``...Experts.forward``) with a Module-class impl + # (e.g. liger's LigerExperts): setattr'ing a class onto ``forward`` would + # fail on first call -> replace the owner class instead + _replace_class(model, owner, impl) else: setattr(owner, final_attr, impl) return True @@ -211,8 +220,9 @@ def kernelize(model: nn.Module, mapping: dict | None = None) -> nn.Module: - ``type[nn.Module]``: replace ``m.__class__`` for every module of the exact type (no subclass walking). - ``str`` dotted path: resolved by the default installer — an - ``nn.Module`` subclass resolves to class replacement, anything else - (module function / class method) to ``setattr``. Unresolvable + ``nn.Module`` subclass resolves to class replacement; a Module-class + impl on a function target (e.g. ``...Experts.forward``) replaces the + owner class; anything else resolves to ``setattr``. Unresolvable ``transformers.*`` family paths are skipped: DEBUG on the default path, WARNING (with typo hint) for explicit mappings. diff --git a/src/twinkle/kernel/ops/rms_norm/liger.py b/src/twinkle/kernel/ops/rms_norm/liger.py index f9d064709..5460b44b0 100644 --- a/src/twinkle/kernel/ops/rms_norm/liger.py +++ b/src/twinkle/kernel/ops/rms_norm/liger.py @@ -88,7 +88,7 @@ class LigerRMSNormGemma4Replacement(_LigerRMSNormBase): Note: Gemma4RMSNorm has a ``with_scale=False`` variant (no weight, used for attention ``v_norm``). That path falls back to a plain torch RMSNorm; this - adapter is only wired onto the scale-bearing variants — see ``liger_builtin``. + adapter is only wired onto the scale-bearing variants — see ``config.py``. """ _liger_offset = 0.0 diff --git a/tests/kernel/test_registry.py b/tests/kernel/test_registry.py index 897394d3d..62819925f 100644 --- a/tests/kernel/test_registry.py +++ b/tests/kernel/test_registry.py @@ -316,6 +316,34 @@ def forward(self, x): sys.modules.pop(mod_name, None) +def test_dotted_function_target_with_class_impl_replaces_owner_class(fake_ops): + """Module-class impl on a ``...Experts.forward``-style target -> owner class + replacement, NOT setattr (moe_experts liger regression: LigerExperts must + not end up assigned to ``Experts.forward``).""" + mod_name = 'tests.kernel._tmp_class_impl' + mod = types.ModuleType(mod_name) + + class FakeExperts(nn.Module): + def forward(self, x): + return x + + class LigerLikeExperts(nn.Module): + def forward(self, x): + return x + 1 + + mod.FakeExperts = FakeExperts + original_forward = FakeExperts.forward + sys.modules[mod_name] = mod + try: + parent = nn.Sequential(FakeExperts()) + kernelize(parent, {f'{mod_name}.FakeExperts.forward': LigerLikeExperts}) + assert type(parent[0]) is LigerLikeExperts + # the class attribute must stay untouched (no class assigned to forward) + assert FakeExperts.forward is original_forward + finally: + sys.modules.pop(mod_name, None) + + def test_family_skip_default_path_stays_debug(twinkle_log): """Default config path (warn=False): a missing family stays DEBUG, not escalated to WARNING.""" default_installer(nn.Linear(1, 1), From ff28d32282fcc0af4357a581e687a6334ad0fdcb Mon Sep 17 00:00:00 2001 From: clc Date: Tue, 4 Aug 2026 09:34:43 +0800 Subject: [PATCH 5/6] test(transformers): point qwen3.5 FLA precision test at ops.fla.npu The test still imported the retired twinkle.kernel.npu_impls.fla module, failing with ModuleNotFoundError on NPU CI and silently dropping FLA precision coverage. Also updates a stale npu_builtin reference in the class docstring. --- tests/transformers/test_qwen35_fla_bwd_precision.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/transformers/test_qwen35_fla_bwd_precision.py b/tests/transformers/test_qwen35_fla_bwd_precision.py index 6bdd83c87..d0fbd5775 100644 --- a/tests/transformers/test_qwen35_fla_bwd_precision.py +++ b/tests/transformers/test_qwen35_fla_bwd_precision.py @@ -126,7 +126,7 @@ def _force_fla_off(model: torch.nn.Module) -> None: def _force_fla_on(model: torch.nn.Module) -> int: import importlib - import twinkle.kernel.npu_impls.fla as fla_mod + import twinkle.kernel.ops.fla.npu as fla_mod importlib.reload(fla_mod) return fla_mod.apply_qwen3_5_fla(model) @@ -275,8 +275,8 @@ class TestQwen35FlaBwdPrecision: """Compare FLA ON (bf16 Triton kernel) vs FLA OFF (fp32 torch reference). These tests do NOT require multiple devices and are the precise analogue of - the FLA toggle used by ``cookbook/transformers/fsdp2.sh`` (which sets the - same ``TWINKLE_NPU_FLA`` env var via ``npu_builtin(model)``). + the FLA toggle used by ``cookbook/transformers/fsdp2.sh`` (the NPU default + kernel config sets the same ``TWINKLE_NPU_FLA`` env var via ``kernelize(model)``). """ @pytest.mark.parametrize('seed', [1234, 2024]) From 3300c02b06638192fc4878a3e70c1260020ca668 Mon Sep 17 00:00:00 2001 From: clc Date: Tue, 4 Aug 2026 09:34:44 +0800 Subject: [PATCH 6/6] docs(cookbook): correct --enable-liger scope in fsdp2 comments The flag only gates the Liger fused-linear-CE loss; per-layer kernels always come from kernelize(model)'s default config (NPU: CANN-first chains) regardless of the flag. The old comments (fsdp2.sh, fsdp2.py, cli.py) described the pre-refactor behavior where the flag swapped in liger_builtin's per-layer kernels. fsdp2.sh: comment block only, no parameter changes. --- cookbook/transformers/fsdp2.py | 16 +++++++++------- cookbook/transformers/fsdp2.sh | 12 +++++++----- src/twinkle/cli/cli.py | 11 ++++++----- 3 files changed, 22 insertions(+), 17 deletions(-) diff --git a/cookbook/transformers/fsdp2.py b/cookbook/transformers/fsdp2.py index c501f6132..621e1bf7d 100644 --- a/cookbook/transformers/fsdp2.py +++ b/cookbook/transformers/fsdp2.py @@ -76,13 +76,15 @@ def train(): if not _torch_baseline: model = kernelize(model) - # `--enable-liger` turns on BOTH the per-layer Liger/CANN kernels (above) - # AND, by default (`enable_fused_ce=True`), the LigerFusedLinearCrossEntropyLoss - # which skips the lm_head GEMM so the (B,T,V) logits tensor is never materialised. - # The forward then runs under task='fused_lm_ce' (TransformersFusedCEPatch). - # Pass `--no-fused-ce` to keep only the per-layer kernels (standard CE loss). - # The loss is device-agnostic: on NPU/CUDA it auto-falls-back to materialised - # CE if the fused kernel raises for a given shape (defensive). + # `--enable-liger` gates ONLY the fused-CE loss here — per-layer kernels + # above come from kernelize(model)'s default config regardless of the flag. + # With `enable_fused_ce=True` (default) it enables the + # LigerFusedLinearCrossEntropyLoss, which skips the lm_head GEMM so the + # (B,T,V) logits tensor is never materialised; the forward then runs under + # task='fused_lm_ce' (TransformersFusedCEPatch). Pass `--no-fused-ce` to + # keep the standard CE loss. The loss is device-agnostic: on NPU/CUDA it + # auto-falls-back to materialised CE if the fused kernel raises for a + # given shape (defensive). _use_fused_ce = args.model.enable_liger and args.model.enable_fused_ce _task = 'fused_lm_ce' if _use_fused_ce else 'causal_lm' diff --git a/cookbook/transformers/fsdp2.sh b/cookbook/transformers/fsdp2.sh index 5a9a77722..c05923c2a 100644 --- a/cookbook/transformers/fsdp2.sh +++ b/cookbook/transformers/fsdp2.sh @@ -2,11 +2,13 @@ # All training config passed as CLI flags. Override at invocation, e.g.: # sh fsdp2.sh --batch-size 16 --lr 5e-5 # -# Liger Kernel: pass --enable-liger to swap in Liger's Triton (CUDA) / Ascend -# (NPU) kernels via `kernelize(model, liger_builtin(model))`. Without the flag -# the NPU path uses `npu_builtin()` and the CUDA path is untouched. -# sh fsdp2.sh --enable-liger # enable Liger -# sh fsdp2.sh --no-enable-liger # explicitly disable (default) +# Liger fused linear cross-entropy: --enable-liger turns on the fused-CE loss +# (pair with --no-fused-ce to make it a no-op). Per-layer kernels are NOT +# gated by this flag — they always come from `kernelize(model)`'s default +# config (NPU: CANN-first chains). To force Liger per-layer kernels, pass a +# custom mapping with liger-first KernelChoice chains (see Kernel.md). +# sh fsdp2.sh --enable-liger # enable Liger fused-CE loss +# sh fsdp2.sh --no-enable-liger # disable fused-CE (default) CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \ torchrun --nproc_per_node=8 fsdp2.py \ diff --git a/src/twinkle/cli/cli.py b/src/twinkle/cli/cli.py index c09264658..1467a68de 100644 --- a/src/twinkle/cli/cli.py +++ b/src/twinkle/cli/cli.py @@ -38,11 +38,12 @@ class ModelArgs: # Off by default — opt in with --enable-liger / TWINKLE_ENABLE_LIGER. enable_liger: bool = False # Fused-linear-CE loss toggle. Only meaningful when `enable_liger` is True. - # Defaults True so `--enable-liger` turns on BOTH the per-layer Liger/CANN - # kernels AND the LigerFusedLinearCrossEntropyLoss (skip-lm_head-GEMM, no - # (B,T,V) logits). Pass `--no-fused-ce` to opt out of the fused-CE loss and - # keep only the per-layer kernels (the loss falls back to standard CE). The - # loss itself is device-agnostic: on NPU/CUDA it auto-falls-back to + # Defaults True so `--enable-liger` alone turns on the + # LigerFusedLinearCrossEntropyLoss (skip-lm_head-GEMM, no (B,T,V) logits). + # Pass `--no-fused-ce` to opt out (the loss falls back to standard CE). + # Note: per-layer kernels are NOT gated by these flags — they come from + # kernelize(model)'s default config (NPU: CANN-first chains). The loss + # itself is device-agnostic: on NPU/CUDA it auto-falls-back to # materialised CE if the fused kernel raises for a given shape. enable_fused_ce: bool = True