Skip to content

refactor: separate Ascend TileLang DSL implementations from AOT descriptors. - #2259

Merged
zhang-minchao merged 2 commits into
xLLM-AI:mainfrom
zhang-minchao:refactor/python-tilelang-kernels
Aug 18, 2026
Merged

refactor: separate Ascend TileLang DSL implementations from AOT descriptors.#2259
zhang-minchao merged 2 commits into
xLLM-AI:mainfrom
zhang-minchao:refactor/python-tilelang-kernels

Conversation

@zhang-minchao

@zhang-minchao zhang-minchao commented Aug 18, 2026

Copy link
Copy Markdown
Collaborator

Summary

  • Move the reusable Ascend TileLang DSL and JIT implementations from the compiler package to xllm/python/kernels_npu/tilelang.
  • Keep compiler-owned AOT metadata and lowering logic in thin descriptors under xllm/compiler/tilelang/targets/ascend/aot.
  • Initialize the active Python kernel package explicitly after native torch.ops.xllm_ops registration.
  • Update AOT cache dependencies, kernel development documentation, and the existing embedded-Python tests.

Motivation

The previous compiler/tilelang/targets/ascend/kernels modules mixed two responsibilities:

  1. reusable TileLang program builders and Python JIT launchers;
  2. AOT-only registration, specialization, ABI, and code-generation metadata.

This made Python JIT callers depend on the compiler package. Moving the DSL implementation directly under kernels_npu would previously also trigger the NPU semantic package and require the native xLLM operators while the C++ binary was still being built.

This refactor separates those responsibilities and keeps the dependency direction one-way:

Python model runtime -> kernels_npu semantic API -> kernels_npu/tilelang
Ascend AOT compiler -----------------------------> kernels_npu/tilelang

The TileLang leaf package does not depend on the NPU semantic API, torch.ops.xllm_ops, the AOT compiler, or the native xLLM extension.

Key Changes

  • xllm/python/kernels_npu/tilelang becomes the owner of TileLang DSL builders, JIT launchers, and implementation-local helpers.
  • xllm/compiler/tilelang/targets/ascend/aot owns only AOT descriptors and imports the shared TileLang builders.
  • xllm.python.initialize_runtime() selects and initializes the active platform kernel package after native operator registration.
  • AOT cache keys include the imported DSL implementation modules so source changes invalidate generated artifacts.
  • The obsolete xllm/compiler/tilelang/targets/ascend/kernels package is removed.

Validation

  • Pre-commit checks pass: Ruff, codespell, and clang-format.
  • The NPU build compiles all 10 migrated Ascend TileLang AOT kernel families.
  • AOT registry output matches the pre-refactor registry for all migrated families and specializations.
  • mtp_async_input_builder_test passes all 8 tests.
  • npu_xllm_ops_test passes all 6 tests.

Reviewer Notes

Please focus on the package dependency direction, the native-operator/runtime initialization order, and the split between reusable DSL implementations and compiler-only AOT descriptors.

- initialize platform kernel packages explicitly after native operator registration
- move TileLang DSL/JIT implementations under kernels_npu and keep compiler-owned AOT descriptors thin
- preserve AOT registry behavior while updating cache dependencies, documentation, and tests
- drop standalone runtime and layout tests that are not part of the test suite
- keep existing C++ tests aligned with explicit Python runtime initialization
@zhang-minchao
zhang-minchao force-pushed the refactor/python-tilelang-kernels branch from b42b65c to 1990b12 Compare August 18, 2026 05:46
@zhang-minchao zhang-minchao changed the title refactor: separate python tilelang kernels from aot descriptors. refactor: separate Ascend TileLang DSL implementations from AOT descriptors. Aug 18, 2026
@zhang-minchao
zhang-minchao merged commit 6b9f852 into xLLM-AI:main Aug 18, 2026
16 of 25 checks passed
AmpereKowk pushed a commit to AmpereKowk/xllm that referenced this pull request Aug 20, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants