functional BHSD fix

aoti_attention.py

@@ -82,18 +82,17 @@ def sparse_attention_functional(
     untiling, then the compression branch scaled by the trained per-row gate. `gate` is
     `[B, H, S_real, D]` (the eager path's `gate.transpose(1, 2)`), applied after untililing.
     """
-    batch, heads, padded_len, dim = query.shape
+    batch, heads, seq_len, dim = query.shape  # inputs are `[B, H, S_real, D]` (BHSD)
     n_tiles = variable_block_sizes.numel()
-    n_real = untile_index.numel()
-    num_q_video = n_tiles - num_prefix_tiles
+    padded_len = n_tiles * 64
 
     # Tile: scatter the packed rows into the kernels' `[B, H, S_pad, D]` layout, pad slots zero.
     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))
+    query_tiled.index_copy_(2, untile_index, query)
     key_tiled = torch.zeros_like(query_tiled)
-    key_tiled.index_copy_(2, untile_index, key.transpose(1, 2))
+    key_tiled.index_copy_(2, untile_index, key)
     value_tiled = torch.zeros_like(query_tiled)
-    value_tiled.index_copy_(2, untile_index, value.transpose(1, 2))
+    value_tiled.index_copy_(2, untile_index, value)
 
     # 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