[BUG] Constexpr missing in dynamic control-flow regions

Open
  • Is something wrong?
#3,632 1 comment 0 reactions 0 assignees View on GitHub

Maintainers usually reply within 4 days

@nodeeeeee is already working on this.

Since Sep 3, 2026.

  • #3580 by @nodeeeeee — open

Assessment

Difficulty
3/5
Estimated time
1-2 days
Newbie friendliness
68/100
Issue type
Bug
Clarity
Mostly clear
Activity status
Active
Tech stack
python
Domain
compilers

Research direction

Start by tracing the AST preprocessor's handling of Constexpr parameter annotations and RegionAnalyzer.visit_Call(), using the reduced receiver.store example as the reproduction. The receiver of a Constexpr parameter should not be added to invoked_args or carried through the dynamic if; verify that the example lowers without the plain Python Meta value error.

Written by the indexing model from the issue text.

Description

? - Needs Triage bug CuTe DSL inactive-30d
Which component has the problem?

CuTe DSL

Bug Report

Motivation

I encountered this while implementing FlashAttention. The implementation naturally groups the load and MMA paths into JIT methods on a compile-time configuration object, and selects between them using the runtime warp index. A reduced version looks like this:

class FlashAttention:
    @cute.jit
    def load(self, tensor, pipeline):
        ...

    @cute.jit
    def mma(self, tensor, pipeline):
        ...


@cute.jit
def fa4_device_body(
    fa: cutlass.Constexpr[FlashAttention],
    tensor: cute.Tensor,
    storage,
):
    warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx())

    pipeline = make_pipeline(
        storage=storage,
        num_stages=fa.num_stages,
    )

    if warp_idx < 4:
        cute.arch.setmaxregister_decrease(fa.num_producer_regs)
        fa.load(tensor, pipeline)
    else:
        cute.arch.setmaxregister_increase(fa.num_mma_regs)
        fa.mma(tensor, pipeline)

Here, self in FlashAttention.load() and FlashAttention.mma() is a compile-time method receiver. The fa parameter in fa4_device_body() refers to the same kind of compile-time object, made explicit by its Constexpr annotation.

Writing the device body as an actual instance method makes this code compile, but only because the compiler has a special case that excludes a receiver whose name is literally self. Passing the same object as a named Constexpr[FlashAttention] parameter should have equivalent staging semantics, but currently does not.

Problem

CuTe DSL treats Constexpr[T] parameters as compile-time Python meta values. They do not have an MLIR representation and must not become arguments or results of runtime control-flow regions.

However, when a method is called through a Constexpr receiver inside a dynamic if, the receiver is incorrectly captured as a region argument:

class Receiver:
    @cute.jit
    def store(self, output):
        ...


@cute.kernel
def kernel(receiver: cutlass.Constexpr[Receiver], output: cute.Tensor):
    tidx, _, _ = cute.arch.thread_idx()

    if tidx == Int32(0):
        receiver.store(output)

The region analyzer currently treats the base object of every method call as a mutable runtime value, except when its name is literally self.

Consequently, receiver is added to the values carried through the dynamic if. Lowering then attempts to flatten the plain Python Receiver instance into MLIR values and fails with:

error[TYPE_DYNAMIC_EXPR_UNSUPPORTED]:
A value carried through this `if` is a plain Python value (Meta value)
(a `Receiver`) that cannot be turned into a Runtime value (Staged value)

This means self.method() works because of the existing self special case, while an explicitly annotated receiver: Constexpr[Receiver] does not, despite having the same staging semantics.

Root cause

Constexpr is a property of the original function parameter binding. The AST preprocessor removes parameter annotations while transforming the function, so the later control-flow region analysis no longer knows which names refer to Constexpr parameters.

RegionAnalyzer.visit_Call() therefore sees only a method call on a Python object and adds its receiver to invoked_args.

Dominant language
C++
Stars
10.7k
Forks
2.2k
Avg merge
3d 20h
Merged PRs (30d)
4

Getting set up

This project ships no dev container, Dockerfile or contributing guide, so setting up is up to you: start from its README, and see our first-contribution guide for the general steps.

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from NVIDIA/cutlass

All issues in NVIDIA/cutlass

Similar issues

More C++ issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.