This operator wraps find_algo + gemm. This is because we don't want find_algo to be visible by torch compile, otherwise it will remove it fro

Converts dense 2d matrix to a csr sparse matrix.

This kernel can accept non-quantized or int4-quantized keys/values. PACKED_PER_VAL determines the quantization type: - PACKED_PER_VAL == 1

Each letter in this diagram is a whole row of length dim. INPUT xq xk xv head_dim ─► batch qqqqqq

This version of reduce kernel takes attention and LSE of chunks as lists of tensors, as opposed to _splitK_reduce, which takes each as a stac

Backward for _splitK_reduce_varargs. Similar to forward, it takes attention and LSE of chunks as lists of tensors, and outputs the corres

r""" Use the additive bias computation from ALiBi_ to generate a mask. Note that this mask can in turn be used to generate a blocksparse atten

Gather a tensor of shape (n, m) into a tensor of shape (n, mp_size * m).

Note that inp can be of type InputsFp8, in which case K/V are assumed to be row-wise FP8-quantized. This is different from int4 quant

Block sparsify a tensor, given a mask and block size