[CuTeDSL] Isolate internal MLIR/NVVM dependencies to prevent cross-version breakage

[icavan]

Summary

This is a repository-wide CuTeDSL compatibility issue, not an SM100-only issue.

cuLA currently imports implementation-detail modules under cutlass._mlir from shared helpers and kernels targeting multiple architectures. These imports include generated NVVM bindings as well as MLIR LLVM, vector, arithmetic, IR, and CuTe dialect wrappers.

The SM100 tcgen05_ld/st failure after upgrading CutDSL 4.5.2 to 4.5.3 is the first confirmed incident, but the root cause is general: private/generated compiler bindings do not provide a stable cross-version contract.

The current dependency range makes this especially risky:

nvidia-cutlass-dsl>=4.4.2,<4.7,!=4.5.0

A resolver can select a syntactically valid but incompatible CutDSL patch release, and pip check still succeeds. Failures are delayed until the affected kernel is imported or JIT-compiled.

Repository-wide internal API surface

At least the following areas directly import cutlass._mlir modules:

Area	Representative files	Internal dependency	Usage
Shared PTX/IR helpers	cula/ops/ptx.py	ir, arith, llvm, vector	inline PTX, multi-register results, vector bitcasts/slices, wide stores
SM90 KDA	cula/ops/kda/sm90/_common.py	llvm	inline PTX primitives
SM90 Lightning	cula/ops/lightning/sm90/schedule.py, prefill_kernel.py	llvm, internal CuTe dialect	scheduling/control instructions and inline PTX
Decode/verification	cula/ops/kda/decode/mtp_kvbuffer.py, cula/lightning/la_verify_kvbuffer.py	arith, llvm	mma.sync wrappers and multi-register extraction
SM100 Lightning/KDA	cula/ops/lightning/prefill_sm100.py, cula/ops/kda/sm100/delta_h.py	llvm	inline PTX helpers
SM100 tcgen05	cula/ops/sm100/ptx.py	ir, arith, llvm, generated nvvm	TMEM copy/load/store/fence and MMA/MMA-WS

The exact file list will evolve, so the fix should include an automated rule preventing new scattered imports rather than relying only on this table.

Current generated _nvvm tcgen05 call surface

cula/ops/sm100/ptx.py is currently the largest direct generated-NVVM dependency:

Internal operation	Logical wrappers	Purpose
_nvvm.tcgen05_ld	1	TMEM to register vector<Nxi32> load
_nvvm.tcgen05_st	1	Register vector<Nxi32> to TMEM store
_nvvm.tcgen05_cp	2	SMEM to TMEM 128x256b and 128x128b copies
_nvvm.tcgen05_fence	2	before_thread_sync / after_thread_sync ordering
_nvvm.tcgen05_mma	2	TF32 SS/TS MMA with a four-word write-disable mask
_nvvm.tcgen05_mma_ws	4	TF32/F16 weight-stationary SS/TS MMA with collector controls

It also directly depends on generated enums including Tcgen05LdStShape, Tcgen05CpShape, Tcgen05FenceKind, Tcgen05GroupKind, Tcgen05MMAKind, Tcgen05MMACollectorBBuffer, and Tcgen05MMACollectorOp.

Confirmed case study: CutDSL 4.5.3 breaks SM100 KDA backward

CutDSL 4.5.2 generated bindings:

tcgen05_ld(res, shape, num, tmem_addr, *, pack=None, half_split_offset=None, ...)
tcgen05_st(shape, num, tmem_addr, r, *, unpack=None, half_split_offset=None, ...)

CutDSL 4.5.3 generated bindings:

tcgen05_ld(res, shape, tmem_addr, *, pack=None, offset=None, ...)
tcgen05_st(shape, tmem_addr, val, *, unpack=None, offset=None, ...)

In 4.5.3, the .xN variant is inferred from the result/value vector type. The explicit num argument was removed and the store keyword changed from r to val.

The first observed failure is:

TypeError: tcgen05_ld() got an unexpected keyword argument 'num'

