diff --git a/.buildkite/test_areas/model_executor.yaml b/.buildkite/test_areas/model_executor.yaml index 1e54538cc3..689d19488b 100644 --- a/.buildkite/test_areas/model_executor.yaml +++ b/.buildkite/test_areas/model_executor.yaml @@ -10,7 +10,9 @@ steps: - vllm/engine/arg_utils.py - vllm/config/model.py - vllm/model_executor + - vllm/model_executor/warmup - tests/model_executor + - tests/model_executor/test_jit_warmup.py - tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py commands: - apt-get update && apt-get install -y curl libsodium23 @@ -34,7 +36,9 @@ steps: - vllm/engine/arg_utils.py - vllm/config/model.py - vllm/model_executor + - vllm/model_executor/warmup - tests/model_executor + - tests/model_executor/test_jit_warmup.py - tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py - vllm/_aiter_ops.py - vllm/platforms/rocm.py diff --git a/tests/model_executor/test_jit_warmup.py b/tests/model_executor/test_jit_warmup.py new file mode 100644 index 0000000000..da6b21091b --- /dev/null +++ b/tests/model_executor/test_jit_warmup.py @@ -0,0 +1,306 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import ast +from dataclasses import dataclass +from types import SimpleNamespace +from typing import Any, cast + +import pytest + +from vllm.model_executor.warmup.jit_warmup import ( + VllmJitKernel, + WarmupIntRange, + get_ast_full_name, + zip_inputs, +) + + +def _next_power_of_2(value: int) -> int: + return 1 << max(0, value - 1).bit_length() + + +def _round_up(value: int, *, multiple: int) -> int: + return ((value + multiple - 1) // multiple) * multiple + + +def _config( + *, + bias: int = 0, + disabled: bool = False, + name: str = "base", + vectorized: bool = False, +) -> SimpleNamespace: + return SimpleNamespace( + bias=bias, + disabled=disabled, + name=name, + vectorized=vectorized, + ) + + +class ToyKernel(VllmJitKernel["ToyKernel.CompileKey"]): + @dataclass(frozen=True) + class CompileKey: + block_size: int + work: int + vector_width: int + descriptor: tuple[object, ...] + enabled: bool + + def dispatch( # type: ignore[override] + self, + *, + tokens: int, + cfg: Any, + lanes: int = 1, + mode: str = "default", + debug: int = 0, + ) -> CompileKey: + block_size = _next_power_of_2(tokens) + work: int = block_size * lanes + cfg.bias + return self.CompileKey( + block_size=block_size, + work=work, + vector_width=4 if cfg.vectorized and block_size >= 4 else 1, + descriptor=( + cfg.name, + mode, + -block_size, + block_size % 3, + block_size**2, + ), + enabled=not cfg.disabled, + ) + + def get_warmup_keys(self, max_tokens: int, cfg: Any) -> list[CompileKey]: + return self._trace_dispatch(self.dispatch)( + tokens=WarmupIntRange(1, max_tokens + 1), + cfg=cfg, + # This argument is intentionally unused by dispatch expressions. + debug=WarmupIntRange(0, 100), + ) + + def compile(self, compile_key: CompileKey) -> None: + pass + + +class RecordingToyKernel(ToyKernel): + def __init__(self) -> None: + self.compiled: list[ToyKernel.CompileKey] = [] + super().__init__() + + def compile(self, compile_key: ToyKernel.CompileKey) -> None: + self.compiled.append(compile_key) + + +def test_trace_dispatch_expands_ranges_dedupes_and_ignores_unused_inputs() -> None: + cfg = _config() + + assert ToyKernel().get_warmup_keys(5, cfg) == [ + ToyKernel.CompileKey(1, 1, 1, ("base", "default", -1, 1, 1), True), + ToyKernel.CompileKey(2, 2, 1, ("base", "default", -2, 2, 4), True), + ToyKernel.CompileKey(4, 4, 1, ("base", "default", -4, 1, 16), True), + ToyKernel.CompileKey(8, 8, 1, ("base", "default", -8, 2, 64), True), + ] + + +def test_compile_key_uses_defaults_locals_attributes_and_expressions() -> None: + cfg = _config(bias=3, disabled=True, name="cfg", vectorized=True) + + assert ToyKernel().compile_key( + { + "tokens": 4, + "cfg": cfg, + "lanes": 2, + } + ) == ToyKernel.CompileKey( + block_size=4, + work=11, + vector_width=4, + descriptor=("cfg", "default", -4, 1, 16), + enabled=False, + ) + + +def test_trace_dispatch_combines_zipped_rows_with_independent_values() -> None: + cfg = _config(vectorized=True) + + keys = ToyKernel()._trace_dispatch(ToyKernel().dispatch)( + zip_inputs( + dict(tokens=1, mode="small"), + dict(tokens=4, mode="wide"), + ), + cfg=cfg, + lanes=(1, 2), + ) + + assert keys == [ + ToyKernel.CompileKey(1, 1, 1, ("base", "small", -1, 1, 1), True), + ToyKernel.CompileKey(1, 2, 1, ("base", "small", -1, 1, 1), True), + ToyKernel.CompileKey(4, 4, 4, ("base", "wide", -4, 1, 16), True), + ToyKernel.CompileKey(4, 8, 4, ("base", "wide", -4, 1, 16), True), + ] + + +def test_zip_inputs_validates_input_rows() -> None: + with pytest.raises(ValueError, match="requires at least one"): + zip_inputs() + with pytest.raises(ValueError, match="rows must be mappings"): + zip_inputs(cast(Any, ("tokens", 1))) + with pytest.raises(ValueError, match="at least one dispatch input name"): + zip_inputs({}) + with pytest.raises(ValueError, match="dispatch input names must be strings"): + zip_inputs(cast(Any, {1: 2})) + with pytest.raises(ValueError, match="same dispatch input names"): + zip_inputs({"tokens": 1}, {"mode": "small"}) + + +def test_trace_dispatch_rejects_bad_positional_groups_and_duplicates() -> None: + kernel = ToyKernel() + + with pytest.raises(TypeError, match="zip_inputs"): + kernel._trace_dispatch(kernel.dispatch)( + cast(Any, {"tokens": 1}), + cfg=_config(), + ) + + with pytest.raises(ValueError, match="specified more than once"): + kernel._trace_dispatch(kernel.dispatch)( + zip_inputs(dict(tokens=1, mode="small")), + tokens=2, + cfg=_config(), + ) + + +def test_helper_calls_support_keywords_and_reject_star_kwargs() -> None: + class HelperKernel(VllmJitKernel["HelperKernel.CompileKey"]): + @dataclass(frozen=True) + class CompileKey: + value: int + + def dispatch( # type: ignore[override] + self, + *, + tokens: int, + block_size: int, + ) -> CompileKey: + return self.CompileKey(value=_round_up(tokens, multiple=block_size)) + + def get_warmup_keys(self) -> list[CompileKey]: + return [] + + def compile(self, compile_key: CompileKey) -> None: + pass + + class StarKwargsKernel(VllmJitKernel["StarKwargsKernel.CompileKey"]): + @dataclass(frozen=True) + class CompileKey: + value: int + + def dispatch( # type: ignore[override] + self, + *, + tokens: int, + block_size: int, + ) -> CompileKey: + return self.CompileKey(value=_round_up(tokens, **{"multiple": block_size})) + + def get_warmup_keys(self) -> list[CompileKey]: + return [] + + def compile(self, compile_key: CompileKey) -> None: + pass + + assert HelperKernel().compile_key( + { + "tokens": 5, + "block_size": 4, + } + ) == HelperKernel.CompileKey(value=8) + with pytest.raises(ValueError, match=r"cannot use \*\*kwargs"): + StarKwargsKernel().compile_key({"tokens": 5, "block_size": 4}) + + +def test_dispatch_body_must_be_local_assignments_then_compile_key_return() -> None: + class BranchKernel(VllmJitKernel["BranchKernel.CompileKey"]): + @dataclass(frozen=True) + class CompileKey: + value: int + + def dispatch(self, *, value: int) -> CompileKey: # type: ignore[override] + if value > 0: + value = 1 + return self.CompileKey(value=value) + + def get_warmup_keys(self) -> list[CompileKey]: + return [] + + def compile(self, compile_key: CompileKey) -> None: + pass + + class KwargsReturnKernel(VllmJitKernel["KwargsReturnKernel.CompileKey"]): + @dataclass(frozen=True) + class CompileKey: + value: int + + def dispatch(self, *, value: int) -> CompileKey: # type: ignore[override] + return self.CompileKey(**{"value": value}) + + def get_warmup_keys(self) -> list[CompileKey]: + return [] + + def compile(self, compile_key: CompileKey) -> None: + pass + + with pytest.raises(ValueError, match="local assignments"): + BranchKernel() + with pytest.raises(ValueError, match=r"cannot use \*\*kwargs in CompileKey"): + KwargsReturnKernel() + + +def test_dispatch_reports_unsupported_expression_with_context() -> None: + class UnsupportedKernel(VllmJitKernel["UnsupportedKernel.CompileKey"]): + @dataclass(frozen=True) + class CompileKey: + value: object + + def dispatch(self, *, value: int) -> CompileKey: # type: ignore[override] + return self.CompileKey(value={value}) + + def get_warmup_keys(self) -> list[CompileKey]: + return [] + + def compile(self, compile_key: CompileKey) -> None: + pass + + with pytest.raises(ValueError) as exc_info: + UnsupportedKernel().compile_key({"value": 1}) + + message = str(exc_info.value) + assert "Unsupported dispatch expression" in message + assert "{value}" in message + assert "Supported dispatch expressions" in message + + +def test_warmup_compiles_all_returned_keys_in_order() -> None: + kernel = RecordingToyKernel() + cfg = _config() + + kernel.warmup(3, cfg) + + assert kernel.compiled == [ + ToyKernel.CompileKey(1, 1, 1, ("base", "default", -1, 1, 1), True), + ToyKernel.CompileKey(2, 2, 1, ("base", "default", -2, 2, 4), True), + ToyKernel.CompileKey(4, 4, 1, ("base", "default", -4, 1, 16), True), + ] + + +def test_get_ast_full_name_handles_names_attributes_and_other_nodes() -> None: + dotted_expr = ast.parse("foo.bar.baz").body[0] + call_expr = ast.parse("foo()").body[0] + assert isinstance(dotted_expr, ast.Expr) + assert isinstance(call_expr, ast.Expr) + + assert get_ast_full_name(dotted_expr.value) == "foo.bar.baz" + assert get_ast_full_name(call_expr.value) is None diff --git a/vllm/config/kernel.py b/vllm/config/kernel.py index f13562f59c..16856dc358 100644 --- a/vllm/config/kernel.py +++ b/vllm/config/kernel.py @@ -175,9 +175,15 @@ class KernelConfig: enable_flashinfer_autotune: bool = None # type: ignore[assignment] """If True, run FlashInfer autotuning during kernel warmup.""" + # TODO(roberto): Remove after registered CuTeDSL warmups are migrated + # to the shared JIT warmup infrastructure. + # https://github.com/vllm-project/vllm/pull/47451 enable_cutedsl_warmup: bool = True """If True, run CuTeDSL compile warmup during kernel warmup.""" + enable_jit_warmup: bool = True + """If True, run JIT compile warmup during kernel warmup.""" + enable_bf16x3_router_gemm: bool = False """If True, use the experimental SM100 BF16x3 CuteDSL router GEMM.""" @@ -250,6 +256,7 @@ class KernelConfig: """ ignored_factors = { "enable_cutedsl_warmup", + "enable_jit_warmup", "enable_flashinfer_autotune", "ir_op_priority", # handled separately below } @@ -260,6 +267,7 @@ class KernelConfig: @field_validator( "enable_flashinfer_autotune", "enable_cutedsl_warmup", + "enable_jit_warmup", mode="wrap", ) @classmethod diff --git a/vllm/model_executor/warmup/cutedsl_warmup.py b/vllm/model_executor/warmup/cutedsl_warmup.py index 991eab22b8..198f0e183e 100644 --- a/vllm/model_executor/warmup/cutedsl_warmup.py +++ b/vllm/model_executor/warmup/cutedsl_warmup.py @@ -2,6 +2,10 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Run registered CuTeDSL warmup compile units.""" +# TODO(roberto): Remove this compatibility registry after registered CuTeDSL +# warmups are migrated to the shared JIT warmup infrastructure. +# https://github.com/vllm-project/vllm/pull/47451 + from __future__ import annotations import time diff --git a/vllm/model_executor/warmup/fa4_cutedsl_config.py b/vllm/model_executor/warmup/fa4_cutedsl_config.py deleted file mode 100644 index 916ffd8a67..0000000000 --- a/vllm/model_executor/warmup/fa4_cutedsl_config.py +++ /dev/null @@ -1,204 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""FA4 MLA prefill CuTeDSL compile warmup config.""" - -from __future__ import annotations - -from collections.abc import Hashable, Iterator -from dataclasses import dataclass -from typing import TYPE_CHECKING, Literal - -import torch - -if TYPE_CHECKING: - from vllm.v1.attention.backends.fa_utils import ( - FlashAttentionCuTeDSLCompileSpec, - ) - -FA4ArchitectureFamily = Literal["sm90", "sm100f", "sm120"] - -FA4_STANDARD_DTYPES = (torch.bfloat16, torch.float16) - -# Current vLLM MLA prefill expands K/V to num_heads before FA4, so this plan -# covers qhead_per_kvhead=1. -# Batch is not a current FA4 MLA-prefill key field. Use b1 for compile-only -# specs because it is the conservative case for Split-KV shape heuristics. -# TODO(roberto): FA4 also has direct-GQA and qv/top-k absorbed-MLA paths, but vLLM -# does not use them in this backend yet; they need a separate -# num_kv_heads/qv/top-k-aware warmup plan if wired in later. -FA4_MLA_PREFILL_COMPILE_BATCH_SIZE = 1 -FA4_MLA_PREFILL_Q_TILE = 128 -FA4_MLA_PREFILL_K_TILE = 128 -FA4_MLA_PREFILL_LONG_K_BLOCKS = 32 -FA4_MLA_PREFILL_VERY_LONG_K_BLOCKS = 64 -FA4_MLA_PREFILL_CAUSAL_OPTIONS = (False, True) -FA4_MLA_PREFILL_LSE_OPTIONS = (False, True) - - -@dataclass(frozen=True) -class FA4MLAPrefillCompileContext: - dtype: torch.dtype - num_heads: int - qk_head_dim: int - v_head_dim: int - kv_nope_head_dim: int - requires_v_padding: bool - scale: float - num_splits: int - fa_version: int - - # Return the V head dim FA4 sees. - @property - def effective_v_head_dim(self) -> int: - if self.requires_v_padding: - return self.qk_head_dim - return self.v_head_dim - - -@dataclass(frozen=True) -class FA4MLAPrefillCompileRequest: - """One compile-only FA4 MLA prefill request.""" - - key: Hashable - compile_spec: FlashAttentionCuTeDSLCompileSpec - - # Compile this request. - def compile(self) -> None: - self.compile_spec.compile() - - -# Yield deduped compile requests. -def iter_fa4_mla_prefill_compile_requests( - ctx: FA4MLAPrefillCompileContext, -) -> Iterator[FA4MLAPrefillCompileRequest]: - """Yield compile requests for this fixed MLA backend. - - FA4 dedupes duplicate atomic kernel selections in its own JIT cache. - """ - seen: set[Hashable] = set() - for compile_spec in iter_fa4_mla_prefill_compile_specs(ctx): - key = compile_spec.request_key() - if key in seen: - continue - seen.add(key) - yield FA4MLAPrefillCompileRequest( - key=key, - compile_spec=compile_spec, - ) - - -# Build compile specs for this setup. -def iter_fa4_mla_prefill_compile_specs( - ctx: FA4MLAPrefillCompileContext, -) -> Iterator[FlashAttentionCuTeDSLCompileSpec]: - """Yield compile-only FA4 MLA prefill requests for this fixed setup.""" - - arch_family = _fa4_architecture_family_from_compute_capability( - *torch.cuda.get_device_capability() - ) - if not _supports_fa4_mla_prefill(ctx, arch_family): - return - - from vllm.v1.attention.backends.fa_utils import ( - FlashAttentionCuTeDSLCompileSpec, - ) - - batch_size = FA4_MLA_PREFILL_COMPILE_BATCH_SIZE - v_stride = None - if not ctx.requires_v_padding: - v_stride = ( - ctx.num_heads * ctx.kv_nope_head_dim, - ctx.kv_nope_head_dim, - 1, - ) - - for _, max_seqlen_q, max_seqlen_k in _shape_probes_for_context(ctx, arch_family): - total_q_tokens = batch_size * max_seqlen_q - total_kv_tokens = batch_size * max_seqlen_k - for causal in FA4_MLA_PREFILL_CAUSAL_OPTIONS: - for return_lse in FA4_MLA_PREFILL_LSE_OPTIONS: - yield FlashAttentionCuTeDSLCompileSpec( - q_shape=(total_q_tokens, ctx.num_heads, ctx.qk_head_dim), - k_shape=(total_kv_tokens, ctx.num_heads, ctx.qk_head_dim), - v_shape=( - total_kv_tokens, - ctx.num_heads, - ctx.effective_v_head_dim, - ), - v_stride=v_stride, - q_dtype=ctx.dtype, - cu_seqlens_q_shape=(batch_size + 1,), - cu_seqlens_k_shape=(batch_size + 1,), - max_seqlen_q=max_seqlen_q, - max_seqlen_k=max_seqlen_k, - softmax_scale=ctx.scale, - causal=causal, - return_softmax_lse=return_lse, - num_splits=ctx.num_splits, - fa_version=ctx.fa_version, - ) - - -# Pick one q/k point per current FA4 MLA-prefill shape regime. -def _shape_probes_for_context( - ctx: FA4MLAPrefillCompileContext, - arch_family: FA4ArchitectureFamily, -) -> tuple[tuple[str, int, int], ...]: - q_stage1_q = 1 - q_stage2_q = FA4_MLA_PREFILL_Q_TILE + 1 - # FA4 never auto-splits when ceil(max_seqlen_k / tile_n) <= 4. - no_split_k = 4 * FA4_MLA_PREFILL_K_TILE - long_k = FA4_MLA_PREFILL_LONG_K_BLOCKS * FA4_MLA_PREFILL_K_TILE - # Diff-head-dim Blackwell Split-KV switches tile_n at 64 K blocks. - very_long_k = FA4_MLA_PREFILL_VERY_LONG_K_BLOCKS * FA4_MLA_PREFILL_K_TILE - - base_probes = ( - ("q_stage1", q_stage1_q, FA4_MLA_PREFILL_K_TILE), - ("q_stage2", q_stage2_q, no_split_k), - ) - # SM120 currently rejects Split-KV in FA4; num_splits=1 also has no split - # shape regimes on any architecture. - if ctx.num_splits == 1 or arch_family == "sm120": - return base_probes - - long_k_probes = ( - ("q_stage1_long_k", q_stage1_q, long_k), - ("q_stage2_long_k", q_stage2_q, long_k), - ) - - # SM90 does not have the SM100 q_stage or diff-head-dim tile_n=64 branch. - # Same-dim SM100-family MLA also does not need the very-long-K probe. - if arch_family == "sm90" or ctx.qk_head_dim == ctx.effective_v_head_dim: - return (*base_probes, *long_k_probes) - - very_long_k_probes = ( - ("q_stage1_very_long_k", q_stage1_q, very_long_k), - ("q_stage2_very_long_k", q_stage2_q, very_long_k), - ) - return (*base_probes, *long_k_probes, *very_long_k_probes) - - -# Check whether this setup can use FA4 MLA prefill. -def _supports_fa4_mla_prefill( - ctx: FA4MLAPrefillCompileContext, - arch_family: FA4ArchitectureFamily, -) -> bool: - return ( - ctx.dtype in FA4_STANDARD_DTYPES - and ctx.num_heads > 0 - and (arch_family != "sm120" or ctx.num_splits == 1) - ) - - -# Map CUDA capability to the FA4 arch family used by warmup checks. -def _fa4_architecture_family_from_compute_capability( - major: int, - minor: int, -) -> FA4ArchitectureFamily: - if (major, minor) == (9, 0): - return "sm90" - if major == 10: - return "sm100f" - if (major, minor) == (12, 0): - return "sm120" - raise ValueError(f"FA4 warmup does not know CUDA capability {major}.{minor}") diff --git a/vllm/model_executor/warmup/fa4_cutedsl_warmup.py b/vllm/model_executor/warmup/fa4_cutedsl_warmup.py new file mode 100644 index 0000000000..cbcca1f38c --- /dev/null +++ b/vllm/model_executor/warmup/fa4_cutedsl_warmup.py @@ -0,0 +1,30 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Warm up FA4 CuTeDSL MLA prefill compile keys.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from vllm.v1.attention.backends.mla.prefill import get_mla_prefill_backend + +if TYPE_CHECKING: + from vllm.v1.worker.gpu_worker import Worker + + +def fa4_cutedsl_warmup(worker: Worker) -> None: + runner = worker.model_runner + if runner.is_pooling_model: + return + + vllm_config = runner.vllm_config + if not vllm_config.model_config.use_mla: + return + + backend_cls = get_mla_prefill_backend(vllm_config) + if backend_cls.get_name() != "FLASH_ATTN": + return + + from vllm.v1.attention.backends.mla.prefill import flash_attn + + flash_attn.FA4_MLA_PREFILL_KERNEL.warmup(vllm_config) diff --git a/vllm/model_executor/warmup/jit_warmup.py b/vllm/model_executor/warmup/jit_warmup.py new file mode 100644 index 0000000000..c5411b5efb --- /dev/null +++ b/vllm/model_executor/warmup/jit_warmup.py @@ -0,0 +1,474 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Shared interfaces and tracing helpers for explicit JIT warmup keys.""" + +from __future__ import annotations + +import ast +import inspect +import itertools +import operator +import textwrap +from abc import ABC, abstractmethod +from collections.abc import Callable, Iterable, Mapping +from dataclasses import dataclass +from typing import Any, Generic, TypeVar + +__all__ = [ + "VllmJitKernel", + "WarmupIntRange", + "get_ast_full_name", + "get_function_source_node", + "zip_inputs", +] + + +CompileKeyT = TypeVar("CompileKeyT") + + +@dataclass(frozen=True) +class WarmupIntRange: + start: int + stop: int + step: int = 1 + + +WarmupValues = Any +CompileKeyDispatchFn = Callable[..., CompileKeyT] + + +@dataclass(frozen=True) +class _WarmupInputRows: + """Warmup dispatch inputs expanded in lockstep.""" + + rows: tuple[Mapping[str, WarmupValues], ...] + + +def _expand_warmup_values(values: WarmupValues) -> tuple[Any, ...]: + if isinstance(values, WarmupIntRange): + return tuple(range(values.start, values.stop, values.step)) + if isinstance(values, (list, tuple)): + return tuple(values) + return (values,) + + +def zip_inputs(*rows: Mapping[str, WarmupValues]) -> _WarmupInputRows: + """Group row-wise dispatch inputs that should be expanded in lockstep.""" + if not rows: + raise ValueError("zip_inputs requires at least one dispatch input row") + if not all(isinstance(row, Mapping) for row in rows): + raise ValueError("zip_inputs rows must be mappings") + + first_names = frozenset(rows[0]) + if not first_names: + raise ValueError("zip_inputs rows require at least one dispatch input name") + if not all(isinstance(name, str) for name in first_names): + raise ValueError("zip_inputs dispatch input names must be strings") + + input_rows: list[Mapping[str, WarmupValues]] = [] + for row in rows: + names = frozenset(row) + if names != first_names: + raise ValueError("zip_inputs rows must use the same dispatch input names") + input_rows.append(dict(row)) + + return _WarmupInputRows(rows=tuple(input_rows)) + + +def _expand_warmup_value_grid( + values: Mapping[str, WarmupValues], + input_names: frozenset[str], +) -> tuple[dict[str, Any], ...]: + names = tuple(name for name in values if name in input_names) + if not names: + return ({},) + + expanded_values = tuple(_expand_warmup_values(values[name]) for name in names) + return tuple( + dict(zip(names, value_set)) for value_set in itertools.product(*expanded_values) + ) + + +def _expand_warmup_input_rows( + rows: tuple[Mapping[str, WarmupValues], ...], + input_names: frozenset[str], +) -> tuple[dict[str, Any], ...]: + active_names = frozenset(name for name in rows[0] if name in input_names) + if not active_names: + return ({},) + + return tuple( + {name: value for name, value in row.items() if name in active_names} + for row in rows + ) + + +def _merge_warmup_kwargs(parts: Iterable[Mapping[str, Any]]) -> dict[str, Any]: + merged: dict[str, Any] = {} + for part in parts: + for name, value in part.items(): + if name in merged: + raise ValueError( + f"Warmup dispatch input '{name}' is specified more than once" + ) + merged[name] = value + return merged + + +@dataclass(frozen=True) +class _CompileKeyDispatchTrace: + local_exprs: tuple[tuple[str, ast.AST], ...] + field_exprs: tuple[tuple[str, ast.AST], ...] + globals: Mapping[str, Any] + input_names: frozenset[str] + defaults: Mapping[str, Any] + + def compile_key( + self, + compile_key_type: type[CompileKeyT], + kwargs: Mapping[str, Any], + ) -> CompileKeyT: + dispatch_values = {**self.defaults, **kwargs} + for name, expr in self.local_exprs: + dispatch_values[name] = _eval_dispatch_expr( + expr, dispatch_values, self.globals + ) + return compile_key_type( + **{ + field: _eval_dispatch_expr(expr, dispatch_values, self.globals) + for field, expr in self.field_exprs + } + ) + + +_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, + ast.Pow: operator.pow, +} +_CMP_OPS: dict[type[ast.cmpop], Callable[[Any, Any], bool]] = { + ast.Eq: operator.eq, + ast.NotEq: operator.ne, + ast.Lt: operator.lt, + ast.LtE: operator.le, + ast.Gt: operator.gt, + ast.GtE: operator.ge, +} + + +def _dispatch_expr_source(node: ast.AST) -> str: + try: + return ast.unparse(node) + except Exception: + return ast.dump(node) + + +def _dispatch_expr_error(node: ast.AST, reason: str) -> ValueError: + return ValueError( + f"{reason}: {_dispatch_expr_source(node)}. " + "Supported dispatch expressions are names, constants, attributes, " + "tuple/list literals, conditional expressions, comparisons, boolean " + "operators, unary not/minus, arithmetic, and calls without **kwargs." + ) + + +class _DispatchExprEvaluator(ast.NodeVisitor): + def __init__( + self, + values: Mapping[str, Any], + globals_: Mapping[str, Any], + ) -> None: + self.values = values + self.globals = globals_ + + def eval(self, node: ast.AST) -> Any: + return self.visit(node) + + def generic_visit(self, node: ast.AST) -> Any: + raise _dispatch_expr_error(node, "Unsupported dispatch expression") + + def visit_Name(self, node: ast.Name) -> Any: + if node.id in self.values: + return self.values[node.id] + if node.id in self.globals: + return self.globals[node.id] + raise _dispatch_expr_error(node, f"Unknown dispatch name '{node.id}'") + + def visit_Constant(self, node: ast.Constant) -> Any: + return node.value + + def visit_IfExp(self, node: ast.IfExp) -> Any: + return self.visit(node.body if self.visit(node.test) else node.orelse) + + def visit_Tuple(self, node: ast.Tuple) -> tuple[Any, ...]: + return tuple(self.visit(elt) for elt in node.elts) + + def visit_List(self, node: ast.List) -> list[Any]: + return [self.visit(elt) for elt in node.elts] + + def visit_BoolOp(self, node: ast.BoolOp) -> Any: + if isinstance(node.op, ast.And): + result = None + for value in node.values: + result = self.visit(value) + if not result: + return result + return result + if isinstance(node.op, ast.Or): + result = None + for value in node.values: + result = self.visit(value) + if result: + return result + return result + raise _dispatch_expr_error(node, "Unsupported dispatch boolean operator") + + def visit_Compare(self, node: ast.Compare) -> bool: + left = self.visit(node.left) + for op_node, comparator in zip(node.ops, node.comparators): + right = self.visit(comparator) + op = _CMP_OPS.get(type(op_node)) + if op is None: + raise _dispatch_expr_error( + node, "Unsupported dispatch comparison operator" + ) + if not op(left, right): + return False + left = right + return True + + def visit_UnaryOp(self, node: ast.UnaryOp) -> Any: + operand = self.visit(node.operand) + if isinstance(node.op, ast.Not): + return not operand + if isinstance(node.op, ast.USub): + return -operand + raise _dispatch_expr_error(node, "Unsupported dispatch unary operator") + + def visit_BinOp(self, node: ast.BinOp) -> Any: + op = _BIN_OPS.get(type(node.op)) + if op is None: + raise _dispatch_expr_error(node, "Unsupported dispatch binary operator") + return op(self.visit(node.left), self.visit(node.right)) + + def visit_Call(self, node: ast.Call) -> Any: + args = [self.visit(arg) for arg in node.args] + fn = self.visit(node.func) + call_kwargs: dict[str, Any] = {} + for keyword in node.keywords: + if keyword.arg is None: + raise _dispatch_expr_error( + node, "Dispatch helper calls cannot use **kwargs" + ) + call_kwargs[keyword.arg] = self.visit(keyword.value) + return fn(*args, **call_kwargs) + + def visit_Attribute(self, node: ast.Attribute) -> Any: + return getattr(self.visit(node.value), node.attr) + + +def get_ast_full_name(node: ast.AST) -> str | None: + if isinstance(node, ast.Name): + return node.id + if isinstance(node, ast.Attribute): + parent = get_ast_full_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: + return _DispatchExprEvaluator(kwargs, globals_).eval(node) + + +def _collect_input_names( + node: ast.AST, + candidate_names: set[str], + local_names: set[str] | None = None, +) -> set[str]: + if local_names is None: + local_names = set() + return { + child.id + for child in ast.walk(node) + if ( + isinstance(child, ast.Name) + and child.id in candidate_names + and child.id not in local_names + ) + } + + +def _collect_dispatch_body( + fn: CompileKeyDispatchFn[Any], + function_def: ast.FunctionDef, +) -> tuple[list[tuple[str, ast.AST]], ast.Call]: + local_exprs: list[tuple[str, ast.AST]] = [] + for statement in function_def.body: + if ( + isinstance(statement, ast.Expr) + and isinstance(statement.value, ast.Constant) + and isinstance(statement.value.value, str) + ): + continue + + if isinstance(statement, ast.Assign): + if len(statement.targets) != 1 or not isinstance( + statement.targets[0], ast.Name + ): + raise _dispatch_expr_error( + statement, "Dispatch assignments must target one local name" + ) + local_exprs.append((statement.targets[0].id, statement.value)) + continue + + if isinstance(statement, ast.AnnAssign): + if statement.value is None: + raise _dispatch_expr_error( + statement, "Dispatch annotations must assign a value" + ) + if not isinstance(statement.target, ast.Name): + raise _dispatch_expr_error( + statement, "Dispatch assignments must target one local name" + ) + local_exprs.append((statement.target.id, statement.value)) + continue + + if isinstance(statement, ast.Return) and isinstance(statement.value, ast.Call): + return local_exprs, statement.value + + if isinstance(statement, ast.Return): + raise _dispatch_expr_error( + statement, "Dispatch must return one CompileKey(...) call" + ) + + raise _dispatch_expr_error( + statement, + "Dispatch may only contain local assignments before CompileKey return", + ) + + raise ValueError(f"Expected {fn.__name__} to return one CompileKey(...) call") + + +def _trace_compile_key_dispatch( + fn: CompileKeyDispatchFn[Any], +) -> _CompileKeyDispatchTrace: + source_fn = getattr(fn, "__func__", fn) + globals_ = source_fn.__globals__ + function_def = get_function_source_node(fn) + + local_exprs, return_call = _collect_dispatch_body(fn, function_def) + + field_exprs: list[tuple[str, ast.AST]] = [] + signature = inspect.signature(fn) + defaults = { + name: parameter.default + for name, parameter in signature.parameters.items() + if parameter.default is not inspect.Parameter.empty + } + candidate_names = set(signature.parameters) + input_names: set[str] = set() + local_names = {name for name, _ in local_exprs} + for _, expr in local_exprs: + input_names.update(_collect_input_names(expr, candidate_names)) + for keyword in return_call.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, local_names) + ) + + return _CompileKeyDispatchTrace( + tuple(local_exprs), + tuple(field_exprs), + globals_, + frozenset(input_names), + defaults, + ) + + +class VllmJitKernel(Generic[CompileKeyT], ABC): + """Kernel wrapper that owns dispatch, warmup keys, and compilation.""" + + CompileKey: type[CompileKeyT] + + def __init__(self) -> None: + self.compile_key_dispatch_trace = _trace_compile_key_dispatch(self.dispatch) + + def compile_key(self, kwargs: Mapping[str, Any]) -> CompileKeyT: + return self.compile_key_dispatch_trace.compile_key(self.CompileKey, kwargs) + + def _trace_dispatch( + self, dispatch: CompileKeyDispatchFn[CompileKeyT] + ) -> Callable[..., list[CompileKeyT]]: + compile_key_dispatch_trace = _trace_compile_key_dispatch(dispatch) + + def traced( + *input_groups: _WarmupInputRows, + **kwargs: WarmupValues, + ) -> list[CompileKeyT]: + for group in input_groups: + if not isinstance(group, _WarmupInputRows): + raise TypeError( + "_trace_dispatch positional arguments must be " + "zip_inputs(...) groups" + ) + expanded_input_groups = tuple( + _expand_warmup_input_rows( + group.rows, compile_key_dispatch_trace.input_names + ) + for group in input_groups + ) + expanded_kwargs = _expand_warmup_value_grid( + kwargs, compile_key_dispatch_trace.input_names + ) + dispatch_value_groups = (*expanded_input_groups, expanded_kwargs) + return list( + dict.fromkeys( + compile_key_dispatch_trace.compile_key( + self.CompileKey, _merge_warmup_kwargs(dispatch_values) + ) + for dispatch_values in itertools.product(*dispatch_value_groups) + ) + ) + + return traced + + @abstractmethod + def dispatch(self, **kwargs: Any) -> CompileKeyT: + """Build one compile key from one concrete dispatch point.""" + raise NotImplementedError + + @abstractmethod + def get_warmup_keys(self, *args: Any, **kwargs: Any) -> list[CompileKeyT]: + """Return compile keys that should be warmed for this kernel.""" + raise NotImplementedError + + @abstractmethod + def compile(self, compile_key: CompileKeyT) -> None: + """Compile one warmup key.""" + raise NotImplementedError + + def warmup(self, *args: Any, **kwargs: Any) -> None: + """Compile this kernel's warmup keys.""" + for compile_key in self.get_warmup_keys(*args, **kwargs): + self.compile(compile_key) diff --git a/vllm/model_executor/warmup/jit_warmup_triton_helper.py b/vllm/model_executor/warmup/jit_warmup_triton_helper.py new file mode 100644 index 0000000000..b909757621 --- /dev/null +++ b/vllm/model_executor/warmup/jit_warmup_triton_helper.py @@ -0,0 +1,183 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import ast +import inspect +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any + +from vllm.model_executor.warmup.jit_warmup import ( + get_ast_full_name, + get_function_source_node, +) + + +@dataclass(frozen=True) +class TritonWarmupTensor: + # Compile-only tensor descriptor for Triton pointer specialization. + dtype: Any + aligned: bool = True + shape: tuple[int, ...] = (1,) + + def data_ptr(self) -> int: + return 0 if self.aligned else 1 + + def ptr_range(self) -> int: + return 0 + + def stride(self) -> tuple[int, ...]: + strides: list[int] = [] + stride = 1 + for size in reversed(self.shape): + strides.append(stride) + stride *= size + return tuple(reversed(strides)) + + +@dataclass(frozen=True) +class TritonPointerInputVariant: + # Named pointer-alignment variant for compile-only Triton warmup. + alignments: tuple[tuple[str, bool], ...] + + @classmethod + def from_alignment(cls, **aligned: bool) -> "TritonPointerInputVariant": + return cls(tuple(aligned.items())) + + def is_aligned(self, name: str) -> bool: + for alignment_name, aligned in self.alignments: + if alignment_name == name: + return aligned + raise KeyError(f"Unknown Triton pointer input variant: {name}") + + def pointer( + self, + name: str, + dtype: Any, + shape: tuple[int, ...] = (1,), + ) -> TritonWarmupTensor: + return TritonWarmupTensor(dtype, aligned=self.is_aligned(name), shape=shape) + + +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 = get_ast_full_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 = get_ast_full_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 get_ast_full_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) + ) diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index d4b9a735d2..e21ab112e7 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/vllm/model_executor/warmup/kernel_warmup.py @@ -17,6 +17,9 @@ from vllm.model_executor.warmup.deep_gemm_warmup import deep_gemm_warmup from vllm.model_executor.warmup.deepseek_v4_mhc_warmup import ( deepseek_v4_mhc_warmup, ) +from vllm.model_executor.warmup.fa4_cutedsl_warmup import ( + fa4_cutedsl_warmup, +) from vllm.model_executor.warmup.flashinfer_autotune_cache import ( resolve_flashinfer_autotune_file, write_flashinfer_autotune_cache, @@ -27,7 +30,7 @@ from vllm.model_executor.warmup.flashinfer_sparse_mla_warmup import ( ) from vllm.model_executor.warmup.qwen_triton_warmup import qwen_triton_warmup from vllm.model_executor.warmup.sparse_mla_triton_warmup import ( - sparse_mla_triton_warmup_if_needed, + sparse_mla_triton_warmup, ) from vllm.model_executor.warmup.v1_block_table_warmup import ( warm_v1_block_table_kernels, @@ -94,7 +97,6 @@ def kernel_warmup(worker: "Worker"): ) # Run next so input-prep kernels JIT against pristine runner state. - sparse_mla_triton_warmup_if_needed(worker) flashinfer_sparse_mla_decode_autotune_warmup(worker) deepseek_v4_sparse_mla_attention_warmup(worker) @@ -156,8 +158,15 @@ def kernel_warmup(worker: "Worker"): ) if worker.vllm_config.kernel_config.enable_cutedsl_warmup: + # TODO(roberto): Remove after registered CuTeDSL warmups are migrated + # to the shared JIT warmup infrastructure. + # https://github.com/vllm-project/vllm/pull/47451 cutedsl_warmup() + if worker.vllm_config.kernel_config.enable_jit_warmup: + fa4_cutedsl_warmup(worker) + sparse_mla_triton_warmup(worker) + def _flashinfer_autotune_skip_ops(runner: "GPUModelRunner") -> set[str] | None: if envs.VLLM_FLASHINFER_AUTOTUNE_SKIP_OPS is not None: diff --git a/vllm/model_executor/warmup/sparse_mla_triton_warmup.py b/vllm/model_executor/warmup/sparse_mla_triton_warmup.py index 185449cb8e..cf35bf40e8 100644 --- a/vllm/model_executor/warmup/sparse_mla_triton_warmup.py +++ b/vllm/model_executor/warmup/sparse_mla_triton_warmup.py @@ -4,11 +4,10 @@ from typing import TYPE_CHECKING -import torch - from vllm.logger import init_logger if TYPE_CHECKING: + from vllm.config import VllmConfig from vllm.v1.worker.gpu_model_runner import GPUModelRunner from vllm.v1.worker.gpu_worker import Worker @@ -31,47 +30,7 @@ _GENERIC_SPARSE_MLA_BACKENDS = frozenset( ) _INDEXER_PREFILL_CHUNK_METADATA_BACKENDS = frozenset({"DEEPSEEK_V32_INDEXER"}) -_SPARSE_PREFILL_METADATA_NUM_PREFILLS = (1, 2, 4, 8) -_SPARSE_PREFILL_METADATA_NUM_DECODES = (0, 1, 2) -_DSV4_PREFILL_CHUNK_METADATA_COMPRESS_RATIOS = (4, 128) -_PREFILL_CHUNK_METADATA_SEQ_LEN_MULTIPLIERS = (2, 3) -_PREFILL_CHUNK_METADATA_QUERY_SLICE_OFFSETS = ( - # query_slice_start offset, query_slice_stop offset - (0, 0), - (0, -1), - (1, 0), - (1, -1), -) -_COMBINE_TOPK_SWA_INPUT_VARIANTS = ( - # offset_topk, offset_query_and_seq, offset_gather - (False, False, False), - (False, True, False), - (True, True, True), -) -_DSV4_COMBINE_TOPK_SWA_WARMUP_CASES = ( - # compress_ratio, topk, topk_width, N - (1, 0, 512, 512), - (4, 512, 512, 512 * 4), - # DSv4-Pro C4A traffic uses top-k 1024 with N=1024. - (4, 1024, 1024, 1024), - (128, 8192, 8192, 8192 * 128), - # Real C128A traffic also specializes N=1 in one call path. - (128, 8192, 8192, 1), -) - - -def _clamp_warmup_tokens(num_tokens: int, max_tokens: int) -> int: - return max(0, min(num_tokens, max_tokens)) - - -def _next_power_of_2(x: int) -> int: - return 1 << (x - 1).bit_length() - - -def _hf_config_int(runner: "GPUModelRunner", name: str, default: int) -> int: - model_config = getattr(runner.vllm_config, "model_config", None) - hf_config = getattr(model_config, "hf_config", None) - return int(getattr(hf_config, name, default) or default) +_INDEXER_PREFILL_CHUNK_METADATA_BACKENDS = frozenset({"DEEPSEEK_V32_INDEXER"}) def _attention_backend_name(backend: object) -> str | None: @@ -96,242 +55,57 @@ def _has_attention_backend( return False -def _warm_sparse_swa_prefill_metadata_kernel( - device: torch.device, - window_size: int, - prefill_tokens: int, +def _compile_sparse_swa_prefill_metadata_kernel( + vllm_config: "VllmConfig", ) -> 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: - for num_decodes in _SPARSE_PREFILL_METADATA_NUM_DECODES: - query_lens = [1] * num_decodes - query_lens += [prefill_tokens] * num_prefills - query_start_locs = [0] - for query_len in query_lens: - query_start_locs.append(query_start_locs[-1] + query_len) - query_start_loc = torch.tensor( - query_start_locs, - dtype=torch.int32, - device=device, - ) - seq_lens = torch.tensor( - [1] * num_decodes + [window_size + q for q in query_lens[num_decodes:]], - dtype=torch.int32, - device=device, - ) - 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.warmup(vllm_config) -def _warm_prefill_chunk_metadata_kernel( - device: torch.device, - compress_ratio: int, - query_len: int, +def _compile_prefill_chunk_metadata_kernel( + vllm_config: "VllmConfig", ) -> None: - from vllm.v1.attention.backends.mla.indexer import build_prefill_chunk_metadata - - num_reqs = 2 - query_start_loc_cpu = torch.arange( - 0, (num_reqs + 1) * query_len, query_len, dtype=torch.int32 - ) - query_start_loc = query_start_loc_cpu.to(device=device) - - uncompressed_seq_lens_cpu = torch.tensor( - [ - compress_ratio * multiplier + query_len - for multiplier in _PREFILL_CHUNK_METADATA_SEQ_LEN_MULTIPLIERS - ], - dtype=torch.int32, - ) - compressed_seq_lens_cpu = uncompressed_seq_lens_cpu // compress_ratio - uncompressed_seq_lens = uncompressed_seq_lens_cpu.to(device=device) - compressed_seq_lens = compressed_seq_lens_cpu.to(device=device) - block_table = torch.zeros( - (num_reqs, int(compressed_seq_lens_cpu.max().item())), - dtype=torch.int32, - device=device, + from vllm.v1.attention.backends.mla.indexer import ( + _BUILD_PREFILL_CHUNK_METADATA_KERNEL, ) - 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 ( - 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, - ) + _BUILD_PREFILL_CHUNK_METADATA_KERNEL.warmup(vllm_config) -def _warm_combine_topk_swa_indices_kernel( - device: torch.device, - num_tokens: int, - window_size: int, - compress_ratio: int, - topk: int, - topk_width: int, - n: int, +def _compile_combine_topk_swa_indices_kernel( + vllm_config: "VllmConfig", ) -> None: - from vllm.models.deepseek_v4.common.ops.cache_utils import combine_topk_swa_indices - - if num_tokens <= 0: - return - - def _make_topk_indices(*, offset: bool) -> torch.Tensor: - if offset: - topk_storage = torch.full( - (num_tokens * topk_width + 1,), - -1, - dtype=torch.int32, - device=device, - ) - topk_indices = topk_storage[1:].reshape(num_tokens, topk_width) - else: - topk_indices = torch.full( - (num_tokens, topk_width), -1, dtype=torch.int32, device=device - ) - if topk > 0: - topk_indices.copy_( - torch.arange(num_tokens * topk_width, dtype=torch.int32, device=device) - .reshape(num_tokens, topk_width) - .remainder(topk_width) - ) - return topk_indices - - query_start_loc = torch.tensor([0, num_tokens], dtype=torch.int32, device=device) - seq_lens = torch.tensor( - [window_size + num_tokens], dtype=torch.int32, device=device - ) - gather_lens = torch.tensor( - [min(window_size + num_tokens, window_size + num_tokens - 1)], - dtype=torch.int32, - device=device, - ) - offset_query_start_loc = torch.empty(3, dtype=torch.int32, device=device)[1:] - offset_query_start_loc.copy_(query_start_loc) - offset_seq_lens = torch.empty(2, dtype=torch.int32, device=device)[1:] - offset_seq_lens.copy_(seq_lens) - offset_gather_lens = torch.empty(2, dtype=torch.int32, device=device)[1:] - offset_gather_lens.copy_(gather_lens) - - for ( - offset_topk, - offset_query_and_seq, - offset_gather, - ) in _COMBINE_TOPK_SWA_INPUT_VARIANTS: - warmup_topk_indices = _make_topk_indices(offset=offset_topk) - warmup_query_start_loc = ( - offset_query_start_loc if offset_query_and_seq else query_start_loc - ) - warmup_seq_lens = offset_seq_lens if offset_query_and_seq else seq_lens - warmup_gather_lens = offset_gather_lens if offset_gather else gather_lens - n_values = (n,) if n == 1 else (n, n + 1) - for m in (window_size + num_tokens, topk_width): - for n_value in n_values: - combine_topk_swa_indices( - warmup_topk_indices, - warmup_query_start_loc, - warmup_seq_lens, - warmup_gather_lens, - window_size, - compress_ratio, - topk, - M=m, - N=n_value, - ) - - -@torch.inference_mode() -def sparse_mla_triton_warmup( - runner: "GPUModelRunner", - num_tokens: int, - *, - compress_ratios: tuple[int, ...], - combine_topk_swa_cases: tuple[tuple[int, int, int, int], ...] = (), -) -> None: - 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) - for compress_ratio, topk, topk_width, n in combine_topk_swa_cases: - _warm_combine_topk_swa_indices_kernel( - device, - num_tokens, - window_size, - compress_ratio, - topk, - topk_width, - n, - ) - - -def deepseek_v4_sparse_triton_warmup( - runner: "GPUModelRunner", - num_tokens: int, -) -> None: - sparse_mla_triton_warmup( - runner, - num_tokens, - compress_ratios=_DSV4_PREFILL_CHUNK_METADATA_COMPRESS_RATIOS, - combine_topk_swa_cases=_DSV4_COMBINE_TOPK_SWA_WARMUP_CASES, + from vllm.models.deepseek_v4.common.ops.cache_utils import ( + _COMBINE_TOPK_SWA_INDICES_KERNEL, ) + _COMBINE_TOPK_SWA_INDICES_KERNEL.warmup(vllm_config) -def sparse_mla_triton_warmup_if_needed(worker: "Worker") -> None: + +def sparse_mla_triton_warmup(worker: "Worker") -> None: runner = worker.model_runner if runner.is_pooling_model: return max_tokens = worker.scheduler_config.max_num_batched_tokens - num_tokens = _clamp_warmup_tokens(8, max_tokens) - if num_tokens <= 0: + max_num_prefills = min(worker.scheduler_config.max_num_seqs, max_tokens) + if max_tokens <= 0 or max_num_prefills <= 0: return + vllm_config = runner.vllm_config try: if _has_attention_backend(runner, _DEEPSEEK_V4_SPARSE_MLA_BACKENDS): - deepseek_v4_sparse_triton_warmup(runner, num_tokens) + _compile_sparse_swa_prefill_metadata_kernel(vllm_config) + _compile_prefill_chunk_metadata_kernel(vllm_config) + _compile_combine_topk_swa_indices_kernel(vllm_config) elif _has_attention_backend(runner, _GENERIC_SPARSE_MLA_BACKENDS): - sparse_mla_triton_warmup( - runner, - num_tokens, - compress_ratios=(1,), - ) + _compile_sparse_swa_prefill_metadata_kernel(vllm_config) + _compile_prefill_chunk_metadata_kernel(vllm_config) elif _has_attention_backend(runner, _INDEXER_PREFILL_CHUNK_METADATA_BACKENDS): - _warm_prefill_chunk_metadata_kernel( - getattr(runner, "device", torch.device("cuda")), - compress_ratio=1, - query_len=num_tokens, - ) + _compile_prefill_chunk_metadata_kernel(vllm_config) + except Exception: logger.warning("Skipping sparse MLA Triton warmup.", exc_info=True) diff --git a/vllm/models/deepseek_v4/common/ops/cache_utils.py b/vllm/models/deepseek_v4/common/ops/cache_utils.py index 55106d06af..fc6e4ee330 100644 --- a/vllm/models/deepseek_v4/common/ops/cache_utils.py +++ b/vllm/models/deepseek_v4/common/ops/cache_utils.py @@ -14,14 +14,23 @@ preparation. window indices for sparse prefill. """ +from dataclasses import dataclass +from typing import Any + import torch from vllm.model_executor.layers.quantization.utils.quant_utils import ( get_fp8_min_max, ) +from vllm.model_executor.warmup.jit_warmup import VllmJitKernel, zip_inputs +from vllm.model_executor.warmup.jit_warmup_triton_helper import ( + TritonPointerInputVariant, + TritonWarmupTensor, +) from vllm.platforms import current_platform from vllm.triton_utils import tl, triton from vllm.utils.import_utils import has_cutedsl +from vllm.utils.math_utils import next_power_of_2 @triton.jit @@ -528,7 +537,6 @@ def combine_topk_swa_indices( N: int, ) -> tuple[torch.Tensor, torch.Tensor]: num_tokens = topk_indices.shape[0] - num_reqs = seq_lens.shape[0] combined_topk = ( (topk + window_size + _SPARSE_PREFILL_TOPK_ALIGNMENT - 1) // _SPARSE_PREFILL_TOPK_ALIGNMENT @@ -544,13 +552,10 @@ def combine_topk_swa_indices( num_tokens, dtype=torch.int32, device=topk_indices.device ) - NUM_WORKERS = 128 - _combine_topk_swa_indices_kernel[(num_reqs, NUM_WORKERS)]( + _COMBINE_TOPK_SWA_INDICES_KERNEL( combined_indices, - combined_indices.stride(0), combined_lens, topk_indices, - topk_indices.stride(0), query_start_loc, seq_lens, gather_lens, @@ -559,82 +564,245 @@ def combine_topk_swa_indices( TOP_K=topk, COMPRESS_RATIO=compress_ratio, WINDOW_SIZE=window_size, - PADDED_TOP_K=triton.next_power_of_2(topk_indices.shape[-1]), ) return combined_indices, combined_lens -@triton.jit -def _combine_topk_swa_indices_kernel( - combined_indices_ptr, - combined_indices_stride, - combined_lens_ptr, - topk_indices_ptr, - topk_indices_stride, - query_start_loc_ptr, - seq_lens_ptr, - gather_lens_ptr, - M, - N, - TOP_K: tl.constexpr, - COMPRESS_RATIO: tl.constexpr, - WINDOW_SIZE: tl.constexpr, - PADDED_TOP_K: tl.constexpr, +_COMBINE_TOPK_SWA_NUM_WORKERS = 128 + + +# Representative pointer alignment variants for Triton pointer specialization. +_COMBINE_TOPK_SWA_POINTER_INPUTS = zip_inputs( + dict( + topk_indices=True, + query_start_loc=True, + seq_lens=True, + gather_lens=True, + ), + dict( + topk_indices=True, + query_start_loc=False, + seq_lens=False, + gather_lens=True, + ), + dict( + topk_indices=False, + query_start_loc=False, + seq_lens=False, + gather_lens=False, + ), +) + + +_DSV4_COMBINE_TOPK_SWA_WARMUP_INPUTS = zip_inputs( + # DSv4-Flash / SWA-only and C4A. + dict(compress_ratio=1, topk=0, topk_width=512), + dict(compress_ratio=4, topk=512, topk_width=512), + # DSv4-Pro C4A. + dict(compress_ratio=4, topk=1024, topk_width=1024), + # DSv4-Pro C128A. + dict(compress_ratio=128, topk=8192, topk_width=8192), +) + + +def _hf_config_int(vllm_config: Any, name: str, default: int) -> int: + model_config = getattr(vllm_config, "model_config", None) + hf_config = getattr(model_config, "hf_config", None) + return int(getattr(hf_config, name, default) or default) + + +def _scheduler_config_int(vllm_config: Any, name: str, default: int) -> int: + scheduler_config = getattr(vllm_config, "scheduler_config", None) + return int(getattr(scheduler_config, name, default) or default) + + +class CombineTopkSwaIndicesKernel( + VllmJitKernel["CombineTopkSwaIndicesKernel.CompileKey"] ): - batch_idx = tl.program_id(0) - worker_id = tl.program_id(1) - num_workers = tl.num_programs(1) + @dataclass(frozen=True) + class CompileKey: + TOP_K: int + COMPRESS_RATIO: int + WINDOW_SIZE: int + PADDED_TOP_K: int + input_variant: TritonPointerInputVariant - # query_start_loc is a global tensor; rebase to chunk-local offsets - # by subtracting the chunk's starting value. - base = tl.load(query_start_loc_ptr) - query_start = tl.load(query_start_loc_ptr + batch_idx) - base - query_end = tl.load(query_start_loc_ptr + batch_idx + 1) - base - query_len = query_end - query_start - seq_len = tl.load(seq_lens_ptr + batch_idx) - gather_len = tl.load(gather_lens_ptr + batch_idx) - start_pos = seq_len - query_len - # The SWA portion of the gathered buffer starts from position - # (seq_len - gather_len), not position 0. We need this offset - # to correctly index into the gathered buffer. - gather_start = seq_len - gather_len + @staticmethod + @triton.jit( + do_not_specialize=[ + "combined_indices_stride", + "topk_indices_stride", + "M", + "N", + ] + ) + def kernel( + combined_indices_ptr, + combined_indices_stride, + combined_lens_ptr, + topk_indices_ptr, + topk_indices_stride, + query_start_loc_ptr, + seq_lens_ptr, + gather_lens_ptr, + M, + N, + TOP_K: tl.constexpr, + COMPRESS_RATIO: tl.constexpr, + WINDOW_SIZE: tl.constexpr, + PADDED_TOP_K: tl.constexpr, + ): + batch_idx = tl.program_id(0) + worker_id = tl.program_id(1) + num_workers = tl.num_programs(1) - for token_idx in range(query_start + worker_id, query_end, num_workers): - # topk_len is fully determined by the query token's absolute position: - # both the C4A indexer and the C128A metadata builder emit - # min((pos + 1) // compress_ratio, topk_tokens) valid entries. - # Caller passes TOP_K=0 for SWA-only layers to zero this out. - token_idx_in_query = token_idx - query_start - pos = start_pos + token_idx_in_query - topk_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOP_K) - swa_len = tl.minimum(pos + 1, WINDOW_SIZE) + # query_start_loc is a global tensor; rebase to chunk-local offsets + # by subtracting the chunk's starting value. + base = tl.load(query_start_loc_ptr) + query_start = tl.load(query_start_loc_ptr + batch_idx) - base + query_end = tl.load(query_start_loc_ptr + batch_idx + 1) - base + query_len = query_end - query_start + seq_len = tl.load(seq_lens_ptr + batch_idx) + gather_len = tl.load(gather_lens_ptr + batch_idx) + start_pos = seq_len - query_len + # The SWA portion of the gathered buffer starts from position + # (seq_len - gather_len), not position 0. We need this offset + # to correctly index into the gathered buffer. + gather_start = seq_len - gather_len - offset = tl.arange(0, PADDED_TOP_K) - mask = offset < topk_len - topk_indices = tl.load( - topk_indices_ptr + token_idx * topk_indices_stride + offset, - mask=mask, + for token_idx in range(query_start + worker_id, query_end, num_workers): + # topk_len is fully determined by the query token's absolute position: + # both the C4A indexer and the C128A metadata builder emit + # min((pos + 1) // compress_ratio, topk_tokens) valid entries. + # Caller passes TOP_K=0 for SWA-only layers to zero this out. + token_idx_in_query = token_idx - query_start + pos = start_pos + token_idx_in_query + topk_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOP_K) + swa_len = tl.minimum(pos + 1, WINDOW_SIZE) + + offset = tl.arange(0, PADDED_TOP_K) + mask = offset < topk_len + topk_indices = tl.load( + topk_indices_ptr + token_idx * topk_indices_stride + offset, + mask=mask, + ) + tl.store( + combined_indices_ptr + token_idx * combined_indices_stride + offset, + topk_indices + M * batch_idx, + mask=mask, + ) + offset = tl.arange(0, WINDOW_SIZE) + # Index into gathered buffer: N + (position - gather_start) + # For positions [pos - swa_len + 1, pos], the buffer indices are: + # [N + pos - swa_len + 1 - gather_start, N + pos - gather_start] + tl.store( + combined_indices_ptr + + token_idx * combined_indices_stride + + topk_len + + offset, + M * batch_idx + N + offset + pos - swa_len + 1 - gather_start, + mask=offset < swa_len, + ) + + combined_len = topk_len + swa_len + tl.store(combined_lens_ptr + token_idx, combined_len) + + def dispatch( # type: ignore[override] + self, + *, + topk_width: int, + topk_indices: bool, + query_start_loc: bool, + seq_lens: bool, + gather_lens: bool, + topk: int, + compress_ratio: int, + WINDOW_SIZE: int, + ) -> CompileKey: + padded_topk = next_power_of_2(topk_width) + input_variant = TritonPointerInputVariant.from_alignment( + topk_indices=topk_indices, + query_start_loc=query_start_loc, + seq_lens=seq_lens, + gather_lens=gather_lens, ) - tl.store( - combined_indices_ptr + token_idx * combined_indices_stride + offset, - topk_indices + M * batch_idx, - mask=mask, - ) - offset = tl.arange(0, WINDOW_SIZE) - # Index into gathered buffer: N + (position - gather_start) - # For positions [pos - swa_len + 1, pos], the buffer indices are: - # [N + pos - swa_len + 1 - gather_start, N + pos - gather_start] - tl.store( - combined_indices_ptr - + token_idx * combined_indices_stride - + topk_len - + offset, - M * batch_idx + N + offset + pos - swa_len + 1 - gather_start, - mask=offset < swa_len, + return self.CompileKey( + TOP_K=topk, + COMPRESS_RATIO=compress_ratio, + WINDOW_SIZE=WINDOW_SIZE, + PADDED_TOP_K=padded_topk, + input_variant=input_variant, ) - combined_len = topk_len + swa_len - tl.store(combined_lens_ptr + token_idx, combined_len) + def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]: + if _scheduler_config_int(vllm_config, "max_num_batched_tokens", 0) <= 0: + return [] + + window_size = _hf_config_int(vllm_config, "sliding_window", 128) + return self._trace_dispatch(self.dispatch)( + _DSV4_COMBINE_TOPK_SWA_WARMUP_INPUTS, + _COMBINE_TOPK_SWA_POINTER_INPUTS, + WINDOW_SIZE=window_size, + ) + + def compile(self, compile_key: CompileKey) -> None: + warmup = getattr(self.kernel, "warmup", None) + assert warmup is not None + int32_ptr = TritonWarmupTensor(torch.int32) + input_variant = compile_key.input_variant + warmup( + int32_ptr, + 1, # do not specialize combined_indices_stride + int32_ptr, + input_variant.pointer("topk_indices", torch.int32), + 1, # do not specialize topk_indices_stride + input_variant.pointer("query_start_loc", torch.int32), + input_variant.pointer("seq_lens", torch.int32), + input_variant.pointer("gather_lens", torch.int32), + 1, # do not specialize M + 1, # do not specialize N + TOP_K=compile_key.TOP_K, + COMPRESS_RATIO=compile_key.COMPRESS_RATIO, + WINDOW_SIZE=compile_key.WINDOW_SIZE, + PADDED_TOP_K=compile_key.PADDED_TOP_K, + grid=(1, _COMBINE_TOPK_SWA_NUM_WORKERS), + ) + + def __call__( + self, + combined_indices: torch.Tensor, + combined_lens: torch.Tensor, + topk_indices: torch.Tensor, + query_start_loc: torch.Tensor, + seq_lens: torch.Tensor, + gather_lens: torch.Tensor, + M: int, + N: int, + *, + TOP_K: int, + COMPRESS_RATIO: int, + WINDOW_SIZE: int, + ) -> None: + num_reqs = seq_lens.shape[0] + self.kernel[(num_reqs, _COMBINE_TOPK_SWA_NUM_WORKERS)]( + combined_indices, + combined_indices.stride(0), + combined_lens, + topk_indices, + topk_indices.stride(0), + query_start_loc, + seq_lens, + gather_lens, + M, + N, + TOP_K=TOP_K, + COMPRESS_RATIO=COMPRESS_RATIO, + WINDOW_SIZE=WINDOW_SIZE, + PADDED_TOP_K=next_power_of_2(topk_indices.shape[-1]), + ) + + +_COMBINE_TOPK_SWA_INDICES_KERNEL = CombineTopkSwaIndicesKernel() def build_flashinfer_mixed_sparse_indices( diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index bf76cbeda8..fe8bfef530 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -8,6 +8,14 @@ import vllm.envs as envs from vllm.config import VllmConfig from vllm.distributed import get_dcp_group, get_pcp_group from vllm.logger import init_logger +from vllm.model_executor.warmup.jit_warmup import ( + VllmJitKernel, + WarmupIntRange, +) +from vllm.model_executor.warmup.jit_warmup_triton_helper import ( + TritonPointerInputVariant, + TritonWarmupTensor, +) from vllm.platforms import current_platform from vllm.triton_utils import tl, triton from vllm.utils.deep_gemm import ( @@ -190,6 +198,211 @@ class DeepseekV32IndexerPrefillChunkMetadata: max_local_total_seq_lens: int = 0 +_BUILD_PREFILL_CHUNK_METADATA_INPUT_VARIANTS = ( + TritonPointerInputVariant.from_alignment(uncompressed_seq_lens=True), + TritonPointerInputVariant.from_alignment(uncompressed_seq_lens=False), +) + + +class BuildPrefillChunkMetadataKernel( + VllmJitKernel["BuildPrefillChunkMetadataKernel.CompileKey"] +): + BLOCK_SIZE = 1024 + + @dataclass(frozen=True) + class CompileKey: + query_slice_start: int + query_slice_stop: int + DCP_RANK: int + DCP_WORLD: int + DCP_INTERLEAVE: int + BLOCK_SIZE: int + COMPRESS_RATIO: int + input_variant: TritonPointerInputVariant + + @staticmethod + @triton.jit + def kernel( + # Inputs + query_start_loc_ptr, + uncompressed_seq_lens_ptr, + cu_compressed_seq_lens_ptr, + # Row-start base for cu_seq_len_ks/ke: local cumulative lens under DCP, + # aliases cu_compressed_seq_lens_ptr otherwise. + row_start_cu_compressed_seq_lens_ptr, + # Outputs + token_to_seq_ptr, + cu_compressed_seq_len_ks_ptr, + cu_compressed_seq_len_ke_ptr, + query_slice_start, + query_slice_stop, + DCP_RANK, + DCP_WORLD, + DCP_INTERLEAVE, + BLOCK_SIZE: tl.constexpr, + COMPRESS_RATIO: tl.constexpr, + ): + batch_idx = tl.program_id(0) + + query_start = tl.load(query_start_loc_ptr + batch_idx) + query_end = tl.load(query_start_loc_ptr + batch_idx + 1) + query_len = query_end - query_start + + seq_start = tl.load(cu_compressed_seq_lens_ptr + batch_idx) + seq_end = tl.load(cu_compressed_seq_lens_ptr + batch_idx + 1) + compressed_seq_len = seq_end - seq_start + + # Row start for the (possibly localized) cu_seq_len_ks/ke. Equals seq_start + # when DCP is disabled (the pointer aliases cu_compressed_seq_lens_ptr). + row_start = tl.load(row_start_cu_compressed_seq_lens_ptr + batch_idx) + + uncompressed_seq_len = tl.load(uncompressed_seq_lens_ptr + batch_idx) + start_pos = uncompressed_seq_len - query_len + + for i in range(0, query_len, BLOCK_SIZE): + offset = i + tl.arange(0, BLOCK_SIZE) + abs_pos = query_start + offset + mask = ( + (offset < query_len) + & (abs_pos >= query_slice_start) + & (abs_pos < query_slice_stop) + ) + out_pos = abs_pos - query_slice_start + + # cu_seq_len_ks: row start in the gathered K buffer. + tl.store(cu_compressed_seq_len_ks_ptr + out_pos, row_start, mask=mask) + + # cu_seq_len_ke: row start + per-token context length. Under DCP the + # global per-token length is sharded across ranks. + global_ctx = start_pos + 1 + offset + len_per_token = global_ctx // COMPRESS_RATIO + if DCP_WORLD > 1: + # Per-rank local context length under interleave-aware DCP, matching + # get_dcp_local_seq_lens. K == 1 reduces to (len + world-1-rank)//world. + base = (len_per_token // DCP_INTERLEAVE // DCP_WORLD) * DCP_INTERLEAVE + remainder = len_per_token - base * DCP_WORLD + remainder = tl.minimum( + tl.maximum(remainder - DCP_RANK * DCP_INTERLEAVE, 0), DCP_INTERLEAVE + ) + len_per_token = base + remainder + tl.store( + cu_compressed_seq_len_ke_ptr + out_pos, + row_start + len_per_token, + mask=mask, + ) + + # Compute token_to_seq + for i in range(0, compressed_seq_len, BLOCK_SIZE): + 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) + + def dispatch( # type: ignore[override] + self, + *, + query_slice_start: int, + query_slice_stop: int, + DCP_RANK: int, + DCP_WORLD: int, + DCP_INTERLEAVE: int, + BLOCK_SIZE: int, + COMPRESS_RATIO: int, + input_variant: TritonPointerInputVariant, + ) -> CompileKey: + return self.CompileKey( + query_slice_start=query_slice_start, + query_slice_stop=query_slice_stop, + DCP_RANK=DCP_RANK, + DCP_WORLD=DCP_WORLD, + DCP_INTERLEAVE=DCP_INTERLEAVE, + BLOCK_SIZE=BLOCK_SIZE, + COMPRESS_RATIO=COMPRESS_RATIO, + input_variant=input_variant, + ) + + 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 + parallel_config = vllm_config.parallel_config + dcp_world = parallel_config.decode_context_parallel_size + dcp_interleave = parallel_config.cp_kv_cache_interleave_size + dcp_rank = get_dcp_group().rank_in_group if dcp_world > 1 else 0 + 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=(1, 2 * max_tokens - 1, 2 * max_tokens), + DCP_RANK=dcp_rank, + DCP_WORLD=dcp_world, + DCP_INTERLEAVE=dcp_interleave, + BLOCK_SIZE=self.BLOCK_SIZE, + COMPRESS_RATIO=list(compress_ratios), + input_variant=_BUILD_PREFILL_CHUNK_METADATA_INPUT_VARIANTS, + ) + + def compile(self, compile_key: CompileKey) -> None: + warmup = getattr(self.kernel, "warmup", None) + assert warmup is not None + int32_ptr = TritonWarmupTensor(torch.int32) + warmup( + int32_ptr, + compile_key.input_variant.pointer("uncompressed_seq_lens", torch.int32), + int32_ptr, + int32_ptr, + int32_ptr, + int32_ptr, + int32_ptr, + compile_key.query_slice_start, + compile_key.query_slice_stop, + compile_key.DCP_RANK, + compile_key.DCP_WORLD, + compile_key.DCP_INTERLEAVE, + BLOCK_SIZE=compile_key.BLOCK_SIZE, + COMPRESS_RATIO=compile_key.COMPRESS_RATIO, + grid=(1,), + ) + + def __call__( + self, + query_start_loc: torch.Tensor, + uncompressed_seq_lens: torch.Tensor, + cu_compressed_seq_lens: torch.Tensor, + row_start_cu_compressed_seq_lens: torch.Tensor, + token_to_seq: torch.Tensor, + cu_compressed_seq_len_ks: torch.Tensor, + cu_compressed_seq_len_ke: torch.Tensor, + query_slice_start: int, + query_slice_stop: int, + DCP_RANK: int, + DCP_WORLD: int, + DCP_INTERLEAVE: int, + *, + num_reqs: int, + COMPRESS_RATIO: int, + ) -> None: + self.kernel[(num_reqs,)]( + query_start_loc, + uncompressed_seq_lens, + cu_compressed_seq_lens, + row_start_cu_compressed_seq_lens, + token_to_seq, + cu_compressed_seq_len_ks, + cu_compressed_seq_len_ke, + query_slice_start, + query_slice_stop, + DCP_RANK, + DCP_WORLD, + DCP_INTERLEAVE, + BLOCK_SIZE=self.BLOCK_SIZE, + COMPRESS_RATIO=COMPRESS_RATIO, + ) + + +_BUILD_PREFILL_CHUNK_METADATA_KERNEL = BuildPrefillChunkMetadataKernel() + + @dataclass class DeepseekV32IndexerPrefillMetadata: chunks: list[DeepseekV32IndexerPrefillChunkMetadata] @@ -841,7 +1054,7 @@ def build_prefill_chunk_metadata( # Under DCP the kernel writes this rank's local row bounds into # cu_seq_len_ks/ke; otherwise local_cu_seq_lens aliases cu_seq_lens. - _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, @@ -854,7 +1067,7 @@ def build_prefill_chunk_metadata( dcp_rank, dcp_world_size, cp_kv_cache_interleave_size, - BLOCK_SIZE=1024, + num_reqs=num_reqs, COMPRESS_RATIO=compress_ratio, ) @@ -881,80 +1094,3 @@ def build_prefill_chunk_metadata( local_total_seq_lens=local_total_seq_lens, max_local_total_seq_lens=max_local_total_seq_lens, ) - - -@triton.jit -def _build_prefill_chunk_metadata_kernel( - # Inputs - query_start_loc_ptr, - uncompressed_seq_lens_ptr, - cu_compressed_seq_lens_ptr, - # Row-start base for cu_seq_len_ks/ke: local cumulative lens under DCP, - # aliases cu_compressed_seq_lens_ptr otherwise. - row_start_cu_compressed_seq_lens_ptr, - # Outputs - token_to_seq_ptr, - cu_compressed_seq_len_ks_ptr, - cu_compressed_seq_len_ke_ptr, - query_slice_start, - query_slice_stop, - DCP_RANK, - DCP_WORLD, - DCP_INTERLEAVE, - BLOCK_SIZE: tl.constexpr, - COMPRESS_RATIO: tl.constexpr, -): - batch_idx = tl.program_id(0) - - query_start = tl.load(query_start_loc_ptr + batch_idx) - query_end = tl.load(query_start_loc_ptr + batch_idx + 1) - query_len = query_end - query_start - - seq_start = tl.load(cu_compressed_seq_lens_ptr + batch_idx) - seq_end = tl.load(cu_compressed_seq_lens_ptr + batch_idx + 1) - compressed_seq_len = seq_end - seq_start - - # Row start for the (possibly localized) cu_seq_len_ks/ke. Equals seq_start - # when DCP is disabled (the pointer aliases cu_compressed_seq_lens_ptr). - row_start = tl.load(row_start_cu_compressed_seq_lens_ptr + batch_idx) - - uncompressed_seq_len = tl.load(uncompressed_seq_lens_ptr + batch_idx) - start_pos = uncompressed_seq_len - query_len - - for i in range(0, query_len, BLOCK_SIZE): - offset = i + tl.arange(0, BLOCK_SIZE) - abs_pos = query_start + offset - mask = ( - (offset < query_len) - & (abs_pos >= query_slice_start) - & (abs_pos < query_slice_stop) - ) - out_pos = abs_pos - query_slice_start - - # cu_seq_len_ks: row start in the gathered K buffer. - tl.store(cu_compressed_seq_len_ks_ptr + out_pos, row_start, mask=mask) - - # cu_seq_len_ke: row start + per-token context length. Under DCP the - # global per-token length is sharded across ranks. - global_ctx = start_pos + 1 + offset - len_per_token = global_ctx // COMPRESS_RATIO - if DCP_WORLD > 1: - # Per-rank local context length under interleave-aware DCP, matching - # get_dcp_local_seq_lens. K == 1 reduces to (len + world-1-rank)//world. - base = (len_per_token // DCP_INTERLEAVE // DCP_WORLD) * DCP_INTERLEAVE - remainder = len_per_token - base * DCP_WORLD - remainder = tl.minimum( - tl.maximum(remainder - DCP_RANK * DCP_INTERLEAVE, 0), DCP_INTERLEAVE - ) - len_per_token = base + remainder - tl.store( - cu_compressed_seq_len_ke_ptr + out_pos, - row_start + len_per_token, - mask=mask, - ) - - # Compute token_to_seq - for i in range(0, compressed_seq_len, BLOCK_SIZE): - 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) diff --git a/vllm/v1/attention/backends/mla/prefill/flash_attn.py b/vllm/v1/attention/backends/mla/prefill/flash_attn.py index 30c0c6d8a6..c5fdbaebc0 100644 --- a/vllm/v1/attention/backends/mla/prefill/flash_attn.py +++ b/vllm/v1/attention/backends/mla/prefill/flash_attn.py @@ -3,7 +3,9 @@ """FlashAttention backend for MLA prefill.""" import functools -from typing import TYPE_CHECKING +from collections.abc import Callable +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any import torch @@ -11,14 +13,7 @@ import vllm.envs as envs from vllm.model_executor.layers.quantization.utils.quant_utils import ( kFp8StaticTensorSym, ) -from vllm.model_executor.warmup.cutedsl_warmup import ( - CuTeDSLCompileUnit, - register_cutedsl_warmup_provider, -) -from vllm.model_executor.warmup.fa4_cutedsl_config import ( - FA4MLAPrefillCompileContext, - iter_fa4_mla_prefill_compile_requests, -) +from vllm.model_executor.warmup.jit_warmup import VllmJitKernel from vllm.platforms import current_platform from vllm.v1.attention.backends.fa_utils import ( compile_flash_attn_varlen_func_from_specs, @@ -32,6 +27,7 @@ from vllm.v1.attention.backends.mla.prefill.base import ( if TYPE_CHECKING: from vllm.config import VllmConfig + from vllm.model_executor.layers.attention.mla_attention import MLADims from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey if is_flash_attn_varlen_func_available(): @@ -40,6 +36,245 @@ else: flash_attn_varlen_func = None # type: ignore[assignment] +FA4_STANDARD_DTYPES = (torch.bfloat16, torch.float16) + +# Current vLLM MLA prefill expands K/V to num_heads before FA4, so this plan +# covers qhead_per_kvhead=1. Batch is not a current FA4 MLA-prefill key field. +# Use b1 for compile-only specs because it is the conservative case for +# Split-KV shape heuristics. +# TODO(roberto): FA4 also has direct-GQA and qv/top-k absorbed-MLA paths, but +# vLLM does not use them in this backend yet; they need a separate +# num_kv_heads/qv/top-k-aware warmup plan if wired in later. +FA4_MLA_PREFILL_COMPILE_BATCH_SIZE = 1 +FA4_MLA_PREFILL_Q_TILE = 128 +FA4_MLA_PREFILL_K_TILE = 128 +FA4_MLA_PREFILL_LONG_K_BLOCKS = 32 +FA4_MLA_PREFILL_VERY_LONG_K_BLOCKS = 64 +FA4_MLA_PREFILL_CAUSAL_OPTIONS = (False, True) +FA4_MLA_PREFILL_LSE_OPTIONS = (False, True) + + +@dataclass(frozen=True) +class _FA4MLAPrefillShapeProbe: + max_seqlen_q: int + max_seqlen_k: int + + +class FA4MLAPrefillKernel(VllmJitKernel["FA4MLAPrefillKernel.CompileKey"]): + """FA4 MLA prefill compile-key wrapper used by generic JIT warmup.""" + + @dataclass(frozen=True) + class CompileKey: + """High-level FA4 compile-only key used by vLLM warmup. + + This is not the CuTeDSL cache key. FA4 owns the selector that maps these + serving inputs to the actual compile-static fields: tile sizes, q_stage, + Split-KV, scheduler choice, layout-presence booleans, dtype/head dims, + arch, and related fields. + """ + + q_shape: tuple[int, ...] + k_shape: tuple[int, ...] + v_shape: tuple[int, ...] + q_dtype: torch.dtype + max_seqlen_q: int + max_seqlen_k: int + softmax_scale: float + causal: bool + fa_version: int + v_stride: tuple[int, ...] | None = None + cu_seqlens_q_shape: tuple[int, ...] | None = None + cu_seqlens_k_shape: tuple[int, ...] | None = None + window_size: tuple[int, int] | None = None + return_softmax_lse: bool = False + num_splits: int = 0 + + def dispatch( # type: ignore[override] + self, + *, + batch_size: int, + dtype: torch.dtype, + num_heads: int, + mla_dims: "MLADims", + shape_probe: _FA4MLAPrefillShapeProbe, + requires_v_padding: bool, + causal: bool, + fa_version: int, + window_size: tuple[int, int] | None = None, + return_lse: bool = False, + num_splits: int = 0, + ) -> CompileKey: + return self.CompileKey( + q_shape=( + batch_size * shape_probe.max_seqlen_q, + num_heads, + mla_dims.qk_nope_head_dim + mla_dims.qk_rope_head_dim, + ), + k_shape=( + batch_size * shape_probe.max_seqlen_k, + num_heads, + mla_dims.qk_nope_head_dim + mla_dims.qk_rope_head_dim, + ), + v_shape=( + batch_size * shape_probe.max_seqlen_k, + num_heads, + ( + mla_dims.qk_nope_head_dim + mla_dims.qk_rope_head_dim + if requires_v_padding + else mla_dims.v_head_dim + ), + ), + q_dtype=dtype, + max_seqlen_q=shape_probe.max_seqlen_q, + max_seqlen_k=shape_probe.max_seqlen_k, + softmax_scale=(mla_dims.qk_nope_head_dim + mla_dims.qk_rope_head_dim) + ** -0.5, + causal=causal, + fa_version=fa_version, + v_stride=( + None + if requires_v_padding + else ( + num_heads * (mla_dims.qk_nope_head_dim + mla_dims.v_head_dim), + mla_dims.qk_nope_head_dim + mla_dims.v_head_dim, + 1, + ) + ), + cu_seqlens_q_shape=(batch_size + 1,), + cu_seqlens_k_shape=(batch_size + 1,), + window_size=window_size, + return_softmax_lse=return_lse, + num_splits=num_splits, + ) + + def get_warmup_keys(self, vllm_config: "VllmConfig") -> list[CompileKey]: + from vllm.model_executor.layers.attention.mla_attention import ( + get_mla_dims, + ) + + mla_dims = get_mla_dims(vllm_config.model_config) + qk_head_dim = mla_dims.qk_nope_head_dim + mla_dims.qk_rope_head_dim + fa_version = get_flash_attn_version(head_size=qk_head_dim) + if fa_version != 4: + return [] + + dtype = vllm_config.model_config.dtype + if dtype not in FA4_STANDARD_DTYPES: + dtype = torch.bfloat16 + + is_sm90 = current_platform.is_device_capability(90) + is_sm100_family = current_platform.is_device_capability_family(100) + is_sm120 = current_platform.is_device_capability_family(120) + if not (is_sm90 or is_sm100_family or is_sm120): + return [] + + num_splits = 1 if envs.VLLM_BATCH_INVARIANT else 0 + if is_sm120 and num_splits != 1: + return [] + + num_heads = vllm_config.model_config.get_num_attention_heads( + vllm_config.parallel_config + ) + if num_heads <= 0: + return [] + + requires_v_padding = False + effective_v_head_dim = ( + qk_head_dim if requires_v_padding else mla_dims.v_head_dim + ) + shape_probes = [ + _FA4MLAPrefillShapeProbe(1, FA4_MLA_PREFILL_K_TILE), + _FA4MLAPrefillShapeProbe( + FA4_MLA_PREFILL_Q_TILE + 1, + 4 * FA4_MLA_PREFILL_K_TILE, + ), + ] + if num_splits != 1: + shape_probes.extend( + ( + _FA4MLAPrefillShapeProbe( + 1, + FA4_MLA_PREFILL_LONG_K_BLOCKS * FA4_MLA_PREFILL_K_TILE, + ), + _FA4MLAPrefillShapeProbe( + FA4_MLA_PREFILL_Q_TILE + 1, + FA4_MLA_PREFILL_LONG_K_BLOCKS * FA4_MLA_PREFILL_K_TILE, + ), + ) + ) + if not is_sm90 and qk_head_dim != effective_v_head_dim: + shape_probes.extend( + ( + _FA4MLAPrefillShapeProbe( + 1, + FA4_MLA_PREFILL_VERY_LONG_K_BLOCKS * FA4_MLA_PREFILL_K_TILE, + ), + _FA4MLAPrefillShapeProbe( + FA4_MLA_PREFILL_Q_TILE + 1, + FA4_MLA_PREFILL_VERY_LONG_K_BLOCKS * FA4_MLA_PREFILL_K_TILE, + ), + ) + ) + + return self._trace_dispatch(self.dispatch)( + batch_size=FA4_MLA_PREFILL_COMPILE_BATCH_SIZE, + dtype=dtype, + num_heads=num_heads, + mla_dims=mla_dims, + shape_probe=shape_probes, + requires_v_padding=requires_v_padding, + causal=FA4_MLA_PREFILL_CAUSAL_OPTIONS, + return_lse=FA4_MLA_PREFILL_LSE_OPTIONS, + num_splits=num_splits, + fa_version=fa_version, + ) + + @staticmethod + def kernel(*args: Any, **kwargs: Any) -> Any: + assert flash_attn_varlen_func is not None + return flash_attn_varlen_func(*args, **kwargs) + + def compile(self, compile_key: CompileKey) -> None: + assert compile_flash_attn_varlen_func_from_specs is not None + window_size = ( + list(compile_key.window_size) + if compile_key.window_size is not None + else None + ) + compile_flash_attn_varlen_func_from_specs( + q_shape=compile_key.q_shape, + k_shape=compile_key.k_shape, + v_shape=compile_key.v_shape, + q_dtype=compile_key.q_dtype, + v_stride=compile_key.v_stride, + cu_seqlens_q_shape=compile_key.cu_seqlens_q_shape, + cu_seqlens_k_shape=compile_key.cu_seqlens_k_shape, + max_seqlen_q=compile_key.max_seqlen_q, + max_seqlen_k=compile_key.max_seqlen_k, + softmax_scale=compile_key.softmax_scale, + causal=compile_key.causal, + window_size=window_size, + return_softmax_lse=compile_key.return_softmax_lse, + fa_version=compile_key.fa_version, + num_splits=compile_key.num_splits, + ) + + def __call__( + self, + *, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + runtime_kernel: Callable[..., Any] | None = None, + **kwargs: Any, + ) -> Any: + kernel = self.kernel if runtime_kernel is None else runtime_kernel + return kernel(q=q, k=k, v=v, **kwargs) + + +FA4_MLA_PREFILL_KERNEL = FA4MLAPrefillKernel() + + class FlashAttnPrefillBackend(MLAPrefillBackend): """FlashAttention backend for MLA prefill.""" @@ -127,46 +362,6 @@ class FlashAttnPrefillBackend(MLAPrefillBackend): # Track whether we're using vllm's FA or upstream (for ROCm) self._is_vllm_fa = current_platform.is_cuda() or current_platform.is_xpu() - if self.vllm_flash_attn_version == 4: - register_cutedsl_warmup_provider(self) - - def get_cutedsl_warmup_compile_units(self) -> tuple[CuTeDSLCompileUnit, ...]: - if self.vllm_flash_attn_version != 4: - return () - if compile_flash_attn_varlen_func_from_specs is None: - raise RuntimeError( - "FA4 compile-only API is unavailable; CuTeDSL warmup does not " - "fall back to synthetic forward passes." - ) - - dtype = self.vllm_config.model_config.dtype - if dtype not in self.supported_dtypes: - dtype = torch.bfloat16 - - qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim - ctx = FA4MLAPrefillCompileContext( - dtype=dtype, - num_heads=self.num_heads, - qk_head_dim=qk_head_dim, - v_head_dim=self.v_head_dim, - kv_nope_head_dim=self.qk_nope_head_dim + self.v_head_dim, - requires_v_padding=self.requires_v_padding, - scale=self.scale, - num_splits=1 if envs.VLLM_BATCH_INVARIANT else 0, - fa_version=self.vllm_flash_attn_version, - ) - compile_requests = tuple(iter_fa4_mla_prefill_compile_requests(ctx)) - if not compile_requests: - return () - - return tuple( - CuTeDSLCompileUnit( - name="fa4_mla_prefill", - key=request.key, - compile=request.compile, - ) - for request in compile_requests - ) def supports_quant_output(self, quant_key: "QuantKey") -> bool: device_capability = current_platform.get_device_capability() @@ -207,10 +402,11 @@ class FlashAttnPrefillBackend(MLAPrefillBackend): if envs.VLLM_BATCH_INVARIANT: kwargs["num_splits"] = 1 - attn_out = self.flash_attn_varlen_func( + attn_out = FA4_MLA_PREFILL_KERNEL( q=q, k=k, v=maybe_padded_v, + runtime_kernel=self.flash_attn_varlen_func, softmax_scale=softmax_scale, **kwargs, ) diff --git a/vllm/v1/attention/backends/mla/sparse_swa.py b/vllm/v1/attention/backends/mla/sparse_swa.py index 9faf56f639..c823f82be2 100644 --- a/vllm/v1/attention/backends/mla/sparse_swa.py +++ b/vllm/v1/attention/backends/mla/sparse_swa.py @@ -7,9 +7,14 @@ import torch from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase +from vllm.model_executor.warmup.jit_warmup import ( + VllmJitKernel, + WarmupIntRange, +) +from vllm.model_executor.warmup.jit_warmup_triton_helper import TritonWarmupTensor from vllm.platforms import current_platform from vllm.triton_utils import tl, triton -from vllm.utils.math_utils import cdiv +from vllm.utils.math_utils import cdiv, next_power_of_2 from vllm.v1.attention.backend import ( AttentionBackend, AttentionCGSupport, @@ -279,6 +284,107 @@ class DeepseekSparseSWAMetadata: return chunk_plan +class ComputePrefillMetadataKernel( + VllmJitKernel["ComputePrefillMetadataKernel.CompileKey"] +): + @dataclass(frozen=True) + class CompileKey: + BLOCK_SIZE: int + + @staticmethod + @triton.jit(do_not_specialize=["num_prefills", "num_decodes", "window_size"]) + def kernel( + # Outputs + prefill_gather_lens_ptr, + # Inputs + seq_lens_ptr, + query_start_loc_ptr, + num_prefills, + num_decodes, + window_size, + BLOCK_SIZE: tl.constexpr, + ): + """Compute prefill gather_lens in a single pass.""" + offset = tl.arange(0, BLOCK_SIZE) + mask = offset < num_prefills + # SM12x + Triton 3.6 raises IMA on out-of-bounds address arithmetic for + # masked-off lanes even though the load mask gates the actual read, so + # clamp the offset. Caller guarantees num_prefills > 0. + safe_offset = tl.minimum(offset, num_prefills - 1) + + seq_len = tl.load(seq_lens_ptr + num_decodes + safe_offset, mask=mask) + qsl_start = tl.load(query_start_loc_ptr + num_decodes + safe_offset, mask=mask) + qsl_end = tl.load( + query_start_loc_ptr + num_decodes + safe_offset + 1, mask=mask + ) + + query_len = qsl_end - qsl_start + prefix_len = seq_len - query_len + gather_len = query_len + tl.minimum(prefix_len, window_size - 1) + + tl.store(prefill_gather_lens_ptr + offset, gather_len, mask=mask) + + def dispatch( # type: ignore[override] + self, + *, + num_prefills: int, + ) -> CompileKey: + return self.CompileKey( + BLOCK_SIZE=next_power_of_2(num_prefills), + ) + + def get_warmup_keys(self, vllm_config: VllmConfig) -> list[CompileKey]: + scheduler_config = vllm_config.scheduler_config + max_prefills = max( + 1, + min( + scheduler_config.max_num_seqs, + scheduler_config.max_num_batched_tokens, + ), + ) + return self._trace_dispatch(self.dispatch)( + num_prefills=WarmupIntRange(1, max_prefills + 1), + ) + + def compile(self, compile_key: CompileKey) -> None: + warmup = getattr(self.kernel, "warmup", None) + assert warmup is not None + int32_ptr = TritonWarmupTensor(torch.int32) + warmup( + int32_ptr, + int32_ptr, + int32_ptr, + compile_key.BLOCK_SIZE, + 0, + 1, + BLOCK_SIZE=compile_key.BLOCK_SIZE, + grid=(1,), + ) + + def __call__( + self, + prefill_gather_lens: torch.Tensor, + seq_lens: torch.Tensor, + query_start_loc: torch.Tensor, + num_prefills: int, + num_decodes: int, + window_size: int, + ) -> None: + compile_key = self.dispatch(num_prefills=num_prefills) + self.kernel[(1,)]( + prefill_gather_lens, + seq_lens, + query_start_loc, + num_prefills, + num_decodes, + window_size, + BLOCK_SIZE=compile_key.BLOCK_SIZE, + ) + + +_COMPUTE_PREFILL_METADATA_KERNEL = ComputePrefillMetadataKernel() + + class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder): """Builds metadata for DeepseekV4 SWA cache. @@ -602,14 +708,13 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder): pfx_gather_lens = torch.empty( num_prefills, dtype=torch.int32, device=seq_lens.device ) - _compute_prefill_metadata_kernel[(1,)]( + _COMPUTE_PREFILL_METADATA_KERNEL( pfx_gather_lens, seq_lens, query_start_loc, num_prefills, num_decodes, self.window_size, - BLOCK_SIZE=triton.next_power_of_2(num_prefills), ) result["prefill_seq_lens"] = seq_lens[num_decodes:] @@ -626,37 +731,6 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder): return result -@triton.jit -def _compute_prefill_metadata_kernel( - # Outputs - prefill_gather_lens_ptr, - # Inputs - seq_lens_ptr, - query_start_loc_ptr, - num_prefills, - num_decodes, - window_size, - BLOCK_SIZE: tl.constexpr, -): - """Compute prefill gather_lens in a single pass.""" - offset = tl.arange(0, BLOCK_SIZE) - mask = offset < num_prefills - # SM12x + Triton 3.6 raises IMA on out-of-bounds address arithmetic for - # masked-off lanes even though the load mask gates the actual read, so - # clamp the offset. Caller guarantees num_prefills > 0. - safe_offset = tl.minimum(offset, num_prefills - 1) - - seq_len = tl.load(seq_lens_ptr + num_decodes + safe_offset, mask=mask) - qsl_start = tl.load(query_start_loc_ptr + num_decodes + safe_offset, mask=mask) - qsl_end = tl.load(query_start_loc_ptr + num_decodes + safe_offset + 1, mask=mask) - - query_len = qsl_end - qsl_start - prefix_len = seq_len - query_len - gather_len = query_len + tl.minimum(prefix_len, window_size - 1) - - tl.store(prefill_gather_lens_ptr + offset, gather_len, mask=mask) - - @triton.jit(do_not_specialize=["token_offset"]) def _compute_swa_indices_and_lens_kernel( swa_indices_ptr,