iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Feat][Perf] Add new warmup infrastructure for JITs (#47451)

Browse Source 
Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>
Signed-off-by: Roberto L. Castro <38211239+LopezCastroRoberto@users.noreply.github.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
This commit is contained in:
Roberto L. Castro
2026-07-20 13:21:55 -07:00
committed by GitHub
co-authored by mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
parent fbfe58133d
commit 7ca017778f
14 changed files with 1855 additions and 693 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.buildkite/test_areas/model_executor.yaml
+4
View File
@@ -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
tests/model_executor/test_jit_warmup.py
+306
View File
@@ -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
vllm/config/kernel.py
+8
View File
@@ -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
vllm/model_executor/warmup/cutedsl_warmup.py
+4
View File
@@ -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
vllm/model_executor/warmup/fa4_cutedsl_config.py
-204
View File
@@ -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}")
vllm/model_executor/warmup/fa4_cutedsl_warmup.py
+30
View File
@@ -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)
vllm/model_executor/warmup/jit_warmup.py
+474
View File
@@ -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)
vllm/model_executor/warmup/jit_warmup_triton_helper.py
+183
View File
@@ -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)
)
vllm/model_executor/warmup/kernel_warmup.py
+11 -2
View File
@@ -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:
vllm/model_executor/warmup/sparse_mla_triton_warmup.py
+28 -254
View File
@@ -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)
vllm/models/deepseek_v4/common/ops/cache_utils.py
+238 -70
View File
@@ -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(
vllm/v1/attention/backends/mla/indexer.py
+215 -79
View File
@@ -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)
vllm/v1/attention/backends/mla/prefill/flash_attn.py
+246 -50
View File
@@ -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,
)
vllm/v1/attention/backends/mla/sparse_swa.py
+108 -34
View File
@@ -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,