| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Flash-MaxSim rerank: one query vs variable-length docs in a packed tensor. |
| 4 | |
| 5 | The kernel reads doc embeddings directly from a packed [total_tokens, d] |
| 6 | tensor using per-doc (offset, length) pairs — no padding, no per-doc |
| 7 | tensors, no materialised [B, Lq, Ld] similarity tensor. |
| 8 | |
| 9 | Entry point: |
| 10 | flash_maxsim_rerank_direct(Q, batch_tensor, doc_offsets, doc_lengths, max_ld) |
| 11 | doc_offsets[i] = start token index of doc i in batch_tensor. |
| 12 | doc_lengths[i] = number of tokens in doc i. |
| 13 | Docs can be scattered (non-contiguous) in the batch tensor. |
| 14 | """ |
| 15 | |
| 16 | import torch |
| 17 | |
| 18 | from vllm.triton_utils import tl, triton |
| 19 | |
| 20 | from ._common import _get_configs, _next_pow2, _prune_configs |
| 21 | |
| 22 | |
| 23 | def _round_to_bucket(x): |
| 24 | """Round up to nearest bucket: 32, 64, 128, 256, 512, 1024, 2048, 4096.""" |
| 25 | for b in [32, 64, 128, 256, 512, 1024, 2048, 4096]: |
| 26 | if x <= b: |
| 27 | return b |
| 28 | return _next_pow2(x) |
| 29 | |
| 30 | |
| 31 | # Kernel: shared Q, scattered D via (offset, length) per doc |
| 32 | |
| 33 | |
| 34 | @triton.autotune( |
| 35 | configs=_get_configs(), |
| 36 | key=["Lq_bucket", "max_Ld_bucket", "d_pad"], |
| 37 | prune_configs_by={"early_config_prune": _prune_configs}, |
| 38 | ) |
| 39 | @triton.jit |
| 40 | def _maxsim_rerank_kernel( |
| 41 | Q_ptr, |
| 42 | D_ptr, |
| 43 | doc_offsets_ptr, |
| 44 | doc_lengths_ptr, |
| 45 | scores_ptr, |
| 46 | B, |
| 47 | Lq, |
| 48 | Lq_bucket: tl.constexpr, |
| 49 | max_Ld_bucket: tl.constexpr, |
| 50 | d: tl.constexpr, |
| 51 | d_pad: tl.constexpr, |
| 52 | stride_q_t, |
| 53 | stride_q_d, |
| 54 | stride_d_t, |
| 55 | stride_d_d, |
| 56 | BLOCK_Q: tl.constexpr, |
| 57 | BLOCK_D: tl.constexpr, |
| 58 | ): |
| 59 | """One program per doc. Q shared (read from same location). D scattered.""" |
| 60 | pid = tl.program_id(0) |
| 61 | if pid >= B: |
| 62 | return |
| 63 | |
| 64 | # Each doc has its own offset and length in the batch tensor. |
| 65 | # int64: (d_start + d_off) * stride overflows int32 past ~16.7M batch |
| 66 | # tokens (same widen-on-load as the pairs kernel below; the host-side |
| 67 | # offsets tensor stays int32, which indexes any realistic batch). |
| 68 | d_start = tl.load(doc_offsets_ptr + pid).to(tl.int64) |
| 69 | d_len = tl.load(doc_lengths_ptr + pid).to(tl.int32) |
| 70 | |
| 71 | k_off = tl.arange(0, d_pad) |
| 72 | k_mask = k_off < d |
| 73 | score_acc = tl.zeros([], dtype=tl.float32) |
| 74 | |
| 75 | # Q is shared — always read from start of Q_ptr |
| 76 | for q_blk in tl.static_range(0, Lq_bucket, BLOCK_Q): |
| 77 | q_off = q_blk + tl.arange(0, BLOCK_Q) |
| 78 | q_valid = q_off < Lq |
| 79 | |
| 80 | Q_block = tl.load( |
| 81 | Q_ptr + q_off[:, None] * stride_q_t + k_off[None, :] * stride_q_d, |
| 82 | mask=q_valid[:, None] & k_mask[None, :], |
| 83 | other=0.0, |
| 84 | ).to(tl.float16) |
| 85 | |
| 86 | m = tl.full([BLOCK_Q], float("-inf"), dtype=tl.float32) |
| 87 | |
| 88 | for d_blk in range(0, max_Ld_bucket, BLOCK_D): |
| 89 | d_off = d_blk + tl.arange(0, BLOCK_D) |
| 90 | d_valid = d_off < d_len |
| 91 | |
| 92 | D_block = tl.load( |
| 93 | D_ptr |
| 94 | + (d_start + d_off[:, None]) * stride_d_t |
| 95 | + k_off[None, :] * stride_d_d, |
| 96 | mask=d_valid[:, None] & k_mask[None, :], |
| 97 | other=0.0, |
| 98 | ).to(tl.float16) |
| 99 | |
| 100 | S = tl.dot(Q_block, tl.trans(D_block)) |
| 101 | S = tl.where(d_valid[None, :], S, float("-inf")) |
| 102 | m = tl.maximum(m, tl.max(S, axis=1)) |
| 103 | |
| 104 | m = tl.where(q_valid, m, 0.0) |
| 105 | score_acc += tl.sum(m) |
| 106 | |
| 107 | tl.store(scores_ptr + pid, score_acc) |
| 108 | |
| 109 | |
| 110 | # Internal helper |
| 111 | |
| 112 | |
| 113 | def _run_rerank_kernel(Q, D_tensor, doc_offsets, doc_lengths, max_seqlen_d): |
| 114 | """Common kernel launch for both rerank and rerank_direct.""" |
| 115 | assert Q.is_cuda, f"Q must be on CUDA, got {Q.device}" |
| 116 | assert D_tensor.is_cuda, f"D must be on CUDA, got {D_tensor.device}" |
| 117 | Lq, d = Q.shape |
| 118 | B = doc_offsets.shape[0] |
| 119 | d_pad = _next_pow2(max(d, 16)) |
| 120 | |
| 121 | Q = Q.contiguous().half() |
| 122 | # D_tensor is NOT copied — kernel reads from it directly |
| 123 | if not D_tensor.is_contiguous(): |
| 124 | D_tensor = D_tensor.contiguous() |
| 125 | if D_tensor.dtype != torch.float16: |
| 126 | D_tensor = D_tensor.half() |
| 127 | |
| 128 | doc_offsets = doc_offsets.to(torch.int32).contiguous() |
| 129 | doc_lengths = doc_lengths.to(torch.int32).contiguous() |
| 130 | scores = torch.empty(B, device=Q.device, dtype=torch.float32) |
| 131 | |
| 132 | if B == 0: |
| 133 | return scores |
| 134 | |
| 135 | # Pad Q dimension if needed (Q is small, safe to copy). |
| 136 | # D is NOT padded — the kernel uses k_mask to handle d < d_pad. |
| 137 | # This preserves zero-copy: D_tensor is read in-place. |
| 138 | if d < d_pad: |
| 139 | Q_pad = torch.zeros(Lq, d_pad, device=Q.device, dtype=Q.dtype) |
| 140 | Q_pad[:, :d] = Q |
| 141 | else: |
| 142 | Q_pad = Q |
| 143 | D_pad = D_tensor |
| 144 | |
| 145 | Lq_bucket = _round_to_bucket(Lq) |
| 146 | max_ld_bucket = _round_to_bucket(max_seqlen_d) |
| 147 | |
| 148 | _maxsim_rerank_kernel[(B,)]( |
| 149 | Q_pad, |
| 150 | D_pad, |
| 151 | doc_offsets, |
| 152 | doc_lengths, |
| 153 | scores, |
| 154 | B, |
| 155 | Lq, |
| 156 | Lq_bucket, |
| 157 | max_ld_bucket, |
| 158 | d, |
| 159 | d_pad, |
| 160 | Q_pad.stride(0), |
| 161 | Q_pad.stride(1), |
| 162 | D_pad.stride(0), |
| 163 | D_pad.stride(1), |
| 164 | ) |
| 165 | return scores |
| 166 | |
| 167 | |
| 168 | # Public API |
| 169 | |
| 170 | |
| 171 | def flash_maxsim_rerank_direct( |
| 172 | Q: torch.Tensor, |
| 173 | batch_tensor: torch.Tensor, |
| 174 | doc_offsets: torch.Tensor, |
| 175 | doc_lengths: torch.Tensor, |
| 176 | max_seqlen_d: int, |
| 177 | ) -> torch.Tensor: |
| 178 | """TRUE zero-copy: score query against docs scattered in a batch tensor. |
| 179 | |
| 180 | The kernel reads doc embeddings directly from batch_tensor at the |
| 181 | positions specified by doc_offsets. No torch.stack, no torch.cat, |
| 182 | no copy of any kind. The batch tensor is the model's output. |
| 183 | |
| 184 | Memory for doc scoring: 0 bytes additional. |
| 185 | |
| 186 | Args: |
| 187 | Q: [Lq, d] — single query embedding (from cache) |
| 188 | batch_tensor: [total_tokens, d] — the model's projected output tensor. |
| 189 | Contains ALL requests' tokens (queries + docs + others). |
| 190 | doc_offsets: [B] int32 — start token index of each doc in batch_tensor |
| 191 | doc_lengths: [B] int32 — number of tokens per doc |
| 192 | max_seqlen_d: int — max(doc_lengths) |
| 193 | |
| 194 | Returns: |
| 195 | scores: [B] float32 — one MaxSim score per document |
| 196 | |
| 197 | """ |
| 198 | assert Q.dim() == 2, f"Q must be 2D [Lq, d], got {Q.dim()}D" |
| 199 | assert batch_tensor.dim() == 2, ( |
| 200 | f"batch_tensor must be 2D, got {batch_tensor.dim()}D" |
| 201 | ) |
| 202 | assert Q.shape[1] == batch_tensor.shape[1] |
| 203 | |
| 204 | return _run_rerank_kernel(Q, batch_tensor, doc_offsets, doc_lengths, max_seqlen_d) |