Fixing only ld would expose the corresponding st break later in the same backward path. This demonstrates why patching individual call sites is insufficient.

Risks

• Patch-level CutDSL changes can break imports, keyword arguments, enums, result types, or lowering behavior.
• Errors occur only when a particular JIT path is reached, making them difficult to detect with normal packaging checks.
• The same pattern exists across SM90, SM100, shared decode code, and architecture-independent helpers.
• Duplicated inline-assembly and result-extraction wrappers make fixes inconsistent.
• Future architectures will multiply the maintenance burden if private bindings remain scattered.
• Even when code still compiles, type/lowering changes can silently affect PTX/SASS, register pressure, synchronization, or performance.

Proposed solution

Phase 0 — restore affected versions and define support policy

• Add a centralized version adapter for confirmed breaks such as the CutDSL 4.5.3 tcgen05_ld/st change.
• Keep version decisions at Python/JIT compile time; they must not become runtime device branches.
• Constrain dependency ranges to versions actually covered by CI.
• Fail early with a clear error when a CutDSL version is outside the supported matrix.

Phase 1 — create a repository-wide low-level compatibility layer

Create one documented compatibility package for all unavoidable compiler/ISA-level operations, for example:

cula/ops/cutedsl_compat/
  version.py
  llvm.py
  nvvm.py
  sm90.py
  sm100.py

The exact layout is flexible, but kernel modules should not import cutlass._mlir directly. The compatibility layer should own:

• CutDSL version detection and feature gates;
• internal/public symbol mapping;
• inline-assembly construction and operand constraints;
• MLIR vector/struct creation and result extraction;
• architecture and PTX-version checks;
• consistent diagnostics for unsupported combinations.

Add a lint or source-level CI check that rejects new cutlass._mlir imports outside this package.

Phase 2 — use the most stable available emission interface

Choose the backend per operation in this priority order.

1. Public CuTe/CutDSL abstractions

Use public APIs whenever they model the complete semantics:

• typed tensors, copy atoms, and cute.copy for data movement;
• Ld32x32bOp / St32x32bOp for TMEM load/store;
• Cp128x256bOp / Cp128x128bOp for SMEM-to-TMEM copy;
• public MMA atoms and cute.gemm when masks/collector behavior are representable;
• public cute.arch synchronization and architecture primitives.

This is the preferred high-level route, but migration must preserve layouts, register counts, instruction selection, and scheduling.

2. CutDSL 4.7 NVVM primitives

CutDSL 4.7 releases a supported NVVM primitives layer. Use it as the preferred formal low-level interface when public CuTe atoms do not expose the complete instruction form.

Likely candidates include output-lane masks, weight-stationary collector controls, exact fence forms, and other architecture-specific operations. Before raising the current <4.7 bound:

• make the 4.7 artifact available in the package index used by cuLA CI;
• map the existing SM90/SM100 low-level operations to official primitives;
• verify generated PTX/SASS, register usage, synchronization, and performance;
• retain pre-4.7 adapters only for versions cuLA intentionally supports.

3. Centralized llvm.inline_asm fallback

For instructions not covered by public abstractions or the 4.7 primitive layer, emit PTX through one centralized llvm.inline_asm wrapper.

Compared with generated _nvvm Python bindings, inline assembly is more stable in practice because cuLA controls the PTX spelling, operands, register constraints, predicates, and version adaptation. It avoids Python signature drift such as the 4.5.3 num / r removal.

TileLang provides a useful tcgen05 reference:

https://github.com/tile-ai/tilelang/blob/main/tilelang/contrib/cutedsl/gemm_tcgen05.py

The fallback must explicitly handle:

• PTX/architecture requirements;
• input/output constraints and multi-register result types;
• elect.sync and single-thread issue semantics;
• side-effect annotations, predicates, and fences;
• large .xN operation splitting where LLVM operand counts become problematic.

The current LLVM MLIR bridge may itself remain an internal dependency on older CutDSL versions; the goal is to quarantine that dependency behind one stable cuLA API rather than scatter it across kernels.

