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..621e1bf7d 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,27 +66,25 @@ 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) - - # `--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). + if not _torch_baseline: + model = kernelize(model) + + # `--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/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..1467a68de 100644 --- a/src/twinkle/cli/cli.py +++ b/src/twinkle/cli/cli.py @@ -34,15 +34,16 @@ 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 - # 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 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..67d623e7a --- /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..f7f728b7b 100644 --- a/src/twinkle/kernel/core.py +++ b/src/twinkle/kernel/core.py @@ -1,17 +1,25 @@ # 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 +57,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 +68,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 +95,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 +124,156 @@ 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) + - 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`` + (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) + 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 + + +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; 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. 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..14ce9cd1e --- /dev/null +++ b/src/twinkle/kernel/ops/__init__.py @@ -0,0 +1,21 @@ +# 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). +""" +# 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', + '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..5a8c7c0a8 --- /dev/null +++ b/src/twinkle/kernel/ops/geglu/__init__.py @@ -0,0 +1,16 @@ +# 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..0233849fe --- /dev/null +++ b/src/twinkle/kernel/ops/moe/__init__.py @@ -0,0 +1,40 @@ +# 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..97b9d263d --- /dev/null +++ b/src/twinkle/kernel/ops/rms_norm/__init__.py @@ -0,0 +1,54 @@ +# 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 99% rename from src/twinkle/kernel/liger_impls/rms_norm.py rename to src/twinkle/kernel/ops/rms_norm/liger.py index f9d064709..5460b44b0 100644 --- a/src/twinkle/kernel/liger_impls/rms_norm.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/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..f47a31910 --- /dev/null +++ b/src/twinkle/kernel/ops/rotary/__init__.py @@ -0,0 +1,39 @@ +# 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..386d6b463 --- /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..de9ca3af5 --- /dev/null +++ b/src/twinkle/kernel/ops/swiglu/__init__.py @@ -0,0 +1,21 @@ +# 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..adb52a4de --- /dev/null +++ b/src/twinkle/kernel/registry.py @@ -0,0 +1,175 @@ +# 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 +import torch.nn as nn +from dataclasses import dataclass +from typing import Any, Callable + +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 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..62819925f --- /dev/null +++ b/tests/kernel/test_registry.py @@ -0,0 +1,447 @@ +"""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_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), + '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 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])