NVIDIA / NVIDIA/cutlass

[BUG] [cuteDSL] mark_compact_shape_dynamic fails to pass shape args to JIT kernel (resolves to 0 at runtime)

Open
#2,794 3 comments 0 reactions 1 assignee View on GitHub

@Flink-ddd is already working on this.

Since Apr 4, 2026.

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

Description

Issue Description

Hi, I've encountered a bug where dynamic tensor shapes marked with mark_compact_shape_dynamic are not correctly marshalled to the device kernel. When using static shapes, the kernel receives the correct values.

Behavior:

  • With mark_layout_dynamic & mark_compact_shape_dynamic: Inside the @cute.kernel, the dynamic dimensions resolve to 0 and strides contain garbage values, even when passing a valid concrete runtime tensor.
  • Without these dynamic markers (Static Shape): The kernel correctly receives the shape values (e.g., (128, 128)).

Steps to Reproduce:

  1. Create a template GPU tensor.
  2. (Bug Case) Mark layout and dimensions as dynamic using mark_layout_dynamic and mark_compact_shape_dynamic.
  3. Compile a JIT kernel that launches a device kernel.
  4. Pass a concrete runtime tensor.
  5. Inspect tensor.shape inside the device kernel via cute.printf.

Reproduction Script:

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

# Device Kernel
@cute.kernel
def device_kernel(tensor: cute.Tensor):
    tidx, _, _ = cute.arch.thread_idx()
    if tidx == 0:
        # BUG: These prints show (0, 0) when using dynamic shape
        # CORRECT: Shows (128, 128) when using static shape
        cute.printf("Device Kernel >> Tensor: {}\n", tensor)
        cute.printf("Device Kernel >> Shape: ({}, {})\n", tensor.shape[0], tensor.shape[1])

# Host Launcher
@cute.jit
def launcher_func(tensor: cute.Tensor):
    # Launch kernel with 1 thread to inspect values
    device_kernel(tensor).launch(grid=(1,1,1), block=(1,1,1))

def main():
    M, N = 128, 128
    
    # 1. Prepare Template (GPU Tensor -> gmem)
    fake_input = torch.zeros(M, N, device='cuda', dtype=torch.bfloat16)
    template = from_dlpack(fake_input, assumed_align=8)
    
    # [BUG REPRODUCTION]
    # Case A (Bug): Dynamic Shape -> Device sees (0, 0)
    # Case B (Working): Static Shape -> Device sees (128, 128)
    # Comment out the following lines to switch to Case B
    template = template.mark_layout_dynamic(leading_dim=1)
    template = template.mark_compact_shape_dynamic(
        mode=1,
        stride_order=(0, 1),
        divisibility=16
    )
    
    print(f"Compiling with template: {template}")
    
    compiled_func = cute.compile(launcher_func, template)
    
    print("\nRunning compiled function...")
    input_real = torch.randn(M, N, device='cuda', dtype=torch.bfloat16)
    input_runtime = from_dlpack(input_real, assumed_align=8)
    
    compiled_func(input_runtime)
    torch.cuda.synchronize()

if __name__ == "__main__":
    main()

Actual Output (Dynamic Shape - Bug):

Device Kernel >> Shape: (0, 0)

Expected Output (Static Shape - Reference):

Device Kernel >> Shape: (128, 128)

Thanks.

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.

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.