feat: add fused_gdn_gating kernel in tilelang-ascend. (#1267)

Author: zhang-minchao

third_party/tilelang-ascend

@@ -46,19 +46,26 @@ def main(argv: list[str] | None = None) -> None:
 
     if args.target == "ascend":
         from ..bootstrap import prepare_ascend
+
         prepare_ascend()
         from ..targets.ascend.build import build_kernels
+
+        manifests = build_kernels(
+            output_root=output_root,
+            kernel_names=args.kernels,
+            force=args.force,
+            device=args.device,
+        )
     elif args.target == "cuda":
         from ..targets.cuda.build import build_kernels
+
+        manifests = build_kernels(
+            output_root=output_root,
+            kernel_names=args.kernels,
+            force=args.force,
+        )
     else:
         raise ValueError(f"Unsupported target: {args.target}")
-
-    manifests = build_kernels(
-        output_root=output_root,
-        kernel_names=args.kernels,
-        force=args.force,
-        device=args.device,
-    )
     for manifest in manifests:
         print(f"[INFO] built {manifest.target}:{manifest.kernel_name}")
         print(f"[INFO] manifest: {Path(manifest.output_dir) / 'manifest.json'}")

xllm/compiler/tilelang/cli/compile_kernels.py

@@ -34,9 +34,9 @@ def require_env(name: str) -> str:
 def prepend_pythonpath(env: dict[str, str], path: str) -> None:
     current = env.get("PYTHONPATH", "")
     items = [item for item in current.split(os.pathsep) if item]
-    if path not in items:
-        items.insert(0, path)
-        env["PYTHONPATH"] = os.pathsep.join(items)
+    items = [item for item in items if item != path]
+    items.insert(0, path)
+    env["PYTHONPATH"] = os.pathsep.join(items)
 
 
 def prepare_tilelang_import(tilelang_root: str | Path | None = None) -> Path:
@@ -45,8 +45,11 @@ def prepare_tilelang_import(tilelang_root: str | Path | None = None) -> Path:
     )
     os.environ["TL_ROOT"] = str(tl_root)
     prepend_pythonpath(os.environ, str(tl_root))
-    if str(tl_root) not in sys.path:
-        sys.path.insert(0, str(tl_root))
+    tl_root_str = str(tl_root)
+    # Keep TL_ROOT at sys.path front to avoid resolving the sibling
+    # package xllm/compiler/tilelang as top-level `tilelang`.
+    sys.path = [p for p in sys.path if p != tl_root_str]
+    sys.path.insert(0, tl_root_str)
     os.environ.setdefault("ACL_OP_INIT_MODE", "1")
     return tl_root
 

xllm/compiler/tilelang/common/toolchain.py

@@ -1,5 +1,8 @@
-from __future__ import annotations
-
+import importlib
+import os
+import pkgutil
+import re
+from dataclasses import dataclass
 from pathlib import Path
 
 from ...common.manifest import KernelFamilyManifest

xllm/compiler/tilelang/targets/ascend/build.py

@@ -1,10 +1,13 @@
 from __future__ import annotations
 
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from dataclasses import dataclass
+import os
 from pathlib import Path
 
 from ...common.cache import compute_cache_key, is_cache_hit
 from ...common.manifest import KernelAbi, KernelFamilyManifest, KernelVariantManifest
-from ...common.spec import DispatchField, KernelCompileSpec, TilelangKernel
+from ...common.spec import DispatchField, KernelCompileSpec, KernelSpec, TilelangKernel
 from ...common.toolchain import repo_root, run_checked
 from . import abi_entry, kernel_registry, toolchain
 from .kernel_registry import RegisteredKernelFamily
@@ -13,6 +16,21 @@
 from .toolchain import AscendBuildContext, TILELANG_BISHENG_COMMON_FLAGS
 
 
