Add Triton warmup compile-key tracing

Co-authored-by: Codex <codex@openai.com>

Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

tests/test_triton_utils.py

@@ -92,3 +92,226 @@ def test_no_triton_fallback():
         assert triton.__class__.__name__ == "TritonPlaceholder"
         assert triton.language.__class__.__name__ == "TritonLanguagePlaceholder"
         assert tl.__class__.__name__ == "TritonLanguagePlaceholder"
+
+
+def test_prefill_chunk_metadata_compile_keys_trace_specializations():
+    for module_name in (
+        "triton",
+        "triton.language",
+        "vllm.triton_utils",
+        "vllm.triton_utils.importing",
+        "vllm.triton_utils.warmup",
+        "vllm.v1.attention.backends.mla.indexer",
+        "vllm.v1.attention.backends.mla.sparse_swa",
+    ):
+        sys.modules.pop(module_name, None)
+
+    with mock.patch.dict(sys.modules, {"triton": None}):
+        from vllm.triton_utils.warmup import (
+            WarmupIntRange,
+            trace_triton_kernel_specialization_args,
+        )
+        from vllm.v1.attention.backends.mla.indexer import (
+            BUILD_PREFILL_CHUNK_METADATA_KERNEL,
+            PrefillChunkMetadataKernelCompileKey,
+            _build_prefill_chunk_metadata_kernel,
+        )
+
+        assert trace_triton_kernel_specialization_args(
+            _build_prefill_chunk_metadata_kernel
+        ) == (
+            "query_slice_start",
+            "query_slice_stop",
+            "BLOCK_SIZE",
+            "COMPRESS_RATIO",
+        )
+
+        keys = tuple(
+            dict.fromkeys(
+                BUILD_PREFILL_CHUNK_METADATA_KERNEL.compile_key(
+                    {
+                        "query_slice_start": query_slice_start,
+                        "query_slice_stop": query_slice_stop,
+                        "BLOCK_SIZE": 1024,
+                        "COMPRESS_RATIO": compress_ratio,
+                    }
+                )
+                for compress_ratio in (1, 4)
+                for query_slice_start, query_slice_stop in ((0, 16), (1, 15))
+            )
+        )
+
+        assert keys == (
+            PrefillChunkMetadataKernelCompileKey(
+                query_slice_start=0,
+                query_slice_stop=16,
+                BLOCK_SIZE=1024,
+                COMPRESS_RATIO=1,
+            ),
+            PrefillChunkMetadataKernelCompileKey(
+                query_slice_start=1,
+                query_slice_stop=15,
+                BLOCK_SIZE=1024,
+                COMPRESS_RATIO=1,
+            ),
+            PrefillChunkMetadataKernelCompileKey(
+                query_slice_start=0,
+                query_slice_stop=16,
+                BLOCK_SIZE=1024,
+                COMPRESS_RATIO=4,
+            ),
+            PrefillChunkMetadataKernelCompileKey(
+                query_slice_start=1,
+                query_slice_stop=15,
+                BLOCK_SIZE=1024,
+                COMPRESS_RATIO=4,
+            ),
+        )
+
+        from vllm.v1.attention.backends.mla.sparse_swa import (
+            ComputePrefillMetadataKernel,
+            _compute_prefill_metadata_kernel,
+        )
+
+        assert trace_triton_kernel_specialization_args(
+            _compute_prefill_metadata_kernel
+        ) == ("BLOCK_SIZE",)
+
+        kernel = ComputePrefillMetadataKernel()
+        assert kernel.compile_key(
+            {
+                "num_prefills": 3,
+            }
+        ) == ComputePrefillMetadataKernel.CompileKey(
+            BLOCK_SIZE=4,
+        )
+        assert kernel._trace_dispatch(kernel.dispatch)(
+            num_prefills=WarmupIntRange(1, 5),
+        ) == [
+            ComputePrefillMetadataKernel.CompileKey(
+                BLOCK_SIZE=1,
+            ),
+            ComputePrefillMetadataKernel.CompileKey(
+                BLOCK_SIZE=2,
+            ),
+            ComputePrefillMetadataKernel.CompileKey(
+                BLOCK_SIZE=4,
+            ),
+        ]
+
+        import torch
+
+        from vllm.v1.attention.backends.mla import indexer, sparse_swa
+
+        class FakeKernel:
+            def __init__(self):
+                self.warmup_calls = []
+                self.launch_calls = []
+
+            def warmup(self, *args, grid, **kwargs):
+                self.warmup_calls.append((args, grid, kwargs))
+
+            def __getitem__(self, grid):
+                def launch(*args, **kwargs):
+                    self.launch_calls.append((grid, args, kwargs))
+
+                return launch
+
+        fake_build_kernel = FakeKernel()
+        build_key = PrefillChunkMetadataKernelCompileKey(
+            query_slice_start=0,
+            query_slice_stop=16,
+            BLOCK_SIZE=1024,
+            COMPRESS_RATIO=4,
+        )
+        with mock.patch.object(
+            indexer, "_build_prefill_chunk_metadata_kernel", fake_build_kernel
+        ):
+            build_call = BUILD_PREFILL_CHUNK_METADATA_KERNEL.compile(build_key)
+            build_call(
+                num_reqs=2,
+                query_start_loc_ptr="query_start_loc",
+                uncompressed_seq_lens_ptr="uncompressed_seq_lens",
+                cu_compressed_seq_lens_ptr="cu_seq_lens",
+                token_to_seq_ptr="token_to_seq",
+                cu_compressed_seq_len_ks_ptr="cu_seq_len_ks",
+                cu_compressed_seq_len_ke_ptr="cu_seq_len_ke",
+                query_slice_start=0,
+                query_slice_stop=16,
+            )
+
+        assert fake_build_kernel.warmup_calls == [
+            (
+                (
+                    torch.int32,
+                    torch.int32,
+                    torch.int32,
+                    torch.int32,
+                    torch.int32,
+                    torch.int32,
+                    0,
+                    16,
+                ),
+                (1,),
+                {
+                    "BLOCK_SIZE": 1024,
+                    "COMPRESS_RATIO": 4,
+                },
+            )
+        ]
+        assert fake_build_kernel.launch_calls == [
+            (
+                (2,),
+                (),
+                {
+                    "query_start_loc_ptr": "query_start_loc",
+                    "uncompressed_seq_lens_ptr": "uncompressed_seq_lens",
+                    "cu_compressed_seq_lens_ptr": "cu_seq_lens",
+                    "token_to_seq_ptr": "token_to_seq",
+                    "cu_compressed_seq_len_ks_ptr": "cu_seq_len_ks",
+                    "cu_compressed_seq_len_ke_ptr": "cu_seq_len_ke",
+                    "query_slice_start": 0,
+                    "query_slice_stop": 16,
+                    "BLOCK_SIZE": 1024,
+                    "COMPRESS_RATIO": 4,
+                },
+            )
+        ]
+
+        fake_compute_kernel = FakeKernel()
+        compute_key = ComputePrefillMetadataKernel.CompileKey(BLOCK_SIZE=4)
+        with mock.patch.object(
+            sparse_swa, "_compute_prefill_metadata_kernel", fake_compute_kernel
+        ):
+            compute_call = kernel.compile(compute_key)
+            compute_call(
+                prefill_gather_lens_ptr="prefill_gather_lens",
+                seq_lens_ptr="seq_lens",
+                query_start_loc_ptr="query_start_loc",
+                num_prefills=3,
+                num_decodes=2,
+                window_size=128,
+            )
+
+        assert fake_compute_kernel.warmup_calls == [
+            (
+                (torch.int32, torch.int32, torch.int32, 4, 0, 1),
+                (1,),
+                {"BLOCK_SIZE": 4},
+            )
+        ]
+        assert fake_compute_kernel.launch_calls == [
+            (
+                (1,),
+                (),
+                {
+                    "prefill_gather_lens_ptr": "prefill_gather_lens",
+                    "seq_lens_ptr": "seq_lens",
+                    "query_start_loc_ptr": "query_start_loc",
+                    "num_prefills": 3,
+                    "num_decodes": 2,
+                    "window_size": 128,
+                    "BLOCK_SIZE": 4,
+                },
+            )
+        ]

