Skip to content

[BACKEND] Add Metax McTLE support#742

Open
cjh9368 wants to merge 4 commits into
flagos-ai:triton_v3.6.xfrom
cjh9368:triton_v3.6.x_mctle
Open

[BACKEND] Add Metax McTLE support#742
cjh9368 wants to merge 4 commits into
flagos-ai:triton_v3.6.xfrom
cjh9368:triton_v3.6.x_mctle

Conversation

@cjh9368

@cjh9368 cjh9368 commented Jul 3, 2026

Copy link
Copy Markdown

METAX backend support for TLE Lite tile primitives in this patch

This patch adds METAX/MACA support for the main TLE Lite tile primitives:

  • tle.extract_tile(x, index, tile_shape)

    • Extracts a tile of shape tile_shape from a register tensor at a tile-grid index.
    • Supported index forms:
      • Multi-dimensional static index: tuple/list of int/constexpr, linearized at compile time.
      • Scalar static index: int/constexpr, treated as the linearized tile index.
      • Scalar dynamic index: scalar i32/i64 tl.tensor.
      • Multi-dimensional dynamic index: tuple/list containing tl.tensor, row-major linearized in the frontend.
    • tile_shape must be compile-time constants, match the source rank, and divide the source shape evenly along each dimension.
  • tle.insert_tile(x, tile, index)

    • Inserts a register tile into a source tensor at a tile-grid index.
    • The result type matches the source tensor type.
    • Index forms and constraints are the same as extract_tile.
    • The tile shape is derived from the tile operand and must divide the source shape evenly.

Lowering path

Key points for the METAX implementation:

  • The Python frontend is moved under triton.experimental.tle for the METAX backend, matching the shared TLE frontend import path.
  • The old third_party/metax/language/maca/mctle Python location is removed.
  • Tile operations are emitted through the mctle pybind builder interfaces:
    • create_extract_tile
    • create_insert_tile
  • The generated IR uses METAX-local mctle dialect ops:
    • mctle.extract_tile
    • mctle.insert_tile
  • The METAX backend continues to load the mctle dialect and related lowering passes through triton._C.libtriton.mctle.

Scope

This PR only enables:

  • tle.extract_tile
  • tle.insert_tile

@cjh9368 cjh9368 requested a review from zhzhcookie as a code owner July 3, 2026 07:11
@CLAassistant

CLAassistant commented Jul 3, 2026

Copy link
Copy Markdown

CLA assistant check
All committers have signed the CLA.

@cjh9368 cjh9368 force-pushed the triton_v3.6.x_mctle branch 2 times, most recently from 179bb7e to a38cd82 Compare July 6, 2026 07:44
@cjh9368 cjh9368 force-pushed the triton_v3.6.x_mctle branch from a38cd82 to a6af50f Compare July 6, 2026 07:48
@cjh9368 cjh9368 force-pushed the triton_v3.6.x_mctle branch from a6af50f to 04c4a40 Compare July 6, 2026 08:29
@zhzhcookie

Copy link
Copy Markdown
Collaborator

Please format the code by pre-commit, and add tle tests in .github/workflows/metax3.6-build-and-test.yml.
Please update md5 in https://github.com/flagos-ai/FlagTree/blob/triton_v3.6.x/python/setup_tools/setup_helper.py#L698, and change url to metaxTritonPlugin-cpython3.12-x86_64_v0.6.1.tar.gz.

@cjh9368 cjh9368 requested a review from Galaxy1458 as a code owner July 7, 2026 10:28
@cjh9368

cjh9368 commented Jul 8, 2026

Copy link
Copy Markdown
Author

build

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants