From de439bcd75709ee2d9cbf44ab0eb6494610d612e Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sat, 23 May 2026 06:40:05 +0000 Subject: [PATCH] fix: cuda.CUstream import --- tests/unit/test_tmem_budget.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/unit/test_tmem_budget.py b/tests/unit/test_tmem_budget.py index a03eff32..3fb5d759 100644 --- a/tests/unit/test_tmem_budget.py +++ b/tests/unit/test_tmem_budget.py @@ -4,6 +4,7 @@ Uses @cute.jit to construct MMA objects inside a compiled context, where OperandMajorMode values are valid MLIR operands. """ import torch, math, sys +import cuda.bindings.driver as cuda import cutlass, cutlass.cute as cute, cutlass.utils as utils from cutlass.cute.nvgpu import tcgen05 from cutlass import Float32, BFloat16, Int32 @@ -97,7 +98,7 @@ def probe_hd(hd): mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k)) - stream = cutlass.cuda.CUstream(torch.cuda.current_stream().cuda_stream) + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) probe = BudgetProbe(head_dim=hd) print(f' Compiling hd={hd}...', flush=True)