[Cute] Update remaining make_fragment_like to make_rmem_tensor_like

Author: henrylhtsang

flash_attn/cute/flash_bwd_sm90.py

@@ -1344,7 +1344,7 @@ def mma_one_m_block(
         acc_dP = mma_dov_fn(A_idx=smem_idx_Q, wg_wait=1)
 
         if const_expr(self.score_mod_bwd is not None):
-            acc_S_pre = cute.make_fragment_like(acc_S)
+            acc_S_pre = cute.make_rmem_tensor_like(acc_S)
             cute.autovec_copy(acc_S, acc_S_pre)
 
         if const_expr(self.score_mod is not None):
@@ -1498,7 +1498,7 @@ def epilogue_dKV(
         warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx())
 
         if const_expr(self.qhead_per_kvhead == 1):
-            rdV = cute.make_fragment_like(acc_dV, self.dtype)
+            rdV = cute.make_rmem_tensor_like(acc_dV, self.dtype)
             rdV.store(acc_dV.load().to(self.dtype))
             rdK = utils.cvt_f16(acc_dK, self.dtype)
 

flash_attn/cute/paged_kv.py

@@ -153,7 +153,7 @@ def load_KV(self, n_block: Int32, sX: cute.Tensor, K_or_V: str):
         seqlenk_row_limit = self.seqlen_k - n_block * self.n_block_size if n_block >= 0 else 0
         for m in cutlass.range_constexpr(cute.size(tXsX, mode=[1])):
             row_valid = tXcX[0, m, 0][0] < seqlenk_row_limit
-            should_load = cute.make_fragment_like(tXsX[None, m, 0], cute.Boolean)
+            should_load = cute.make_rmem_tensor_like(tXsX[None, m, 0], cute.Boolean)
             should_load.fill(row_valid)
 
             page = self.tPrPage[m]