+@dataclass(frozen=True)
+class _VariantBuildPlan:
+    compile_spec: KernelCompileSpec
+    kernel_spec: KernelSpec
+    generated_source: Path
+    compiled_binary: Path
+    cache_key: str
+
+
+@dataclass(frozen=True)
+class _VariantBuildResult:
+    manifest: KernelVariantManifest
+    kernel_abi: KernelAbi
+
+
 def _variant_entry_symbol(spec: KernelCompileSpec) -> str:
     kernel_entry_name = spec.entry_name or spec.kernel_name
     return f"{kernel_entry_name}__{spec.variant_key}_call"
@@ -116,7 +134,8 @@ def build_kernel_family(
     existing_manifest = _read_family_manifest(manifest_path)
     dependency_files = _build_dependency_files(family)
 
-    variant_manifests: list[KernelVariantManifest] = []
+    variant_manifest_by_key: dict[str, KernelVariantManifest] = {}
+    uncached_plans: list[_VariantBuildPlan] = []
     family_kernel_abi: KernelAbi | None = None
 
     for compile_spec, kernel_spec in family.spec_pairs:
@@ -167,24 +186,37 @@ def build_kernel_family(
                     "All variants in a TileLang kernel must share the same exported "
                     f"C ABI. Mismatch found in variant {compile_spec.variant_key!r}."
                 )
-            variant_manifests.append(
-                KernelVariantManifest(
-                    variant_key=compile_spec.variant_key,
-                    specialization=dict(compile_spec.specialization),
-                    dispatch_values=dict(compile_spec.dispatch_values),
-                    generated_source=cached_variant.generated_source,
-                    compiled_binary=cached_variant.compiled_binary,
-                    entry_symbol=cached_variant.entry_symbol,
-                    cache_key=cached_variant.cache_key,
-                    toolchain_options=dict(context.toolchain_options),
-                    fingerprint=dict(context.fingerprint),
-                    compile_definitions=kernel_spec.render_compile_definitions(
-                        entry_symbol=cached_variant.entry_symbol
-                    ),
-                )
+            variant_manifest_by_key[compile_spec.variant_key] = KernelVariantManifest(
+                variant_key=compile_spec.variant_key,
+                specialization=dict(compile_spec.specialization),
+                dispatch_values=dict(compile_spec.dispatch_values),
+                generated_source=cached_variant.generated_source,
+                compiled_binary=cached_variant.compiled_binary,
+                entry_symbol=cached_variant.entry_symbol,
+                cache_key=cached_variant.cache_key,
+                toolchain_options=dict(context.toolchain_options),
+                fingerprint=dict(context.fingerprint),
+                compile_definitions=kernel_spec.render_compile_definitions(
+                    entry_symbol=cached_variant.entry_symbol
+                ),
             )
             continue
 
+        uncached_plans.append(
+            _VariantBuildPlan(
+                compile_spec=compile_spec,
+                kernel_spec=kernel_spec,
+                generated_source=generated_source,
+                compiled_binary=compiled_binary,
+                cache_key=cache_key,
+            )
+        )
+
+    compile_cwd = repo_root()
+
+    def _run_variant_job(plan: _VariantBuildPlan) -> _VariantBuildResult:
+        compile_spec = plan.compile_spec
+        kernel_spec = plan.kernel_spec
         source = family.kernel_cls.generate_source(**compile_spec.specialization)
         entry_symbol = _variant_entry_symbol(compile_spec)
         rendered_source = abi_entry.rename_variant_internal_symbols(
@@ -194,44 +226,79 @@ def build_kernel_family(
             compile_spec.variant_key,
         )
         kernel_abi = abi_entry.parse_kernel_abi(rendered_source, entry_symbol)
-        if family_kernel_abi is None:
-            family_kernel_abi = kernel_abi
-        elif kernel_abi != family_kernel_abi:
-            raise ValueError(
-                "All variants in a TileLang kernel must share the same exported "
-                f"C ABI. Mismatch found in variant {compile_spec.variant_key!r}."
-            )
-        generated_source.write_text(rendered_source, encoding="utf-8")
+        plan.generated_source.write_text(rendered_source, encoding="utf-8")
 
         compile_cmd = [
             context.bisheng_executable,
             f"--cce-aicore-arch={context.bisheng_arch}",
             *TILELANG_BISHENG_COMMON_FLAGS,
             f"-Dg_tilingKey=g_tilingKey__{compile_spec.variant_key}",
             *[f"-I{include_dir}" for include_dir in context.include_dirs],
-            str(generated_source),
+            str(plan.generated_source),
             "-c",
             "-o",
-            str(compiled_binary),
+            str(plan.compiled_binary),
         ]
-        run_checked(compile_cmd, cwd=repo_root())
-
-        variant_manifests.append(
-            KernelVariantManifest(
-                variant_key=compile_spec.variant_key,
-                specialization=compile_spec.specialization,
-                dispatch_values=compile_spec.dispatch_values,
-                generated_source=str(generated_source),
-                compiled_binary=str(compiled_binary),
-                entry_symbol=entry_symbol,
-                cache_key=cache_key,
-                toolchain_options=dict(context.toolchain_options),
-                fingerprint=dict(context.fingerprint),
-                compile_definitions=kernel_spec.render_compile_definitions(
-                    entry_symbol=entry_symbol
-                ),
-            )
+        run_checked(compile_cmd, cwd=compile_cwd)
+        manifest = KernelVariantManifest(
+            variant_key=compile_spec.variant_key,
+            specialization=dict(compile_spec.specialization),
+            dispatch_values=dict(compile_spec.dispatch_values),
+            generated_source=str(plan.generated_source),
+            compiled_binary=str(plan.compiled_binary),
+            entry_symbol=entry_symbol,
+            cache_key=plan.cache_key,
+            toolchain_options=dict(context.toolchain_options),
+            fingerprint=dict(context.fingerprint),
+            compile_definitions=kernel_spec.render_compile_definitions(
+                entry_symbol=entry_symbol
+            ),
         )
+        return _VariantBuildResult(manifest=manifest, kernel_abi=kernel_abi)
+
+    if uncached_plans:
+        max_workers = max(1, os.cpu_count() or 1)
+        if max_workers == 1:
+            for plan in uncached_plans:
+                result = _run_variant_job(plan)
+                if family_kernel_abi is None:
+                    family_kernel_abi = result.kernel_abi
+                elif result.kernel_abi != family_kernel_abi:
+                    raise ValueError(
+                        "All variants in a TileLang kernel must share the same exported "
+                        "C ABI. Mismatch found in variant "
+                        f"{result.manifest.variant_key!r}."
+                    )
+                variant_manifest_by_key[result.manifest.variant_key] = result.manifest
+        else:
+            with ThreadPoolExecutor(max_workers=max_workers) as executor:
+                future_to_plan = {
+                    executor.submit(_run_variant_job, plan): plan
+                    for plan in uncached_plans
+                }
+                for future in as_completed(future_to_plan):
+                    plan = future_to_plan[future]
+                    try:
+                        result = future.result()
+                    except Exception as exc:
+                        raise RuntimeError(
+                            "Ascend variant build failed for variant "
+                            f"{plan.compile_spec.variant_key!r}"
+                        ) from exc
+                    if family_kernel_abi is None:
+                        family_kernel_abi = result.kernel_abi
+                    elif result.kernel_abi != family_kernel_abi:
+                        raise ValueError(
+                            "All variants in a TileLang kernel must share the same exported "
+                            "C ABI. Mismatch found in variant "
+                            f"{result.manifest.variant_key!r}."
+                        )
+                    variant_manifest_by_key[result.manifest.variant_key] = result.manifest
+
+    variant_manifests: list[KernelVariantManifest] = [
+        variant_manifest_by_key[compile_spec.variant_key]
+        for compile_spec, _ in family.spec_pairs
+    ]
 
     if family_kernel_abi is None:
         raise ValueError(