[BACKEND][MTHREADS] Replace monkey-patch with spec_path dynamic lookup by Stardep-lmc · Pull Request #687 · flagos-ai/FlagTree

Stardep-lmc · 2026-06-10T09:56:14Z

Move 48 mthreads-specialized python files from third_party/mthreads/python/triton/ to third_party/mthreads/backend/spec/triton/ (only files that differ from main)
Add spec_path(__path__) to __init__.py in runtime/, compiler/, language/, tools/ for runtime path specialization via flagtree_spec.py
Remove monkey-patch code from mthreads.py, keep build-time package merging logic and expose apply_mthreads_setup_args()
Update setup.py to call apply_mthreads_setup_args for non-editable installs when FLAGTREE_BACKEND=mthreads

- git mv 12 files with genuine mthreads-specific changes from third_party/mthreads/python/triton/ to backend/spec/triton/ - Add spec_path(__path__) to triton/, compiler/, runtime/, language/, tools/ __init__.py for runtime path specialization - Add max_shared_mem to main compiler/__init__.py and triton/__init__.py - Remove monkey-patch code from mthreads.py (_wrap_setup, _patch_setup_for_mthreads_python_root), keep build-time package merging logic and expose apply_mthreads_setup_args() - Update setup.py to call apply_mthreads_setup_args for non-editable installs when FLAGTREE_BACKEND=mthreads

- Remove runtime/adjust_kernel_param.py (just an older upstream version) - Remove experimental/gluon/__init__.py (missing spec_path setup from upstream) - Restore both files to third_party/mthreads/python/triton/

…only - Remove merge_mthreads_packages, merge_mthreads_package_dir, patch_mthreads_cmdclass, and apply_mthreads_setup_args - get_package_dir() now returns empty dict (no root package override) - Remove apply_mthreads_setup_args call from setup.py - Both editable and non-editable installs now use the standard spec_path mechanism via backends/mthreads/spec/triton/

…/__init__.py - git mv language/core.py and language/standard.py to spec (have mthreads-specific functions: squeeze, unsqueeze, to_tensor, _experimental_descriptor_load, _experimental_descriptor_store) - Add these symbols to main python/triton/language/__init__.py imports and __all__ - Remove async_task import (not in mthreads version) - Revert setup.py to direct setup() call (no intermediate variable)

Spec core.py imports _tuple_create from triton._utils. The mthreads _utils.py has apply_with_path, is_iterable, and _tuple_create that are not in the main version.

Spec core.py is missing upstream async_task, causing ImportError. These files have too much upstream divergence to spec cleanly. Instead, mthreads-specific additions will be merged to main files.

- Add squeeze and unsqueeze functions to python/triton/language/standard.py (mthreads-specific functions now merged to main) - Remove async_task, _experimental_descriptor_load, _experimental_descriptor_store from language/__init__.py (not in main core.py)

…_init__.py - Revert squeeze/unsqueeze from standard.py (stay in mthreads fork) - Revert _experimental_descriptor_* from core.py (stay in mthreads fork) - Keep only spec_path(__path__) in language/__init__.py - Restore language/__init__.py imports to match main upstream

…_language - Remove third_party/mthreads/python/triton/runtime/adjust_kernel_param.py (AI hallucination, not needed) - Remove third_party/mthreads/python/triton/experimental/gluon/__init__.py (identical to main, not callable from third_party) - Add third_party/mthreads/backend/spec/__init__.py with init_language() for mthreads-specific symbols - Update python/triton/__init__.py to explicitly call spec('init_language') Co-Authored-By: Claude <noreply@anthropic.com>

These two files contain mthreads-specific changes that init_language() alone cannot cover. By placing them at backend/spec/triton/language/, spec_path() will automatically use them when FLAGTREE_BACKEND=mthreads. Co-Authored-By: Claude <noreply@anthropic.com>

flagcx_wrapper.py and flagcx/include may not exist in all flagcx repo versions. Check existence before copying to avoid FileNotFoundError. Co-Authored-By: Claude <noreply@anthropic.com>

Move code_generator, semantic, jit, autotuner, interpreter, cache, and _async_compile from the old monkey-patch location into third_party/mthreads/backend/spec/triton/ so the spec_path import mechanism finds the correct mthreads-customized versions instead of falling back to the generic upstream originals. This fixes near-universal test failures caused by the baseline state where the old monkey-patch was disabled but critical files were missing from the spec directory. Co-Authored-By: Claude <noreply@anthropic.com>

Co-Authored-By: Claude <noreply@anthropic.com>

_libtriton.so_ is a compiled C extension installed under the main triton package root (triton/_C/), not under the mthreads backend directory. Use sys.modules['triton'].__file__ to locate the correct path for _C, backends, compiler, and language base files. Co-Authored-By: Claude <noreply@anthropic.com>