4. Generated _nvvm compatibility shim for legacy versions only

Direct generated bindings should be the last resort and should exist only inside the compatibility layer for intentionally supported pre-4.7 releases.

Phase 3 — migrate incrementally

Suggested order:

1. Consolidate duplicated shared inline-assembly/vector helpers.
2. Migrate public-copy and synchronization operations.
3. Move SM90/decode mma.sync and control helpers behind the common interface.
4. Migrate SM100 TMEM load/store/copy.
5. Migrate or isolate masked MMA, MMA-WS, collector, and exact fence operations.
6. Remove legacy version paths after the support window ends.

Each step should compare generated PTX/SASS and performance before deleting the old path.

Testing and CI

The compatibility matrix should cover both version and architecture dimensions.

CutDSL versions

• 4.4.2
• 4.5.2
• 4.5.3
• 4.6.1
• 4.7 using official NVVM primitives
• the newest version allowed by pyproject.toml

Architecture paths

• SM90 KDA and Lightning compile/execute smoke tests;
• shared decode and mma.sync wrappers;
• SM100 KDA forward/backward;
• SM100 masked SS/TS MMA;
• SM100 weight-stationary TF32/F16 collector paths;
• import/compile checks for compatibility helpers that do not require a GPU.

Validation

• numerical correctness;
• PTX/SASS instruction selection;
• register usage and spills;
• synchronization correctness;
• compile time and runtime performance.

Acceptance criteria

• The issue is treated as a repository-wide CuTeDSL compatibility effort, with the SM100 failure retained as a regression case.
• Kernel modules do not add new direct imports from cutlass._mlir.
• Existing direct imports are removed or isolated in the compatibility package.
• Public CuTe/CutDSL APIs are used wherever they express the required semantics.
• CutDSL 4.7+ uses official NVVM primitives for supported low-level operations.
• Unsupported operations use one documented and tested llvm.inline_asm fallback.
• Generated _nvvm bindings are limited to explicit legacy-version shims.
• SM90 and SM100 paths pass the supported CutDSL version matrix.
• Dependency bounds match the versions covered by CI.
• Migration introduces no numerical, synchronization, PTX/SASS, or performance regression.

[icavan]

Status after #119

Confirmed against main at merge commit 5370b14637c3bb4bc4242b334e6b7b00a05c6e41.

Resolved by #119

• The confirmed CuTeDSL 4.5.2/4.5.3 tcgen05_ld / tcgen05_st Python-signature break is handled by capability detection rather than a package-version gate.
• Load and store signatures are detected independently, including the optional num parameter and the r versus val store-value keyword.
• Unknown or ambiguous generated bindings fail early with an explicit diagnostic.
• Unit coverage was added for legacy, inferred-shape, mixed, and unsupported signatures.
• The stale SM100 intracard CP test calls were migrated from the removed private _no_cp keyword to use_intracard_cp=False.

Why #118 remains open

#119 fixes the immediate regression and completes only part of Phase 0. The repository-wide acceptance criteria in this issue are not yet satisfied:

• main still contains direct cutlass._mlir imports in 10 Python files (19 import lines), spanning shared helpers, SM90, SM100, KDA, Lightning, and decode paths.
• There is no source-level CI rule preventing new direct internal imports outside a compatibility layer.
• The dependency range remains nvidia-cutlass-dsl>=4.4.2,<4.7,!=4.5.0; the proposed supported-version CI matrix is not implemented.
• CuTeDSL 4.7 public NVVM primitives have not been adopted.
• Remaining generated _nvvm operations (cp, fence, mma, and mma_ws) are not yet isolated or migrated to public primitives / centralized llvm.inline_asm fallbacks.

Accordingly, #119 should be recorded as the immediate 4.5.3 compatibility fix, but it should not close this broader tracking issue. The next focused PR should add the import-boundary enforcement and move the remaining low-level imports behind the compatibility package before attempting operation-by-operation primitive migration.