@tduncan / Nvidia-Optimization-Stack-VLLM / commits / e9be5323f0

[Perf] Integrate flash-maxsim Triton kernels for late-interaction scoring (#40337)

Signed-off-by: roi.pony <roi.pony@ibm.com> Co-authored-by: roi.pony <roi.pony@ibm.com> Co-authored-by: Taneem Ibrahim <taneem.ibrahim@gmail.com>

roipony committed Oct 1, 2026 at 21:06 UTC e9be5323f0f864c1dc5a22a138530f3e106ebbb6
10 files changed +508 -4
tests/v1/worker/test_late_interaction_runner.py
+105
@@ -10,6 +10,7 @@ from vllm.v1.pool.late_interaction import (
10 LATE_INTERACTION_MODE_CACHE_QUERY,
11 build_late_interaction_doc_params,
12 build_late_interaction_query_params,
13 + compute_maxsim_score_batched,
14 )
15 from vllm.v1.pool.late_interaction_runner import LateInteractionRunner
16
@@ -152,3 +153,107 @@ def test_invalid_query_uses_raises():
153 req_ids=["query-req"],
154 finished_mask=[True],
155 )
156 +
157 +
158 +# ---------------------------------------------------------------------------
159 +# Fused Triton scoring path (PR #40337).
160 +# ---------------------------------------------------------------------------
161 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
162 +def test_score_flash_matches_fp64_oracle():
163 + """Shared-query scoring through the fused kernel must match a per-pair
164 + fp64 MaxSim oracle on ragged doc lengths, including length-1 docs."""
165 + torch.manual_seed(0)
166 + d = 128
167 + q = torch.randn(32, d, device="cuda", dtype=torch.float16)
168 + doc_lengths = [180, 37, 512, 1, 300, 64, 1030, 256]
169 + docs = [
170 + torch.randn(ld, d, device="cuda", dtype=torch.float16) for ld in doc_lengths
171 + ]
172 +
173 + runner = LateInteractionRunner()
174 + assert runner._flash_enabled
175 + scores = runner._score([q] * len(docs), docs)
176 + assert runner._flash_enabled, "kernel path must not have fallen back"
177 +
178 + for i, doc in enumerate(docs):
179 + ref = (q.double() @ doc.double().T).max(dim=1).values.sum().float()
180 + torch.testing.assert_close(
181 + scores[i].to(torch.float32), ref, atol=5e-2, rtol=1e-3
182 + )
183 +
184 +
185 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
186 +def test_score_multi_query_uses_reference_path():
187 + """Distinct queries per pair route to the reference scorer and agree
188 + with it exactly (no kernel dispatch for the N:N pattern)."""
189 + torch.manual_seed(1)
190 + d = 128
191 + queries = [torch.randn(16, d, device="cuda") for _ in range(3)]
192 + docs = [torch.randn(ld, d, device="cuda") for ld in (50, 3, 77)]
193 +
194 + runner = LateInteractionRunner()
195 + scores = runner._score(queries, docs)
196 + ref = compute_maxsim_score_batched(queries, docs)
197 + for s, r in zip(scores, ref):
198 + torch.testing.assert_close(s, r)
199 +
200 +
201 +def test_score_flash_failure_disables_and_falls_back(monkeypatch):
202 + """A kernel failure must serve the batch via the reference scorer and
203 + disable the kernel path for the rest of the process."""
204 + runner = LateInteractionRunner()
205 + runner._flash_enabled = True
206 +
207 + def _boom(query, docs):
208 + raise RuntimeError("synthetic kernel failure")
209 +
210 + monkeypatch.setattr(runner, "_score_flash_shared_query", _boom)
211 + q = torch.randn(8, 32)
212 + docs = [torch.randn(5, 32), torch.randn(7, 32)]
213 + # CPU tensors already skip the kernel; force the dispatch condition.
214 + monkeypatch.setattr(torch.Tensor, "is_cuda", property(lambda self: True))
215 + scores = runner._score([q, q], docs)
216 + assert not runner._flash_enabled
217 + ref = compute_maxsim_score_batched([q, q], docs)
218 + for s, r in zip(scores, ref):
219 + torch.testing.assert_close(s, r)
220 +
221 +
222 +def test_disabled_flag_uses_reference_path():
223 + runner = LateInteractionRunner(enable_flash=False)
224 + assert not runner._flash_enabled
225 +
226 +
227 +# ---------------------------------------------------------------------------
228 +# Autotune config pruning.
229 +# ---------------------------------------------------------------------------
230 +def _prune(monkeypatch, budget: int, named_args: dict):
231 + from vllm.v1.pool.flash_maxsim import _common
232 +
233 + monkeypatch.setattr(_common, "_smem_budget", lambda: budget)
234 + return _common._prune_configs(_common._get_configs(), named_args)
235 +
236 +
237 +def _est(cfg, d_pad: int) -> int:
238 + bq, bd = cfg.kwargs["BLOCK_Q"], cfg.kwargs["BLOCK_D"]
239 + return (bq * d_pad + bd * d_pad) * 2 + bq * bd * 4
240 +
241 +
242 +def test_prune_configs_uses_padded_dim(monkeypatch):
243 + """d=513 pads to 1024: survivors must fit the budget at d_pad, not at
244 + the un-padded d (which would admit ~2x-oversized tiles)."""
245 + budget = 166_912 # A100 opt-in
246 + survivors = _prune(monkeypatch, budget, {"Lq": 1024, "d": 513, "d_pad": 1024})
247 + assert survivors
248 + assert all(_est(c, 1024) <= budget for c in survivors)
249 +
250 +
251 +def test_prune_configs_no_rejected_fallback(monkeypatch):
252 + """When nothing fits, the fallback is the single smallest-footprint
253 + config — never a slice of rejected configs."""
254 + from vllm.v1.pool.flash_maxsim import _common
255 +
256 + survivors = _prune(monkeypatch, 10_000, {"Lq": 1024, "d_pad": 1024})
257 + assert len(survivors) == 1
258 + smallest = min(_common._get_configs(), key=lambda c: _est(c, 1024))
259 + assert survivors[0].kwargs == smallest.kwargs
vllm/config/pooler.py
+8
@@ -63,6 +63,14 @@ class PoolerConfig:
63 `None` uses the pooler's default, which is `True` in most cases.
64 """
65
66 + enable_flash_late_interaction: bool = True
67 + """
68 + Whether the engine-side late-interaction scorer may use the fused
69 + flash-maxsim Triton kernel. Disabled automatically when the API server
70 + is started with `--no-enable-flash-late-interaction`; the reference
71 + scorer is used instead.
72 + """
73 +
74 ## for embedding models
75 dimensions: int | None = None
76 """
vllm/entrypoints/cli/serve.py
+3
@@ -13,6 +13,7 @@ from vllm.entrypoints.cli.types import CLISubcommand
13 from vllm.entrypoints.launchers.api_server.entry import run_server, setup_server
14 from vllm.entrypoints.launchers.cli_args import (
15 make_arg_parser,
16 + propagate_flash_late_interaction,
17 validate_parsed_serve_args,
18 )
19 from vllm.entrypoints.launchers.dp_supervisor import run_dp_supervisor
@@ -185,6 +186,7 @@ def run_headless(args: argparse.Namespace):
186
187 # Create the EngineConfig.
188 engine_args = vllm.AsyncEngineArgs.from_cli_args(args)
189 + propagate_flash_late_interaction(args, engine_args)
190 usage_context = UsageContext.OPENAI_API_SERVER
191 vllm_config = engine_args.create_engine_config(
192 usage_context=usage_context, headless=True
@@ -295,6 +297,7 @@ def run_multi_api_server(args: argparse.Namespace):
297 listen_address, sock = setup_server(args, reuse_port=num_api_servers > 1)
298
299 engine_args = vllm.AsyncEngineArgs.from_cli_args(args)
300 + propagate_flash_late_interaction(args, engine_args)
301 engine_args._api_process_count = num_api_servers
302 engine_args._api_process_rank = -1
303
vllm/entrypoints/launchers/api_server/entry.py
+3
@@ -51,6 +51,9 @@ async def build_async_engine_client(
51 # Context manager to handle engine_client lifecycle
52 # Ensures everything is shutdown and cleaned up on error/exit
53 engine_args = AsyncEngineArgs.from_cli_args(args)
54 + from ..cli_args import propagate_flash_late_interaction
55 +
56 + propagate_flash_late_interaction(args, engine_args)
57 if client_config:
58 engine_args._api_process_count = client_config.get("client_count", 1)
59 engine_args._api_process_rank = client_config.get("client_index", 0)
vllm/entrypoints/launchers/cli_args.py
+21 -1
@@ -363,7 +363,10 @@ class FrontendArgs(BaseFrontendArgs):
363 """
364 enable_flash_late_interaction: bool = True
365 """If set, run pooling score MaxSim on GPU in the API server process.
366 - Can significantly improve late-interaction scoring performance."""
366 + Can significantly improve late-interaction scoring performance.
367 + When disabled, the setting is also propagated to the engine's
368 + `PoolerConfig`, so the engine-side scorer uses the reference MaxSim
369 + path instead of the fused Triton kernel."""
370
371 @classmethod
372 def _customize_cli_kwargs(
@@ -498,3 +501,20 @@ def create_parser_for_docs() -> FlexibleArgumentParser:
501 prog="-m vllm.entrypoints.launchers.api_server.entry"
502 )
503 return make_arg_parser(parser_for_docs)
504 +
505 +
506 +def propagate_flash_late_interaction(args, engine_args) -> None:
507 + """Propagate `--no-enable-flash-late-interaction` into the engine config.
508 +
509 + The frontend flag alone only disables the API-server scoring path;
510 + mirroring it into `PoolerConfig.enable_flash_late_interaction` lets the
511 + engine-side scorer fall back to the reference MaxSim path too.
512 + """
513 + if getattr(args, "enable_flash_late_interaction", True):
514 + return
515 + from vllm.config.pooler import PoolerConfig
516 +
517 + if engine_args.pooler_config is None:
518 + engine_args.pooler_config = PoolerConfig(enable_flash_late_interaction=False)
519 + else:
520 + engine_args.pooler_config.enable_flash_late_interaction = False
vllm/v1/pool/flash_maxsim/__init__.py new
+7
@@ -0,0 +1,7 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Fused Triton kernels for late-interaction (MaxSim) scoring."""
4 +
5 +from .flash_maxsim_rerank import flash_maxsim_rerank_direct
6 +
7 +__all__ = ["flash_maxsim_rerank_direct"]
vllm/v1/pool/flash_maxsim/_common.py new
+85
@@ -0,0 +1,85 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Shared utilities for flash-maxsim kernels: padding, autotune configs,
4 +and pruning."""
5 +
6 +from vllm.triton_utils import triton
7 +
8 +
9 +def _next_pow2(x: int) -> int:
10 + return 1 << (x - 1).bit_length()
11 +
12 +
13 +def _get_configs() -> list:
14 + """Autotune grid over (BLOCK_Q, BLOCK_D); `_prune_configs` filters by
15 + SMEM budget and Lq at runtime, so one list is safe across GPUs."""
16 + return [
17 + triton.Config({"BLOCK_Q": 16, "BLOCK_D": 16}, num_warps=2, num_stages=2),
18 + triton.Config({"BLOCK_Q": 16, "BLOCK_D": 32}, num_warps=2, num_stages=2),
19 + triton.Config({"BLOCK_Q": 32, "BLOCK_D": 16}, num_warps=2, num_stages=2),
20 + triton.Config({"BLOCK_Q": 32, "BLOCK_D": 32}, num_warps=4, num_stages=2),
21 + triton.Config({"BLOCK_Q": 32, "BLOCK_D": 64}, num_warps=4, num_stages=2),
22 + triton.Config({"BLOCK_Q": 32, "BLOCK_D": 128}, num_warps=8, num_stages=2),
23 + triton.Config({"BLOCK_Q": 64, "BLOCK_D": 64}, num_warps=4, num_stages=2),
24 + triton.Config({"BLOCK_Q": 64, "BLOCK_D": 128}, num_warps=8, num_stages=1),
25 + triton.Config({"BLOCK_Q": 128, "BLOCK_D": 64}, num_warps=8, num_stages=1),
26 + triton.Config({"BLOCK_Q": 128, "BLOCK_D": 128}, num_warps=8, num_stages=1),
27 + ]
28 +
29 +
30 +def _smem_budget() -> int:
31 + """Per-block opt-in shared-memory budget of the active device (bytes):
32 + Triton driver query, else a conservative capability table."""
33 + try:
34 + driver = triton.runtime.driver.active
35 + props = driver.utils.get_device_properties(driver.get_current_device())
36 + return int(props["max_shared_mem"])
37 + except Exception:
38 + try:
39 + from vllm.platforms import current_platform
40 +
41 + cap = current_platform.get_device_capability()
42 + if cap is not None:
43 + if cap.major >= 9:
44 + return 232_448 # Hopper+: 227 KB opt-in
45 + if cap.major == 8:
46 + return 101_376 # worst sm_8x (86/89); A100 has more
47 + except Exception:
48 + pass
49 + return 98_304 # conservative pre-Ampere default
50 +
51 +
52 +def _prune_configs(configs, named_args, **kwargs):
53 + """Reject configs that exceed the GPU's shared-memory budget.
54 +
55 + BLOCK_Q and BLOCK_D load fp16 tiles of size [BLOCK_*, d_pad] — the
56 + kernels index with ``tl.arange(0, d_pad)``, so the estimate must use
57 + the padded dimension, not the real one (d=513 pads to 1024 and
58 + nearly doubles the tile) — and the [BLOCK_Q, BLOCK_D] score tile is
59 + fp32. Total:
60 + (BLOCK_Q*d_pad + BLOCK_D*d_pad)*2 + BLOCK_Q*BLOCK_D*4 bytes.
61 + """
62 + Lq = named_args.get("Lq", 32)
63 + d_pad = named_args.get("d_pad")
64 + if d_pad is None:
65 + d_pad = _next_pow2(max(named_args.get("d", 128), 16))
66 + smem_limit = _smem_budget()
67 +
68 + def _est(cfg):
69 + bq = cfg.kwargs["BLOCK_Q"]
70 + bd = cfg.kwargs["BLOCK_D"]
71 + return (bq * d_pad + bd * d_pad) * 2 + bq * bd * 4
72 +
73 + pruned = []
74 + for cfg in configs:
75 + if cfg.kwargs["BLOCK_Q"] > Lq * 2:
76 + continue
77 + if _est(cfg) > smem_limit:
78 + continue
79 + pruned.append(cfg)
80 + if pruned:
81 + return pruned
82 + # Nothing fits the estimate (large d_pad on a small-SMEM device).
83 + # Return the single smallest-footprint config — a deterministic
84 + # least-bad launch attempt — never arbitrary rejected configs.
85 + return [min(configs, key=_est)]
vllm/v1/pool/flash_maxsim/flash_maxsim_rerank.py new
+204
@@ -0,0 +1,204 @@
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)
vllm/v1/pool/late_interaction_runner.py
+67 -2
@@ -4,6 +4,7 @@ from collections.abc import Iterable
4
5 import torch
6
7 +from vllm.logger import init_logger
8 from vllm.pooling_params import PoolingParams
9 from vllm.v1.outputs import PoolerOutput
10 from vllm.v1.pool.late_interaction import (
@@ -12,17 +13,29 @@ from vllm.v1.pool.late_interaction import (
13 compute_maxsim_score_batched,
14 )
15
16 +logger = init_logger(__name__)
17 +
18
19 class LateInteractionRunner:
20 """Worker-side state and postprocessing for late-interaction scoring."""
21
19 - def __init__(self) -> None:
22 + def __init__(self, enable_flash: bool = True) -> None:
23 # query_key -> token embeddings for late-interaction scoring.
24 self._query_cache: dict[str, torch.Tensor] = {}
25 # query_key -> remaining number of docs that should use this query.
26 self._query_uses: dict[str, int] = {}
27 # doc request id -> query key.
28 self._doc_query_keys: dict[str, str] = {}
29 + # Fused Triton scoring (PoolerConfig.enable_flash_late_interaction
30 + # gates it; a runtime kernel failure disables it for the process).
31 + if enable_flash:
32 + try:
33 + from vllm.v1.pool.flash_maxsim import ( # noqa: F401
34 + flash_maxsim_rerank_direct,
35 + )
36 + except ImportError:
37 + enable_flash = False
38 + self._flash_enabled = enable_flash
39
40 def clear(self) -> None:
41 self._query_cache.clear()
@@ -116,7 +129,7 @@ class LateInteractionRunner:
129 raise ValueError(f"Unsupported late-interaction mode: {mode!r}")
130
131 if score_indices:
119 - score_values = compute_maxsim_score_batched(score_queries, score_docs)
132 + score_values = self._score(score_queries, score_docs)
133 for i, req_id, query_key, score in zip(
134 score_indices, score_req_ids, score_query_keys, score_values
135 ):
@@ -126,6 +139,58 @@ class LateInteractionRunner:
139
140 return outputs
141
142 + def _score(
143 + self,
144 + queries: list[torch.Tensor],
145 + docs: list[torch.Tensor],
146 + ) -> list[torch.Tensor]:
147 + """Score (query_i, doc_i) pairs: the fused Triton kernel handles the
148 + common rerank pattern (one shared query, CUDA inputs); multiple
149 + distinct queries, CPU tensors, or a kernel failure fall back to the
150 + reference scorer."""
151 + if self._flash_enabled and docs and queries[0].is_cuda:
152 + first = queries[0]
153 + if all(q is first for q in queries):
154 + try:
155 + return self._score_flash_shared_query(first, docs)
156 + except Exception as exc:
157 + # A persistent compile/launch failure must not take down
158 + # requests: serve through the reference scorer and stop
159 + # trying the kernel for the rest of the process.
160 + self._flash_enabled = False
161 + logger.warning(
162 + "flash-maxsim scoring failed (%s); falling back to "
163 + "the reference MaxSim path for this process.",
164 + exc,
165 + )
166 + return compute_maxsim_score_batched(queries, docs)
167 +
168 + @staticmethod
169 + def _score_flash_shared_query(
170 + query: torch.Tensor,
171 + docs: list[torch.Tensor],
172 + ) -> list[torch.Tensor]:
173 + """One kernel launch for all docs against a single shared query:
174 + docs pack into one [total_tokens, d] tensor (single cat, no padding)
175 + and the kernel reads each by (offset, length) — the [B, Lq, Ld]
176 + similarity tensor is never materialised."""
177 + from vllm.v1.pool.flash_maxsim import flash_maxsim_rerank_direct
178 +
179 + device = docs[0].device
180 + lengths = [int(d.shape[0]) for d in docs]
181 + offsets = [0] * len(lengths)
182 + for i in range(1, len(lengths)):
183 + offsets[i] = offsets[i - 1] + lengths[i - 1]
184 + packed = docs[0] if len(docs) == 1 else torch.cat(docs, dim=0)
185 + scores = flash_maxsim_rerank_direct(
186 + query,
187 + packed,
188 + torch.tensor(offsets, device=device, dtype=torch.int32),
189 + torch.tensor(lengths, device=device, dtype=torch.int32),
190 + max(lengths),
191 + )
192 + return list(scores.unbind(0))
193 +
194 def _release_query_use(self, query_key: str) -> None:
195 remaining = self._query_uses.get(query_key, 1) - 1
196 if remaining <= 0:
vllm/v1/worker/gpu/pool/pooling_runner.py
+5 -1
@@ -51,7 +51,11 @@ class PoolingRunner:
51 self.pooling_params: dict[int, PoolingParams] = {}
52 self.pooling_states: dict[int, PoolingStates] = {}
53 self.prompt_token_ids: dict[int, torch.Tensor] = {}
54 - self.late_interaction_runner = LateInteractionRunner()
54 + pooler_config = self.model_config.pooler_config
55 + self.late_interaction_runner = LateInteractionRunner(
56 + enable_flash=pooler_config is None
57 + or getattr(pooler_config, "enable_flash_late_interaction", True)
58 + )
59
60 @staticmethod
61 def get_supported_tasks(model: nn.Module) -> list[PoolingTask]: