| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import cutlass |
| 5 | import cutlass.cute as cute |
| 6 | |
| 7 | from vllm.vllm_flash_attn.cute import utils |
| 8 | |
| 9 | |
| 10 | @cute.jit |
| 11 | def dense_mask_mod( |
| 12 | batch: cute.TensorSSA, |
| 13 | head: cute.TensorSSA, |
| 14 | q_idx: cute.TensorSSA, |
| 15 | kv_idx: cute.TensorSSA, |
| 16 | seqlen_info, |
| 17 | aux_tensors: list, |
| 18 | ) -> cute.TensorSSA: |
| 19 | dense_mask = aux_tensors[0] |
| 20 | batch_idx = utils.ssa_to_scalar(batch) |
| 21 | q_idx = utils.ssa_to_scalar(q_idx) |
| 22 | kv_idx = utils.ssa_to_scalar(kv_idx) |
| 23 | batch_stride, query_stride, _ = dense_mask.stride |
| 24 | aligned_mask = cute.make_tensor( |
| 25 | dense_mask.iterator, |
| 26 | cute.make_layout( |
| 27 | dense_mask.shape, |
| 28 | stride=( |
| 29 | cute.assume(batch_stride, divby=4), |
| 30 | cute.assume(query_stride, divby=4), |
| 31 | 1, |
| 32 | ), |
| 33 | ), |
| 34 | ) |
| 35 | mask_row = aligned_mask[batch_idx, q_idx, None] |
| 36 | mask_chunks = cute.flat_divide(mask_row, (4,)) |
| 37 | mask_chunk = mask_chunks[None, (kv_idx >> 5) >> 2] |
| 38 | loaded = cute.make_rmem_tensor_like(mask_chunk) |
| 39 | cute.autovec_copy(mask_chunk, loaded) |
| 40 | return cute.recast_tensor(loaded, cutlass.Uint32).load() |
| 41 | |
| 42 | |
| 43 | dense_mask_mod.__vec_size__ = 128 |
| 44 | |
| 45 | |
| 46 | @cute.jit |
| 47 | def offset_dense_mask_mod( |
| 48 | batch: cute.TensorSSA, |
| 49 | head: cute.TensorSSA, |
| 50 | q_idx: cute.TensorSSA, |
| 51 | kv_idx: cute.TensorSSA, |
| 52 | seqlen_info, |
| 53 | aux_tensors: list, |
| 54 | ) -> cute.TensorSSA: |
| 55 | dense_mask = aux_tensors[0] |
| 56 | batch_idx = utils.ssa_to_scalar(batch) |
| 57 | q_idx = utils.ssa_to_scalar(q_idx) |
| 58 | key_start = dense_mask[batch_idx, 0, dense_mask.shape[2] - 1] |
| 59 | kv_idx = utils.ssa_to_scalar(kv_idx) + key_start |
| 60 | word_idx = kv_idx >> 5 |
| 61 | bit_idx = cutlass.Uint32(kv_idx & 31) |
| 62 | word = dense_mask[batch_idx, q_idx, word_idx] |
| 63 | result = cute.make_rmem_tensor(1, dtype=cutlass.Uint32) |
| 64 | result[0] = utils.shr_u32(cutlass.Uint32(word), bit_idx) |
| 65 | return result.load() |
| 66 | |
| 67 | |
| 68 | offset_dense_mask_mod.__vec_size__ = 32 |