vllm/model_executor/warmup/sparse_mla_triton_warmup.py

@@ -101,7 +101,7 @@ def _warm_sparse_swa_prefill_metadata_kernel(
     prefill_tokens: int,
 ) -> None:
     from vllm.v1.attention.backends.mla.sparse_swa import (
-        _compute_prefill_metadata_kernel,
+        _COMPUTE_PREFILL_METADATA_KERNEL,
     )
 
     for num_prefills in _SPARSE_PREFILL_METADATA_NUM_PREFILLS:
@@ -124,20 +124,21 @@ def _warm_sparse_swa_prefill_metadata_kernel(
             prefill_gather_lens = torch.empty(
                 num_prefills, dtype=torch.int32, device=device
             )
-            _compute_prefill_metadata_kernel[(1,)](
-                prefill_gather_lens,
-                seq_lens,
-                query_start_loc,
-                num_prefills,
-                num_decodes,
-                window_size,
-                BLOCK_SIZE=_next_power_of_2(num_prefills),
+            _COMPUTE_PREFILL_METADATA_KERNEL(
+                prefill_gather_lens_ptr=prefill_gather_lens,
+                seq_lens_ptr=seq_lens,
+                query_start_loc_ptr=query_start_loc,
+                num_prefills=num_prefills,
+                num_decodes=num_decodes,
+                window_size=window_size,
             )
 
 
 def _warm_prefill_chunk_metadata_kernel(
     device: torch.device,
     compress_ratio: int,
+    query_slice_start: int,
+    query_slice_stop: int,
     query_len: int,
 ) -> None:
     from vllm.v1.attention.backends.mla.indexer import build_prefill_chunk_metadata
@@ -164,31 +165,18 @@ def _warm_prefill_chunk_metadata_kernel(
         device=device,
     )
 
-    offset_uncompressed_seq_lens = torch.empty(
-        num_reqs + 1, dtype=torch.int32, device=device
-    )[1:]
-    offset_uncompressed_seq_lens.copy_(uncompressed_seq_lens)
-    query_slices = tuple(
-        slice(start, num_reqs * query_len + stop)
-        for start, stop in _PREFILL_CHUNK_METADATA_QUERY_SLICE_OFFSETS
-    )
-    for warmup_uncompressed_seq_lens in (
+    build_prefill_chunk_metadata(
+        0,
+        num_reqs,
+        query_start_loc,
+        query_start_loc_cpu,
         uncompressed_seq_lens,
-        offset_uncompressed_seq_lens,
-    ):
-        for query_slice in query_slices:
-            build_prefill_chunk_metadata(
-                0,
-                num_reqs,
-                query_start_loc,
-                query_start_loc_cpu,
-                warmup_uncompressed_seq_lens,
-                compressed_seq_lens,
-                compressed_seq_lens_cpu,
-                block_table,
-                compress_ratio,
-                query_slice=query_slice,
-            )
+        compressed_seq_lens,
+        compressed_seq_lens_cpu,
+        block_table,
+        compress_ratio,
+        query_slice=slice(query_slice_start, query_slice_stop),
+    )
 
 
 def _warm_combine_topk_swa_indices_kernel(
@@ -277,12 +265,40 @@ def sparse_mla_triton_warmup(
     compress_ratios: tuple[int, ...],
     combine_topk_swa_cases: tuple[tuple[int, int, int, int], ...] = (),
 ) -> None:
+    from vllm.v1.attention.backends.mla.indexer import (
+        BUILD_PREFILL_CHUNK_METADATA_KERNEL,
+    )
+
     device = getattr(runner, "device", torch.device("cuda"))
     window_size = _hf_config_int(runner, "sliding_window", 128)
 
     _warm_sparse_swa_prefill_metadata_kernel(device, window_size, num_tokens)
-    for compress_ratio in compress_ratios:
-        _warm_prefill_chunk_metadata_kernel(device, compress_ratio, num_tokens)
+    query_slice_bounds = tuple(
+        (start, 2 * num_tokens + stop)
+        for start, stop in _PREFILL_CHUNK_METADATA_QUERY_SLICE_OFFSETS
+    )
+    compile_keys = tuple(
+        dict.fromkeys(
+            BUILD_PREFILL_CHUNK_METADATA_KERNEL.compile_key(
+                {
+                    "query_slice_start": query_slice_start,
+                    "query_slice_stop": query_slice_stop,
+                    "BLOCK_SIZE": 1024,
+                    "COMPRESS_RATIO": compress_ratio,
+                }
+            )
+            for compress_ratio in compress_ratios
+            for query_slice_start, query_slice_stop in query_slice_bounds
+        )
+    )
+    for key in compile_keys:
+        _warm_prefill_chunk_metadata_kernel(
+            device,
+            key.COMPRESS_RATIO,
+            key.query_slice_start,
+            key.query_slice_stop,
+            num_tokens,
+        )
     for compress_ratio, topk, topk_width, n in combine_topk_swa_cases:
         _warm_combine_topk_swa_indices_kernel(
             device,

vllm/triton_utils/warmup.py

@@ -0,0 +1,377 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import ast
+import inspect
+import itertools
+import operator
+import textwrap
+from abc import ABC, abstractmethod
+from collections.abc import Callable, Mapping
+from dataclasses import dataclass, fields, is_dataclass
+from typing import Any, Generic, TypeVar
+
+from vllm.config import VllmConfig
+
+CompileKeyT = TypeVar("CompileKeyT")
+
+
+@dataclass(frozen=True)
+class WarmupIntRange:
+    start: int
+    stop: int
+    step: int = 1
+
+
+WarmupIntValues = int | WarmupIntRange | list[int] | tuple[int, ...]
+
+
+def expand_warmup_int_values(values: WarmupIntValues) -> tuple[int, ...]:
+    if isinstance(values, WarmupIntRange):
+        return tuple(range(values.start, values.stop, values.step))
+    if isinstance(values, int):
+        return (values,)
+    return tuple(values)
+
+
+@dataclass(frozen=True)
+class _DispatchTrace:
+    field_exprs: tuple[tuple[str, ast.AST], ...]
+    globals: Mapping[str, Any]
+    input_names: frozenset[str]
+
+    def compile_key(
+        self,
+        compile_key_type: type[CompileKeyT],
+        kwargs: Mapping[str, Any],
+    ) -> CompileKeyT:
+        return compile_key_type(
+            **{
+                field: int(_eval_dispatch_expr(expr, kwargs, self.globals))
+                for field, expr in self.field_exprs
+            }
+        )
+
+
+def _next_power_of_2(x: int) -> int:
+    return 1 << (x - 1).bit_length()
+
+
+_BIN_OPS: dict[type[ast.operator], Callable[[Any, Any], Any]] = {
+    ast.Add: operator.add,
+    ast.Sub: operator.sub,
+    ast.Mult: operator.mul,
+    ast.FloorDiv: operator.floordiv,
+    ast.Mod: operator.mod,
+}
+
+
+def _full_attr_name(node: ast.AST) -> str | None:
+    if isinstance(node, ast.Name):
+        return node.id
+    if isinstance(node, ast.Attribute):
+        parent = _full_attr_name(node.value)
+        if parent is not None:
+            return f"{parent}.{node.attr}"
+    return None
+
+
+def _get_function_source_node(fn: Callable[..., Any]) -> ast.FunctionDef:
+    source_fn = getattr(fn, "fn", fn)
+    source = textwrap.dedent(inspect.getsource(source_fn))
+    tree = ast.parse(source)
+    function_defs = [node for node in tree.body if isinstance(node, ast.FunctionDef)]
+    if len(function_defs) != 1:
+        name = getattr(source_fn, "__name__", type(source_fn).__name__)
+        raise ValueError(f"Expected one function in {name}, found {len(function_defs)}")
+    return function_defs[0]
+
+
+def _eval_dispatch_expr(
+    node: ast.AST,
+    kwargs: Mapping[str, Any],
+    globals_: Mapping[str, Any],
+) -> Any:
+    if isinstance(node, ast.Name):
+        if node.id in kwargs:
+            return kwargs[node.id]
+        if node.id in globals_:
+            return globals_[node.id]
+        raise ValueError(f"Unknown dispatch name: {node.id}")
+
+    if isinstance(node, ast.Constant):
+        return node.value
+
+    if isinstance(node, ast.UnaryOp) and isinstance(node.op, ast.USub):
+        return -_eval_dispatch_expr(node.operand, kwargs, globals_)
+
+    if isinstance(node, ast.BinOp):
+        op = _BIN_OPS.get(type(node.op))
+        if op is None:
+            raise ValueError(f"Unsupported dispatch binary op: {ast.dump(node.op)}")
+        return op(
+            _eval_dispatch_expr(node.left, kwargs, globals_),
+            _eval_dispatch_expr(node.right, kwargs, globals_),
+        )
+
+    if isinstance(node, ast.Call):
+        full_name = _full_attr_name(node.func)
+        args = [_eval_dispatch_expr(arg, kwargs, globals_) for arg in node.args]
+        if full_name == "triton.next_power_of_2":
+            if len(args) != 1 or node.keywords:
+                raise ValueError("triton.next_power_of_2 dispatch call takes one arg")
+            return _next_power_of_2(int(args[0]))
+
+        fn = _eval_dispatch_expr(node.func, kwargs, globals_)
+        call_kwargs = {
+            keyword.arg: _eval_dispatch_expr(keyword.value, kwargs, globals_)
+            for keyword in node.keywords
+            if keyword.arg is not None
+        }
+        return fn(*args, **call_kwargs)
+
+    if isinstance(node, ast.Attribute):
+        value = _eval_dispatch_expr(node.value, kwargs, globals_)
+        return getattr(value, node.attr)
+
+    raise ValueError(f"Unsupported dispatch expression: {ast.dump(node)}")
+
+
+def _collect_input_names(
+    node: ast.AST,
+    candidate_names: set[str],
+) -> set[str]:
+    return {
+        child.id
+        for child in ast.walk(node)
+        if isinstance(child, ast.Name) and child.id in candidate_names
+    }
+
+
+def trace_dispatch(fn: Callable[..., Any]) -> _DispatchTrace:
+    source_fn = getattr(fn, "__func__", fn)
+    globals_ = source_fn.__globals__
+    function_def = _get_function_source_node(fn)
+
+    returns = [
+        node.value
+        for node in ast.walk(function_def)
+        if isinstance(node, ast.Return) and node.value is not None
+    ]
+    if len(returns) != 1 or not isinstance(returns[0], ast.Call):
+        raise ValueError(
+            f"Expected {fn.__name__} to return exactly one CompileKey(...) call"
+        )
+
+    field_exprs: list[tuple[str, ast.AST]] = []
+    signature = inspect.signature(fn)
+    candidate_names = set(signature.parameters)
+    input_names: set[str] = set()
+    for keyword in returns[0].keywords:
+        if keyword.arg is None:
+            raise ValueError(f"{fn.__name__} cannot use **kwargs in CompileKey")
+        field_exprs.append((keyword.arg, keyword.value))
+        input_names.update(_collect_input_names(keyword.value, candidate_names))
+
+    return _DispatchTrace(tuple(field_exprs), globals_, frozenset(input_names))
+
+
+def _literal_str_refs(node: ast.AST) -> tuple[str | int, ...]:
+    if isinstance(node, ast.Constant) and isinstance(node.value, str | int):
+        return (node.value,)
+    if isinstance(node, ast.List | ast.Tuple):
+        refs: list[str | int] = []
+        for elt in node.elts:
+            if isinstance(elt, ast.Constant) and isinstance(elt.value, str | int):
+                refs.append(elt.value)
+            else:
+                raise ValueError(
+                    f"Unsupported Triton specialization ref: {ast.dump(elt)}"
+                )
+        return tuple(refs)
+    raise ValueError(f"Unsupported Triton specialization refs: {ast.dump(node)}")
+
+
+def _normalize_arg_refs(
+    refs: tuple[str | int, ...],
+    arg_names: tuple[str, ...],
+) -> frozenset[str]:
+    names: set[str] = set()
+    for ref in refs:
+        if isinstance(ref, int):
+            names.add(arg_names[ref])
+        else:
+            names.add(ref)
+    return frozenset(names)
+
+
+def _decorator_keyword_refs(
+    function_def: ast.FunctionDef,
+    keyword_name: str,
+) -> tuple[str | int, ...]:
+    for decorator in function_def.decorator_list:
+        if not isinstance(decorator, ast.Call):
+            continue
+        decorator_name = _full_attr_name(decorator.func)
+        if decorator_name not in ("triton.jit", "jit"):
+            continue
+        for keyword in decorator.keywords:
+            if keyword.arg == keyword_name:
+                return _literal_str_refs(keyword.value)
+    return ()
+
+
+def _triton_do_not_specialize_args(
+    kernel: Callable[..., Any],
+    function_def: ast.FunctionDef,
+    arg_names: tuple[str, ...],
+) -> frozenset[str]:
+    refs = getattr(kernel, "do_not_specialize", None)
+    if refs is not None:
+        return _normalize_arg_refs(tuple(refs), arg_names)
+    return _normalize_arg_refs(
+        _decorator_keyword_refs(function_def, "do_not_specialize"),
+        arg_names,
+    )
+
+
+def _triton_constexpr_arg_names(
+    kernel: Callable[..., Any],
+    function_def: ast.FunctionDef,
+    arg_names: tuple[str, ...],
+) -> frozenset[str]:
+    constexprs = getattr(kernel, "constexprs", None)
+    if constexprs is not None:
+        return frozenset(arg_names[index] for index in constexprs)
+
+    names: set[str] = set()
+    for arg in function_def.args.args + function_def.args.kwonlyargs:
+        if arg.annotation is None:
+            continue
+        annotation = _full_attr_name(arg.annotation)
+        if annotation in ("tl.constexpr", "triton.language.constexpr", "constexpr"):
+            names.add(arg.arg)
+    return frozenset(names)
+
+
+def _leftmost_name(node: ast.AST) -> str | None:
+    if isinstance(node, ast.Name):
+        return node.id
+    if isinstance(node, ast.BinOp):
+        return _leftmost_name(node.left)
+    return None
+
+
+def _pointer_arg_names(
+    function_def: ast.FunctionDef,
+    arg_names: tuple[str, ...],
+) -> frozenset[str]:
+    candidate_names = set(arg_names)
+    pointer_names = {name for name in arg_names if name.endswith("_ptr")}
+    for node in ast.walk(function_def):
+        if not isinstance(node, ast.Call):
+            continue
+        if _full_attr_name(node.func) not in ("tl.load", "tl.store"):
+            continue
+        if not node.args:
+            continue
+        name = _leftmost_name(node.args[0])
+        if name in candidate_names:
+            pointer_names.add(name)
+    return frozenset(pointer_names)
+
+
+def trace_triton_kernel_specialization_args(
+    kernel: Callable[..., Any],
+) -> tuple[str, ...]:
+    function_def = _get_function_source_node(kernel)
+    source_fn = getattr(kernel, "fn", kernel)
+    arg_names = tuple(inspect.signature(source_fn).parameters)
+    constexpr_args = _triton_constexpr_arg_names(kernel, function_def, arg_names)
+    do_not_specialize_args = _triton_do_not_specialize_args(
+        kernel, function_def, arg_names
+    )
+    pointer_args = _pointer_arg_names(function_def, arg_names)
+
+    return tuple(
+        name
+        for name in arg_names
+        if name in constexpr_args
+        or (name not in pointer_args and name not in do_not_specialize_args)
+    )
+
+
+class VllmJitKernelWithWarmup(Generic[CompileKeyT], ABC):
+    CompileKey: type[CompileKeyT]
+
+    def __init__(self) -> None:
+        self.traced_dispatch_args = trace_dispatch(self.dispatch)
+        self.callables: dict[CompileKeyT, Callable[..., Any]] = {}
+
+    def compile_key(self, kwargs: Mapping[str, Any]) -> CompileKeyT:
+        return self.traced_dispatch_args.compile_key(self.CompileKey, kwargs)
+
+    def assert_compile_key_matches_triton(
+        self,
+        kernel: Callable[..., Any],
+    ) -> None:
+        if not is_dataclass(self.CompileKey):
+            raise TypeError(f"{type(self).__name__}.CompileKey must be a dataclass")
+
+        compile_key_args = tuple(field.name for field in fields(self.CompileKey))
+        triton_args = trace_triton_kernel_specialization_args(kernel)
+        if compile_key_args != triton_args:
+            raise ValueError(
+                f"{type(self).__name__}.CompileKey fields {compile_key_args} "
+                f"do not match Triton specialization args {triton_args}"
+            )
+
+    def _trace_dispatch(
+        self, dispatch: Callable[..., CompileKeyT]
+    ) -> Callable[..., list[CompileKeyT]]:
+        traced_dispatch_args = trace_dispatch(dispatch)
+
+        def traced(**kwargs: WarmupIntValues) -> list[CompileKeyT]:
+            kwarg_names = tuple(
+                name for name in kwargs if name in traced_dispatch_args.input_names
+            )
+            kwarg_values = tuple(
+                expand_warmup_int_values(kwargs[name]) for name in kwarg_names
+            )
+            return list(
+                dict.fromkeys(
+                    traced_dispatch_args.compile_key(
+                        self.CompileKey, dict(zip(kwarg_names, values))
+                    )
+                    for values in itertools.product(*kwarg_values)
+                )
+            )
+
+        return traced
+
+    @abstractmethod
+    def dispatch(self, **kwargs: Any) -> CompileKeyT:
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_warmup_keys(self, vllm_config: VllmConfig) -> list[CompileKeyT]:
+        raise NotImplementedError
+
+    @abstractmethod
+    def compile(self, compile_key: CompileKeyT) -> Callable[..., Any]:
+        raise NotImplementedError
+
+    def warmup(self, vllm_config: VllmConfig) -> None:
+        for compile_key in self.get_warmup_keys(vllm_config):
+            if compile_key not in self.callables:
+                self.callables[compile_key] = self.compile(compile_key)
+
+    def __call__(self, **kwargs: Any) -> Any:
+        compile_key = self.compile_key(kwargs)
+        fn = self.callables.get(compile_key)
+        if fn is None:
+            fn = self.compile(compile_key)
+            self.callables[compile_key] = fn
+        return fn(**kwargs)
+
+    def saved_compile_keys(self) -> tuple[CompileKeyT, ...]:
+        return tuple(self.callables)

vllm/v1/attention/backends/mla/indexer.py

@@ -1,6 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from collections.abc import Callable
 from dataclasses import dataclass
+from typing import Any
 
 import torch
 
@@ -9,6 +11,7 @@ from vllm.config import VllmConfig
 from vllm.logger import init_logger
 from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
+from vllm.triton_utils.warmup import VllmJitKernelWithWarmup, WarmupIntRange
 from vllm.utils.deep_gemm import (
     get_paged_mqa_logits_metadata,
     has_deep_gemm,
@@ -179,6 +182,113 @@ class DeepseekV32IndexerPrefillChunkMetadata:
     skip_kv_gather: bool = False
 
 
+class BuildPrefillChunkMetadataKernel(
+    VllmJitKernelWithWarmup["BuildPrefillChunkMetadataKernel.CompileKey"]
+):
+    @dataclass(frozen=True)
+    class CompileKey:
+        query_slice_start: int
+        query_slice_stop: int
+        BLOCK_SIZE: int
+        COMPRESS_RATIO: int
+
+    def dispatch(  # type: ignore[override]
+        self,
+        *,
+        query_slice_start: int,
+        query_slice_stop: int,
+        BLOCK_SIZE: int,
+        COMPRESS_RATIO: int,
+    ) -> CompileKey:
+        return self.CompileKey(
+            query_slice_start=query_slice_start,
+            query_slice_stop=query_slice_stop,
+            BLOCK_SIZE=BLOCK_SIZE,
+            COMPRESS_RATIO=COMPRESS_RATIO,
+        )
+
+    def get_warmup_keys(self, vllm_config: VllmConfig) -> list[CompileKey]:
+        max_tokens = max(1, min(vllm_config.scheduler_config.max_num_batched_tokens, 8))
+        hf_config = vllm_config.model_config.hf_config
+        compress_ratios = tuple(
+            int(ratio)
+            for ratio in (getattr(hf_config, "compress_ratios", None) or (1,))
+        )
+        return self._trace_dispatch(self.dispatch)(
+            query_slice_start=WarmupIntRange(0, 2),
+            query_slice_stop=WarmupIntRange(2 * max_tokens - 1, 2 * max_tokens + 1),
+            BLOCK_SIZE=1024,
+            COMPRESS_RATIO=list(compress_ratios),
+        )
+
+    def compile(self, compile_key: CompileKey) -> Callable[..., Any]:
+        warmup = getattr(_build_prefill_chunk_metadata_kernel, "warmup", None)
+        if warmup is not None:
+            warmup(
+                torch.int32,
+                torch.int32,
+                torch.int32,
+                torch.int32,
+                torch.int32,
+                torch.int32,
+                compile_key.query_slice_start,
+                compile_key.query_slice_stop,
+                BLOCK_SIZE=compile_key.BLOCK_SIZE,
+                COMPRESS_RATIO=compile_key.COMPRESS_RATIO,
+                grid=(1,),
+            )
+
+        def call(num_reqs: int, **kwargs: Any) -> None:
+            _build_prefill_chunk_metadata_kernel[(num_reqs,)](
+                **kwargs,
+                BLOCK_SIZE=compile_key.BLOCK_SIZE,
+                COMPRESS_RATIO=compile_key.COMPRESS_RATIO,
+            )
+
+        return call
+
+    def __call__(  # type: ignore[override]
+        self,
+        query_start_loc_ptr: Any,
+        uncompressed_seq_lens_ptr: Any,
+        cu_compressed_seq_lens_ptr: Any,
+        token_to_seq_ptr: Any,
+        cu_compressed_seq_len_ks_ptr: Any,
+        cu_compressed_seq_len_ke_ptr: Any,
+        query_slice_start: int,
+        query_slice_stop: int,
+        *,
+        num_reqs: int,
+        BLOCK_SIZE: int,
+        COMPRESS_RATIO: int,
+    ) -> Any:
+        compile_key = self.dispatch(
+            query_slice_start=query_slice_start,
+            query_slice_stop=query_slice_stop,
+            BLOCK_SIZE=BLOCK_SIZE,
+            COMPRESS_RATIO=COMPRESS_RATIO,
+        )
+        fn = self.callables.get(compile_key)
+        if fn is None:
+            fn = self.compile(compile_key)
+            self.callables[compile_key] = fn
+        return fn(
+            num_reqs=num_reqs,
+            query_start_loc_ptr=query_start_loc_ptr,
+            uncompressed_seq_lens_ptr=uncompressed_seq_lens_ptr,
+            cu_compressed_seq_lens_ptr=cu_compressed_seq_lens_ptr,
+            token_to_seq_ptr=token_to_seq_ptr,
+            cu_compressed_seq_len_ks_ptr=cu_compressed_seq_len_ks_ptr,
+            cu_compressed_seq_len_ke_ptr=cu_compressed_seq_len_ke_ptr,
+            query_slice_start=query_slice_start,
+            query_slice_stop=query_slice_stop,
+        )
+
+
+BUILD_PREFILL_CHUNK_METADATA_KERNEL = BuildPrefillChunkMetadataKernel()
+PrefillChunkMetadataKernelCompileKey = BuildPrefillChunkMetadataKernel.CompileKey
+
+
 @dataclass
 class DeepseekV32IndexerPrefillMetadata:
     chunks: list[DeepseekV32IndexerPrefillChunkMetadata]
@@ -677,8 +787,8 @@ def build_prefill_chunk_metadata(
         (query_start_loc_cpu[end_idx] - query_start_loc_cpu[start_idx]).item()
     )
     if query_slice is not None:
-        qs_start = query_slice.start
-        qs_stop = query_slice.stop
+        qs_start = int(query_slice.start)
+        qs_stop = int(query_slice.stop)
     else:
         qs_start = 0
         qs_stop = total_query_len
@@ -687,7 +797,7 @@ def build_prefill_chunk_metadata(
     cu_seq_len_ks = torch.empty(output_query_len, dtype=torch.int32, device=device)
     cu_seq_len_ke = torch.empty(output_query_len, dtype=torch.int32, device=device)
 
-    _build_prefill_chunk_metadata_kernel[(num_reqs,)](
+    BUILD_PREFILL_CHUNK_METADATA_KERNEL(
         query_start_loc,
         uncompressed_seq_lens[start_idx:end_idx],
         cu_seq_lens,
@@ -696,6 +806,7 @@ def build_prefill_chunk_metadata(
         cu_seq_len_ke,
         qs_start,
         qs_stop,
+        num_reqs=num_reqs,
         BLOCK_SIZE=1024,
         COMPRESS_RATIO=compress_ratio,
     )
@@ -776,3 +887,8 @@ def _build_prefill_chunk_metadata_kernel(
         offset = i + tl.arange(0, BLOCK_SIZE)
         mask = offset < compressed_seq_len
         tl.store(token_to_seq_ptr + seq_start + offset, batch_idx, mask=mask)
+
+
+BUILD_PREFILL_CHUNK_METADATA_KERNEL.assert_compile_key_matches_triton(
+    _build_prefill_chunk_metadata_kernel
+)

vllm/v1/attention/backends/mla/sparse_swa.py

@@ -1,7 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from collections.abc import Callable
 from dataclasses import dataclass, field
-from typing import ClassVar, cast
+from typing import Any, ClassVar, cast
 
 import torch
 
@@ -9,6 +10,7 @@ from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config
 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
 from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
+from vllm.triton_utils.warmup import VllmJitKernelWithWarmup, WarmupIntRange
 from vllm.utils.math_utils import cdiv
 from vllm.v1.attention.backend import (
     AttentionBackend,
@@ -276,6 +278,55 @@ class DeepseekSparseSWAMetadata:
         return chunk_plan
 
 
+class ComputePrefillMetadataKernel(
+    VllmJitKernelWithWarmup["ComputePrefillMetadataKernel.CompileKey"]
+):
+    @dataclass(frozen=True)
+    class CompileKey:
+        BLOCK_SIZE: int
+
+    def dispatch(  # type: ignore[override]
+        self,
+        *,
+        num_prefills: int,
+    ) -> CompileKey:
+        return self.CompileKey(
+            BLOCK_SIZE=triton.next_power_of_2(num_prefills),
+        )
+
+    def get_warmup_keys(self, vllm_config: VllmConfig) -> list[CompileKey]:
+        max_prefills = max(1, min(vllm_config.scheduler_config.max_num_seqs, 8))
+        return self._trace_dispatch(self.dispatch)(
+            num_prefills=WarmupIntRange(1, max_prefills + 1),
+        )
+
+    def compile(self, compile_key: CompileKey) -> Callable[..., Any]:
+        warmup = getattr(_compute_prefill_metadata_kernel, "warmup", None)
+        if warmup is not None:
+            warmup(
+                torch.int32,
+                torch.int32,
+                torch.int32,
+                compile_key.BLOCK_SIZE,
+                0,
+                1,
+                BLOCK_SIZE=compile_key.BLOCK_SIZE,
+                grid=(1,),
+            )
+
+        def call(**kwargs: Any) -> None:
+            _compute_prefill_metadata_kernel[(1,)](
+                **kwargs,
+                BLOCK_SIZE=compile_key.BLOCK_SIZE,
+            )
+
+        return call
+
+
+_COMPUTE_PREFILL_METADATA_KERNEL = ComputePrefillMetadataKernel()
+ComputePrefillMetadataKernelCompileKey = ComputePrefillMetadataKernel.CompileKey
+
+
 class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
     """Builds metadata for DeepseekV4 SWA cache.
 
@@ -553,14 +604,13 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
             pfx_gather_lens = torch.empty(
                 num_prefills, dtype=torch.int32, device=seq_lens.device
             )
-            _compute_prefill_metadata_kernel[(1,)](
-                pfx_gather_lens,
-                seq_lens,
-                query_start_loc,
-                num_prefills,
-                num_decodes,
-                self.window_size,
-                BLOCK_SIZE=triton.next_power_of_2(num_prefills),
+            _COMPUTE_PREFILL_METADATA_KERNEL(
+                prefill_gather_lens_ptr=pfx_gather_lens,
+                seq_lens_ptr=seq_lens,
+                query_start_loc_ptr=query_start_loc,
+                num_prefills=num_prefills,
+                num_decodes=num_decodes,
+                window_size=self.window_size,
             )
 
             result["prefill_seq_lens"] = seq_lens[num_decodes:]
@@ -577,7 +627,7 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
         return result
 
 
-@triton.jit
+@triton.jit(do_not_specialize=["num_prefills", "num_decodes", "window_size"])
 def _compute_prefill_metadata_kernel(
     # Outputs
     prefill_gather_lens_ptr,
@@ -608,6 +658,11 @@ def _compute_prefill_metadata_kernel(
     tl.store(prefill_gather_lens_ptr + offset, gather_len, mask=mask)
 
 
+_COMPUTE_PREFILL_METADATA_KERNEL.assert_compile_key_matches_triton(
+    _compute_prefill_metadata_kernel
+)
+
+
 @triton.jit(do_not_specialize=["token_offset"])
 def _compute_swa_indices_and_lens_kernel(
     swa_indices_ptr,