refactor(kernel): register-dispatch architecture for kernelize (registry/config/ops 三层拆分) - #255
Merged
Merged
Conversation
…try/config/ops split)
Rework kernelize into a register-dispatch model, replacing the
hard-coded builtin/liger module pair:
- registry: ops REGISTER themselves with per-backend impls (lazy load +
availability check); KernelChoice DISPATCHES to the first available
backend in the priority chain
- config: DEFAULT_KERNEL_CONFIG as pure data (dotted-path targets only)
— the mapping declares WHAT to replace and WHICH op; never imports
transformers/torch_npu/liger_kernel
- core: kernelize() walks the mapping and dispatches each entry to an
installer — HOW to install is decided by priority
KernelChoice.installer -> OpDefinition.installer -> default_installer
(class swap for nn.Module targets, setattr for functions/methods);
transformers.* family-missing skips are DEBUG on the default path and
WARNING (typo hint) for explicit mappings
- ops: swiglu / geglu / rms_norm / rotary / moe / sdpa_attention / fla
moved from npu_impls & liger_impls into per-op packages; new ep op
(NPU grouped-matmul) replaces the inlined EP expert loop in
expert_parallel.py
- moe_experts default chain is ('npu',) — LigerExperts is opt-in only,
absorbing the old _prefer_cann_on_npu drop semantics
- hub(): HubRef values resolve lazily via the optional kernels package
Cookbooks: kernelize is gated only by the torch-baseline escape hatch;
--enable-liger is narrowed to gating the fused-linear-CE loss. EP utils:
advanced indexing instead of index_select (aclnnIndexAdd broken on some
torch_npu/CANN), sync D2H copies for split sizes.
Contributor
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
werwrewe
marked this pull request as draft
July 30, 2026 08:27
Formatting only, no behavior change: yapf + isort + double-quote-string -fixer per .pre-commit-config.yaml; fixes E501 in core.py.
kevssim
marked this pull request as ready for review
August 3, 2026 07:33
Contributor
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
- npu_impls/ -> ops/: update imports to twinkle.kernel.ops.*.npu (impl functions moved, not removed; all CPU-importable via lazy torch_npu), plus NPU-skipif numerical parity cases - test_builtin.py / test_liger_builtin.py: dropped, builtin module retired - test_replace.py: drop _replace_attr cases (helper removed), keep _replace_class coverage - test_resolve_value.py: retarget to resolve_direct_value (passthrough + HubRef delegation); device-dict mechanism superseded by KernelChoice chains - test_kernelize.py: drop device-dict cases, npu_builtin auto-apply case now covers DEFAULT_KERNEL_CONFIG path - test_public_api.py: expect new __all__ (DEFAULT_KERNEL_CONFIG, KernelChoice, hub, kernelize) - test_registry.py: new coverage for registry/config resolution, installer priority, dotted targets, KernelChoice fallback chains
kevssim
reviewed
Aug 3, 2026
Collaborator
|
Qwen3.5 FLA 精度测试仍导入已删除的 twinkle.kernel.npu_impls.fla;在 NPU 测试环境中会直接 ModuleNotFoundError,无法覆盖 FLA。应迁移至 twinkle.kernel.ops.fla.npu。见 tests/transformers/test_qwen35_fla_bwd_precision.py:129 |
Collaborator
|
FSDP2 脚本注释仍指导用户调用已删除的 npu_builtin / liger_builtin,且描述与现有默认配置行为相反。见 cookbook/transformers/fsdp2.sh:5 |
added 3 commits
August 4, 2026 09:32
…unction path _install_dotted dispatched any non-class target (e.g. ...Experts.forward) to setattr, so a Module-class impl like liger's LigerExperts was assigned onto Experts.forward and failed on the first call. The dispatcher now replaces the owner class in that case (function impls keep the setattr path); family-missing DEBUG/WARN semantics are unchanged. Also fixes a stale liger_builtin reference in the rms_norm liger adapter docstring, and adds a regression test covering the class-impl dispatch.
The test still imported the retired twinkle.kernel.npu_impls.fla module, failing with ModuleNotFoundError on NPU CI and silently dropping FLA precision coverage. Also updates a stale npu_builtin reference in the class docstring.
The flag only gates the Liger fused-linear-CE loss; per-layer kernels always come from kernelize(model)'s default config (NPU: CANN-first chains) regardless of the flag. The old comments (fsdp2.sh, fsdp2.py, cli.py) described the pre-refactor behavior where the flag swapped in liger_builtin's per-layer kernels. fsdp2.sh: comment block only, no parameter changes.
kevssim
approved these changes
Aug 4, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
主要工作
将 kernelize 重构为 register-dispatch 模型,替换掉硬编码的builtin/liger 双模块:
cookbook:kernelize 仅由 torch-baseline 逃生门门控;--enable-liger 语义收窄为只门控 fused-linear-CE loss。EP utils:用高级索引替代 index_select(部分 torch_npu/CANN 版本 aclnnIndexAdd 损坏),split sizes 改为同步 D2H 拷贝。
docs:中英文档同步。
特性
kernelize(model, mapping)的核心模型:key 表示目标,value 表示如何替换KernelChoice)ops/<op>/下管理分发流程(register-dispatch)
kernelize 接口变化
新接口用法举例:
目录结构
旧
builtin.py/liger.py/npu_impls//liger_impls/全部删除(impl 本体 git mv 进ops/<op>/<backend>.py,内容不改)。关键设计点
config.py一张纯数据 mapping;新增家族 = 家族表加一行 tupleimplementationsdict 的字符串 key;平台/依赖判断统一收进KernelImpl.available(),返回(bool, reason)进日志lazy_import('pkg.mod:attr')字符串工厂和find_spec探测,import twinkle.kernel不触碰 torch_npu / liger / transformersmoe_experts两个 backend 都注册,但默认链只写('npu',)(liger 类替换会挡住 CANN grouped-matmul 快路径,吸收旧_prefer_cann_on_npudrop 语义);qwen3_5 的 rope 数值不兼容 liger,由 config 链只写('npu',)排除KernelImpl.load(target):liger RMSNorm 是一族实现(gemma/gemma4/qwen3_5 变体,选错 → NaN),load 接收 mapping target 按家族分派行为变化
kernelize(model)对这些家族不动);Qwen 家族替换结果与旧版完全一致(golden 证明)--enable-liger语义收窄:不再切逐层 kernel,现在只门控 fused-CE loss;强制 Liger 逐层需自建('liger',...)链(见上面例 ②)kernelize(model):从 no-op 变为应用默认配置(装了 liger 即生效,未装 ≈ no-op 且全 DEBUG 无噪音);cookbook 入口仅保留TWINKLE_TORCH_BASELINE=1逃生门Test plan