@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 1713b9866a

[2/N][warmup][DSv4] Migrate sequence and DCP kernels (#53564)

Signed-off-by: LopezCastroRoberto <rocastro@redhat.com> Co-authored-by: Codex <codex@openai.com> Co-authored-by: OpenAI Codex <noreply@openai.com> Co-authored-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Roberto L. Castro committed Sep 7, 2026 at 15:35 UTC 1713b9866ae5c380d97da5796217064ac7ab57b8
12 files changed +1461 -583
.buildkite/test_areas/model_executor.yaml
+4
@@ -13,6 +13,8 @@ steps:
13 - vllm/model_executor/warmup
14 - tests/model_executor
15 - tests/model_executor/test_jit_warmup.py
16 + - tests/model_executor/test_jit_warmup_cutedsl_launcher.py
17 + - tests/model_executor/test_jit_warmup_triton_launcher.py
18 - tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py
19 commands:
20 - apt-get update && apt-get install -y curl libsodium23
@@ -41,6 +43,8 @@ steps:
43 - vllm/model_executor/warmup
44 - tests/model_executor
45 - tests/model_executor/test_jit_warmup.py
46 + - tests/model_executor/test_jit_warmup_cutedsl_launcher.py
47 + - tests/model_executor/test_jit_warmup_triton_launcher.py
48 - tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py
49 - vllm/_aiter_ops.py
50 - vllm/platforms/rocm.py
tests/model_executor/test_jit_warmup.py
+3 -1
@@ -532,10 +532,12 @@ def test_runtime_cache_miss_compiles_and_caches_executor() -> None:
532 def test_registry_records_only_inside_model_setup_context() -> None:
533 registry = JitWarmupRegistry(_config())
534 kernel = RecordingToyKernel()
535 + config = _config()
536
537 kernel.register_warmup(3, _config())
538 with registry.activate():
538 - kernel.register_warmup(3, _config())
539 + kernel.register_warmup(3, config)
540 + kernel.register_warmup(3, config)
541
542 assert len(registry) == 1
543 assert kernel.compiled == []
tests/model_executor/test_jit_warmup_cutedsl_launcher.py new
+76
@@ -0,0 +1,76 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from dataclasses import dataclass
5 +from typing import Any
6 +
7 +from vllm.model_executor.warmup import jit_warmup_cutedsl_helper
8 +from vllm.model_executor.warmup.jit_warmup_cutedsl_helper import (
9 + CuTeDSLLaunchSpec,
10 + VllmCuTeDSLJitKernel,
11 + cutedsl_kernel_launcher,
12 +)
13 +
14 +
15 +class _TestCuTeDSLKernel(VllmCuTeDSLJitKernel["_TestCuTeDSLKernel.CompileKey"]):
16 + @dataclass(frozen=True)
17 + class CompileKey:
18 + variant: int
19 +
20 + @staticmethod
21 + def kernel(compile_key: CompileKey) -> tuple[str, int]:
22 + return "entry", compile_key.variant
23 +
24 + def dispatch(self, *, variant: int) -> CompileKey:
25 + return self.CompileKey(variant=variant)
26 +
27 + def get_warmup_keys(self) -> list[CompileKey]:
28 + return self._trace_dispatch(self.dispatch)(variant=(1, 2))
29 +
30 + def warmup_inputs(self, compile_key: CompileKey) -> tuple[Any, ...]:
31 + return (f"fake-{compile_key.variant}",)
32 +
33 + @cutedsl_kernel_launcher
34 + def __call__(
35 + self,
36 + payload: str,
37 + *,
38 + variant: int,
39 + ) -> CuTeDSLLaunchSpec[CompileKey]:
40 + compile_key = self.dispatch(variant=variant)
41 + return compile_key, (payload,), {"variant": variant}
42 +
43 +
44 +def test_cutedsl_launcher_reuses_compiled_executor(monkeypatch) -> None:
45 + compile_calls: list[tuple[Any, tuple[Any, ...]]] = []
46 + launch_calls: list[tuple[int, tuple[Any, ...]]] = []
47 +
48 + def compile_cutedsl(entry: tuple[str, int], *args: Any) -> Any:
49 + compile_calls.append((entry, args))
50 + variant = entry[1]
51 +
52 + def executor(*runtime_args: Any) -> tuple[int, tuple[Any, ...]]:
53 + launch_calls.append((variant, runtime_args))
54 + return variant, runtime_args
55 +
56 + return executor
57 +
58 + monkeypatch.setattr(
59 + jit_warmup_cutedsl_helper,
60 + "compile_cutedsl",
61 + compile_cutedsl,
62 + )
63 + kernel = _TestCuTeDSLKernel()
64 + compile_key = kernel.CompileKey(variant=1)
65 +
66 + kernel.compile(compile_key)
67 + kernel.compile(compile_key)
68 + assert compile_calls == [(("entry", 1), ("fake-1",))]
69 +
70 + assert kernel("runtime", variant=1) == (1, ("runtime",))
71 + assert kernel("other", variant=2) == (2, ("other",))
72 + assert compile_calls == [
73 + (("entry", 1), ("fake-1",)),
74 + (("entry", 2), ("fake-2",)),
75 + ]
76 + assert launch_calls == [(1, ("runtime",)), (2, ("other",))]
tests/model_executor/test_jit_warmup_triton_launcher.py new
+128
@@ -0,0 +1,128 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from dataclasses import dataclass
5 +from typing import Any
6 +
7 +from vllm.model_executor.warmup.jit_warmup_triton_helper import (
8 + LaunchSpec,
9 + VllmTritonJitKernel,
10 + kernel_launcher,
11 +)
12 +
13 +
14 +class _FakeTritonKernel:
15 + arg_names = ("first", "second", "CONST")
16 +
17 + def __init__(self) -> None:
18 + self.warmup_calls: list[dict[str, Any]] = []
19 +
20 + def warmup(self, **kwargs: Any) -> None:
21 + self.warmup_calls.append(kwargs)
22 +
23 +
24 +class _TestTritonKernel(VllmTritonJitKernel["_TestTritonKernel.CompileKey"]):
25 + kernel = _FakeTritonKernel()
26 +
27 + @dataclass(frozen=True)
28 + class CompileKey:
29 + value: int
30 +
31 + def dispatch(self, *, value: int) -> CompileKey:
32 + return self.CompileKey(value=value)
33 +
34 + def get_warmup_keys(self) -> list[CompileKey]:
35 + return self._trace_dispatch(self.dispatch)(value=1)
36 +
37 + def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
38 + return dict(first="warmup", second=compile_key.value, runtime_launcher=None)
39 +
40 + @kernel_launcher
41 + def __call__(
42 + self,
43 + first: str,
44 + second: int,
45 + runtime_launcher: Any,
46 + ) -> LaunchSpec:
47 + return (2,), dict(
48 + CONST=7,
49 + _runtime_launcher=runtime_launcher,
50 + _runtime_launcher_arg_count=2,
51 + )
52 +
53 +
54 +def test_triton_launcher_supports_compile_and_runtime_adapters() -> None:
55 + owner = _TestTritonKernel()
56 + owner.kernel.warmup_calls.clear()
57 +
58 + owner.compile(owner.CompileKey(value=1))
59 + assert owner.kernel.warmup_calls == [
60 + {"grid": (1,), "first": "warmup", "second": 1, "CONST": 7}
61 + ]
62 +
63 + runtime_calls: list[tuple[Any, ...]] = []
64 +
65 + def runtime_launcher(kernel: Any, grid: Any, *args: Any, **kwargs: Any) -> None:
66 + runtime_calls.append((kernel, grid, args, kwargs))
67 +
68 + owner("runtime", 2, runtime_launcher)
69 + assert runtime_calls == [(owner.kernel, (2,), ("runtime", 2), {"CONST": 7})]
70 +
71 +
72 +def test_triton_launcher_supports_cpu_function_wrappers() -> None:
73 + calls: list[tuple[Any, ...]] = []
74 +
75 + def kernel(first: str, second: int, CONST: int) -> None:
76 + calls.append((first, second, CONST))
77 +
78 + class FuncWrapper:
79 + def __init__(self) -> None:
80 + self.func = kernel
81 +
82 + def __getitem__(self, _grid: Any) -> Any:
83 + return self.func
84 +
85 + class TestCpuKernel(_TestTritonKernel):
86 + kernel: Any = FuncWrapper()
87 +
88 + owner = TestCpuKernel()
89 +
90 + owner("runtime", 2, None)
91 + assert calls == [("runtime", 2, 7)]
92 +
93 +
94 +def test_compute_slot_mapping_uses_named_launcher_inputs(monkeypatch) -> None:
95 + from vllm.v1.attention.backends.utils import PAD_SLOT_ID
96 + from vllm.v1.worker.block_table import ComputeSlotMappingKernel
97 +
98 + owner = ComputeSlotMappingKernel()
99 + compile_key = owner.CompileKey(
100 + kv_cache_block_size=16,
101 + blocks_per_kv_block=1,
102 + total_cp_world_size=2,
103 + total_cp_rank=1,
104 + cp_kv_cache_interleave_size=1,
105 + block_table_stride=128,
106 + block_size=16,
107 + )
108 + launches: list[tuple[Any, ...]] = []
109 +
110 + def launch(grid: Any, inputs: Any, **kwargs: Any) -> None:
111 + launches.append((grid, inputs, kwargs))
112 +
113 + monkeypatch.setattr(owner, "launch", launch)
114 + owner.compile(compile_key)
115 +
116 + grid, inputs, kwargs = launches[0]
117 + assert grid == (2,)
118 + assert inputs["num_tokens"] == 2
119 + assert inputs["block_table_stride"] == 128
120 + assert kwargs == {
121 + "KV_CACHE_BLOCK_SIZE": 16,
122 + "BLOCKS_PER_KV_BLOCK": 1,
123 + "TOTAL_CP_WORLD_SIZE": 2,
124 + "TOTAL_CP_RANK": 1,
125 + "CP_KV_CACHE_INTERLEAVE_SIZE": 1,
126 + "PAD_ID": PAD_SLOT_ID,
127 + "BLOCK_SIZE": owner.triton_block_size,
128 + }
vllm/model_executor/kernels/attention/dsa/dcp_indexer_cutedsl.py
+501 -328
@@ -1,7 +1,8 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 -from functools import cache
4 +from dataclasses import dataclass
5 +from typing import Any
6
7 import cutlass
8 import cutlass.cute as cute
@@ -11,6 +12,17 @@ from cutlass import Float32, Int32, Uint32, Uint64
12 from quack.compile_utils import make_fake_tensor
13
14 from vllm.cute_utils import recast_val
15 +from vllm.model_executor.warmup.jit_warmup_cutedsl_helper import (
16 + CuTeDSLLaunchSpec,
17 + VllmCuTeDSLJitKernel,
18 + cutedsl_kernel_launcher,
19 +)
20 +from vllm.model_executor.warmup.jit_warmup_triton_helper import (
21 + LaunchSpec,
22 + TritonWarmupTensor,
23 + VllmTritonJitKernel,
24 + kernel_launcher,
25 +)
26 from vllm.triton_utils import tl, triton
27
28
@@ -25,9 +37,7 @@ def stable_topk_from_gathered_candidates_cutedsl(
37 dtype=torch.int32,
38 device=gathered.device,
39 )
28 - StableTopKFromGatheredCandidatesKernel.compile(topk, gathered.shape[1])(
29 - gathered, out
30 - )
40 + _STABLE_TOPK_FROM_GATHERED_CANDIDATES_KERNEL(gathered, out, topk=topk)
41 return out
42
43
@@ -41,380 +51,543 @@ def pack_dcp_topk_candidates_cutedsl(
51 row_starts: torch.Tensor | None,
52 ) -> None:
53 topk = topk_indices.shape[1]
44 - grid = (topk_indices.shape[0], triton.cdiv(topk, 512))
54 row_starts_arg = row_starts if row_starts is not None else topk_indices
46 - _pack_dcp_topk_candidates_triton_kernel[grid](
55 + _PACK_DCP_TOPK_CANDIDATES_KERNEL(
56 logits,
57 topk_indices,
58 packed,
59 row_starts_arg,
51 - logits.stride(0),
52 - logits.stride(1),
53 - topk_indices.stride(0),
54 - topk_indices.stride(1),
55 - packed.stride(0),
56 - packed.stride(1),
57 - packed.stride(2),
58 - logits.shape[1],
59 - DCP_RANK=dcp_rank,
60 - DCP_WORLD_SIZE=dcp_world_size,
61 - CP_INTERLEAVE=cp_interleave,
62 - HAS_ROW_STARTS=row_starts is not None,
63 - TOPK=topk,
64 - BLOCK_SIZE=512,
65 - num_warps=8,
60 + logits_stride0=logits.stride(0),
61 + logits_stride1=logits.stride(1),
62 + topk_stride0=topk_indices.stride(0),
63 + topk_stride1=topk_indices.stride(1),
64 + packed_stride0=packed.stride(0),
65 + packed_stride1=packed.stride(1),
66 + packed_stride2=packed.stride(2),
67 + num_cols=logits.shape[1],
68 + dcp_rank=dcp_rank,
69 + dcp_world_size=dcp_world_size,
70 + cp_interleave=cp_interleave,
71 + has_row_starts=row_starts is not None,
72 + topk=topk,
73 + block_size=512,
74 )
75
76
69 -@triton.jit
70 -def _pack_dcp_topk_candidates_triton_kernel(
71 - logits,
72 - topk_indices,
73 - packed,
74 - row_starts,
75 - logits_stride0: tl.constexpr,
76 - logits_stride1: tl.constexpr,
77 - topk_stride0: tl.constexpr,
78 - topk_stride1: tl.constexpr,
79 - packed_stride0: tl.constexpr,
80 - packed_stride1: tl.constexpr,
81 - packed_stride2: tl.constexpr,
82 - num_cols,
83 - DCP_RANK: tl.constexpr,
84 - DCP_WORLD_SIZE: tl.constexpr,
85 - CP_INTERLEAVE: tl.constexpr,
86 - HAS_ROW_STARTS: tl.constexpr,
87 - TOPK: tl.constexpr,
88 - BLOCK_SIZE: tl.constexpr,
77 +class PackDCPTopkCandidatesKernel(
78 + VllmTritonJitKernel["PackDCPTopkCandidatesKernel.CompileKey"]
79 ):
90 - row = tl.program_id(0)
91 - tile = tl.program_id(1)
92 - cols = tile * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
93 - mask = cols < TOPK
94 -
95 - local_idx = tl.load(
96 - topk_indices + row * topk_stride0 + cols * topk_stride1,
97 - mask=mask,
98 - other=-1,
99 - )
100 - valid = local_idx >= 0
101 - safe_local_idx = tl.maximum(local_idx, 0)
102 -
103 - row_start = 0
104 - if HAS_ROW_STARTS:
105 - row_start = tl.load(row_starts + row)
106 -
107 - score_col = safe_local_idx + row_start
108 - score_col = tl.minimum(score_col, tl.maximum(num_cols - 1, 0))
109 - score = tl.load(
110 - logits + row * logits_stride0 + score_col * logits_stride1,
111 - mask=mask & valid,
112 - other=-float("inf"),
80 + @dataclass(frozen=True)
81 + class CompileKey:
82 + dcp_rank: int
83 + dcp_world_size: int
84 + cp_interleave: int
85 + has_row_starts: bool
86 + topk: int
87 + block_size: int
88 +
89 + # These scalars only describe runtime layouts and bounds; the constexpr
90 + # fields below own the launch geometry and algorithmic specialization.
91 + @staticmethod
92 + @triton.jit(
93 + do_not_specialize=[
94 + "logits_stride0",
95 + "logits_stride1",
96 + "topk_stride0",
97 + "topk_stride1",
98 + "packed_stride0",
99 + "packed_stride1",
100 + "packed_stride2",
101 + "num_cols",
102 + ]
103 )
104 + def kernel(
105 + logits,
106 + topk_indices,
107 + packed,
108 + row_starts,
109 + logits_stride0,
110 + logits_stride1,
111 + topk_stride0,
112 + topk_stride1,
113 + packed_stride0,
114 + packed_stride1,
115 + packed_stride2,
116 + num_cols,
117 + DCP_RANK: tl.constexpr,
118 + DCP_WORLD_SIZE: tl.constexpr,
119 + CP_INTERLEAVE: tl.constexpr,
120 + HAS_ROW_STARTS: tl.constexpr,
121 + TOPK: tl.constexpr,
122 + BLOCK_SIZE: tl.constexpr,
123 + ):
124 + row = tl.program_id(0)
125 + tile = tl.program_id(1)
126 + cols = tile * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
127 + mask = cols < TOPK
128 +
129 + local_idx = tl.load(
130 + topk_indices + row * topk_stride0 + cols * topk_stride1,
131 + mask=mask,
132 + other=-1,
133 + )
134 + valid = local_idx >= 0
135 + safe_local_idx = tl.maximum(local_idx, 0)
136 +
137 + row_start = 0
138 + if HAS_ROW_STARTS:
139 + row_start = tl.load(row_starts + row)
140 +
141 + score_col = safe_local_idx + row_start
142 + score_col = tl.minimum(score_col, tl.maximum(num_cols - 1, 0))
143 + score = tl.load(
144 + logits + row * logits_stride0 + score_col * logits_stride1,
145 + mask=mask & valid,
146 + other=-float("inf"),
147 + )
148
115 - global_id = (
116 - (safe_local_idx // CP_INTERLEAVE) * (DCP_WORLD_SIZE * CP_INTERLEAVE)
117 - + DCP_RANK * CP_INTERLEAVE
118 - + safe_local_idx % CP_INTERLEAVE
119 - )
120 - global_id = tl.where(valid, global_id, -1)
121 -
122 - packed_base = packed + row * packed_stride0 + cols * packed_stride1
123 - tl.store(packed_base, score, mask=mask)
124 - tl.store(packed_base + packed_stride2, global_id.to(tl.float32), mask=mask)
125 -
126 -
127 -@cute.jit
128 -def _warp_scan_inclusive_i32(val: Int32, lane: Int32) -> Int32:
129 - for i in cutlass.range_constexpr(cute.arch.WARP_SIZE.bit_length() - 1):
130 - offset = 1 << i
131 - partial = cute.arch.shuffle_sync_up(val, offset=offset, mask_and_clamp=0)
132 - if lane >= offset:
133 - val += partial
134 - return val
135 -
136 -
137 -@cute.jit
138 -def _block_scan_inclusive_i32(
139 - val: Int32,
140 - lane: Int32,
141 - warp_id: Int32,
142 - warp_scratch: cute.Tensor,
143 - warps_per_block: int,
144 -) -> Int32:
145 - prefix = _warp_scan_inclusive_i32(val, lane)
146 - if lane == Int32(cute.arch.WARP_SIZE - 1):
147 - warp_scratch[0, warp_id] = prefix
148 - cute.arch.sync_threads()
149 -
150 - if warp_id == Int32(0):
151 - warp_total = Int32(0)
152 - if lane < Int32(warps_per_block):
153 - warp_total = warp_scratch[0, lane]
154 - warp_prefix = _warp_scan_inclusive_i32(warp_total, lane)
155 - if lane < Int32(warps_per_block):
156 - warp_scratch[0, lane] = warp_prefix - warp_total
157 - cute.arch.sync_threads()
158 -
159 - return prefix + warp_scratch[0, warp_id]
160 -
161 -
162 -class StableTopKFromGatheredCandidatesKernel:
149 + global_id = (
150 + (safe_local_idx // CP_INTERLEAVE) * (DCP_WORLD_SIZE * CP_INTERLEAVE)
151 + + DCP_RANK * CP_INTERLEAVE
152 + + safe_local_idx % CP_INTERLEAVE
153 + )
154 + global_id = tl.where(valid, global_id, -1)
155 +
156 + packed_base = packed + row * packed_stride0 + cols * packed_stride1
157 + tl.store(packed_base, score, mask=mask)
158 + tl.store(packed_base + packed_stride2, global_id.to(tl.float32), mask=mask)
159 +
160 + def dispatch( # type: ignore[override]
161 + self,
162 + *,
163 + has_row_starts: bool,
164 + **compile_key_fields: int,
165 + ) -> CompileKey:
166 + return self.CompileKey(
167 + **compile_key_fields,
168 + has_row_starts=has_row_starts,
169 + )
170 +
171 + def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]:
172 + dcp_world_size = vllm_config.parallel_config.decode_context_parallel_size
173 + if dcp_world_size <= 1:
174 + return []
175 + cp_interleave = vllm_config.parallel_config.cp_kv_cache_interleave_size
176 + topk = vllm_config.model_config.hf_config.index_topk
177 + if topk <= 0:
178 + return []
179 +
180 + try:
181 + from vllm.distributed.parallel_state import get_dcp_group
182 +
183 + dcp_rank = get_dcp_group().rank_in_group
184 + except Exception:
185 + dcp_rank = 0
186 +
187 + return self._trace_dispatch(self.dispatch)(
188 + dcp_rank=dcp_rank,
189 + dcp_world_size=dcp_world_size,
190 + cp_interleave=cp_interleave,
191 + has_row_starts=(False, True),
192 + topk=topk,
193 + block_size=512,
194 + )
195 +
196 + def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
197 + fp32_ptr = TritonWarmupTensor(torch.float32)
198 + int32_ptr = TritonWarmupTensor(torch.int32)
199 + return dict(
200 + logits=fp32_ptr,
201 + topk_indices=TritonWarmupTensor(torch.int32, shape=(1, compile_key.topk)),
202 + packed=fp32_ptr,
203 + row_starts_arg=int32_ptr,
204 + logits_stride0=1,
205 + logits_stride1=1,
206 + topk_stride0=1,
207 + topk_stride1=1,
208 + packed_stride0=1,
209 + packed_stride1=1,
210 + packed_stride2=1,
211 + num_cols=1,
212 + dcp_rank=compile_key.dcp_rank,
213 + dcp_world_size=compile_key.dcp_world_size,
214 + cp_interleave=compile_key.cp_interleave,
215 + has_row_starts=compile_key.has_row_starts,
216 + topk=compile_key.topk,
217 + block_size=compile_key.block_size,
218 + )
219 +
220 + @kernel_launcher
221 + def __call__(
222 + self,
223 + logits: torch.Tensor,
224 + topk_indices: torch.Tensor,
225 + packed: torch.Tensor,
226 + row_starts_arg: torch.Tensor,
227 + *,
228 + logits_stride0: int,
229 + logits_stride1: int,
230 + topk_stride0: int,
231 + topk_stride1: int,
232 + packed_stride0: int,
233 + packed_stride1: int,
234 + packed_stride2: int,
235 + num_cols: int,
236 + dcp_rank: int,
237 + dcp_world_size: int,
238 + cp_interleave: int,
239 + has_row_starts: bool,
240 + topk: int,
241 + block_size: int,
242 + ) -> LaunchSpec:
243 + grid = (topk_indices.shape[0], triton.cdiv(topk, block_size))
244 + return grid, dict(
245 + row_starts=row_starts_arg,
246 + DCP_RANK=dcp_rank,
247 + DCP_WORLD_SIZE=dcp_world_size,
248 + CP_INTERLEAVE=cp_interleave,
249 + HAS_ROW_STARTS=has_row_starts,
250 + TOPK=topk,
251 + BLOCK_SIZE=block_size,
252 + num_warps=8,
253 + )
254 +
255 +
256 +class StableTopKFromGatheredCandidatesKernel(
257 + VllmCuTeDSLJitKernel["StableTopKFromGatheredCandidatesKernel.CompileKey"]
258 +):
259 tb_size = 512
260 hist_bins = 2048
261 radix_bits = (hist_bins - 1).bit_length()
166 - assert hist_bins == 1 << radix_bits
262 key_bits = Uint64.width
263 radix_passes = (key_bits + radix_bits - 1) // radix_bits
264 final_radix_bits = key_bits - radix_bits * (radix_passes - 1)
265 hist_chunks = (hist_bins + tb_size - 1) // tb_size
266 warps_per_block = tb_size // cute.arch.WARP_SIZE
267
173 - def __init__(self, topk: int, num_candidates: int):
174 - assert num_candidates % self.tb_size == 0, (
268 + @dataclass(frozen=True)
269 + class CompileKey:
270 + topk: int
271 + num_candidates: int
272 +
273 + @staticmethod
274 + def kernel(compile_key: CompileKey) -> Any:
275 + tb_size = StableTopKFromGatheredCandidatesKernel.tb_size
276 + hist_bins = StableTopKFromGatheredCandidatesKernel.hist_bins
277 + radix_bits = StableTopKFromGatheredCandidatesKernel.radix_bits
278 + key_bits = StableTopKFromGatheredCandidatesKernel.key_bits
279 + radix_passes = StableTopKFromGatheredCandidatesKernel.radix_passes
280 + final_radix_bits = StableTopKFromGatheredCandidatesKernel.final_radix_bits
281 + hist_chunks = StableTopKFromGatheredCandidatesKernel.hist_chunks
282 + warps_per_block = StableTopKFromGatheredCandidatesKernel.warps_per_block
283 + topk = compile_key.topk
284 + assert hist_bins == 1 << radix_bits
285 + assert compile_key.num_candidates % tb_size == 0, (
286 "StableTopKFromGatheredCandidatesKernel requires candidate count "
176 - f"to be a multiple of {self.tb_size}, got {num_candidates}"
287 + f"to be a multiple of {tb_size}, got {compile_key.num_candidates}"
288 )
178 - self.topk = topk
179 - self.keys_per_thread = num_candidates // self.tb_size
289 + keys_per_thread = compile_key.num_candidates // tb_size
290
291 @cute.struct
292 class SharedStorage:
183 - hist: cute.struct.MemRange[Int32, self.hist_bins]
293 + hist: cute.struct.MemRange[Int32, hist_bins]
294 committed_count: cute.struct.MemRange[Int32, 1]
295 running_count: cute.struct.MemRange[Int32, 1]
296 threshold_bin: cute.struct.MemRange[Int32, 1]
297 threshold_found: cute.struct.MemRange[Int32, 1]
298 include_threshold_bin: cute.struct.MemRange[Int32, 1]
299 prefix_s: cute.struct.Align[cute.struct.MemRange[Uint64, 1], 8]
190 - warp_totals: cute.struct.MemRange[Int32, self.warps_per_block]
300 + warp_totals: cute.struct.MemRange[Int32, warps_per_block]
301
192 - self.shared_storage = SharedStorage
302 + shared_storage = SharedStorage
303
194 - @cute.jit
195 - def __call__(
196 - self,
197 - gathered: cute.Tensor,
198 - out: cute.Tensor,
199 - stream: CUstream,
200 - ):
201 - grid = (gathered.shape[0], 1, 1)
202 - self.kernel(gathered, out).launch(
203 - grid=grid,
204 - block=(self.tb_size, 1, 1),
205 - stream=stream,
206 - )
304 + @cute.jit
305 + def warp_scan_inclusive_i32(val: Int32, lane: Int32) -> Int32:
306 + for i in cutlass.range_constexpr(cute.arch.WARP_SIZE.bit_length() - 1):
307 + offset = 1 << i
308 + partial = cute.arch.shuffle_sync_up(
309 + val, offset=offset, mask_and_clamp=0
310 + )
311 + if lane >= offset:
312 + val += partial
313 + return val
314 +
315 + @cute.jit
316 + def block_scan_inclusive_i32(
317 + val: Int32,
318 + lane: Int32,
319 + warp_id: Int32,
320 + warp_scratch: cute.Tensor,
321 + warps_per_block: int,
322 + ) -> Int32:
323 + prefix = warp_scan_inclusive_i32(val, lane)
324 + if lane == Int32(cute.arch.WARP_SIZE - 1):
325 + warp_scratch[0, warp_id] = prefix
326 + cute.arch.sync_threads()
327
208 - @cute.jit
209 - def _stable_key(self, score: Float32, token_id: Int32) -> Uint64:
210 - bits = recast_val(score, Uint32)
211 - mask = Uint32(0x80000000)
212 - if (bits & Uint32(0x80000000)) != Uint32(0):
213 - mask = Uint32(0xFFFFFFFF)
214 - score_key = Uint64(bits ^ mask) << Uint64(32)
215 - id_key = Uint64(~Uint32(token_id))
216 - key = score_key | id_key
217 - if token_id < Int32(0):
218 - key = Uint64(0)
219 - return key
220 -
221 - @cute.jit
222 - def _prefix_matches(
223 - self,
224 - key: Uint64,
225 - prefix: Uint64,
226 - prefix_bits: Int32,
227 - ):
228 - matches = prefix_bits == Int32(0)
229 - if prefix_bits != Int32(0):
230 - shift = Int32(self.key_bits) - prefix_bits
231 - matches = (key >> Uint64(shift)) == (prefix >> Uint64(shift))
232 - return matches
233 -
234 - @cute.jit
235 - def _radix_pass(
236 - self,
237 - keys: cute.Tensor,
238 - output: cute.Tensor,
239 - storage,
240 - tid: Int32,
241 - step: Int32,
242 - bits: int,
243 - is_final_pass: bool,
244 - ):
245 - hist_smem = storage.hist.get_tensor(cute.make_layout((self.hist_bins,)))
246 - committed_count_smem = storage.committed_count.data_ptr()
247 - running_count_smem = storage.running_count.data_ptr()
248 - threshold_bin_smem = storage.threshold_bin.data_ptr()
249 - threshold_found_smem = storage.threshold_found.data_ptr()
250 - include_threshold_bin_smem = storage.include_threshold_bin.data_ptr()
251 - prefix_smem = storage.prefix_s.data_ptr()
252 - warp_totals_smem = storage.warp_totals.get_tensor(
253 - cute.make_layout((1, self.warps_per_block))
254 - )
328 + if warp_id == Int32(0):
329 + warp_total = Int32(0)
330 + if lane < Int32(warps_per_block):
331 + warp_total = warp_scratch[0, lane]
332 + warp_prefix = warp_scan_inclusive_i32(warp_total, lane)
333 + if lane < Int32(warps_per_block):
334 + warp_scratch[0, lane] = warp_prefix - warp_total
335 + cute.arch.sync_threads()
336
256 - prefix_bits = step * Int32(self.radix_bits)
257 - num_bins = 1 << bits
258 - block_scan_iterations = (num_bins + self.tb_size - 1) // self.tb_size
259 - shift = Int32(self.key_bits) - prefix_bits - Int32(bits)
260 - bin_mask = Uint64(num_bins - 1)
261 - prefix = prefix_smem.load()
262 -
263 - for chunk in cutlass.range_constexpr(self.hist_chunks):
264 - hist_smem[tid + Int32(chunk * self.tb_size)] = Int32(0)
265 - if tid == Int32(0):
266 - running_count_smem.store(committed_count_smem.load())
267 - include_threshold_bin_smem.store(Int32(0))
268 - threshold_found_smem.store(Int32(0))
269 - cute.arch.sync_threads()
270 -
271 - for key_idx in cutlass.range_constexpr(self.keys_per_thread):
272 - key = keys[key_idx]
273 - if self._prefix_matches(key, prefix, prefix_bits):
274 - bin_idx = Int32((key >> Uint64(shift)) & bin_mask)
275 - cute.arch.atomic_add(
276 - hist_smem.iterator + bin_idx,
277 - Int32(1),
278 - sem="relaxed",
279 - scope="cta",
280 - )
281 - cute.arch.sync_threads()
282 -
283 - lane = cute.arch.lane_idx()
284 - warp_id = cute.arch.warp_idx()
285 - # Each iteration scans one tb_size-wide slice of bins, high to low.
286 - iter = Int32(0)
287 - threshold_found = threshold_found_smem.load()
288 - while threshold_found == Int32(0) and iter < Int32(block_scan_iterations):
289 - bin_idx = Int32(num_bins - 1) - (iter * Int32(self.tb_size) + tid)
290 - count = hist_smem[bin_idx]
291 - chunk_inclusive = _block_scan_inclusive_i32(
292 - count,
293 - lane,
294 - warp_id,
295 - warp_totals_smem,
296 - self.warps_per_block,
337 + return prefix + warp_scratch[0, warp_id]
338 +
339 + @cute.jit
340 + def stable_key(score: Float32, token_id: Int32) -> Uint64:
341 + bits = recast_val(score, Uint32)
342 + mask = Uint32(0x80000000)
343 + if (bits & Uint32(0x80000000)) != Uint32(0):
344 + mask = Uint32(0xFFFFFFFF)
345 + score_key = Uint64(bits ^ mask) << Uint64(32)
346 + id_key = Uint64(~Uint32(token_id))
347 + key = score_key | id_key
348 + if token_id < Int32(0):
349 + key = Uint64(0)
350 + return key
351 +
352 + @cute.jit
353 + def prefix_matches(
354 + key: Uint64,
355 + prefix: Uint64,
356 + prefix_bits: Int32,
357 + ):
358 + matches = prefix_bits == Int32(0)
359 + if prefix_bits != Int32(0):
360 + shift = Int32(key_bits) - prefix_bits
361 + matches = (key >> Uint64(shift)) == (prefix >> Uint64(shift))
362 + return matches
363 +
364 + @cute.jit
365 + def radix_pass(
366 + keys: cute.Tensor,
367 + output: cute.Tensor,
368 + storage,
369 + tid: Int32,
370 + step: Int32,
371 + bits: int,
372 + is_final_pass: bool,
373 + ):
374 + hist_smem = storage.hist.get_tensor(cute.make_layout((hist_bins,)))
375 + committed_count_smem = storage.committed_count.data_ptr()
376 + running_count_smem = storage.running_count.data_ptr()
377 + threshold_bin_smem = storage.threshold_bin.data_ptr()
378 + threshold_found_smem = storage.threshold_found.data_ptr()
379 + include_threshold_bin_smem = storage.include_threshold_bin.data_ptr()
380 + prefix_smem = storage.prefix_s.data_ptr()
381 + warp_totals_smem = storage.warp_totals.get_tensor(
382 + cute.make_layout((1, warps_per_block))
383 )
298 - running_count = running_count_smem.load()
299 - prior_in_scan_slice = chunk_inclusive - count
300 - remaining = Int32(self.topk) - running_count - prior_in_scan_slice
301 - if count > Int32(0) and remaining > Int32(0) and remaining <= count:
302 - threshold_bin_smem.store(bin_idx)
303 - if count <= remaining or cutlass.const_expr(is_final_pass):
304 - include_threshold_bin_smem.store(Int32(1))
305 - threshold_found_smem.store(Int32(1))
306 - # Barrier: every thread must finish reading running_count for this
307 - # slice before tb_size-1 advances it, else a warp racing ahead to
308 - # the store makes a lagging thread double-count the slice total
309 - # (-> remaining too small -> threshold too high -> under-fill).
310 - cute.arch.sync_threads()
311 - if tid == Int32(self.tb_size - 1):
312 - running_count_smem.store(running_count + chunk_inclusive)
384 +
385 + prefix_bits = step * Int32(radix_bits)
386 + num_bins = 1 << bits
387 + block_scan_iterations = (num_bins + tb_size - 1) // tb_size
388 + shift = Int32(key_bits) - prefix_bits - Int32(bits)
389 + bin_mask = Uint64(num_bins - 1)
390 + prefix = prefix_smem.load()
391 +
392 + for chunk in cutlass.range_constexpr(hist_chunks):
393 + hist_smem[tid + Int32(chunk * tb_size)] = Int32(0)
394 + if tid == Int32(0):
395 + running_count_smem.store(committed_count_smem.load())
396 + include_threshold_bin_smem.store(Int32(0))
397 + threshold_found_smem.store(Int32(0))
398 cute.arch.sync_threads()
399
315 - threshold_found = threshold_found_smem.load()
316 - iter += Int32(1)
317 -
318 - threshold = threshold_bin_smem.load()
319 - should_include_threshold = include_threshold_bin_smem.load() != Int32(0)
320 - for key_idx in cutlass.range_constexpr(self.keys_per_thread):
321 - key = keys[key_idx]
322 - if self._prefix_matches(key, prefix, prefix_bits):
323 - bin_idx = Int32((key >> Uint64(shift)) & bin_mask)
324 - selected = bin_idx > threshold
325 - if should_include_threshold:
326 - selected = selected or bin_idx == threshold
327 - if selected:
328 - dst = cute.arch.atomic_add(
329 - committed_count_smem,
400 + for key_idx in cutlass.range_constexpr(keys_per_thread):
401 + key = keys[key_idx]
402 + if prefix_matches(key, prefix, prefix_bits):
403 + bin_idx = Int32((key >> Uint64(shift)) & bin_mask)
404 + cute.arch.atomic_add(
405 + hist_smem.iterator + bin_idx,
406 Int32(1),
407 sem="relaxed",
408 scope="cta",
409 )
334 - if dst < Int32(self.topk):
335 - output[dst] = recast_val(~Uint32(key), Int32)
336 - cute.arch.sync_threads()
410 + cute.arch.sync_threads()
411
338 - pass_finished = include_threshold_bin_smem.load()
339 - if tid == Int32(0) and pass_finished == Int32(0):
340 - prefix_smem.store(prefix | (Uint64(threshold) << Uint64(shift)))
341 - cute.arch.sync_threads()
342 - return pass_finished
412 + lane = cute.arch.lane_idx()
413 + warp_id = cute.arch.warp_idx()
414 + # Each iteration scans one tb_size-wide slice of bins, high to low.
415 + iter = Int32(0)
416 + threshold_found = threshold_found_smem.load()
417 + while threshold_found == Int32(0) and iter < Int32(block_scan_iterations):
418 + bin_idx = Int32(num_bins - 1) - (iter * Int32(tb_size) + tid)
419 + count = hist_smem[bin_idx]
420 + chunk_inclusive = block_scan_inclusive_i32(
421 + count,
422 + lane,
423 + warp_id,
424 + warp_totals_smem,
425 + warps_per_block,
426 + )
427 + running_count = running_count_smem.load()
428 + prior_in_scan_slice = chunk_inclusive - count
429 + remaining = Int32(topk) - running_count - prior_in_scan_slice
430 + if count > Int32(0) and remaining > Int32(0) and remaining <= count:
431 + threshold_bin_smem.store(bin_idx)
432 + if count <= remaining or cutlass.const_expr(is_final_pass):
433 + include_threshold_bin_smem.store(Int32(1))
434 + threshold_found_smem.store(Int32(1))
435 + # Barrier: every thread must finish reading running_count for this
436 + # slice before tb_size-1 advances it, else a warp racing ahead to
437 + # the store makes a lagging thread double-count the slice total
438 + # (-> remaining too small -> threshold too high -> under-fill).
439 + cute.arch.sync_threads()
440 + if tid == Int32(tb_size - 1):
441 + running_count_smem.store(running_count + chunk_inclusive)
442 + cute.arch.sync_threads()
443 +
444 + threshold_found = threshold_found_smem.load()
445 + iter += Int32(1)
446 +
447 + threshold = threshold_bin_smem.load()
448 + should_include_threshold = include_threshold_bin_smem.load() != Int32(0)
449 + for key_idx in cutlass.range_constexpr(keys_per_thread):
450 + key = keys[key_idx]
451 + if prefix_matches(key, prefix, prefix_bits):
452 + bin_idx = Int32((key >> Uint64(shift)) & bin_mask)
453 + selected = bin_idx > threshold
454 + if should_include_threshold:
455 + selected = selected or bin_idx == threshold
456 + if selected:
457 + dst = cute.arch.atomic_add(
458 + committed_count_smem,
459 + Int32(1),
460 + sem="relaxed",
461 + scope="cta",
462 + )
463 + if dst < Int32(topk):
464 + output[dst] = recast_val(~Uint32(key), Int32)
465 + cute.arch.sync_threads()
466
344 - @cute.kernel
345 - def kernel(
346 - self,
347 - input: cute.Tensor,
348 - out: cute.Tensor,
349 - ):
350 - row, _, _ = cute.arch.block_idx()
351 - tid, _, _ = cute.arch.thread_idx()
352 - input_row = input[row, None, None]
353 - output_row = out[row, None]
354 - keys = cute.make_rmem_tensor((self.keys_per_thread,), Uint64)
355 -
356 - smem = cutlass.utils.SmemAllocator()
357 - storage = smem.allocate(self.shared_storage, 8)
358 - committed_count_smem = storage.committed_count.data_ptr()
359 - prefix_smem = storage.prefix_s.data_ptr()
360 - for i in range(tid, self.topk, self.tb_size):
361 - output_row[i] = Int32(-1)
362 -
363 - for key_idx in cutlass.range_constexpr(self.keys_per_thread):
364 - col = tid + Int32(key_idx * self.tb_size)
365 - score = Float32(input_row[col, 0])
366 - token_id = Int32(input_row[col, 1])
367 - keys[key_idx] = self._stable_key(score, token_id)
368 -
369 - if tid == Int32(0):
370 - committed_count_smem.store(Int32(0))
371 - prefix_smem.store(Uint64(0))
372 - cute.arch.sync_threads()
373 -
374 - step = Int32(0)
375 - finished = Int32(0)
376 - while finished == Int32(0) and step < Int32(self.radix_passes - 1):
377 - finished = self._radix_pass(
378 - keys,
379 - output_row,
380 - storage,
381 - tid,
382 - step,
383 - self.radix_bits,
384 - False,
385 - )
386 - step += Int32(1)
387 -
388 - if finished == Int32(0):
389 - self._radix_pass(
390 - keys,
391 - output_row,
392 - storage,
393 - tid,
394 - Int32(self.radix_passes - 1),
395 - self.final_radix_bits,
396 - True,
467 + pass_finished = include_threshold_bin_smem.load()
468 + if tid == Int32(0) and pass_finished == Int32(0):
469 + prefix_smem.store(prefix | (Uint64(threshold) << Uint64(shift)))
470 + cute.arch.sync_threads()
471 + return pass_finished
472 +
473 + @cute.kernel
474 + def device_kernel(input: cute.Tensor, out: cute.Tensor):
475 + row, _, _ = cute.arch.block_idx()
476 + tid, _, _ = cute.arch.thread_idx()
477 + input_row = input[row, None, None]
478 + output_row = out[row, None]
479 + keys = cute.make_rmem_tensor((keys_per_thread,), Uint64)
480 +
481 + smem = cutlass.utils.SmemAllocator()
482 + storage = smem.allocate(shared_storage, 8)
483 + committed_count_smem = storage.committed_count.data_ptr()
484 + prefix_smem = storage.prefix_s.data_ptr()
485 + for i in range(tid, topk, tb_size):
486 + output_row[i] = Int32(-1)
487 +
488 + for key_idx in cutlass.range_constexpr(keys_per_thread):
489 + col = tid + Int32(key_idx * tb_size)
490 + score = Float32(input_row[col, 0])
491 + token_id = Int32(input_row[col, 1])
492 + keys[key_idx] = stable_key(score, token_id)
493 +
494 + if tid == Int32(0):
495 + committed_count_smem.store(Int32(0))
496 + prefix_smem.store(Uint64(0))
497 + cute.arch.sync_threads()
498 +
499 + step = Int32(0)
500 + finished = Int32(0)
501 + while finished == Int32(0) and step < Int32(radix_passes - 1):
502 + finished = radix_pass(
503 + keys,
504 + output_row,
505 + storage,
506 + tid,
507 + step,
508 + radix_bits,
509 + False,
510 + )
511 + step += Int32(1)
512 +
513 + if finished == Int32(0):
514 + radix_pass(
515 + keys,
516 + output_row,
517 + storage,
518 + tid,
519 + Int32(radix_passes - 1),
520 + final_radix_bits,
521 + True,
522 + )
523 +
524 + @cute.jit
525 + def host_entrypoint(
526 + gathered: cute.Tensor,
527 + out: cute.Tensor,
528 + stream: CUstream,
529 + ):
530 + grid = (gathered.shape[0], 1, 1)
531 + device_kernel(gathered, out).launch(
532 + grid=grid,
533 + block=(tb_size, 1, 1),
534 + stream=stream,
535 )
536
399 - @cache
400 - @staticmethod
401 - def compile(topk: int, num_candidates: int):
402 - num_rows = cute.sym_int()
537 + return host_entrypoint
538
539 + def dispatch( # type: ignore[override]
540 + self,
541 + **compile_key_fields: int,
542 + ) -> CompileKey:
543 + return self.CompileKey(**compile_key_fields)
544 +
545 + def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]:
546 + dcp_world_size = vllm_config.parallel_config.decode_context_parallel_size
547 + if dcp_world_size <= 1:
548 + return []
549 + topk = vllm_config.model_config.hf_config.index_topk
550 + if topk <= 0:
551 + return []
552 + return self._trace_dispatch(self.dispatch)(
553 + topk=topk,
554 + num_candidates=topk * dcp_world_size,
555 + )
556 +
557 + def warmup_inputs(self, compile_key: CompileKey) -> tuple[Any, ...]:
558 + num_rows = cute.sym_int()
559 gathered = cute.runtime.make_fake_tensor(
560 Float32,
406 - (num_rows, num_candidates, 2),
561 + (num_rows, compile_key.num_candidates, 2),
562 stride=(cute.sym_int64(divisibility=2), 2, 1),
563 assumed_align=8,
564 )
410 - out = make_fake_tensor(Int32, (num_rows, topk), divisibility=1)
411 -
412 - kernel = StableTopKFromGatheredCandidatesKernel(topk, num_candidates)
413 - stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
414 - return cute.compile(
415 - kernel,
416 - gathered,
417 - out,
418 - stream,
419 - options="--enable-tvm-ffi",
565 + out = make_fake_tensor(
566 + Int32,
567 + (num_rows, compile_key.topk),
568 + divisibility=1,
569 + )
570 + return gathered, out
571 +
572 + @cutedsl_kernel_launcher
573 + def __call__(
574 + self,
575 + gathered: torch.Tensor,
576 + out: torch.Tensor,
577 + *,
578 + topk: int,
579 + ) -> CuTeDSLLaunchSpec[CompileKey]:
580 + compile_key = self.dispatch(topk=topk, num_candidates=gathered.shape[1])
581 + return (
582 + compile_key,
583 + (gathered, out),
584 + {
585 + "gathered_shape": tuple(gathered.shape),
586 + "out_shape": tuple(out.shape),
587 + "topk": topk,
588 + },
589 )
590 +
591 +
592 +_PACK_DCP_TOPK_CANDIDATES_KERNEL = PackDCPTopkCandidatesKernel()
593 +_STABLE_TOPK_FROM_GATHERED_CANDIDATES_KERNEL = StableTopKFromGatheredCandidatesKernel()
vllm/model_executor/layers/sparse_attn_indexer.py
+24 -1
@@ -768,7 +768,8 @@ class SparseAttnIndexer(CustomOp):
768 # DCP scalars are constant for the run; resolve them here (config is set
769 # during model construction) and pass them into the custom op, rather
770 # than threading them through per-step metadata.
771 - parallel_config = get_current_vllm_config().parallel_config
771 + vllm_config = get_current_vllm_config()
772 + parallel_config = vllm_config.parallel_config
773 self._parallel_config = parallel_config
774 self.dcp_world_size = parallel_config.decode_context_parallel_size
775 self.dcp_rank = get_dcp_group().rank_in_group if self.dcp_world_size > 1 else 0
@@ -780,6 +781,28 @@ class SparseAttnIndexer(CustomOp):
781 "the current vLLM environment."
782 )
783
784 + if vllm_config.kernel_config.enable_jit_warmup:
785 + from vllm.v1.attention.ops.common import (
786 + _PACK_SEQ_TRITON_KERNEL,
787 + _UNPACK_SEQ_TRITON_KERNEL,
788 + )
789 +
790 + pack_dtype = torch.uint8 if use_fp4_cache else current_platform.fp8_dtype()
791 + _PACK_SEQ_TRITON_KERNEL.register_warmup(
792 + dtype=pack_dtype,
793 + pad_value=0 if use_fp4_cache else -float("inf"),
794 + )
795 + _UNPACK_SEQ_TRITON_KERNEL.register_warmup()
796 +
797 + if self.dcp_world_size > 1 and current_platform.is_cuda() and has_cutedsl():
798 + from vllm.model_executor.kernels.attention.dsa.dcp_indexer_cutedsl import ( # noqa: E501
799 + _PACK_DCP_TOPK_CANDIDATES_KERNEL,
800 + _STABLE_TOPK_FROM_GATHERED_CANDIDATES_KERNEL,
801 + )
802 +
803 + _PACK_DCP_TOPK_CANDIDATES_KERNEL.register_warmup()
804 + _STABLE_TOPK_FROM_GATHERED_CANDIDATES_KERNEL.register_warmup()
805 +
806 @property
807 def cp_kv_cache_interleave_size(self) -> int:
808 """With PD+DCP, the real value isn't known until block_size is finalized,
vllm/model_executor/warmup/jit_warmup.py
+37 -7
@@ -670,6 +670,34 @@ class VllmJitKernel(Generic[CompileKeyT], ABC):
670 self.compile(compile_key)
671
672
673 +def _same_value(left: Any, right: Any) -> bool:
674 + """True if two registration values are the same object or compare equal.
675 +
676 + Identity is checked first so shared singletons (e.g. ``vllm_config``, torch
677 + dtypes) short-circuit before any potentially deep or non-boolean ``__eq__``.
678 + """
679 + if left is right:
680 + return True
681 + try:
682 + return bool(left == right)
683 + except (TypeError, ValueError, RuntimeError):
684 + return False
685 +
686 +
687 +def _same_registration(
688 + left: tuple[tuple[Any, ...], dict[str, Any]],
689 + right: tuple[tuple[Any, ...], dict[str, Any]],
690 +) -> bool:
691 + """True if two ``(args, kwargs)`` registrations are equivalent."""
692 + left_args, left_kwargs = left
693 + right_args, right_kwargs = right
694 + if len(left_args) != len(right_args) or left_kwargs.keys() != right_kwargs.keys():
695 + return False
696 + return all(
697 + _same_value(x, y) for x, y in zip(left_args, right_args, strict=True)
698 + ) and all(_same_value(left_kwargs[k], right_kwargs[k]) for k in left_kwargs)
699 +
700 +
701 class JitWarmupRegistry:
702 """Collect and compile JIT kernels selected during runner setup."""
703
@@ -713,13 +741,15 @@ class JitWarmupRegistry:
741 kwargs: dict[str, Any],
742 ) -> None:
743 registrations = self._registrations.setdefault(kernel, [])
716 - if (
717 - not args
718 - and not kwargs
719 - and any(
720 - not registered_args and not registered_kwargs
721 - for registered_args, registered_kwargs in registrations
722 - )
744 + # Every layer of a deep model appends an identical (args, kwargs) here
745 + # (e.g. a 61-layer DSA model registers the same pack (dtype, pad_value)
746 + # 61 times); each expands to the same compile keys, so tracing
747 + # get_warmup_keys once per distinct registration is sufficient. Dedup on
748 + # identity-or-equality -- identity short-circuits shared singletons like
749 + # vllm_config before any deep __eq__.
750 + if any(
751 + _same_registration(registered, (args, kwargs))
752 + for registered in registrations
753 ):
754 return
755 registrations.append((args, kwargs))
vllm/model_executor/warmup/jit_warmup_cutedsl_helper.py new
+80
@@ -0,0 +1,80 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from abc import abstractmethod
5 +from collections.abc import Callable, Mapping
6 +from functools import wraps
7 +from typing import Any, ClassVar, Generic, TypeAlias, TypeVar
8 +
9 +from vllm.model_executor.warmup.jit_warmup import VllmJitKernel
10 +
11 +DEFAULT_CUTEDSL_COMPILE_OPTIONS = "--enable-tvm-ffi"
12 +
13 +CompileKeyT = TypeVar("CompileKeyT")
14 +CuTeDSLLaunchSpec: TypeAlias = tuple[
15 + CompileKeyT,
16 + tuple[Any, ...],
17 + Mapping[str, Any] | None,
18 +]
19 +
20 +
21 +def cutedsl_fake_stream(*, use_tvm_ffi_env_stream: bool = True) -> Any:
22 + import cutlass.cute as cute
23 +
24 + return cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=use_tvm_ffi_env_stream)
25 +
26 +
27 +def compile_cutedsl(
28 + entry: Callable[..., Any],
29 + *args: Any,
30 + options: str = DEFAULT_CUTEDSL_COMPILE_OPTIONS,
31 + use_tvm_ffi_env_stream: bool = True,
32 +) -> Any:
33 + import cutlass.cute as cute
34 +
35 + return cute.compile(
36 + entry,
37 + *args,
38 + cutedsl_fake_stream(use_tvm_ffi_env_stream=use_tvm_ffi_env_stream),
39 + options=options,
40 + )
41 +
42 +
43 +class VllmCuTeDSLJitKernel(VllmJitKernel[CompileKeyT], Generic[CompileKeyT]):
44 + """CuTeDSL owner whose compiled executor is shared by warmup and runtime."""
45 +
46 + kernel: ClassVar[Any]
47 +
48 + @abstractmethod
49 + def warmup_inputs(self, compile_key: CompileKeyT) -> tuple[Any, ...]:
50 + """Return fake arguments that compile one executor specialization."""
51 + raise NotImplementedError
52 +
53 + def compile(self, compile_key: CompileKeyT) -> None:
54 + if compile_key in self._compiled_cache:
55 + return
56 + self._compiled_cache[compile_key] = compile_cutedsl(
57 + self.kernel(compile_key),
58 + *self.warmup_inputs(compile_key),
59 + )
60 +
61 +
62 +def cutedsl_kernel_launcher(
63 + call_fn: Callable[..., CuTeDSLLaunchSpec[CompileKeyT]],
64 +) -> Callable[..., Any]:
65 + """Invoke a cached CuTeDSL executor from a declarative ``__call__``."""
66 +
67 + @wraps(call_fn)
68 + def wrapper(
69 + self: VllmCuTeDSLJitKernel[CompileKeyT],
70 + *args: Any,
71 + **kwargs: Any,
72 + ) -> Any:
73 + compile_key, launch_args, runtime_context = call_fn(self, *args, **kwargs)
74 + executor = self._get_or_compile(
75 + compile_key,
76 + runtime_context=runtime_context,
77 + )
78 + return executor(*launch_args)
79 +
80 + return wrapper
vllm/model_executor/warmup/jit_warmup_triton_helper.py
+20 -4
@@ -106,8 +106,16 @@ class VllmTritonJitKernel(VllmJitKernel[CompileKeyT], Generic[CompileKeyT]):
106 self._warming = False
107
108 @cached_property
109 - def _kernel_param_names(self) -> frozenset[str]:
110 - return frozenset(self.kernel.arg_names)
109 + def _kernel_arg_names(self) -> tuple[str, ...]:
110 + arg_names = getattr(self.kernel, "arg_names", None)
111 + if arg_names is not None:
112 + return tuple(arg_names)
113 + wrapped = getattr(self.kernel, "func", None)
114 + if wrapped is not None:
115 + return tuple(inspect.signature(wrapped).parameters)
116 + raise TypeError(
117 + f"Cannot inspect kernel parameters for {type(self.kernel).__name__}"
118 + )
119
120 def launch(
121 self,
@@ -116,14 +124,22 @@ class VllmTritonJitKernel(VllmJitKernel[CompileKeyT], Generic[CompileKeyT]):
124 /,
125 **kwargs: Any,
126 ) -> Any:
127 + runtime_launcher = kwargs.pop("_runtime_launcher", None)
128 + runtime_launcher_arg_count = kwargs.pop("_runtime_launcher_arg_count", 0)
129 for name, value in inputs.items():
120 - target = name if name in self._kernel_param_names else f"{name}_ptr"
121 - if target in self._kernel_param_names and target not in kwargs:
130 + target = name if name in self._kernel_arg_names else f"{name}_ptr"
131 + if target in self._kernel_arg_names and target not in kwargs:
132 kwargs[target] = value
133 if self._warming:
134 warmup = getattr(self.kernel, "warmup", None)
135 assert warmup is not None
136 return warmup(grid=(1,), **kwargs)
137 + if runtime_launcher is not None:
138 + regular_args = [
139 + kwargs.pop(name)
140 + for name in self._kernel_arg_names[:runtime_launcher_arg_count]
141 + ]
142 + return runtime_launcher(self.kernel, grid, *regular_args, **kwargs)
143 return self.kernel[grid](**kwargs)
144
145
vllm/v1/attention/ops/common.py
+271 -115
@@ -1,63 +1,157 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from dataclasses import dataclass
4 +from typing import Any
5 +
6 import torch
7
8 +from vllm.model_executor.warmup.jit_warmup_triton_helper import (
9 + LaunchSpec,
10 + TritonWarmupTensor,
11 + VllmTritonJitKernel,
12 + kernel_launcher,
13 +)
14 from vllm.triton_utils import tl, triton
15
16
8 -@triton.jit
9 -def _pack_seq_kernel(
10 - x_ptr, # [N, D]
11 - out_ptr, # [B, Lmax, D]
12 - lengths_ptr, # *i32, [B]
13 - N: tl.constexpr,
14 - D: tl.constexpr,
15 - Lmax: tl.constexpr,
16 - PAD_VALUE: tl.constexpr,
17 - PAD_IS_UINT8: tl.constexpr,
18 - BLOCK_T: tl.constexpr, # timesteps per program
19 - BLOCK_D: tl.constexpr, # features per program
20 -):
21 - pid_b = tl.program_id(0) # batch id
22 - pid_t = tl.program_id(1) # block over time dimension
23 - pid_d = tl.program_id(2) # block over feature dimension
24 - off_t = pid_t * BLOCK_T + tl.arange(0, BLOCK_T) # [BLOCK_T]
25 - off_d = pid_d * BLOCK_D + tl.arange(0, BLOCK_D) # [BLOCK_D]
26 -
27 - # Compute start index and sequence length from cumulative lengths
28 - in_start = 0
29 - for i in range(pid_b):
30 - in_start += tl.load(lengths_ptr + i)
31 - seq_len = tl.load(lengths_ptr + pid_b)
32 -
33 - # valid time positions for this block
34 - t_mask = off_t < Lmax
35 -
36 - # compute input row indices for valid (b, t)
37 - in_row = in_start + off_t
38 - valid_row = (off_t < seq_len) & t_mask
39 -
40 - # Pointers
41 - # x_ptr: row-major [N, D]
42 - x_row_ptr = x_ptr + in_row[:, None] * D + off_d[None, :]
43 -
44 - # out_ptr: row-major [B, Lmax, D]
45 - out_row_ptr = out_ptr + (pid_b * Lmax + off_t)[:, None] * D + off_d[None, :]
46 -
47 - # Initialize with PAD. PAD_IS_UINT8 selects the pad tensor's dtype so
48 - # integer-typed outputs (e.g. MXFP4 packed nibbles, ue8m0 scale bytes)
49 - # get an exact-byte pad rather than going through an fp32→uint8 cast
50 - # that's implementation-defined outside of value 0.
51 - d_mask = off_d[None, :] < D
52 - if PAD_IS_UINT8:
53 - pad_vals = tl.full([BLOCK_T, BLOCK_D], PAD_VALUE, tl.uint8)
54 - else:
55 - pad_vals = tl.full([BLOCK_T, BLOCK_D], PAD_VALUE, tl.float32)
56 - tl.store(out_row_ptr, pad_vals, mask=t_mask[:, None] & d_mask)
17 +class PackSeqTritonKernel(VllmTritonJitKernel["PackSeqTritonKernel.CompileKey"]):
18 + @dataclass(frozen=True)
19 + class CompileKey:
20 + dtype: torch.dtype
21 + pad_value: float | int
22 + pad_is_uint8: bool
23 + block_t: int
24 + block_d: int
25 +
26 + # N is unused; D and Lmax only affect runtime addressing and masks. The
27 + # compile-time tile shape remains controlled by BLOCK_T and BLOCK_D.
28 + @staticmethod
29 + @triton.jit(do_not_specialize=["N", "D", "Lmax"])
30 + def kernel(
31 + x_ptr, # [N, D]
32 + out_ptr, # [B, Lmax, D]
33 + lengths_ptr, # *i32, [B]
34 + N,
35 + D,
36 + Lmax,
37 + PAD_VALUE: tl.constexpr,
38 + PAD_IS_UINT8: tl.constexpr,
39 + BLOCK_T: tl.constexpr, # timesteps per program
40 + BLOCK_D: tl.constexpr, # features per program
41 + ):
42 + pid_b = tl.program_id(0) # batch id
43 + pid_t = tl.program_id(1) # block over time dimension
44 + pid_d = tl.program_id(2) # block over feature dimension
45 + off_t = pid_t * BLOCK_T + tl.arange(0, BLOCK_T) # [BLOCK_T]
46 + off_d = pid_d * BLOCK_D + tl.arange(0, BLOCK_D) # [BLOCK_D]
47 +
48 + # Compute start index and sequence length from cumulative lengths
49 + in_start = 0
50 + for i in range(pid_b):
51 + in_start += tl.load(lengths_ptr + i)
52 + seq_len = tl.load(lengths_ptr + pid_b)
53 +
54 + # valid time positions for this block
55 + t_mask = off_t < Lmax
56 +
57 + # compute input row indices for valid (b, t)
58 + in_row = in_start + off_t
59 + valid_row = (off_t < seq_len) & t_mask
60 +
61 + # Pointers
62 + # x_ptr: row-major [N, D]
63 + x_row_ptr = x_ptr + in_row[:, None] * D + off_d[None, :]
64 +
65 + # out_ptr: row-major [B, Lmax, D]
66 + out_row_ptr = out_ptr + (pid_b * Lmax + off_t)[:, None] * D + off_d[None, :]
67 +
68 + # Initialize with PAD. PAD_IS_UINT8 selects the pad tensor's dtype so
69 + # integer-typed outputs (e.g. MXFP4 packed nibbles, ue8m0 scale bytes)
70 + # get an exact-byte pad rather than going through an fp32→uint8 cast
71 + # that's implementation-defined outside of value 0.
72 + d_mask = off_d[None, :] < D
73 + if PAD_IS_UINT8:
74 + pad_vals = tl.full([BLOCK_T, BLOCK_D], PAD_VALUE, tl.uint8)
75 + else:
76 + pad_vals = tl.full([BLOCK_T, BLOCK_D], PAD_VALUE, tl.float32)
77 + tl.store(out_row_ptr, pad_vals, mask=t_mask[:, None] & d_mask)
78 +
79 + # Load & write only where within seq_len
80 + x_vals = tl.load(x_row_ptr, mask=valid_row[:, None] & d_mask)
81 + tl.store(out_row_ptr, x_vals, mask=valid_row[:, None] & d_mask)
82 +
83 + def dispatch( # type: ignore[override]
84 + self,
85 + *,
86 + dtype: torch.dtype,
87 + pad_value: float | int,
88 + **compile_key_fields: int,
89 + ) -> CompileKey:
90 + is_uint8 = dtype == torch.uint8
91 + return self.CompileKey(
92 + **compile_key_fields,
93 + dtype=dtype,
94 + pad_value=int(pad_value) if is_uint8 else float(pad_value),
95 + pad_is_uint8=is_uint8,
96 + )
97
58 - # Load & write only where within seq_len
59 - x_vals = tl.load(x_row_ptr, mask=valid_row[:, None] & d_mask)
60 - tl.store(out_row_ptr, x_vals, mask=valid_row[:, None] & d_mask)
98 + def get_warmup_keys(
99 + self,
100 + *,
101 + dtype: torch.dtype,
102 + pad_value: float | int,
103 + ) -> list[CompileKey]:
104 + return self._trace_dispatch(self.dispatch)(
105 + dtype=dtype,
106 + pad_value=pad_value,
107 + block_t=64,
108 + block_d=64,
109 + )
110 +
111 + def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
112 + return dict(
113 + x_reshaped=TritonWarmupTensor(compile_key.dtype),
114 + out=TritonWarmupTensor(compile_key.dtype),
115 + lengths=TritonWarmupTensor(torch.int32),
116 + N=1,
117 + D=1,
118 + Lmax=1,
119 + pad_value=compile_key.pad_value,
120 + pad_is_uint8=compile_key.pad_is_uint8,
121 + block_t=compile_key.block_t,
122 + block_d=compile_key.block_d,
123 + )
124 +
125 + @kernel_launcher
126 + def __call__(
127 + self,
128 + x_reshaped: torch.Tensor,
129 + out: torch.Tensor,
130 + lengths: torch.Tensor,
131 + *,
132 + N: int,
133 + D: int,
134 + Lmax: int,
135 + pad_value: float | int,
136 + pad_is_uint8: bool,
137 + block_t: int,
138 + block_d: int,
139 + ) -> LaunchSpec:
140 + grid = (
141 + lengths.shape[0],
142 + triton.cdiv(Lmax, block_t),
143 + triton.cdiv(D, block_d),
144 + )
145 + return grid, dict(
146 + x_ptr=x_reshaped,
147 + lengths_ptr=lengths,
148 + PAD_VALUE=pad_value,
149 + PAD_IS_UINT8=pad_is_uint8,
150 + BLOCK_T=block_t,
151 + BLOCK_D=block_d,
152 + num_warps=4,
153 + num_stages=2,
154 + )
155
156
157 def pack_seq_triton(
@@ -109,20 +203,18 @@ def pack_seq_triton(
203
204 out = torch.empty((B, Lmax, D), device=x.device, dtype=x.dtype)
205
112 - grid = (B, triton.cdiv(Lmax, block_t), triton.cdiv(D, block_d))
113 - _pack_seq_kernel[grid](
206 + lengths = lengths.int()
207 + _PACK_SEQ_TRITON_KERNEL(
208 x_reshaped,
209 out,
116 - lengths.int(),
117 - N,
118 - D,
119 - Lmax,
120 - PAD_VALUE=pad_constexpr,
121 - PAD_IS_UINT8=is_uint8,
122 - BLOCK_T=block_t,
123 - BLOCK_D=block_d,
124 - num_warps=4,
125 - num_stages=2,
210 + lengths,
211 + N=N,
212 + D=D,
213 + Lmax=Lmax,
214 + pad_value=pad_constexpr,
215 + pad_is_uint8=is_uint8,
216 + block_t=block_t,
217 + block_d=block_d,
218 )
219
220 if len(original_shape) > 2:
@@ -131,47 +223,109 @@ def pack_seq_triton(
223 return out
224
225
134 -@triton.jit
135 -def _unpack_seq_triton_kernel(
136 - packed_ptr, # [B, Lmax, D]
137 - out_ptr, # [N, D]
138 - lengths_ptr, # *i32, [B]
139 - B: tl.constexpr,
140 - Lmax: tl.constexpr,
141 - D: tl.constexpr,
142 - BLOCK_T: tl.constexpr, # timesteps per program
143 - BLOCK_D: tl.constexpr, # features per program
144 -):
145 - pid_b = tl.program_id(0) # batch id
146 - pid_t = tl.program_id(1) # block over time dimension
147 - pid_d = tl.program_id(2) # block over feature dimension
148 - off_t = pid_t * BLOCK_T + tl.arange(0, BLOCK_T) # [BLOCK_T]
149 - off_d = pid_d * BLOCK_D + tl.arange(0, BLOCK_D) # [BLOCK_D]
150 -
151 - # bounds: compute start from cumulative lengths
152 - in_start = 0
153 - for i in range(pid_b):
154 - in_start += tl.load(lengths_ptr + i)
155 - seq_len = tl.load(lengths_ptr + pid_b)
156 -
157 - # valid time positions for this block
158 - t_mask = off_t < Lmax
159 - valid_row = (off_t < seq_len) & t_mask
160 -
161 - # compute output row indices for valid (b, t)
162 - out_row = in_start + off_t
163 -
164 - # Pointers
165 - # packed_ptr: row-major [B, Lmax, D]
166 - packed_row_ptr = packed_ptr + (pid_b * Lmax + off_t)[:, None] * D + off_d[None, :]
167 -
168 - # out_ptr: row-major [N, D]
169 - out_row_ptr = out_ptr + out_row[:, None] * D + off_d[None, :]
170 -
171 - # Load from packed tensor and store to output
172 - d_mask = off_d[None, :] < D
173 - packed_vals = tl.load(packed_row_ptr, mask=valid_row[:, None] & d_mask)
174 - tl.store(out_row_ptr, packed_vals, mask=valid_row[:, None] & d_mask)
226 +class UnpackSeqTritonKernel(VllmTritonJitKernel["UnpackSeqTritonKernel.CompileKey"]):
227 + @dataclass(frozen=True)
228 + class CompileKey:
229 + dtype: torch.dtype
230 + block_t: int
231 + block_d: int
232 +
233 + # B, Lmax, and D only affect runtime addressing and masks. The compile-time
234 + # tile shape remains controlled by BLOCK_T and BLOCK_D.
235 + @staticmethod
236 + @triton.jit(do_not_specialize=["B", "Lmax", "D"])
237 + def kernel(
238 + packed_ptr, # [B, Lmax, D]
239 + out_ptr, # [N, D]
240 + lengths_ptr, # *i32, [B]
241 + B,
242 + Lmax,
243 + D,
244 + BLOCK_T: tl.constexpr, # timesteps per program
245 + BLOCK_D: tl.constexpr, # features per program
246 + ):
247 + pid_b = tl.program_id(0) # batch id
248 + pid_t = tl.program_id(1) # block over time dimension
249 + pid_d = tl.program_id(2) # block over feature dimension
250 + off_t = pid_t * BLOCK_T + tl.arange(0, BLOCK_T) # [BLOCK_T]
251 + off_d = pid_d * BLOCK_D + tl.arange(0, BLOCK_D) # [BLOCK_D]
252 +
253 + # bounds: compute start from cumulative lengths
254 + in_start = 0
255 + for i in range(pid_b):
256 + in_start += tl.load(lengths_ptr + i)
257 + seq_len = tl.load(lengths_ptr + pid_b)
258 +
259 + # valid time positions for this block
260 + t_mask = off_t < Lmax
261 + valid_row = (off_t < seq_len) & t_mask
262 +
263 + # compute output row indices for valid (b, t)
264 + out_row = in_start + off_t
265 +
266 + # Pointers
267 + # packed_ptr: row-major [B, Lmax, D]
268 + packed_row_ptr = (
269 + packed_ptr + (pid_b * Lmax + off_t)[:, None] * D + off_d[None, :]
270 + )
271 +
272 + # out_ptr: row-major [N, D]
273 + out_row_ptr = out_ptr + out_row[:, None] * D + off_d[None, :]
274 +
275 + # Load from packed tensor and store to output
276 + d_mask = off_d[None, :] < D
277 + packed_vals = tl.load(packed_row_ptr, mask=valid_row[:, None] & d_mask)
278 + tl.store(out_row_ptr, packed_vals, mask=valid_row[:, None] & d_mask)
279 +
280 + def dispatch( # type: ignore[override]
281 + self,
282 + *,
283 + dtype: torch.dtype,
284 + **compile_key_fields: int,
285 + ) -> CompileKey:
286 + return self.CompileKey(**compile_key_fields, dtype=dtype)
287 +
288 + def get_warmup_keys(self) -> list[CompileKey]:
289 + return self._trace_dispatch(self.dispatch)(
290 + dtype=torch.int32,
291 + block_t=64,
292 + block_d=64,
293 + )
294 +
295 + def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
296 + return dict(
297 + packed_reshaped=TritonWarmupTensor(compile_key.dtype),
298 + out=TritonWarmupTensor(compile_key.dtype),
299 + lengths=TritonWarmupTensor(torch.int32),
300 + B=1,
301 + Lmax=1,
302 + D=1,
303 + block_t=compile_key.block_t,
304 + block_d=compile_key.block_d,
305 + )
306 +
307 + @kernel_launcher
308 + def __call__(
309 + self,
310 + packed_reshaped: torch.Tensor,
311 + out: torch.Tensor,
312 + lengths: torch.Tensor,
313 + *,
314 + B: int,
315 + Lmax: int,
316 + D: int,
317 + block_t: int,
318 + block_d: int,
319 + ) -> LaunchSpec:
320 + grid = (B, triton.cdiv(Lmax, block_t), triton.cdiv(D, block_d))
321 + return grid, dict(
322 + packed_ptr=packed_reshaped,
323 + lengths_ptr=lengths,
324 + BLOCK_T=block_t,
325 + BLOCK_D=block_d,
326 + num_warps=4,
327 + num_stages=2,
328 + )
329
330
331 def unpack_seq_triton(
@@ -209,18 +363,16 @@ def unpack_seq_triton(
363
364 out = torch.empty((N, D), device=packed_tensor.device, dtype=packed_tensor.dtype)
365
212 - grid = (B, triton.cdiv(Lmax, block_t), triton.cdiv(D, block_d))
213 - _unpack_seq_triton_kernel[grid](
366 + lengths = lengths.int()
367 + _UNPACK_SEQ_TRITON_KERNEL(
368 packed_reshaped,
369 out,
216 - lengths.int(),
217 - B,
218 - Lmax,
219 - D,
220 - BLOCK_T=block_t,
221 - BLOCK_D=block_d,
222 - num_warps=4,
223 - num_stages=2,
370 + lengths,
371 + B=B,
372 + Lmax=Lmax,
373 + D=D,
374 + block_t=block_t,
375 + block_d=block_d,
376 )
377
378 # Reshape output back to original dimensions (except first dimension)
@@ -229,3 +381,7 @@ def unpack_seq_triton(
381 out = out.reshape(output_shape)
382
383 return out
384 +
385 +
386 +_PACK_SEQ_TRITON_KERNEL = PackSeqTritonKernel()
387 +_UNPACK_SEQ_TRITON_KERNEL = UnpackSeqTritonKernel()
vllm/v1/attention/ops/dcp.py
+271 -99
@@ -6,7 +6,8 @@ from __future__ import annotations
6
7 import functools
8 from collections.abc import Callable
9 -from typing import TYPE_CHECKING, Protocol
9 +from dataclasses import dataclass
10 +from typing import TYPE_CHECKING, Any, Protocol
11
12 import torch
13 import torch.distributed as dist
@@ -15,6 +16,16 @@ import vllm.envs as envs
16 from vllm.config import VllmConfig
17 from vllm.distributed import get_dcp_group
18 from vllm.logger import init_logger
19 +from vllm.model_executor.warmup.jit_warmup import (
20 + WarmupIntRange,
21 +)
22 +from vllm.model_executor.warmup.jit_warmup_triton_helper import (
23 + LaunchSpec,
24 + TritonWarmupTensor,
25 + VllmTritonJitKernel,
26 + kernel_launcher,
27 + triton_scalar_specialization_rep,
28 +)
29 from vllm.triton_utils import tl, triton
30 from vllm.v1.attention.ops.cp_common import (
31 DirectCPWorkspace,
@@ -65,93 +76,254 @@ def mask_dcp_empty_shards_(
76 # AG + RS/AR implementation
77
78
68 -@triton.jit
69 -def _correct_attn_cp_out_kernel(
70 - outputs_ptr,
71 - new_output_ptr,
72 - lses_ptr,
73 - vlse_ptr,
74 - outputs_stride_B,
75 - outputs_stride_H,
76 - outputs_stride_D,
77 - lses_stride_N,
78 - lses_stride_B,
79 - lses_stride_H,
80 - lse_idx,
81 - HEAD_DIM: tl.constexpr,
82 - N_ROUNDED: tl.constexpr,
83 - IS_BASE_E: tl.constexpr,
84 -):
85 - """
86 - Apply the all-gathered lses to correct each local rank's attention
87 - output. we still need perform a cross-rank reduction to obtain the
88 - final attention output.
79 +class CorrectAttnCPOutKernel(VllmTritonJitKernel["CorrectAttnCPOutKernel.CompileKey"]):
80 + @dataclass(frozen=True)
81 + class CompileKey:
82 + output_dtype: torch.dtype
83 + lse_dtype: torch.dtype
84 + outputs_stride_b: int
85 + outputs_stride_h: int
86 + outputs_stride_d: int
87 + lses_stride_n: int
88 + lses_stride_b: int
89 + lses_stride_h: int
90 + lse_idx: int
91 + head_dim: int
92 + n_rounded: int
93 + is_base_e: bool
94 +
95 + @staticmethod
96 + @triton.jit
97 + def kernel(
98 + outputs_ptr,
99 + new_output_ptr,
100 + lses_ptr,
101 + vlse_ptr,
102 + outputs_stride_B,
103 + outputs_stride_H,
104 + outputs_stride_D,
105 + lses_stride_N,
106 + lses_stride_B,
107 + lses_stride_H,
108 + lse_idx,
109 + HEAD_DIM: tl.constexpr,
110 + N_ROUNDED: tl.constexpr,
111 + IS_BASE_E: tl.constexpr,
112 + ):
113 + """
114 + Apply the all-gathered lses to correct each local rank's attention
115 + output. we still need perform a cross-rank reduction to obtain the
116 + final attention output.
117 +
118 + Args:
119 + outputs_ptr (triton.PointerType):
120 + Pointer to input tensor of shape [ B, H, D ]
121 + lses_ptr (triton.PointerType):
122 + Pointer to input tensor of shape [ N, B, H ]
123 + new_output_ptr (triton.PointerType):
124 + Pointer to output tensor of shape [ B, H, D ]
125 + vlse_ptr (triton.PointerType):
126 + Pointer to output tensor of shape [ B, H ]
127 + """
128 + batch_idx = tl.program_id(axis=0).to(tl.int64)
129 + head_idx = tl.program_id(axis=1).to(tl.int64)
130 + d_offsets = tl.arange(0, HEAD_DIM)
131 + num_n_offsets = tl.arange(0, N_ROUNDED)
132 +
133 + # shape = [N]
134 + lse_offsets = (
135 + num_n_offsets * lses_stride_N
136 + + batch_idx * lses_stride_B
137 + + head_idx * lses_stride_H
138 + )
139
90 - Args:
91 - outputs_ptr (triton.PointerType):
92 - Pointer to input tensor of shape [ B, H, D ]
93 - lses_ptr (triton.PointerType):
94 - Pointer to input tensor of shape [ N, B, H ]
95 - new_output_ptr (triton.PointerType):
96 - Pointer to output tensor of shape [ B, H, D ]
97 - vlse_ptr (triton.PointerType):
98 - Pointer to output tensor of shape [ B, H ]
99 - """
100 - batch_idx = tl.program_id(axis=0).to(tl.int64)
101 - head_idx = tl.program_id(axis=1).to(tl.int64)
102 - d_offsets = tl.arange(0, HEAD_DIM)
103 - num_n_offsets = tl.arange(0, N_ROUNDED)
140 + # calc final lse
141 + lse = tl.load(lses_ptr + lse_offsets).to(tl.float32)
142 + lse = tl.where((lse != lse) | (lse == float("inf")), -float("inf"), lse)
143 + lse_max = tl.max(lse, axis=0)
144 + lse_max = tl.where(lse_max == -float("inf"), 0, lse_max)
145 + lse -= lse_max
146 + if IS_BASE_E:
147 + lse_exp = tl.exp(lse)
148 + lse_acc = tl.sum(lse_exp, axis=0)
149 + lse = tl.log(lse_acc)
150 + else:
151 + lse_exp = tl.exp2(lse)
152 + lse_acc = tl.sum(lse_exp, axis=0)
153 + lse = tl.log2(lse_acc)
154 + lse += lse_max
155 +
156 + lse_offsets = batch_idx * lses_stride_B + head_idx * lses_stride_H
157 + tl.store(vlse_ptr + lse_offsets, lse)
158 +
159 + # shape = [D]
160 + output_offsets = (
161 + batch_idx * outputs_stride_B
162 + + head_idx * outputs_stride_H
163 + + d_offsets * outputs_stride_D
164 + )
165
105 - # shape = [N]
106 - lse_offsets = (
107 - num_n_offsets * lses_stride_N
108 - + batch_idx * lses_stride_B
109 - + head_idx * lses_stride_H
110 - )
166 + # correct output
167 + lse_offset = (
168 + lse_idx * lses_stride_N
169 + + batch_idx * lses_stride_B
170 + + head_idx * lses_stride_H
171 + )
172 + lse_tmp = tl.load(lses_ptr + lse_offset).to(tl.float32)
173 + lse_finally = lse_tmp - lse
174 + lse_finally = tl.where(
175 + (lse_finally != lse_finally) | (lse_finally == float("inf")),
176 + -float("inf"),
177 + lse_finally,
178 + )
179 + factor = tl.exp(lse_finally) if IS_BASE_E else tl.exp2(lse_finally)
180 + output = tl.load(outputs_ptr + output_offsets)
181 + output = output * factor
182 + output = tl.where(factor == 0.0, 0.0, output)
183
112 - # calc final lse
113 - lse = tl.load(lses_ptr + lse_offsets).to(tl.float32)
114 - lse = tl.where((lse != lse) | (lse == float("inf")), -float("inf"), lse)
115 - lse_max = tl.max(lse, axis=0)
116 - lse_max = tl.where(lse_max == -float("inf"), 0, lse_max)
117 - lse -= lse_max
118 - if IS_BASE_E:
119 - lse_exp = tl.exp(lse)
120 - lse_acc = tl.sum(lse_exp, axis=0)
121 - lse = tl.log(lse_acc)
122 - else:
123 - lse_exp = tl.exp2(lse)
124 - lse_acc = tl.sum(lse_exp, axis=0)
125 - lse = tl.log2(lse_acc)
126 - lse += lse_max
127 -
128 - lse_offsets = batch_idx * lses_stride_B + head_idx * lses_stride_H
129 - tl.store(vlse_ptr + lse_offsets, lse)
130 -
131 - # shape = [D]
132 - output_offsets = (
133 - batch_idx * outputs_stride_B
134 - + head_idx * outputs_stride_H
135 - + d_offsets * outputs_stride_D
136 - )
184 + tl.store(new_output_ptr + output_offsets, output)
185
138 - # correct output
139 - lse_offset = (
140 - lse_idx * lses_stride_N + batch_idx * lses_stride_B + head_idx * lses_stride_H
141 - )
142 - lse_tmp = tl.load(lses_ptr + lse_offset).to(tl.float32)
143 - lse_finally = lse_tmp - lse
144 - lse_finally = tl.where(
145 - (lse_finally != lse_finally) | (lse_finally == float("inf")),
146 - -float("inf"),
147 - lse_finally,
148 - )
149 - factor = tl.exp(lse_finally) if IS_BASE_E else tl.exp2(lse_finally)
150 - output = tl.load(outputs_ptr + output_offsets)
151 - output = output * factor
152 - output = tl.where(factor == 0.0, 0.0, output)
186 + def dispatch( # type: ignore[override]
187 + self,
188 + *,
189 + output_dtype: torch.dtype,
190 + lse_dtype: torch.dtype,
191 + num_tokens: int,
192 + num_heads: int,
193 + head_dim: int,
194 + n_rounded: int,
195 + lse_idx: int,
196 + is_base_e: bool,
197 + ) -> CompileKey:
198 + # Warm the contiguous [B, H, D] output and [N, B, H] LSE layouts used
199 + # by the production collective path. Runtime still forwards real
200 + # strides so unsupported layouts remain correct, but may compile on
201 + # first use under the JIT monitor.
202 + outputs_stride_d = 1
203 + outputs_stride_h = head_dim
204 + outputs_stride_b = num_heads * head_dim
205 + lses_stride_h = 1
206 + lses_stride_b = num_heads
207 + lses_stride_n = num_tokens * num_heads
208 + return self.CompileKey(
209 + output_dtype=output_dtype,
210 + lse_dtype=lse_dtype,
211 + outputs_stride_b=triton_scalar_specialization_rep(outputs_stride_b),
212 + outputs_stride_h=triton_scalar_specialization_rep(outputs_stride_h),
213 + outputs_stride_d=triton_scalar_specialization_rep(outputs_stride_d),
214 + lses_stride_n=triton_scalar_specialization_rep(lses_stride_n),
215 + lses_stride_b=triton_scalar_specialization_rep(lses_stride_b),
216 + lses_stride_h=triton_scalar_specialization_rep(lses_stride_h),
217 + lse_idx=triton_scalar_specialization_rep(lse_idx),
218 + head_dim=head_dim,
219 + n_rounded=n_rounded,
220 + is_base_e=is_base_e,
221 + )
222 +
223 + def get_warmup_keys(
224 + self,
225 + vllm_config: Any,
226 + *,
227 + output_dtype: torch.dtype | None = None,
228 + num_heads: int | None = None,
229 + head_dim: int | None = None,
230 + is_base_e: bool | tuple[bool, ...] = (False, True),
231 + ) -> list[CompileKey]:
232 + from vllm.model_executor.layers.attention.mla_attention import (
233 + get_mla_dims,
234 + )
235 +
236 + dcp_world_size = vllm_config.parallel_config.decode_context_parallel_size
237 + max_tokens = min(
238 + 16,
239 + vllm_config.scheduler_config.max_num_batched_tokens,
240 + )
241 + hf_config = vllm_config.model_config.hf_config
242 + if num_heads is None:
243 + num_heads = (
244 + hf_config.num_attention_heads
245 + // vllm_config.parallel_config.tensor_parallel_size
246 + )
247 + if head_dim is None:
248 + head_dim = get_mla_dims(vllm_config.model_config).v_head_dim
249 + if output_dtype is None:
250 + output_dtype = vllm_config.model_config.dtype
251 + if dcp_world_size <= 1 or max_tokens <= 0 or num_heads <= 0:
252 + return []
253 + return self._trace_dispatch(self.dispatch)(
254 + output_dtype=output_dtype,
255 + lse_dtype=torch.float32,
256 + num_tokens=WarmupIntRange(1, max_tokens + 1),
257 + num_heads=num_heads,
258 + head_dim=head_dim,
259 + n_rounded=dcp_world_size,
260 + lse_idx=WarmupIntRange(0, dcp_world_size),
261 + is_base_e=is_base_e,
262 + )
263 +
264 + def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
265 + output_ptr = TritonWarmupTensor(
266 + compile_key.output_dtype,
267 + shape=(1, 1, compile_key.head_dim),
268 + strides=(
269 + compile_key.outputs_stride_b,
270 + compile_key.outputs_stride_h,
271 + compile_key.outputs_stride_d,
272 + ),
273 + )
274 + lse_ptr = TritonWarmupTensor(
275 + compile_key.lse_dtype,
276 + shape=(compile_key.n_rounded, 1, 1),
277 + strides=(
278 + compile_key.lses_stride_n,
279 + compile_key.lses_stride_b,
280 + compile_key.lses_stride_h,
281 + ),
282 + )
283 + return dict(
284 + outputs=output_ptr,
285 + new_output=output_ptr,
286 + lses=lse_ptr,
287 + vlse=lse_ptr,
288 + lse_idx=compile_key.lse_idx,
289 + ctx=None,
290 + is_base_e=compile_key.is_base_e,
291 + )
292
154 - tl.store(new_output_ptr + output_offsets, output)
293 + @kernel_launcher
294 + def __call__(
295 + self,
296 + outputs: torch.Tensor,
297 + new_output: torch.Tensor,
298 + lses: torch.Tensor,
299 + vlse: torch.Tensor,
300 + lse_idx: int,
301 + ctx: Any,
302 + *,
303 + is_base_e: bool,
304 + ) -> LaunchSpec:
305 + num_tokens, num_heads, head_dim = outputs.shape
306 + n_rounded = lses.shape[0]
307 + outputs_stride_b, outputs_stride_h, outputs_stride_d = outputs.stride()
308 + lses_stride_n, lses_stride_b, lses_stride_h = lses.stride()
309 + grid = (num_tokens, num_heads, 1)
310 + return grid, dict(
311 + outputs_stride_B=outputs_stride_b,
312 + outputs_stride_H=outputs_stride_h,
313 + outputs_stride_D=outputs_stride_d,
314 + lses_stride_N=lses_stride_n,
315 + lses_stride_B=lses_stride_b,
316 + lses_stride_H=lses_stride_h,
317 + HEAD_DIM=head_dim,
318 + N_ROUNDED=n_rounded,
319 + IS_BASE_E=is_base_e,
320 + _runtime_launcher=None if self._warming else ctx.call_kernel,
321 + # CPTritonContext caches the non-constexpr positional prefix; derive
322 + # its length so adding/reordering a kernel arg cannot silently
323 + # misalign the cached replay path.
324 + _runtime_launcher_arg_count=len(self.kernel.arg_names)
325 + - len(self.kernel.constexprs),
326 + )
327
328
329 class CPTritonContext:
@@ -202,14 +374,12 @@ def correct_attn_out(
374 f"got {tuple(lses.shape)}"
375 )
376
205 - B, H, D = out.shape
206 - N = lses.shape[0]
377 + B, H, _ = out.shape
378
379 # Strides after we normalized shapes to 3-D views. The kernel computes
380 # offsets for `vlse_ptr` using lses_stride_B/H, so the output buffer must
381 # have the same B/H stride layout as a slice of `lses`.
211 - o_sB, o_sH, o_sD = out.stride()
212 - l_sN, l_sB, l_sH = lses.stride()
382 + _, l_sB, l_sH = lses.stride()
383
384 # Allocate LSE with the same B/H strides as `lses` so writes land correctly
385 # even when `lses` is a non-contiguous view (e.g., 4-D to 3-D squeeze).
@@ -217,24 +387,15 @@ def correct_attn_out(
387 (B, H), (l_sB, l_sH), device=lses.device, dtype=lses.dtype
388 )
389
220 - # Kernel launch config
221 - grid = (B, H, 1)
222 -
223 - regular_args = (
390 + _CORRECT_ATTN_CP_OUT_KERNEL(
391 out,
392 out,
393 lses,
394 lse,
228 - o_sB,
229 - o_sH,
230 - o_sD,
231 - l_sN,
232 - l_sB,
233 - l_sH,
395 cp_rank,
396 + ctx,
397 + is_base_e=is_lse_base_on_e,
398 )
236 - const_args = {"HEAD_DIM": D, "N_ROUNDED": N, "IS_BASE_E": is_lse_base_on_e}
237 - ctx.call_kernel(_correct_attn_cp_out_kernel, grid, *regular_args, **const_args)
399 return out, lse
400
401
@@ -1250,6 +1411,14 @@ class MLADCPManager:
1411 is_lse_base_on_e,
1412 use_pcp,
1413 )
1414 + if vllm_config.kernel_config.enable_jit_warmup and not self.use_a2a:
1415 + _CORRECT_ATTN_CP_OUT_KERNEL.register_warmup(
1416 + vllm_config,
1417 + output_dtype=output_dtype,
1418 + num_heads=num_heads,
1419 + head_dim=output_head_dim,
1420 + is_base_e=is_lse_base_on_e,
1421 + )
1422 self.query_gather = (
1423 None
1424 if use_pcp
@@ -1411,3 +1580,6 @@ class MLADCPManager:
1580 local_kv: torch.Tensor,
1581 ) -> object:
1582 return self._kv_gather(gathered_kv, local_kv)
1583 +
1584 +
1585 +_CORRECT_ATTN_CP_OUT_KERNEL = CorrectAttnCPOutKernel()
vllm/v1/worker/block_table.py
+46 -28
@@ -11,11 +11,11 @@ import torch
11
12 from vllm.distributed import get_dcp_group, get_pcp_group
13 from vllm.logger import init_logger
14 -from vllm.model_executor.warmup.jit_warmup import (
15 - VllmJitKernel,
16 -)
14 from vllm.model_executor.warmup.jit_warmup_triton_helper import (
15 + LaunchSpec,
16 TritonWarmupTensor,
17 + VllmTritonJitKernel,
18 + kernel_launcher,
19 triton_scalar_specialization_rep,
20 )
21 from vllm.triton_utils import tl, triton
@@ -394,7 +394,9 @@ class MultiGroupBlockTable:
394 return self.block_tables[idx]
395
396
397 -class ComputeSlotMappingKernel(VllmJitKernel["ComputeSlotMappingKernel.CompileKey"]):
397 +class ComputeSlotMappingKernel(
398 + VllmTritonJitKernel["ComputeSlotMappingKernel.CompileKey"]
399 +):
400 triton_block_size = 1024
401
402 @dataclass(frozen=True)
@@ -490,37 +492,53 @@ class ComputeSlotMappingKernel(VllmJitKernel["ComputeSlotMappingKernel.CompileKe
492 def get_warmup_keys(self, **dispatch_kwargs: int) -> list[CompileKey]:
493 return self._trace_dispatch(self.dispatch)(**dispatch_kwargs)
494
493 - def compile(self, compile_key: CompileKey) -> None:
494 - warmup = getattr(self.kernel, "warmup", None)
495 - assert warmup is not None
495 + def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
496 int32_ptr = TritonWarmupTensor(torch.int32)
497 int64_ptr = TritonWarmupTensor(torch.int64)
498 - warmup(
499 - 2, # arbitrary, num_tokens in do_not_specialize
500 - 2, # arbitrary, max_num_tokens in do_not_specialize
501 - int32_ptr,
502 - int64_ptr,
503 - int32_ptr,
504 - compile_key.block_table_stride,
505 - compile_key.block_size,
506 - int64_ptr,
507 - KV_CACHE_BLOCK_SIZE=compile_key.kv_cache_block_size,
508 - BLOCKS_PER_KV_BLOCK=compile_key.blocks_per_kv_block,
509 - TOTAL_CP_WORLD_SIZE=compile_key.total_cp_world_size,
510 - TOTAL_CP_RANK=compile_key.total_cp_rank,
511 - CP_KV_CACHE_INTERLEAVE_SIZE=compile_key.cp_kv_cache_interleave_size,
512 - PAD_ID=PAD_SLOT_ID,
513 - BLOCK_SIZE=self.triton_block_size,
514 - grid=(2,),
498 + return dict(
499 + num_reqs=1,
500 + num_tokens=2, # arbitrary, in do_not_specialize
501 + max_num_tokens=2, # arbitrary, in do_not_specialize
502 + query_start_loc=int32_ptr,
503 + positions=int64_ptr,
504 + block_table=TritonWarmupTensor(
505 + torch.int32,
506 + shape=(1, compile_key.block_table_stride),
507 + ),
508 + block_table_stride=compile_key.block_table_stride,
509 + block_size=compile_key.block_size,
510 + slot_mapping=int64_ptr,
511 + kv_cache_block_size=compile_key.kv_cache_block_size,
512 + blocks_per_kv_block=compile_key.blocks_per_kv_block,
513 + total_cp_world_size=compile_key.total_cp_world_size,
514 + total_cp_rank=compile_key.total_cp_rank,
515 + cp_kv_cache_interleave_size=compile_key.cp_kv_cache_interleave_size,
516 )
517
518 + @kernel_launcher
519 def __call__(
520 self,
521 num_reqs: int,
520 - *args: Any,
521 - ) -> None:
522 - self.kernel[(num_reqs + 1,)](
523 - *args,
522 + num_tokens: int,
523 + max_num_tokens: int,
524 + query_start_loc: torch.Tensor,
525 + positions: torch.Tensor,
526 + block_table: torch.Tensor,
527 + block_table_stride: int,
528 + block_size: int,
529 + slot_mapping: torch.Tensor,
530 + kv_cache_block_size: int,
531 + blocks_per_kv_block: int,
532 + total_cp_world_size: int,
533 + total_cp_rank: int,
534 + cp_kv_cache_interleave_size: int,
535 + ) -> LaunchSpec:
536 + return (num_reqs + 1,), dict(
537 + KV_CACHE_BLOCK_SIZE=kv_cache_block_size,
538 + BLOCKS_PER_KV_BLOCK=blocks_per_kv_block,
539 + TOTAL_CP_WORLD_SIZE=total_cp_world_size,
540 + TOTAL_CP_RANK=total_cp_rank,
541 + CP_KV_CACHE_INTERLEAVE_SIZE=cp_kv_cache_interleave_size,
542 PAD_ID=PAD_SLOT_ID,
543 BLOCK_SIZE=self.triton_block_size,
544 )