feat: add NPU SAS + LI patch for DeepSeek-V4 without mindspeed

pyproject.toml

@@ -80,3 +80,4 @@ where = ["src"]
 
 [tool.setuptools.package-data]
 "twinkle_client.skills.bundled" = ["*.md"]
+"twinkle.kernel.ops.dsv4_sas_li.aclnn" = ["*.h", "*.cpp", "**/*.cpp"]

src/twinkle/kernel/config.py

@@ -111,6 +111,12 @@ def _build() -> dict[Any, Any]:
     # 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', ))
+
+    # DeepSeek-V4 SAS + LI (env-gated by TWINKLE_NPU_DSV4_SAS, npu only)
+    _dsv4 = 'transformers.models.deepseek_v4.modeling_deepseek_v4'
+    cfg[f'{_dsv4}.DeepseekV4Attention.forward'] = KernelChoice(op='dsv4_attention', backends=('npu', ))
+    cfg[f'{_dsv4}.DeepseekV4Indexer.forward'] = KernelChoice(op='dsv4_indexer', backends=('npu', ))
+    cfg[f'{_dsv4}.DeepseekV4CSACompressor.forward'] = KernelChoice(op='dsv4_csa_compressor', backends=('npu', ))
     return cfg
 
 

src/twinkle/kernel/ops/__init__.py

@@ -12,7 +12,7 @@
 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 . import dsv4_sas_li, fla, geglu, moe, rms_norm, rotary, sdpa_attention, swiglu  # noqa: F401,E402
 from .ep import EpExpertsGmm, ep_forward
 
 __all__ = [

src/twinkle/kernel/ops/dsv4_sas_li/__init__.py

@@ -0,0 +1,65 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+"""DeepSeek-V4 SAS/LI op registration: three forward-level replacements gated
+by the ``TWINKLE_NPU_DSV4_SAS`` env var.
+
+When enabled, the full patch set is applied:
+
+  - ``DeepseekV4Attention.forward``     → NPU sparse attention (SAS)
+  - ``DeepseekV4Indexer.forward``       → Lightning Indexer (LI)
+  - ``DeepseekV4CSACompressor.forward`` → full replacement returning a
+    3-tuple ``(compressed_kv, block_bias, top_k_indices)``
+
+LI is always on under SAS — there is no use case for SAS without LI
+(CSA would fall back to the slower stock indexer) or LI without SAS
+(indices would go unused). The CSA compressor is a **full forward
+replacement** rather than a wrapper (see ``npu.py`` docstring for details).
+"""
+from __future__ import annotations
+
+import os
+
+from ...registry import KernelImpl, is_npu_available, lazy_import, register_op
+
+
+def _dsv4_sas_available() -> tuple[bool, str | None]:
+    env = os.environ.get('TWINKLE_NPU_DSV4_SAS', '').lower().strip()
+    if not env or env in ('0', 'false', 'off', 'no'):
+        return False, 'TWINKLE_NPU_DSV4_SAS not enabled'
+    ok, reason = is_npu_available()
+    if not ok:
+        return ok, reason
+    return True, None
+
+
+_DSV4_BASE = 'twinkle.kernel.ops.dsv4_sas_li.npu'
+
+register_op(
+    'dsv4_attention',
+    implementations={
+        'npu': KernelImpl(
+            load=lazy_import(f'{_DSV4_BASE}:npu_dsv4_attention_forward'),
+            available=_dsv4_sas_available,
+        ),
+    },
+)
+
+register_op(
+    'dsv4_indexer',
+    implementations={
+        'npu': KernelImpl(
+            load=lazy_import(f'{_DSV4_BASE}:npu_dsv4_indexer_forward'),
+            available=_dsv4_sas_available,
+        ),
+    },
+)
+
+register_op(
+    'dsv4_csa_compressor',
+    implementations={
+        'npu':
+        KernelImpl(
+            load=lazy_import(f'{_DSV4_BASE}:npu_dsv4_csa_compressor_forward'),
+            available=_dsv4_sas_available,
+        ),
+    },
+)

src/twinkle/kernel/ops/dsv4_sas_li/aclnn/__init__.py

@@ -0,0 +1,9 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+"""Self-compiled ACLNN C++ extensions for Ascend NPU fusion operators.
+
+Provides JIT-compiled bindings for DeepSeek-V4 SAS (Sparse Attention with
+Shared-KV) and LI (Lightning Indexer) without depending on mindspeed.
+"""
+from .builder import build_op
+
+__all__ = ['build_op']