diff --git a/tests/unit/test_d1_rescale_debug.py b/tests/unit/test_d1_rescale_debug.py index ddb06a6e..af538441 100644 --- a/tests/unit/test_d1_rescale_debug.py +++ b/tests/unit/test_d1_rescale_debug.py @@ -43,7 +43,7 @@ def test_multi_kv_debug(hd=64, s_k=256): pv_n_tile = kernel.pv_n_tile n_pv_tiles = kernel.n_pv_tiles print(f" n_kv_tiles={kernel.n_kv_tiles}, pv_n_tile={pv_n_tile}, n_pv_tiles={n_pv_tiles}") - print(f" tmem_o0_offset={kernel.tmem_o0_offset}, tmem_p0_offset={kernel.tmem_p0_offset}") + # tmem_o0_offset is set in _setup, not __init__ stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)