feat(training): enable LyCORIS 4 fused kernels

docs/adr/0014-lycoris4-fused-kernels.md

@@ -0,0 +1,64 @@
+# 0014 — LyCORIS 4 fused kernels 与 Windows Triton
+
+**状态**：Accepted
+
+**日期**：2026-09-20
+**决策者**：项目维护者
+
+## 背景
+
+训练栈通过 `utils/lycoris_adapter.py` 接入 LyCORIS 3.4。普通 LoRA 已走
+bypass forward，但 LoHa / LoKr / DoRA 与 plain T-LoRA 的 rebuild 路径仍可能
+物化完整 `ΔW`。LyCORIS 4.0 新增 Triton / TileLang fused kernels，并按
+Triton → TileLang → `torch.compile` → eager 的顺序逐调用选择与回退。
+
+Windows 上 PyTorch 不自带可导入的 Triton，缺失时还会让 xformers 在启动期打印
+非致命 traceback。当前项目的目标本机环境是 PyTorch 2.11 + CUDA 12.8；该组合
+按 triton-windows 的兼容表对应 Triton 3.6。
+
+同时，LyCORIS 4.0 移除了旧 patch 所依赖的内部 `make_kron` /
+`rebuild_tucker` import，但 4.0.0 的 LoKr `rank_dropout` 仍会在 CPU 创建 mask。
+
+## 候选方案
+
+1. **保留 LyCORIS 3.4**：风险最低，但无法使用 fused kernels，启动警告保留。
+2. **直接升级并删除本地 adapter / patch**：代码最少，但丢失 family preset、
+   checkpoint metadata、训练 hook、T-LoRA timestep mask，并重新暴露 LoKr device bug。
+3. **升级到 4.0，保留薄 integration adapter，数学运算委托官方 functional API**：
+   需要小范围兼容改动，但保留项目协议并获得官方 kernel dispatch。
+
+## 决策
+
+- 固定 `lycoris-lora==4.0.0`，避免未来 major API 静默漂移。
+- Windows dependency 声明允许 Triton 3.6–3.8；启动器在导入 xformers/flash-attn
+  之前按已安装 PyTorch minor 自动收窄或修复版本（2.10/2.11→3.6、
+  2.12/2.13→3.7、2.14→3.8）。
+- 保留 `LycorisAdapter`：它只负责 AdapterProtocol、模型族 target preset、保存元数据、
+  sample/eval 状态和 T-LoRA step hook；LoRA 数学继续交给 LyCORIS。
+- plain T-LoRA 在两个低秩因子上应用 timestep mask。forward 调用官方
+  `functional.locon.bypass_forward_diff`，不物化 `ΔW`；merge/save 路径调用
+  `functional.locon.diff_weight`。两者都进入 LyCORIS 4 fused/compile/eager dispatch。
+- LoKr device patch 改为包装 upstream `get_weight`：只临时关闭坏的 dropout 分支，
+  调用 upstream（保留 fused dispatch），然后在 `weight.device` 重放 dropout。
+- OrthoLoRA 保持项目自实现。LyCORIS 4 没有与其 Cayley/SVD 参数化等价的官方
+  adapter，不能为了去掉 wrapper 改变算法或 checkpoint 语义。
+
+## 理由
+
+方案 3 把外部算法实现尽量留给上游，同时保留 Anima/Krea2 集成必需的边界。
+显式 pin 能让 device patch 的适用版本可审计。Triton 缺少或某个 shape 不受支持时，
+LyCORIS 自带回退，因此加速失败不会改变训练正确性。
+
+## 后果
+
+- Windows 首次安装多下载约 50 MB，并在首次遇到新 shape 时支付 JIT/tuning 开销。
+- fused kernel 的 microbenchmark 提升不等于整步训练等比例提升；DiT 主干仍占大头。
+- `rank_dropout`、卷积、过大 rank 等超出 fused scope 的调用会自动回退。
+- Triton 与 PyTorch minor 有配对约束；支持表之外的新 PyTorch 会保守回退并要求
+  更新 `_WINDOWS_TRITON_FOR_TORCH`，不会猜测 ABI。
+
+## 参考
+
+- [LyCORIS 4.0 fused kernels](https://github.com/KohakuBlueleaf/LyCORIS/blob/v4.0.0/docs/kernels/README.md)
+- [LyCORIS backend selection](https://github.com/KohakuBlueleaf/LyCORIS/blob/v4.0.0/docs/kernels/backends.md)
+- [triton-windows PyTorch compatibility](https://github.com/triton-lang/triton-windows#3-pytorch)

docs/adr/README.md

@@ -16,6 +16,7 @@
 | 0008 | [studio/ 4 层重构（0.11.0）](0008-studio-restructure-0.11.0.md) | Accepted | 2026-05-28 |
 | 0009 | [统一日志 + 错误体系（0.12.0）](0009-logging-error-system.md) | Accepted | 2026-05-28 |
 | 0010 | [preprocess scope 从项目级 download 下沉到 version 级 train](0010-preprocess-train-scope.md) | Accepted | 2026-06-03 |
+| 0014 | [LyCORIS 4 fused kernels 与 Windows Triton](0014-lycoris4-fused-kernels.md) | Accepted | 2026-09-20 |
 
 ## 状态值
 

docs/user-guide/training-tips.md

@@ -493,6 +493,26 @@ masterpiece, best quality, newest, safe,
 
 ## 硬件优化
 
+### LyCORIS 4 fused kernels（Windows / NVIDIA）
+
+项目使用 LyCORIS 4.0，并在 Windows 自动校正与当前 PyTorch 配对的 Triton；
+例如 PyTorch 2.11 对应 `triton-windows 3.6`。启动时应看到类似：
+
+```text
+[studio] LoRA kernels：LyCORIS 4.0.0 · preferred=triton · fused=triton · Triton 3.6.0.post26
+```
+
+backend 按调用自动选择：Triton → TileLang → `torch.compile` → eager PyTorch。
+某个 shape、dtype 或 dropout 不适合 fused kernel 时会自动回退，不需要手工切换。
+首次遇到新 shape 会有 JIT / tuning 预热，后续从项目内 `.cache/triton` 复用缓存。
+
+普通 LoRA / LoHa / LoKr / DoRA 可使用官方 fused 路径；plain T-LoRA 保留 timestep
+rank mask，并调用不创建独立 `ΔW` 的官方 fused bypass。`T-LoRA + Ortho` 使用项目自己的 Cayley/SVD
+参数化，不会伪装成 LyCORIS 算法，因此目前不走这套 fused kernels。
+
+这些数字只加速 adapter 运算，不代表整个训练步会同倍数提速；模型 attention / MLP、
+数据和采样通常仍是主耗时。
+
 ### RTX 3090/4090 (24GB)
 
 ```yaml

requirements.txt

@@ -13,7 +13,14 @@ wandb>=0.16.0
                         #           总能装上，而 adamw8bit 是可选优化器而非默认。未装时选它
                         #           会在训练启动期报可操作错误（optimizers/adamw8bit.py
                         #           validate），不会跑到一半才炸。装：pip install bitsandbytes
-lycoris-lora>=1.9.0
+# LyCORIS 4 brings fused adapter kernels and keeps the wrapper/module API used
+# by this project.  Pin the major migration so internal compatibility patches
+# remain auditable instead of silently running against a future breaking API.
+lycoris-lora==4.0.0
+# Windows PyTorch does not bring Triton itself.  The Studio bootstrap narrows
+# this broad range to the minor paired with the installed PyTorch (for example,
+# torch 2.10/2.11 -> Triton 3.6) before importing flash-attn/xformers.
+triton-windows>=3.6.0.post26,<3.9; sys_platform == "win32"
 prodigyopt>=1.0          # Prodigy 优化器（anima_train.py optimizer=prodigy 时必需）
 # ProdigyPlusScheduleFree (LoganBooker) — Flux/Qwen-Image LoRA 社区主流，解决 Prodigy 的
 # mutation ep 问题。pin >=2.0.0 因为 v1.9.2 ↔ v2.0.0 的 state_dict 不兼容。

studio/cli.py

@@ -16,6 +16,7 @@
 import argparse
 import hashlib
 import os
+import re
 import shutil
 import signal
 import subprocess
@@ -392,6 +393,121 @@ def _try_enable_flash_attn() -> None:
         )
 
 
+def _report_lycoris_kernels() -> None:
+    """打印 LyCORIS 4 的 adapter-kernel 选择，不触发 JIT 编译。
+
+    ``available_backends`` 只探测模块是否可导入；真正的 shape/device 选择仍由
+    LyCORIS 每次调用完成。这里用 ``preferred`` 而不是 ``active``，避免把尚未
+    首次编译的 Triton kernel 描述成已经执行过。
+    """
+    try:
+        from importlib.metadata import PackageNotFoundError, version  # noqa: PLC0415
+
+        lycoris_version = version("lycoris-lora")
+    except PackageNotFoundError:
+        return
+    except Exception as exc:  # noqa: BLE001
+        print(f"[studio] LyCORIS 状态读取失败（{exc}）", file=sys.stderr)
+        return
+
+    try:
+        from lycoris.kernels import available_backends, resolve_backend  # noqa: PLC0415
+
+        available = available_backends()
+        preferred = resolve_backend()
+    except Exception as exc:  # noqa: BLE001
+        print(
+            f"[studio] 警告：LyCORIS {lycoris_version} kernel backend 初始化失败（{exc}）",
+            file=sys.stderr,
+        )
+        return
+
+    triton_version = None
+    for dist_name in ("triton-windows", "triton"):
+        try:
+            triton_version = version(dist_name)
+            break
+        except PackageNotFoundError:
+            continue
+
+    fused = ",".join(name for name in available if name in {"triton", "tilelang"}) or "none"
+    triton_label = f" · Triton {triton_version}" if triton_version else ""
+    _say(
+        f"LoRA kernels：LyCORIS {lycoris_version} · preferred={preferred} "
+        f"· fused={fused}{triton_label}"
+    )
+
+
+_WINDOWS_TRITON_FOR_TORCH: dict[tuple[int, int], tuple[int, int]] = {
+    (2, 10): (3, 6),
+    (2, 11): (3, 6),
+    (2, 12): (3, 7),
+    (2, 13): (3, 7),
+    (2, 14): (3, 8),
+}
+
+
+def _major_minor(raw: str) -> Optional[tuple[int, int]]:
+    match = re.match(r"^(\d+)\.(\d+)", raw)
+    return (int(match.group(1)), int(match.group(2))) if match else None
+
+
+def _ensure_windows_triton() -> None:
+    """Match triton-windows to PyTorch before xformers imports Triton.
+
+    PyTorch and Triton minor releases have an ABI/toolchain pairing.  A static
+    requirement cannot express that dependency, so the launcher repairs a
+    missing or mismatched Windows wheel from the official compatibility table.
+    Unsupported/older Torch builds keep LyCORIS' compile/eager fallback.
+    """
+    if sys.platform != "win32":
+        return
+
+    try:
+        import torch  # noqa: PLC0415
+
+        torch_minor = _major_minor(torch.__version__)
+    except Exception:  # noqa: BLE001
+        return
+
+    target = _WINDOWS_TRITON_FOR_TORCH.get(torch_minor)
+    if target is None:
+        if torch_minor is not None:
+            print(
+                f"[studio] 警告：暂无 torch {torch_minor[0]}.{torch_minor[1]} "
+                "对应的 Windows Triton pin；LyCORIS 将自动回退",
+                file=sys.stderr,
+            )
+        return
+
+    from importlib.metadata import PackageNotFoundError, version  # noqa: PLC0415
+
+    try:
+        installed_minor = _major_minor(version("triton-windows"))
+    except PackageNotFoundError:
+        installed_minor = None
+    if installed_minor == target:
+        return
+
+    lower = f"{target[0]}.{target[1]}"
+    upper = f"{target[0]}.{target[1] + 1}"
+    requirement = f"triton-windows>={lower},<{upper}"
+    installed = (
+        f"{installed_minor[0]}.{installed_minor[1]}"
+        if installed_minor is not None
+        else "not installed"
+    )
+    _say(
+        f"torch {torch_minor[0]}.{torch_minor[1]} requires Triton {lower} "
+        f"(current: {installed}); installing compatible wheel..."
+    )
+    if _pip_install([requirement]) != 0:
+        print(
+            "[studio] 警告：Triton 安装失败；LyCORIS 将使用 compile/eager fallback",
+            file=sys.stderr,
+        )
+
+
 def _check_torch_cuda() -> None:
     """启动期检查 torch 是否能用 CUDA；CPU-only torch 跑训练 / 出图会极慢。
 
@@ -637,7 +753,9 @@ def cmd_run(args: argparse.Namespace) -> int:
         if not getattr(args, 'skip_pending', False):
             _apply_pending_install()
         _check_torch_cuda()
+        _ensure_windows_triton()
         _try_enable_flash_attn()
+        _report_lycoris_kernels()
         _check_onnxruntime()
         url = f"http://{args.host}:{args.port}/"
         _say(f"启动后端 → {url}")
@@ -692,7 +810,9 @@ def cmd_dev(args: argparse.Namespace) -> int:
     if not getattr(args, 'skip_pending', False):
         _apply_pending_install()
     _check_torch_cuda()
+    _ensure_windows_triton()
     _try_enable_flash_attn()
+    _report_lycoris_kernels()
     _check_onnxruntime()
 
     pg = ProcGroup()

tests/test_lycoris_patch.py

@@ -1,7 +1,7 @@
-"""utils.lycoris_patch — lycoris-lora 3.4.0 LokrModule.get_weight device patch。
+"""utils.lycoris_patch — LyCORIS 3.4/4.0 LoKr rank-dropout device patch。
 
 覆盖：
-- 命中受影响版本（3.4.0）→ patch 真实生效（torch.rand 调用得到 device=weight.device 的 mask）
+- 命中受影响版本 → patch 生效且保留 upstream weight builder
 - 未装 lycoris → skipped_not_installed
 - 未知版本 → skipped_version_unknown + warn
 - 同进程内幂等 → skipped_already_patched
@@ -52,15 +52,24 @@ def fresh_patch_module(monkeypatch: pytest.MonkeyPatch):
             delattr(LokrModule, mod._PATCHED_FLAG)
 
 
+@pytest.mark.parametrize("affected_version", ["3.4.0", "4.0.0"])
 def test_apply_on_known_affected_version_patches_get_weight(
-    fresh_patch_module, monkeypatch: pytest.MonkeyPatch
+    fresh_patch_module, monkeypatch: pytest.MonkeyPatch, affected_version: str,
 ) -> None:
-    """3.4.0 安装时 patch 应替换 LokrModule.get_weight 且生成的 mask 在 weight.device 上。"""
+    """受影响版本的 mask 应跟随 weight device，且 upstream builder 仍被调用。"""
     pytest.importorskip("lycoris.modules.lokr")
     import torch
     from lycoris.modules.lokr import LokrModule
 
-    monkeypatch.setattr(fresh_patch_module, "version", lambda _: "3.4.0")
+    upstream_calls: list[object] = []
+
+    def _fake_upstream(self, shape):
+        upstream_calls.append(shape)
+        assert self.rank_dropout == 0.0
+        return self.weight.clone()
+
+    monkeypatch.setattr(LokrModule, "get_weight", _fake_upstream)
+    monkeypatch.setattr(fresh_patch_module, "version", lambda _: affected_version)
 
     status = fresh_patch_module.apply_lokr_device_patch()
     assert status == "applied"
@@ -77,22 +86,18 @@ def _spy_rand(*args, **kwargs):
 
     monkeypatch.setattr(torch, "rand", _spy_rand)
 
-    # 用 mock self（不构造完整 LokrModule，单测只关心 torch.rand 这一行）
+    # 用 mock self，不依赖不同 LyCORIS 版本的 LoKr constructor 细节。
     class _FakeSelf:
         training = True
         rank_dropout = 0.5
         rank_dropout_scale = False
-        use_w1 = True
-        use_w2 = True
-        tucker = False
-        scale = 1.0
-        lokr_w1 = torch.eye(4)
-        lokr_w2 = torch.eye(4)
+        weight = torch.eye(4)
 
     fake = _FakeSelf()
-    # 直接调被替换后的 get_weight；shape=None 让 weight 保持原状
-    LokrModule.get_weight(fake, None)
-    assert captured["device"] is not None, "patched get_weight 没把 device 传给 torch.rand"
+    LokrModule.get_weight(fake, (4, 4))
+    assert upstream_calls == [(4, 4)]
+    assert fake.rank_dropout == 0.5, "wrapper 必须恢复 module 的 rank_dropout"
+    assert captured["device"] == fake.weight.device
 
 
 def test_apply_when_not_installed_returns_skipped(
@@ -123,6 +128,6 @@ def test_apply_idempotent(
 ) -> None:
     """同进程内重复调用：第一次 applied，之后 skipped_already_patched。"""
     pytest.importorskip("lycoris.modules.lokr")
-    monkeypatch.setattr(fresh_patch_module, "version", lambda _: "3.4.0")
+    monkeypatch.setattr(fresh_patch_module, "version", lambda _: "4.0.0")
     assert fresh_patch_module.apply_lokr_device_patch() == "applied"
     assert fresh_patch_module.apply_lokr_device_patch() == "skipped_already_patched"