- Fix _constexpr annotation error in squeeze/unsqueeze by importing constexpr at module level in backend spec __init__.py, so the JIT code generator can resolve the name from fn.__globals__. - Add shutil.which fallback in both _filecheck.py modules so FileCheck is resolved from system PATH when the local binary is missing, with a clear error message when neither is found. Co-Authored-By: Claude <noreply@anthropic.com>

The module-level raise FileNotFoundError prevented pytest from collecting any tests in test_frontend.py and test_line_info.py. Move the check into a lazy _get_filecheck_path() helper so the module can be imported without FileCheck installed, and only tests that actually invoke FileCheck will fail. Co-Authored-By: Claude <noreply@anthropic.com>

CMake copies FileCheck to ${TRITON_WHEEL_DIR}/FileCheck (the triton package root), but the mthreads _filecheck.py lives 4 levels deeper at triton/backends/mthreads/spec/triton/. Add _filecheck_wheel as a fallback search path between the local directory and system PATH. Co-Authored-By: Claude <noreply@anthropic.com>

This change (5884c75) was unrelated to the mthreads spec_path migration and should not be part of this PR. Co-Authored-By: Claude <noreply@anthropic.com>

All 66 files under third_party/mthreads/python/triton/ were deleted: - 38 byte-identical to python/triton/ (pure duplicates) - 28 with version-skew diffs (older/stale copies or non-mthreads gluon code) The 7 genuinely mthreads-specific files were already git mv'd to backend/spec/triton/ and are loaded via spec_path. The remaining copies served no purpose — this directory is not on sys.path and not referenced by any build system. Retained: third_party/mthreads/python/src/ (C++ sources for CMake) and third_party/mthreads/python/test/ (mthreads-specific tests). Co-Authored-By: Claude <noreply@anthropic.com>

… python/triton/ The previous approach (e9e222f) merged mthreads-specific changes from deleted third_party/mthreads/python/triton/ files directly into shared python/triton/ files. This broke other backends because those changes depend on mthreads' own libtriton modifications. This fix reverts to the correct spec_path approach: - Only add spec_path(__path__) to 4 __init__.py files (safe no-op for other backends — spec_path returns early when spec dir doesn't exist) - Put all 17 mthreads-specific file versions in spec/triton/ where they are only activated when FLAGTREE_BACKEND=mthreads - Add missing to_tensor to triton.language __all__ No shared python/triton/ non-init files are modified. Other backends will use their own versions of _core.py, _layouts.py, _semantic.py, compile.py, ragged_tma.py, etc. exactly as before. Co-Authored-By: Claude <noreply@anthropic.com>

Upstream triton.experimental.gluon.language.__init__.py imports thread_barrier from ._core, but the mthreads spec_path copy only had the older barrier() function. This caused 10 test collection errors across test_frontend and tle tests. Co-Authored-By: Claude <noreply@anthropic.com>

* [SPEC] Fix spec __init__.py, _filecheck.py, _internal_testing.py in triton dir * [SPEC] Fix spec _filecheck.py, _utils.py * Apply code-format changes * [SPEC] Fix spec knobs.py * [SPEC] Fix spec _filecheck.py, testing.py * [SPEC] Fix spec _utils.py * [SPEC] Fix spec _filecheck.py, _internal_testing.py * Apply code-format changes * [SPEC] Fix spec _utils.py, init_language * Apply code-format changes * [SPEC] Fix spec testing.py --------- Co-authored-by: flagtree-bot <flagtree_ai@163.com>

* [SPEC] Fix spec language init * [SPEC] Fix spec compiler & triton init * [SPEC] Fix spec gluon init * [SPEC] Fix spec gluon

Stardep-lmc requested review from Galaxy1458 and zhzhcookie as code owners June 10, 2026 09:56

github-actions Bot added mthreads triton_v3.6.x labels Jun 10, 2026

Stardep-lmc force-pushed the feat/mthreads-spec-path-third-party branch from af6d076 to 57c010b Compare June 10, 2026 11:38

zhzhcookie reviewed Jun 13, 2026

View reviewed changes

Comment thread third_party/mthreads/backend/spec/triton/language/__init__.py

zhzhcookie reviewed Jun 13, 2026

View reviewed changes

Comment thread third_party/mthreads/backend/spec/triton/runtime/__init__.py

zhzhcookie reviewed Jun 13, 2026

View reviewed changes

Comment thread third_party/mthreads/python/triton/runtime/adjust_kernel_param.py

zhzhcookie reviewed Jun 13, 2026

View reviewed changes

Comment thread third_party/mthreads/backend/spec/triton/runtime/errors.py

