Fix: add cutlass import to test_d1_qk512

2026-05-24 14:20:32 +00:00
parent ce267909ad
commit 449a6e7ede
1 changed files with 1 additions and 1 deletions
--- a/tests/unit/test_d1_qk512.py
+++ b/tests/unit/test_d1_qk512.py
@@ -1,6 +1,6 @@
 """Minimal hd=512 test: ONLY QK GEMM, no softmax, no PV.
 Goal: isolate whether the compilation hang is from QK or softmax/PV."""
-import torch, math, time
+import torch, math, time, cutlass
 import cutlass.cute as cute
 import cutlass.torch as ct
 import cuda.bindings.driver as cuda