[BUG][CuTe DSL] TVM-FFI env stream detection fails when all GPU tensors are nested in a tuple

[reubenconducts]

Which component has the problem?

CuTe DSL

Bug Report

Describe the bug
TVMFFIFunctionBuilder.find_env_stream (cutlass/base_dsl/tvm_ffi_builder/tvm_ffi_builder.py) only scans the top-level parameter list for a GPU spec.Tensor to source the implicit environment stream from. If every GPU tensor is nested inside a tuple/NamedTuple argument (a spec.TupleParam), detection fails and compilation raises:

ValueError: EnvStream cannot be detected in `...` we need parameters to contain GPU Tensors

even though the tuple does contain GPU tensors.

Steps/Code to reproduce bug

"""Repro: TVM-FFI env stream detection fails when GPU tensors are nested in a tuple param.

Requirements: nvidia-cutlass-dsl (repro'd on 4.6.0.dev0 and 4.7.0)
Usage:
    python repro_tvm_ffi_nested_env_stream.py               # shows the failure
    python repro_tvm_ffi_nested_env_stream.py --workaround  # flatten patch; compiles & runs
"""

import argparse
import sys
from functools import wraps

import torch

import cuda.bindings.driver as cuda
import cutlass
import cutlass.cute as cute
from cutlass.cute.runtime import from_dlpack, make_fake_stream

N = 128


@cute.kernel
def _add_one_kernel(src: cute.Tensor, dst: cute.Tensor):
    tidx, _, _ = cute.arch.thread_idx()
    dst[tidx] = src[tidx] + 1.0


@cute.jit
def add_one_flat(src: cute.Tensor, dst: cute.Tensor, stream: cuda.CUstream):
    _add_one_kernel(src, dst).launch(grid=[1, 1, 1], block=[N, 1, 1], stream=stream)


@cute.jit
def add_one_nested(tensors: tuple, stream: cuda.CUstream):
    src, dst = tensors
    _add_one_kernel(src, dst).launch(grid=[1, 1, 1], block=[N, 1, 1], stream=stream)


def apply_flatten_workaround():
    """Flatten TupleParams so find_env_stream can see nested GPU tensors."""
    from cutlass.base_dsl.tvm_ffi_builder import spec
    from cutlass.base_dsl.tvm_ffi_builder.tvm_ffi_builder import TVMFFIFunctionBuilder

    original = TVMFFIFunctionBuilder.find_env_stream

    def flatten(params):
        for param in params:
            if isinstance(param, spec.TupleParam):
                yield from flatten(param.params)
            else:
                yield param

    @wraps(original)
    def find_env_stream(self, params):
        stream = original(self, params)
        return stream if stream is not None else original(self, list(flatten(params)))

    TVMFFIFunctionBuilder.find_env_stream = find_env_stream


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--workaround", action="store_true")
    args = parser.parse_args()
    if args.workaround:
        apply_flatten_workaround()

    src_t = torch.arange(N, dtype=torch.float32, device="cuda")
    dst_t = torch.zeros(N, dtype=torch.float32, device="cuda")
    src = from_dlpack(src_t, assumed_align=16, enable_tvm_ffi=True).mark_layout_dynamic()
    dst = from_dlpack(dst_t, assumed_align=16, enable_tvm_ffi=True).mark_layout_dynamic()
    env_stream = make_fake_stream(use_tvm_ffi_env_stream=True)

    # Baseline: tensors as top-level params -> env stream is detected, compiles fine.
    compiled_flat = cute.compile(add_one_flat, src, dst, env_stream, options="--enable-tvm-ffi")
    compiled_flat(src_t, dst_t)
    torch.cuda.synchronize()
    assert torch.equal(dst_t, src_t + 1)
    print("flat params:   compiled and ran OK")

    # Same kernel, tensors nested in a tuple param -> EnvStream detection fails.
    try:
        compiled_nested = cute.compile(
            add_one_nested, (src, dst), env_stream, options="--enable-tvm-ffi"
        )
    except Exception as exc:  # DSLRuntimeError (ICE) wrapping the ValueError
        print(f"nested params: FAILED to compile: {exc.__cause__ or exc}")
        sys.exit(1)

    dst_t.zero_()
    stream = torch.cuda.Stream()
    with torch.cuda.stream(stream):
        compiled_nested((src_t, dst_t))
    torch.cuda.synchronize()
    assert torch.equal(dst_t, src_t + 1)
    print("nested params: compiled and ran OK (on a non-default env stream)")


if __name__ == "__main__":
    main()

Expected behavior
find_env_stream should be able to detect tensors in a tuple.

Environment details (please complete the following information):

• nvidia-cutlass-dsl==4.6.0.dev0 and 4.7.0, GB300.