From 7b872cc41eabcd0320cc0231bd3a8107fd91c18d Mon Sep 17 00:00:00 2001 From: Mook <68294499+Godmook@users.noreply.github.com> Date: Sun, 26 Apr 2026 15:14:37 -0700 Subject: [PATCH] [Perf] Skip bool-mask round-trip in block-sparse VSA attention (#1243) --- .../python/fastvideo_kernel/__init__.py | 7 + .../fastvideo_kernel/block_sparse_attn.py | 281 ++++++++++++------ .../python/fastvideo_kernel/ops.py | 19 +- .../fastvideo_kernel/triton_kernels/index.py | 114 ++++++- 4 files changed, 323 insertions(+), 98 deletions(-) diff --git a/fastvideo-kernel/python/fastvideo_kernel/__init__.py b/fastvideo-kernel/python/fastvideo_kernel/__init__.py index 560717d94..0e6b8da6c 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/__init__.py +++ b/fastvideo-kernel/python/fastvideo_kernel/__init__.py @@ -5,6 +5,11 @@ from fastvideo_kernel.ops import ( video_sparse_attn, ) +from fastvideo_kernel.block_sparse_attn import ( + block_sparse_attn, + block_sparse_attn_from_indices, +) + from fastvideo_kernel.vmoba import ( moba_attn_varlen, process_moba_input, @@ -22,6 +27,8 @@ from fastvideo_kernel.turbodiffusion_ops import ( __all__ = [ "sliding_tile_attention", "video_sparse_attn", + "block_sparse_attn", + "block_sparse_attn_from_indices", "moba_attn_varlen", "process_moba_input", "process_moba_output", diff --git a/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn.py b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn.py index 119349ebc..b5a1217b3 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn.py +++ b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn.py @@ -1,3 +1,5 @@ +"""Autograd-enabled block-sparse attention. Index-native ops with a bool-mask compat shim.""" + from __future__ import annotations import os @@ -6,6 +8,11 @@ from typing import Tuple import torch +# --------------------------------------------------------------------------- +# Backend selection helpers +# --------------------------------------------------------------------------- + + def _get_sm90_ops(): try: from fastvideo_kernel._C import fastvideo_kernel_ops # type: ignore @@ -25,38 +32,66 @@ def _is_sm90() -> bool: def _force_triton() -> bool: - # Force Triton even on SM90 and even if the compiled extension is available. - # Useful for CI / debugging / parity testing. return os.environ.get("FASTVIDEO_KERNEL_VSA_FORCE_TRITON", "0") == "1" -def _map_to_index(block_map: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Preferred map->index conversion used by the wrapper. +# --------------------------------------------------------------------------- +# Index helpers +# --------------------------------------------------------------------------- - This wrapper **requires** the Triton implementation. - If Triton (or the Triton map_to_index module) is not available, it raises. - """ + +def _map_to_index(block_map: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + """Compact a bool block_map to (q2k_idx, q2k_num). Legacy path only.""" if block_map.dim() == 3: block_map = block_map.unsqueeze(0) if block_map.dim() != 4: - raise ValueError(f"block_map must be [B,H,Q,KV] (or [H,Q,KV]), got shape={tuple(block_map.shape)}") + raise ValueError( + f"block_map must be [B,H,Q,KV] (or [H,Q,KV]), " + f"got shape={tuple(block_map.shape)}" + ) if block_map.dtype != torch.bool: block_map = block_map.to(torch.bool) - if not block_map.is_cuda: - raise RuntimeError("block_map must be a CUDA tensor (Triton map_to_index required).") + raise RuntimeError( + "block_map must be a CUDA tensor (Triton map_to_index required)." + ) try: - from fastvideo_kernel.triton_kernels.index import map_to_index as triton_map_to_index # local import - except Exception as e: + from fastvideo_kernel.triton_kernels.index import map_to_index as triton_map_to_index + except Exception as e: # pragma: no cover - environment issue raise ImportError( "Triton map_to_index is required but not available. " - "Ensure Triton is installed and fastvideo_kernel.triton_kernels.index is importable." + "Ensure Triton is installed and " + "fastvideo_kernel.triton_kernels.index is importable." ) from e return triton_map_to_index(block_map) +def _invert_indices_for_backward( + q2k_idx: torch.Tensor, + q2k_num: torch.Tensor, + num_kv_blocks: int, +) -> Tuple[torch.Tensor, torch.Tensor]: + from fastvideo_kernel.triton_kernels.index import invert_indices + return invert_indices(q2k_idx, q2k_num, num_kv_blocks=num_kv_blocks) + + +def _as_int32_contig(t: torch.Tensor, name: str) -> torch.Tensor: + """Return `t` as a contiguous int32 tensor, raising a clear error on CPU input.""" + if not t.is_cuda: + raise RuntimeError(f"{name} must be a CUDA tensor, got device={t.device}") + if t.dtype != torch.int32: + t = t.to(torch.int32) + if not t.is_contiguous(): + t = t.contiguous() + return t + + +# --------------------------------------------------------------------------- +# Triton backend custom ops (index-native) +# --------------------------------------------------------------------------- + + @torch.library.custom_op( "fastvideo_kernel::block_sparse_attn_triton", mutates_args=(), @@ -66,34 +101,40 @@ def block_sparse_attn_triton( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, - block_map: torch.Tensor, + q2k_idx: torch.Tensor, + q2k_num: torch.Tensor, variable_block_sizes: torch.Tensor, ) -> Tuple[torch.Tensor, torch.Tensor]: - q = q.contiguous() - k = k.contiguous() - v = v.contiguous() - block_map = block_map.to(torch.bool) - q2k_idx, q2k_num = _map_to_index(block_map) - - from fastvideo_kernel.triton_kernels.block_sparse_attn_triton import ( # local import + from fastvideo_kernel.triton_kernels.block_sparse_attn_triton import ( triton_block_sparse_attn_forward, ) - o, M = triton_block_sparse_attn_forward(q, k, v, q2k_idx, q2k_num, variable_block_sizes) + o, M = triton_block_sparse_attn_forward( + q.contiguous(), + k.contiguous(), + v.contiguous(), + q2k_idx, + q2k_num, + variable_block_sizes, + ) return o, M - @torch.library.register_fake("fastvideo_kernel::block_sparse_attn_triton") def _block_sparse_attn_triton_fake( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, - block_map: torch.Tensor, + q2k_idx: torch.Tensor, + q2k_num: torch.Tensor, variable_block_sizes: torch.Tensor, ) -> Tuple[torch.Tensor, torch.Tensor]: o = torch.empty_like(q) - M = torch.empty((q.shape[0], q.shape[1], q.shape[2]), device=q.device, dtype=torch.float32) + M = torch.empty( + (q.shape[0], q.shape[1], q.shape[2]), + device=q.device, + dtype=torch.float32, + ) return o, M @@ -109,20 +150,32 @@ def block_sparse_attn_backward_triton( v: torch.Tensor, o: torch.Tensor, M: torch.Tensor, - block_map: torch.Tensor, + q2k_idx: torch.Tensor, + q2k_num: torch.Tensor, variable_block_sizes: torch.Tensor, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - grad_output = grad_output.contiguous() - block_map = block_map.to(torch.bool) - q2k_idx, q2k_num = _map_to_index(block_map) - k2q_idx, k2q_num = _map_to_index(block_map.transpose(-1, -2).contiguous()) - - from fastvideo_kernel.triton_kernels.block_sparse_attn_triton import ( # local import + from fastvideo_kernel.triton_kernels.block_sparse_attn_triton import ( triton_block_sparse_attn_backward, ) + num_kv_blocks = int(variable_block_sizes.numel()) + k2q_idx, k2q_num = _invert_indices_for_backward( + q2k_idx, q2k_num, num_kv_blocks + ) + # q/k/v are saved from the user-facing inputs and may be non-contiguous; + # o/M are kernel outputs so are already contiguous. dq, dk, dv = triton_block_sparse_attn_backward( - grad_output, q, k, v, o, M, q2k_idx, q2k_num, k2q_idx, k2q_num, variable_block_sizes + grad_output.contiguous(), + q.contiguous(), + k.contiguous(), + v.contiguous(), + o, + M, + q2k_idx, + q2k_num, + k2q_idx, + k2q_num, + variable_block_sizes, ) return dq, dk, dv @@ -135,7 +188,8 @@ def _block_sparse_attn_backward_triton_fake( v: torch.Tensor, o: torch.Tensor, M: torch.Tensor, - block_map: torch.Tensor, + q2k_idx: torch.Tensor, + q2k_num: torch.Tensor, variable_block_sizes: torch.Tensor, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: dq = torch.empty_like(q) @@ -144,19 +198,28 @@ def _block_sparse_attn_backward_triton_fake( return dq, dk, dv -def _backward_triton(ctx, grad_o, grad_M): - q, k, v, o, M, block_map, variable_block_sizes = ctx.saved_tensors - dq, dk, dv = block_sparse_attn_backward_triton(grad_o, q, k, v, o, M, block_map, variable_block_sizes) - return dq, dk, dv, None, None - - def _setup_context_triton(ctx, inputs, output): - q, k, v, block_map, variable_block_sizes = inputs + q, k, v, q2k_idx, q2k_num, variable_block_sizes = inputs o, M = output - ctx.save_for_backward(q, k, v, o, M, block_map, variable_block_sizes) + ctx.save_for_backward(q, k, v, o, M, q2k_idx, q2k_num, variable_block_sizes) -block_sparse_attn_triton.register_autograd(_backward_triton, setup_context=_setup_context_triton) +def _backward_triton(ctx, grad_o, grad_M): + q, k, v, o, M, q2k_idx, q2k_num, variable_block_sizes = ctx.saved_tensors + dq, dk, dv = block_sparse_attn_backward_triton( + grad_o, q, k, v, o, M, q2k_idx, q2k_num, variable_block_sizes + ) + return dq, dk, dv, None, None, None + + +block_sparse_attn_triton.register_autograd( + _backward_triton, setup_context=_setup_context_triton +) + + +# --------------------------------------------------------------------------- +# SM90 backend custom ops (index-native) +# --------------------------------------------------------------------------- @torch.library.custom_op( @@ -168,21 +231,21 @@ def block_sparse_attn_sm90( q_padded: torch.Tensor, k_padded: torch.Tensor, v_padded: torch.Tensor, - block_map: torch.Tensor, + q2k_idx: torch.Tensor, + q2k_num: torch.Tensor, variable_block_sizes: torch.Tensor, ) -> Tuple[torch.Tensor, torch.Tensor]: block_sparse_fwd, _ = _get_sm90_ops() if block_sparse_fwd is None: raise ImportError("fastvideo_kernel_ops.block_sparse_fwd is not available") - q_padded = q_padded.contiguous() - k_padded = k_padded.contiguous() - v_padded = v_padded.contiguous() - block_map = block_map.to(torch.bool) - q2k_idx, q2k_num = _map_to_index(block_map) - o_padded, lse_padded = block_sparse_fwd( - q_padded, k_padded, v_padded, q2k_idx, q2k_num, variable_block_sizes.int() + q_padded.contiguous(), + k_padded.contiguous(), + v_padded.contiguous(), + q2k_idx, + q2k_num, + variable_block_sizes, ) return o_padded, lse_padded @@ -192,11 +255,16 @@ def _block_sparse_attn_sm90_fake( q_padded: torch.Tensor, k_padded: torch.Tensor, v_padded: torch.Tensor, - block_map: torch.Tensor, + q2k_idx: torch.Tensor, + q2k_num: torch.Tensor, variable_block_sizes: torch.Tensor, ) -> Tuple[torch.Tensor, torch.Tensor]: o = torch.empty_like(q_padded) - lse = torch.empty((q_padded.shape[0], q_padded.shape[1], q_padded.shape[2], 1), device=q_padded.device, dtype=torch.float32) + lse = torch.empty( + (q_padded.shape[0], q_padded.shape[1], q_padded.shape[2], 1), + device=q_padded.device, + dtype=torch.float32, + ) return o, lse @@ -212,30 +280,34 @@ def block_sparse_attn_backward_sm90( v_padded: torch.Tensor, o_padded: torch.Tensor, lse_padded: torch.Tensor, - block_map: torch.Tensor, + q2k_idx: torch.Tensor, + q2k_num: torch.Tensor, variable_block_sizes: torch.Tensor, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: _, block_sparse_bwd = _get_sm90_ops() if block_sparse_bwd is None: raise ImportError("fastvideo_kernel_ops.block_sparse_bwd is not available") - grad_output_padded = grad_output_padded.contiguous() - block_map = block_map.to(torch.bool) - k2q_idx, k2q_num = _map_to_index(block_map.transpose(-1, -2).contiguous()) + num_kv_blocks = int(variable_block_sizes.numel()) + k2q_idx, k2q_num = _invert_indices_for_backward( + q2k_idx, q2k_num, num_kv_blocks + ) + # q/k/v are saved from user-facing inputs; o/lse are kernel outputs. dq, dk, dv = block_sparse_bwd( - q_padded, - k_padded, - v_padded, + q_padded.contiguous(), + k_padded.contiguous(), + v_padded.contiguous(), o_padded, lse_padded, - grad_output_padded, + grad_output_padded.contiguous(), k2q_idx, k2q_num, - variable_block_sizes.int(), + variable_block_sizes, ) - # C++ kernel returns fp32 grads; cast back to match PyTorch convention if needed - return dq.to(grad_output_padded.dtype), dk.to(grad_output_padded.dtype), dv.to(grad_output_padded.dtype) + # C++ kernel returns fp32 grads; cast back to the input dtype. + out_dtype = grad_output_padded.dtype + return dq.to(out_dtype), dk.to(out_dtype), dv.to(out_dtype) @torch.library.register_fake("fastvideo_kernel::block_sparse_attn_backward_sm90") @@ -246,7 +318,8 @@ def _block_sparse_attn_backward_sm90_fake( v_padded: torch.Tensor, o_padded: torch.Tensor, lse_padded: torch.Tensor, - block_map: torch.Tensor, + q2k_idx: torch.Tensor, + q2k_num: torch.Tensor, variable_block_sizes: torch.Tensor, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: dq = torch.empty_like(q_padded) @@ -255,21 +328,57 @@ def _block_sparse_attn_backward_sm90_fake( return dq, dk, dv -def _backward_sm90(ctx, grad_o, grad_lse): - q, k, v, o, lse, block_map, variable_block_sizes = ctx.saved_tensors - dq, dk, dv = block_sparse_attn_backward_sm90( - grad_o, q, k, v, o, lse, block_map, variable_block_sizes - ) - return dq, dk, dv, None, None - - def _setup_context_sm90(ctx, inputs, output): - q, k, v, block_map, variable_block_sizes = inputs + q, k, v, q2k_idx, q2k_num, variable_block_sizes = inputs o, lse = output - ctx.save_for_backward(q, k, v, o, lse, block_map, variable_block_sizes) + ctx.save_for_backward(q, k, v, o, lse, q2k_idx, q2k_num, variable_block_sizes) -block_sparse_attn_sm90.register_autograd(_backward_sm90, setup_context=_setup_context_sm90) +def _backward_sm90(ctx, grad_o, grad_lse): + q, k, v, o, lse, q2k_idx, q2k_num, variable_block_sizes = ctx.saved_tensors + dq, dk, dv = block_sparse_attn_backward_sm90( + grad_o, q, k, v, o, lse, q2k_idx, q2k_num, variable_block_sizes + ) + return dq, dk, dv, None, None, None + + +block_sparse_attn_sm90.register_autograd( + _backward_sm90, setup_context=_setup_context_sm90 +) + + +# --------------------------------------------------------------------------- +# Public API +# --------------------------------------------------------------------------- + + +def block_sparse_attn_from_indices( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + q2k_idx: torch.Tensor, + q2k_num: torch.Tensor, + variable_block_sizes: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Block-sparse attention with autograd, taking compact per-row KV indices.""" + # Normalize index tensors once at the public boundary so the custom ops + # and their fakes can assume int32/contiguous. No-op on well-formed input. + q2k_idx = _as_int32_contig(q2k_idx, "q2k_idx") + q2k_num = _as_int32_contig(q2k_num, "q2k_num") + variable_block_sizes = _as_int32_contig(variable_block_sizes, "variable_block_sizes") + + block_sparse_fwd, block_sparse_bwd = _get_sm90_ops() + use_sm90 = ( + (not _force_triton()) + and _is_sm90() + and block_sparse_fwd is not None + and block_sparse_bwd is not None + ) + if use_sm90: + return block_sparse_attn_sm90(q, k, v, q2k_idx, q2k_num, variable_block_sizes) + # Triton path: supports q_seq_len != kv_seq_len as long as both are padded + # to a multiple of the block size (64 tokens). + return block_sparse_attn_triton(q, k, v, q2k_idx, q2k_num, variable_block_sizes) def block_sparse_attn( @@ -279,16 +388,8 @@ def block_sparse_attn( block_map: torch.Tensor, variable_block_sizes: torch.Tensor, ) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Unified block-sparse attention op with autograd support. - - On SM90 with compiled extension present: uses fastvideo_kernel_ops.block_sparse_fwd/bwd. - - Otherwise: uses Triton implementation (requires q/k/v to have same padded length today). - """ - block_sparse_fwd, block_sparse_bwd = _get_sm90_ops() - if (not _force_triton()) and _is_sm90() and (block_sparse_fwd is not None) and (block_sparse_bwd is not None): - return block_sparse_attn_sm90(q, k, v, block_map, variable_block_sizes) - # Triton path: supports q_seq_len != kv_seq_len as long as both are padded - # to a multiple of the block size (64 tokens). - return block_sparse_attn_triton(q, k, v, block_map, variable_block_sizes) - - + """Bool-mask compat wrapper; prefer block_sparse_attn_from_indices.""" + q2k_idx, q2k_num = _map_to_index(block_map) + return block_sparse_attn_from_indices( + q, k, v, q2k_idx, q2k_num, variable_block_sizes + ) diff --git a/fastvideo-kernel/python/fastvideo_kernel/ops.py b/fastvideo-kernel/python/fastvideo_kernel/ops.py index 72b49bb3c..398a5c130 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/ops.py +++ b/fastvideo-kernel/python/fastvideo_kernel/ops.py @@ -1,6 +1,6 @@ import math import torch -from .block_sparse_attn import block_sparse_attn +from .block_sparse_attn import block_sparse_attn, block_sparse_attn_from_indices from .triton_kernels.st_attn_triton import sliding_tile_attention_triton # Try to load the C++ extension @@ -125,13 +125,18 @@ def video_sparse_attn( out_c = out_c.repeat(1, 1, 1, block_elements, 1).view(batch, heads, q_seq_len, dim) - # Sparse branch + # Sparse branch: feed top-k indices directly, skipping the bool-mask round-trip. topk_idx = torch.topk(scores, topk, dim=-1).indices - mask = torch.zeros_like(scores, - dtype=torch.bool).scatter_(-1, topk_idx, True) - - # out_s = block_sparse_attn(q, k, v, mask, variable_block_sizes)[0] - out_s = block_sparse_attn(q, k, v, mask, variable_block_sizes)[0] + q2k_idx = topk_idx.to(torch.int32).contiguous() + q2k_num = torch.full( + (batch, heads, q_num_blocks), + topk, + dtype=torch.int32, + device=q.device, + ) + out_s = block_sparse_attn_from_indices( + q, k, v, q2k_idx, q2k_num, variable_block_sizes + )[0] if compress_attn_weight is not None: return out_c * compress_attn_weight + out_s diff --git a/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/index.py b/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/index.py index d84d01650..65e479331 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/index.py +++ b/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/index.py @@ -1,9 +1,10 @@ ## pytorch sdpa version of block sparse ## +from typing import Tuple + import triton import triton.language as tl import torch - @triton.jit def topk_index_to_map_kernel( map_ptr, @@ -153,3 +154,114 @@ def map_to_index(block_map: torch.Tensor): ) return index, index_num + + +@triton.jit +def _invert_indices_kernel( + q2k_idx_ptr, + q2k_num_ptr, + k2q_idx_ptr, + k2q_num_ptr, + q2k_idx_b, q2k_idx_h, q2k_idx_q, q2k_idx_k, + q2k_num_b, q2k_num_h, q2k_num_q, + k2q_idx_b, k2q_idx_h, k2q_idx_k, k2q_idx_q, + k2q_num_b, k2q_num_h, k2q_num_k, + MAX_KV_PER_Q: tl.constexpr, +): + # One program per (b, h, q): reserve a slot in k2q via atomicAdd, write q. + pid_b = tl.program_id(0) + pid_h = tl.program_id(1) + pid_q = tl.program_id(2) + + n = tl.load( + q2k_num_ptr + + pid_b * q2k_num_b + + pid_h * q2k_num_h + + pid_q * q2k_num_q + ) + + q2k_row = ( + q2k_idx_ptr + + pid_b * q2k_idx_b + + pid_h * q2k_idx_h + + pid_q * q2k_idx_q + ) + + for i in tl.range(0, MAX_KV_PER_Q): + if i < n: + kv = tl.load(q2k_row + i * q2k_idx_k) + count_ptr = ( + k2q_num_ptr + + pid_b * k2q_num_b + + pid_h * k2q_num_h + + kv * k2q_num_k + ) + pos = tl.atomic_add(count_ptr, 1) + tl.store( + k2q_idx_ptr + + pid_b * k2q_idx_b + + pid_h * k2q_idx_h + + kv * k2q_idx_k + + pos * k2q_idx_q, + pid_q, + ) + + +def invert_indices( + q2k_idx: torch.Tensor, + q2k_num: torch.Tensor, + num_kv_blocks: int, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Transpose a Q->KV index list into a K->Q one via atomic compaction (GPU).""" + if q2k_idx.dim() != 4: + raise ValueError( + f"q2k_idx must be [B, H, Nq, Mk], got shape={tuple(q2k_idx.shape)}" + ) + if q2k_num.dim() != 3: + raise ValueError( + f"q2k_num must be [B, H, Nq], got shape={tuple(q2k_num.shape)}" + ) + if not q2k_idx.is_cuda or not q2k_num.is_cuda: + raise RuntimeError("invert_indices requires CUDA tensors.") + + B, H, Nq, Mk = q2k_idx.shape + if q2k_num.shape != (B, H, Nq): + raise ValueError( + f"q2k_num shape {tuple(q2k_num.shape)} does not match q2k_idx " + f"[B, H, Nq] = {(B, H, Nq)}" + ) + + q2k_idx = q2k_idx.contiguous() + q2k_num = q2k_num.contiguous() + if q2k_idx.dtype != torch.int32: + q2k_idx = q2k_idx.to(torch.int32) + if q2k_num.dtype != torch.int32: + q2k_num = q2k_num.to(torch.int32) + + # Any KV block is attended by at most Nq Q blocks (one per Q row), so + # `Nq` is a tight upper bound on the compacted K->Q slots. + k2q_idx = torch.empty( + (B, H, num_kv_blocks, Nq), + dtype=torch.int32, + device=q2k_idx.device, + ) + k2q_num = torch.zeros( + (B, H, num_kv_blocks), + dtype=torch.int32, + device=q2k_idx.device, + ) + + grid = (B, H, Nq) + _invert_indices_kernel[grid]( + q2k_idx, + q2k_num, + k2q_idx, + k2q_num, + q2k_idx.stride(0), q2k_idx.stride(1), q2k_idx.stride(2), q2k_idx.stride(3), + q2k_num.stride(0), q2k_num.stride(1), q2k_num.stride(2), + k2q_idx.stride(0), k2q_idx.stride(1), k2q_idx.stride(2), k2q_idx.stride(3), + k2q_num.stride(0), k2q_num.stride(1), k2q_num.stride(2), + MAX_KV_PER_Q=Mk, + ) + + return k2q_idx, k2q_num