functional index_copy fix

aoti_attention.py

@@ -88,16 +88,12 @@ def sparse_attention_functional(
     num_q_video = n_tiles - num_prefix_tiles
 
     # Tile: scatter the packed rows into the kernels' `[B, H, S_pad, D]` layout, pad slots zero.
-    query_tiled = torch.index_copy(
-        torch.zeros((batch, heads, padded_len, dim), dtype=query.dtype, device=query.device),
-        2, untile_index, query.transpose(1, 2),
-    )
-    key_tiled = torch.index_copy(
-        torch.zeros_like(query_tiled), 2, untile_index, key.transpose(1, 2),
-    )
-    value_tiled = torch.index_copy(
-        torch.zeros_like(query_tiled), 2, untile_index, value.transpose(1, 2),
-    )
+    query_tiled = torch.zeros((batch, heads, padded_len, dim), dtype=query.dtype, device=query.device)
+    query_tiled.index_copy_(2, untile_index, query.transpose(1, 2))
+    key_tiled = torch.zeros_like(query_tiled)
+    key_tiled.index_copy_(2, untile_index, key.transpose(1, 2))
+    value_tiled = torch.zeros_like(query_tiled)
+    value_tiled.index_copy_(2, untile_index, value.transpose(1, 2))
 
     # Per-head fp32 pooled tile scores.
     q_pool = query_tiled.view(batch, heads, n_tiles, 64, dim).sum(dim=3, dtype=torch.float32) / tile_divisor