Skip to content

[Metax][MCTLE] Add NVMMAShared encoding builder binding - #1007

Open
cc-c122 wants to merge 1 commit into
flagos-ai:mainfrom
cc-c122:metax-mctle-nvmma-shared-binding
Open

[Metax][MCTLE] Add NVMMAShared encoding builder binding#1007
cc-c122 wants to merge 1 commit into
flagos-ai:mainfrom
cc-c122:metax-mctle-nvmma-shared-binding

Conversation

@cc-c122

@cc-c122 cc-c122 commented Aug 18, 2026

Copy link
Copy Markdown

Summary

MetaX MCTLE already provides the shared-memory allocation and local-pointer
builder APIs required by TLE, including create_local_alloc and
create_local_pointers.

However, it is missing the
make_nv_mma_shared_encoding_attr builder binding required by TLE kernels
using NVMMASharedEncodingAttr.

This change exposes the existing TritonGPU
NVMMASharedEncodingAttr construction through the MCTLE builder, aligning
MCTLE with the corresponding generic TLE builder API.

This change only adds the missing builder binding. It does not introduce a new
IR type, compiler pass, or MMA lowering implementation.

Testing

Tested on MetaX C550 with:

  • PyTorch: 2.8.0+metax3.7.2.0
  • Triton: 3.6.0
  • Target: GPUTarget(backend='maca', arch=80, warp_size=64)

A minimal TLE kernel was used to exercise the following path:

NVMMAShared alloc
    -> local_ptr
    -> shared store/load
    -> global output

Without this binding

The MCTLE builder contains:

make_nv_mma_shared_encoding_attr : False
create_local_alloc               : True
create_local_pointers            : True

The same kernel fails during TLE compilation at tle.gpu.alloc with:

AttributeError:
'triton._C.libtriton.ir.builder' object has no attribute
'make_nv_mma_shared_encoding_attr'

With this binding

The MCTLE builder contains:

make_nv_mma_shared_encoding_attr : True
create_local_alloc               : True
create_local_pointers            : True

The same kernel compiles and runs successfully on MetaX C550:

max_abs_diff : 0.0
equal        : True

NVMMAShared alloc + local_ptr + store/load: PASS

@CLAassistant

CLAassistant commented Aug 18, 2026

Copy link
Copy Markdown

CLA assistant check
All committers have signed the CLA.

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.

2 participants