Stardep-lmc force-pushed the feat/mthreads-spec-path-third-party branch from 048318c to 64f9b09 Compare June 16, 2026 10:59

Stardep-lmc force-pushed the feat/mthreads-spec-path-third-party branch 3 times, most recently from fa6a008 to e2129ef Compare July 2, 2026 10:51

Stardep-lmc and others added 17 commits July 3, 2026 05:09

Apply code-format changes

ff87231

Remove files with no genuine mthreads changes from spec

6b9bd8d

- Remove runtime/adjust_kernel_param.py (just an older upstream version) - Remove experimental/gluon/__init__.py (missing spec_path setup from upstream) - Restore both files to third_party/mthreads/python/triton/

Move _utils.py to spec (needed by core.py for _tuple_create)

5036735

Spec core.py imports _tuple_create from triton._utils. The mthreads _utils.py has apply_with_path, is_iterable, and _tuple_create that are not in the main version.

Revert core.py and standard.py move to spec (version skew)

160e285

Spec core.py is missing upstream async_task, causing ImportError. These files have too much upstream divergence to spec cleanly. Instead, mthreads-specific additions will be merged to main files.

Apply code-format changes

5163f26

Make flagcx _copy_required_files resilient to missing source files

5884c75

flagcx_wrapper.py and flagcx/include may not exist in all flagcx repo versions. Check existence before copying to avoid FileNotFoundError. Co-Authored-By: Claude <noreply@anthropic.com>

Apply code-format changes

d95b3a8

Re-trigger CI for mthreads spec_path file migration

285f798

Co-Authored-By: Claude <noreply@anthropic.com>

chore: re-trigger CI

bec6897

Co-Authored-By: Claude <noreply@anthropic.com>

flagtree-bot and others added 3 commits July 3, 2026 05:09

Apply code-format changes

52655d7

Apply code-format changes

a0eb5a0

Stardep-lmc force-pushed the feat/mthreads-spec-path-third-party branch from 41d53de to a0eb5a0 Compare July 2, 2026 21:10

Stardep-lmc and others added 5 commits July 3, 2026 07:10

Apply code-format changes

6120c17

Apply code-format changes

8758953

zhzhcookie reviewed Jul 3, 2026

View reviewed changes

Comment thread python/setup_tools/utils/default.py

Revert default.py to upstream: remove unrelated flagcx resilience change

d9c87f6

This change (5884c75) was unrelated to the mthreads spec_path migration and should not be part of this PR. Co-Authored-By: Claude <noreply@anthropic.com>

Stardep-lmc requested review from menchunlei and sunnycase as code owners July 3, 2026 13:52

Stardep-lmc and others added 3 commits July 6, 2026 16:32

Merge branch 'triton_v3.6.x' into feat/mthreads-spec-path-third-party

6531fea

Merge branch 'triton_v3.6.x' into feat/mthreads-spec-path-third-party

ce82d8e

Stardep-lmc force-pushed the feat/mthreads-spec-path-third-party branch 2 times, most recently from ff98a93 to b30dd2c Compare July 9, 2026 07:25

Stardep-lmc force-pushed the feat/mthreads-spec-path-third-party branch from a7fc4b4 to dc5b998 Compare July 9, 2026 11:18

Stardep-lmc and others added 2 commits July 11, 2026 20:39

Merge branch 'main' into feat/mthreads-spec-path-third-party

63e6f97

github-actions Bot added the main label Jul 13, 2026

github-actions Bot added the amd label Jul 14, 2026

[SPEC] Fix spec language | compiler | triton | gluon init (#790)

d589da7

* [SPEC] Fix spec language init * [SPEC] Fix spec compiler & triton init * [SPEC] Fix spec gluon init * [SPEC] Fix spec gluon

github-actions Bot added the xpu label Jul 15, 2026

zhzhcookie added 2 commits July 15, 2026 10:03

Merge branch 'main' into feat/mthreads-spec-path-third-party

121d529

[SPEC] Add Triton 3.7 comments

4e6def8

zhzhcookie force-pushed the feat/mthreads-spec-path-third-party branch from 2c9bfb4 to 4e6def8 Compare July 15, 2026 13:12

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[BACKEND][MTHREADS] Replace monkey-patch with spec_path dynamic lookup#687

[BACKEND][MTHREADS] Replace monkey-patch with spec_path dynamic lookup#687
Stardep-lmc wants to merge 38 commits into
mainfrom
feat/mthreads-spec-path-third-party

Stardep-lmc commented Jun 10, 2026 •

edited by zhzhcookie

Loading

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

Stardep-lmc commented Jun 10, 2026 • edited by zhzhcookie Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Stardep-lmc commented Jun 10, 2026 •

edited by zhzhcookie

Loading