NVIDIA / NVIDIA/cutlass

[BUG] [CuTeDSL] cute.composition throws when rhs has a dynamic shape

Open
#3,470 2 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

? - Needs Triage bug CuTe DSL
Dominant language
C++
Stars
10.5k
Forks
2.1k
Avg merge
3d 11h
Merged PRs (30d)
7

Description

Which component has the problem?

CuTe DSL

Bug Report

Describe the bug
Invoking cute.composition on a RHS layout that contains one mode with a dynamic size causes an exception to be raised. At first I thought that this was by design, because if cute.composition were to verify that RHS's codomain falls within LHS's domain, it needs to know RHS's domain. But apparently such a check isn't actually performed (not even for fully-static layouts).

Steps/Code to reproduce bug

import torch
import cutlass.cute as cute
from cutlass.cute.runtime import from_dlpack


@cute.jit
def repro(tensor: cute.Tensor):
    lhs = cute.make_layout((4096, 8192), stride=(1, 4096))
    rhs = cute.make_layout(
        (32, tensor.shape[1] // 8, 8, 128),
        stride=(
            cute.E(0),
            8 * cute.E(1),
            cute.E(1),
            32 * cute.E(0),
        ),
    )
    print(cute.composition(lhs, rhs))


tensor = from_dlpack(torch.empty((4096, 8192), device="cuda"))
tensor = tensor.mark_compact_shape_dynamic(mode=1)
repro(tensor)

Raises:

cutlass._mlir._mlir_libs._site_initialize.<locals>.MLIRError: Operation creation failed:
error: unknown: unable to compute the following composition: '!cute.layout<"(4096,8192):(1,4096)">' o '!cute.layout<"(32,?,8,128):(1@0,8@1,1@1,32@0)">'

Environment details (please complete the following information):
CuTeDSL version 4.7.0

Contributor guide

No contributing guide indexed for this repository

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.

Research direction

Start by running the provided Python reproduction with the stated CuTe DSL 4.7.0 environment and trace the cute.composition entry point with the dynamic RHS layout. Compare it with fully static layouts and confirm that the reproduction no longer raises; add or update a regression check for the corrected behavior.

Written by the indexing model from the issue text.

Assessment

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.