diff --git a/mojo_opset/backends/ttx/kernels/npu/swa.py b/mojo_opset/backends/ttx/kernels/npu/swa.py index a57bb7c2e..6f3337e1f 100644 --- a/mojo_opset/backends/ttx/kernels/npu/swa.py +++ b/mojo_opset/backends/ttx/kernels/npu/swa.py @@ -1930,6 +1930,8 @@ def _sdpa_single_block_bwd_dkdv( # Load (transposed) K block q = tl.load(Q_block_ptr, boundary_check=(0, 1), padding_option="zero") + tl.extra.cann.extension.compile_hint(q, "cv_pipeline_lazy_load", True) + # q_T = tl.trans(q) # qkT = tl.dot(k, q_T) # [BLOCK_N, BLOCK_M] k_T = tl.trans(k) @@ -1949,6 +1951,8 @@ def _sdpa_single_block_bwd_dkdv( # -- Compute dV ---- # dv = pT @ do do = tl.load(DO_block_ptr, boundary_check=(0, 1), padding_option="zero") + tl.extra.cann.extension.compile_hint(do, "cv_pipeline_lazy_load", True) + dv = tl.dot(tl.trans(p_cast), do, dv_ptr) # -- Compute dS ---- @@ -1990,6 +1994,7 @@ def _sdpa_single_block_bwd_dq( # Load (transposed) K block k = tl.load(K_block_ptr, boundary_check=(0, 1), padding_option="zero") + tl.extra.cann.extension.compile_hint(k, "cv_pipeline_lazy_load", True) k_T = tl.trans(k) qk = tl.dot(q, k_T) # [BLOCK_M, BLOCK_N] # tl.extra.cann.extension.compile_hint(qk, "tile_cube_loop")