From 6c95a641e95c0faa6f3aa802d1fd3cce3f3bc3ce Mon Sep 17 00:00:00 2001 From: "Roberto L. Castro" <38211239+LopezCastroRoberto@users.noreply.github.com> Date: Tue, 11 Aug 2026 13:46:07 +0200 Subject: [PATCH] [2/N][Feat][Perf] Add new warmup infrastructure for JITs. Add predicate filtering for JIT warmup, and migrate Inkling FA4 (#49315) Signed-off-by: LopezCastroRoberto Signed-off-by: Roberto L. Castro <38211239+LopezCastroRoberto@users.noreply.github.com> Co-authored-by: Lucas Wilkinson Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- tests/model_executor/test_jit_warmup.py | 65 +++ .../models/inkling/test_fa4_rel_attention.py | 8 +- tests/models/inkling/test_fa4_warmup.py | 113 +++-- vllm/config/kernel.py | 2 +- vllm/model_executor/warmup/cutedsl_warmup.py | 2 +- .../warmup/fa4_cutedsl_warmup.py | 23 +- vllm/model_executor/warmup/jit_warmup.py | 201 ++++++--- vllm/models/inkling/nvidia/attention.py | 24 +- vllm/models/inkling/nvidia/ops/__init__.py | 6 +- .../inkling/nvidia/ops/fa4_rel_attention.py | 392 +++++++++++++++--- vllm/models/inkling/nvidia/ops/fa4_warmup.py | 154 ------- .../backends/mla/prefill/flash_attn.py | 135 +++--- 12 files changed, 739 insertions(+), 386 deletions(-) delete mode 100644 vllm/models/inkling/nvidia/ops/fa4_warmup.py diff --git a/tests/model_executor/test_jit_warmup.py b/tests/model_executor/test_jit_warmup.py index da6b21091b..f93de50a3d 100644 --- a/tests/model_executor/test_jit_warmup.py +++ b/tests/model_executor/test_jit_warmup.py @@ -105,6 +105,39 @@ def test_trace_dispatch_expands_ranges_dedupes_and_ignores_unused_inputs() -> No ] +def test_warmup_range_uses_custom_advancement() -> None: + keys = ToyKernel()._trace_dispatch(ToyKernel().dispatch)( + tokens=WarmupIntRange( + 1, + 10, + advance=lambda value: _next_power_of_2(value) + 1, + ), + cfg=_config(), + ) + + assert [key.block_size for key in keys] == [1, 2, 4, 8, 16] + + +def test_warmup_range_validates_custom_advancement() -> None: + trace_dispatch = ToyKernel()._trace_dispatch(ToyKernel().dispatch) + + with pytest.raises(ValueError, match="both step and advance"): + trace_dispatch( + tokens=WarmupIntRange( + 1, + 10, + step=2, + advance=lambda value: value + 1, + ), + cfg=_config(), + ) + with pytest.raises(ValueError, match="must return a greater value"): + trace_dispatch( + tokens=WarmupIntRange(1, 10, advance=lambda value: value), + cfg=_config(), + ) + + def test_compile_key_uses_defaults_locals_attributes_and_expressions() -> None: cfg = _config(bias=3, disabled=True, name="cfg", vectorized=True) @@ -143,6 +176,38 @@ def test_trace_dispatch_combines_zipped_rows_with_independent_values() -> None: ] +def test_trace_dispatch_filters_with_traced_predicate() -> None: + class PredicateKernel(ToyKernel): + def _is_valid_warmup_dispatch( + self, + *, + tokens: int, + lanes: int, + max_work: int, + ) -> bool: + block_size = _next_power_of_2(tokens) + return block_size * lanes <= max_work + + def get_warmup_keys( + self, max_tokens: int, cfg: Any + ) -> list[ToyKernel.CompileKey]: + return self._trace_dispatch(self.dispatch)( + tokens=WarmupIntRange(1, max_tokens + 1), + lanes=(1, 2), + cfg=cfg, + max_work=4, + _when=self._is_valid_warmup_dispatch, + ) + + assert PredicateKernel().get_warmup_keys(5, _config()) == [ + ToyKernel.CompileKey(1, 1, 1, ("base", "default", -1, 1, 1), True), + ToyKernel.CompileKey(1, 2, 1, ("base", "default", -1, 1, 1), True), + ToyKernel.CompileKey(2, 2, 1, ("base", "default", -2, 2, 4), True), + ToyKernel.CompileKey(2, 4, 1, ("base", "default", -2, 2, 4), True), + ToyKernel.CompileKey(4, 4, 1, ("base", "default", -4, 1, 16), True), + ] + + def test_zip_inputs_validates_input_rows() -> None: with pytest.raises(ValueError, match="requires at least one"): zip_inputs() diff --git a/tests/models/inkling/test_fa4_rel_attention.py b/tests/models/inkling/test_fa4_rel_attention.py index 1a11ac220c..bdd1fbf944 100644 --- a/tests/models/inkling/test_fa4_rel_attention.py +++ b/tests/models/inkling/test_fa4_rel_attention.py @@ -2,8 +2,8 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Correctness test for the Inkling FA4 relative-attention score-mod kernel. -Checks ``inkling_fa4_rel_attention`` against a pure-PyTorch reference that -implements the relative bias exactly as documented in the Inkling architecture +Checks ``INKLING_FA4_REL_ATTENTION_KERNEL`` against a pure-PyTorch reference +that implements the relative bias exactly as documented in the Inkling architecture guide:: logit(i, j, h) = (1 / head_dim) * dot(q[i, h], k[j, h]) + rel_bias(i, j, h) @@ -23,9 +23,9 @@ from vllm.models.inkling.nvidia.attention import ( compute_log_scaling_tau, ) from vllm.models.inkling.nvidia.ops.fa4_rel_attention import ( + INKLING_FA4_REL_ATTENTION_KERNEL, _use_sheared_bias, inkling_fa4_num_splits, - inkling_fa4_rel_attention, ) from vllm.platforms import current_platform from vllm.platforms.interface import DeviceCapability @@ -279,7 +279,7 @@ def _run_case(seq_lens, num_heads, num_kv_heads, rel_extent, window_left, seed=0 num_kv_heads=num_kv_heads, max_kv_len=max(kv_lens), ) - out = inkling_fa4_rel_attention( + out = INKLING_FA4_REL_ATTENTION_KERNEL( q, key_cache, value_cache, diff --git a/tests/models/inkling/test_fa4_warmup.py b/tests/models/inkling/test_fa4_warmup.py index eea9798aa9..8b7a654715 100644 --- a/tests/models/inkling/test_fa4_warmup.py +++ b/tests/models/inkling/test_fa4_warmup.py @@ -1,17 +1,47 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace + import torch +import vllm.models.inkling.nvidia.ops.fa4_rel_attention as fa4_rel_attention +from vllm.models.inkling.configs import InklingMMConfig, InklingModelConfig from vllm.models.inkling.nvidia.ops.fa4_rel_attention import ( + InklingFA4RelAttentionKernel, + _num_warps_bucket, bucket_max_seqlen_q, inkling_fa4_num_splits, ) -from vllm.models.inkling.nvidia.ops.fa4_warmup import ( - InklingFA4WarmupConfig, - _iter_compile_units, - _num_warps_bucket, -) + + +def _vllm_config_from_reference_config(config: SimpleNamespace) -> SimpleNamespace: + return SimpleNamespace( + model_config=SimpleNamespace( + hf_config=InklingMMConfig( + text_config=InklingModelConfig( + num_attention_heads=config.num_heads, + num_key_value_heads=config.num_kv_heads, + head_dim=config.head_dim, + rel_extent=config.rel_extent, + swa_num_attention_heads=config.num_heads, + swa_num_key_value_heads=config.num_kv_heads, + swa_head_dim=config.head_dim, + sliding_window_size=config.rel_extent, + ) + ), + max_model_len=config.max_kv_len, + dtype=config.dtype, + ), + cache_config=SimpleNamespace( + cache_dtype="auto", + block_size=config.block_size, + ), + scheduler_config=SimpleNamespace( + max_num_seqs=config.max_num_reqs, + max_num_batched_tokens=config.max_num_batched_tokens, + ), + ) def test_bucket_max_seqlen_q(): @@ -28,8 +58,13 @@ def test_bucket_max_seqlen_q(): ] -def test_warmup_enumerates_every_runtime_compile_class(): - config = InklingFA4WarmupConfig( +def test_warmup_enumerates_every_runtime_compile_class(monkeypatch): + monkeypatch.setattr( + fa4_rel_attention, + "get_tensor_model_parallel_world_size", + lambda: 1, + ) + config = SimpleNamespace( num_heads=16, num_kv_heads=2, head_dim=128, @@ -43,27 +78,47 @@ def test_warmup_enumerates_every_runtime_compile_class(): max_num_reqs=64, max_num_batched_tokens=192, ) - warmed_keys = {unit.key[-1] for unit in _iter_compile_units(config)} + vllm_config = _vllm_config_from_reference_config(config) + kernel = InklingFA4RelAttentionKernel() + warmed_keys = set(kernel.get_warmup_keys(vllm_config)) + expected_keys: set[InklingFA4RelAttentionKernel.CompileKey] = set() - for query_len in range(1, config.max_num_batched_tokens + 1): - max_seqlen_q = bucket_max_seqlen_q(query_len) - max_num_reqs = min( - config.max_num_reqs, - config.max_num_batched_tokens - query_len + 1, - ) - for num_reqs in range(1, max_num_reqs + 1): - num_splits = inkling_fa4_num_splits( - is_local=config.is_local, - batch_size=num_reqs, - max_query_len=max_seqlen_q, - num_heads=config.num_heads, - num_kv_heads=config.num_kv_heads, - max_kv_len=config.max_kv_len, + for is_local in (False, True): + max_kv_len = config.rel_extent if is_local else config.max_kv_len + window_size = (config.rel_extent - 1, 0) if is_local else (-1, -1) + for query_len in range(1, config.max_num_batched_tokens + 1): + max_seqlen_q = bucket_max_seqlen_q(query_len) + max_num_reqs = min( + config.max_num_reqs, + config.max_num_batched_tokens - query_len + 1, ) - runtime_key = ( - max_seqlen_q, - num_splits, - _num_warps_bucket(num_reqs) if num_splits > 1 else None, - num_reqs > 1024, - ) - assert runtime_key in warmed_keys + for num_reqs in range(1, max_num_reqs + 1): + num_splits = inkling_fa4_num_splits( + is_local=is_local, + batch_size=num_reqs, + max_query_len=max_seqlen_q, + num_heads=config.num_heads, + num_kv_heads=config.num_kv_heads, + max_kv_len=max_kv_len, + ) + expected_keys.add( + InklingFA4RelAttentionKernel.CompileKey( + is_local=is_local, + num_heads=config.num_heads, + num_kv_heads=config.num_kv_heads, + head_dim=config.head_dim, + rel_extent=config.rel_extent, + dtype=config.dtype, + kv_dtype=config.kv_dtype, + block_size=config.block_size, + window_size=window_size, + max_seqlen_q=max_seqlen_q, + num_splits=num_splits, + num_warps_bucket=( + _num_warps_bucket(num_reqs) if num_splits > 1 else None + ), + large_num_reqs=num_reqs > 1024, + ) + ) + + assert warmed_keys == expected_keys diff --git a/vllm/config/kernel.py b/vllm/config/kernel.py index 59195f307b..74e959f33a 100644 --- a/vllm/config/kernel.py +++ b/vllm/config/kernel.py @@ -180,7 +180,7 @@ class KernelConfig: # 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.""" + """Deprecated: run legacy CuTeDSL warmup providers.""" enable_jit_warmup: bool = True """If True, run JIT compile warmup during kernel warmup.""" diff --git a/vllm/model_executor/warmup/cutedsl_warmup.py b/vllm/model_executor/warmup/cutedsl_warmup.py index d2b8142f3b..3abd486e3c 100644 --- a/vllm/model_executor/warmup/cutedsl_warmup.py +++ b/vllm/model_executor/warmup/cutedsl_warmup.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Run registered CuTeDSL warmup compile units.""" +"""Deprecated compatibility registry for legacy CuTeDSL warmups.""" # TODO(roberto): Remove this compatibility registry after registered CuTeDSL # warmups are migrated to the shared JIT warmup infrastructure. diff --git a/vllm/model_executor/warmup/fa4_cutedsl_warmup.py b/vllm/model_executor/warmup/fa4_cutedsl_warmup.py index 7f1042338e..5470da5034 100644 --- a/vllm/model_executor/warmup/fa4_cutedsl_warmup.py +++ b/vllm/model_executor/warmup/fa4_cutedsl_warmup.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Warm up FA4 CuTeDSL MLA prefill compile keys.""" +"""Warm up FA4 CuTeDSL kernels.""" from __future__ import annotations @@ -12,7 +12,7 @@ if TYPE_CHECKING: from vllm.v1.worker.gpu_worker import Worker -def fa4_cutedsl_warmup(worker: Worker) -> None: +def _warm_fa4_mla_prefill(worker: Worker) -> None: runner = worker.model_runner if runner.is_pooling_model: return @@ -32,3 +32,22 @@ def fa4_cutedsl_warmup(worker: Worker) -> None: from vllm.v1.attention.backends.mla.prefill import flash_attn flash_attn.FA4_MLA_PREFILL_KERNEL.warmup(vllm_config) + + +def _warm_inkling_fa4_rel_attention(worker: Worker) -> None: + from vllm.models.inkling.configs import InklingMMConfig, InklingModelConfig + from vllm.models.inkling.nvidia.ops.fa4_rel_attention import ( + INKLING_FA4_REL_ATTENTION_KERNEL, + ) + + vllm_config = worker.vllm_config + hf_config = vllm_config.model_config.hf_config + if not isinstance(hf_config, (InklingMMConfig, InklingModelConfig)): + return + + INKLING_FA4_REL_ATTENTION_KERNEL.warmup(vllm_config) + + +def fa4_cutedsl_warmup(worker: Worker) -> None: + _warm_fa4_mla_prefill(worker) + _warm_inkling_fa4_rel_attention(worker) diff --git a/vllm/model_executor/warmup/jit_warmup.py b/vllm/model_executor/warmup/jit_warmup.py index c5411b5efb..d8c5140950 100644 --- a/vllm/model_executor/warmup/jit_warmup.py +++ b/vllm/model_executor/warmup/jit_warmup.py @@ -28,13 +28,29 @@ CompileKeyT = TypeVar("CompileKeyT") @dataclass(frozen=True) class WarmupIntRange: + """Expand integers with range semantics or a custom monotonic progression.""" + start: int stop: int step: int = 1 + advance: Callable[[int], int] | None = None WarmupValues = Any CompileKeyDispatchFn = Callable[..., CompileKeyT] +WarmupPredicateFn = Callable[..., bool] +_LocalExprs = tuple[tuple[str, ast.AST], ...] + + +def _eval_local_exprs( + local_exprs: _LocalExprs, + values: Mapping[str, Any], + globals_: Mapping[str, Any], +) -> dict[str, Any]: + evaluated_values = dict(values) + for name, expr in local_exprs: + evaluated_values[name] = _eval_dispatch_expr(expr, evaluated_values, globals_) + return evaluated_values @dataclass(frozen=True) @@ -46,7 +62,20 @@ class _WarmupInputRows: def _expand_warmup_values(values: WarmupValues) -> tuple[Any, ...]: if isinstance(values, WarmupIntRange): - return tuple(range(values.start, values.stop, values.step)) + if values.advance is None: + return tuple(range(values.start, values.stop, values.step)) + if values.step != 1: + raise ValueError("WarmupIntRange cannot set both step and advance") + + expanded: list[int] = [] + value = values.start + while value < values.stop: + expanded.append(value) + next_value = values.advance(value) + if next_value <= value: + raise ValueError("WarmupIntRange.advance must return a greater value") + value = next_value + return tuple(expanded) if isinstance(values, (list, tuple)): return tuple(values) return (values,) @@ -75,20 +104,6 @@ def zip_inputs(*rows: Mapping[str, WarmupValues]) -> _WarmupInputRows: 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], @@ -117,7 +132,7 @@ def _merge_warmup_kwargs(parts: Iterable[Mapping[str, Any]]) -> dict[str, Any]: @dataclass(frozen=True) class _CompileKeyDispatchTrace: - local_exprs: tuple[tuple[str, ast.AST], ...] + local_exprs: _LocalExprs field_exprs: tuple[tuple[str, ast.AST], ...] globals: Mapping[str, Any] input_names: frozenset[str] @@ -128,11 +143,9 @@ class _CompileKeyDispatchTrace: 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 - ) + dispatch_values = _eval_local_exprs( + self.local_exprs, {**self.defaults, **kwargs}, self.globals + ) return compile_key_type( **{ field: _eval_dispatch_expr(expr, dispatch_values, self.globals) @@ -141,6 +154,23 @@ class _CompileKeyDispatchTrace: ) +@dataclass(frozen=True) +class _WarmupPredicateTrace: + local_exprs: _LocalExprs + return_expr: ast.AST + globals: Mapping[str, Any] + input_names: frozenset[str] + defaults: Mapping[str, Any] + + def matches(self, kwargs: Mapping[str, Any]) -> bool: + dispatch_values = _eval_local_exprs( + self.local_exprs, {**self.defaults, **kwargs}, self.globals + ) + return bool( + _eval_dispatch_expr(self.return_expr, dispatch_values, self.globals) + ) + + _BIN_OPS: dict[type[ast.operator], Callable[[Any, Any], Any]] = { ast.Add: operator.add, ast.Sub: operator.sub, @@ -317,10 +347,10 @@ def _collect_input_names( } -def _collect_dispatch_body( - fn: CompileKeyDispatchFn[Any], +def _collect_expression_body( + fn: Callable[..., Any], function_def: ast.FunctionDef, -) -> tuple[list[tuple[str, ast.AST]], ast.Call]: +) -> tuple[list[tuple[str, ast.AST]], ast.AST]: local_exprs: list[tuple[str, ast.AST]] = [] for statement in function_def.body: if ( @@ -335,7 +365,8 @@ def _collect_dispatch_body( statement.targets[0], ast.Name ): raise _dispatch_expr_error( - statement, "Dispatch assignments must target one local name" + statement, + "Traced warmup helper assignments must target one local name", ) local_exprs.append((statement.targets[0].id, statement.value)) continue @@ -343,29 +374,45 @@ def _collect_dispatch_body( if isinstance(statement, ast.AnnAssign): if statement.value is None: raise _dispatch_expr_error( - statement, "Dispatch annotations must assign a value" + statement, "Traced warmup helper annotations must assign a value" ) if not isinstance(statement.target, ast.Name): raise _dispatch_expr_error( - statement, "Dispatch assignments must target one local name" + statement, + "Traced warmup helper 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): + if isinstance(statement, ast.Return) and statement.value is not None: return local_exprs, statement.value if isinstance(statement, ast.Return): raise _dispatch_expr_error( - statement, "Dispatch must return one CompileKey(...) call" + statement, "Traced warmup helper must return an expression" ) raise _dispatch_expr_error( statement, - "Dispatch may only contain local assignments before CompileKey return", + "Traced warmup helper may only contain local assignments before return", ) - raise ValueError(f"Expected {fn.__name__} to return one CompileKey(...) call") + raise ValueError(f"Expected {fn.__name__} to return an expression") + + +def _function_trace_inputs( + fn: Callable[..., Any], +) -> tuple[dict[str, Any], set[str]]: + signature = inspect.signature(fn) + defaults = { + name: parameter.default + for name, parameter in signature.parameters.items() + if parameter.default is not inspect.Parameter.empty + } + bound_self = getattr(fn, "__self__", None) + if bound_self is not None: + defaults["self"] = bound_self + return defaults, set(signature.parameters) def _trace_compile_key_dispatch( @@ -375,21 +422,20 @@ def _trace_compile_key_dispatch( globals_ = source_fn.__globals__ function_def = get_function_source_node(fn) - local_exprs, return_call = _collect_dispatch_body(fn, function_def) + local_exprs, return_expr = _collect_expression_body(fn, function_def) + if not isinstance(return_expr, ast.Call): + raise _dispatch_expr_error( + return_expr, + "Dispatch must return one CompileKey(...) call", + ) 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) + defaults, candidate_names = _function_trace_inputs(fn) 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: + for keyword in return_expr.keywords: if keyword.arg is None: raise ValueError(f"{fn.__name__} cannot use **kwargs in CompileKey") field_exprs.append((keyword.arg, keyword.value)) @@ -406,6 +452,30 @@ def _trace_compile_key_dispatch( ) +def _trace_warmup_predicate( + fn: WarmupPredicateFn, +) -> _WarmupPredicateTrace: + source_fn = getattr(fn, "__func__", fn) + globals_ = source_fn.__globals__ + function_def = get_function_source_node(fn) + + local_exprs, return_expr = _collect_expression_body(fn, function_def) + defaults, candidate_names = _function_trace_inputs(fn) + 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)) + input_names.update(_collect_input_names(return_expr, candidate_names, local_names)) + + return _WarmupPredicateTrace( + tuple(local_exprs), + return_expr, + globals_, + frozenset(input_names), + defaults, + ) + + class VllmJitKernel(Generic[CompileKeyT], ABC): """Kernel wrapper that owns dispatch, warmup keys, and compilation.""" @@ -424,6 +494,7 @@ class VllmJitKernel(Generic[CompileKeyT], ABC): def traced( *input_groups: _WarmupInputRows, + _when: WarmupPredicateFn | None = None, **kwargs: WarmupValues, ) -> list[CompileKeyT]: for group in input_groups: @@ -432,24 +503,50 @@ class VllmJitKernel(Generic[CompileKeyT], ABC): "_trace_dispatch positional arguments must be " "zip_inputs(...) groups" ) + predicate_trace = ( + _trace_warmup_predicate(_when) if _when is not None else None + ) + input_names = compile_key_dispatch_trace.input_names + if predicate_trace is not None: + input_names = input_names | predicate_trace.input_names expanded_input_groups = tuple( - _expand_warmup_input_rows( - group.rows, compile_key_dispatch_trace.input_names - ) + _expand_warmup_input_rows(group.rows, input_names) for group in input_groups ) - expanded_kwargs = _expand_warmup_value_grid( - kwargs, compile_key_dispatch_trace.input_names + # Expand independent keyword inputs into cartesian-product axes. + expanded_kwarg_axes = tuple( + (name, _expand_warmup_values(value)) + for name, value in kwargs.items() + if name in 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) + dispatch_value_axes = ( + *expanded_input_groups, + *(values for _, values in expanded_kwarg_axes), + ) + input_group_count = len(expanded_input_groups) + kwarg_names = tuple(name for name, _ in expanded_kwarg_axes) + compile_keys: dict[CompileKeyT, None] = {} + for dispatch_value_set in itertools.product(*dispatch_value_axes): + dispatch_values = _merge_warmup_kwargs( + ( + *dispatch_value_set[:input_group_count], + dict( + zip( + kwarg_names, + dispatch_value_set[input_group_count:], + ) + ), ) - for dispatch_values in itertools.product(*dispatch_value_groups) ) - ) + if predicate_trace is not None and not predicate_trace.matches( + dispatch_values + ): + continue + compile_key = compile_key_dispatch_trace.compile_key( + self.CompileKey, dispatch_values + ) + compile_keys[compile_key] = None + return list(compile_keys) return traced diff --git a/vllm/models/inkling/nvidia/attention.py b/vllm/models/inkling/nvidia/attention.py index fc617b89ad..b18fa81f7e 100644 --- a/vllm/models/inkling/nvidia/attention.py +++ b/vllm/models/inkling/nvidia/attention.py @@ -35,11 +35,10 @@ from vllm.v1.kv_cache_interface import ( from ..configs import InklingModelConfig from .layernorm import InklingRMSNorm from .ops.fa4_rel_attention import ( + INKLING_FA4_REL_ATTENTION_KERNEL, bucket_max_seqlen_q, inkling_fa4_num_splits, - inkling_fa4_rel_attention, ) -from .ops.fa4_warmup import InklingFA4WarmupConfig, register_fa4_warmup from .ops.qkvr_prep import fused_qkvr_prep from .sconv_swa_attn import _K, _V, InklingConvState, InklingSconvMetadata from .short_conv import InklingShortConv @@ -175,25 +174,6 @@ class InklingAttention(nn.Module, AttentionLayerBase): compilation_config.static_forward_context[prefix] = self self.kv_cache = torch.tensor([]) # replaced by bind_kv_cache - register_fa4_warmup( - InklingFA4WarmupConfig( - num_heads=self.num_heads, - num_kv_heads=self.num_kv_heads, - head_dim=self.head_dim, - rel_extent=self.rel_extent, - window_size=self.window_size, - is_local=self.is_local, - max_kv_len=self._max_kv_len, - dtype=vllm_config.model_config.dtype, - kv_dtype=self.kv_cache_torch_dtype, - block_size=vllm_config.cache_config.block_size, - max_num_reqs=vllm_config.scheduler_config.max_num_seqs, - max_num_batched_tokens=( - vllm_config.scheduler_config.max_num_batched_tokens - ), - ) - ) - def get_attn_backend(self) -> type[AttentionBackend]: return FlashAttentionBackend @@ -309,7 +289,7 @@ class InklingAttention(nn.Module, AttentionLayerBase): num_kv_heads=self.num_kv_heads, max_kv_len=self._max_kv_len, ) - inkling_fa4_rel_attention( + INKLING_FA4_REL_ATTENTION_KERNEL( q[:nt], key_cache, value_cache, diff --git a/vllm/models/inkling/nvidia/ops/__init__.py b/vllm/models/inkling/nvidia/ops/__init__.py index 2d364fc614..bb7f066ddb 100644 --- a/vllm/models/inkling/nvidia/ops/__init__.py +++ b/vllm/models/inkling/nvidia/ops/__init__.py @@ -15,11 +15,13 @@ from .sconv import fused_sconv _LAZY_EXPORTS = { "silu_and_mul_triton": "silu_and_mul", "sink_silu_mul_epilogue": "silu_and_mul", - "inkling_fa4_rel_attention": "fa4_rel_attention", + "INKLING_FA4_REL_ATTENTION_KERNEL": "fa4_rel_attention", } if TYPE_CHECKING: - from .fa4_rel_attention import inkling_fa4_rel_attention # noqa: F401 + from .fa4_rel_attention import ( # noqa: F401 + INKLING_FA4_REL_ATTENTION_KERNEL, + ) from .silu_and_mul import ( # noqa: F401 silu_and_mul_triton, sink_silu_mul_epilogue, diff --git a/vllm/models/inkling/nvidia/ops/fa4_rel_attention.py b/vllm/models/inkling/nvidia/ops/fa4_rel_attention.py index 883545e3f2..4385786135 100644 --- a/vllm/models/inkling/nvidia/ops/fa4_rel_attention.py +++ b/vllm/models/inkling/nvidia/ops/fa4_rel_attention.py @@ -3,12 +3,23 @@ from __future__ import annotations from collections.abc import Callable +from dataclasses import dataclass from functools import cache -from typing import Any +from typing import TYPE_CHECKING, Any import torch +from vllm.distributed import get_tensor_model_parallel_world_size +from vllm.model_executor.warmup.jit_warmup import ( + VllmJitKernel, + WarmupIntRange, + zip_inputs, +) from vllm.platforms import current_platform +from vllm.utils.torch_utils import kv_cache_dtype_str_to_dtype + +if TYPE_CHECKING: + from vllm.config import VllmConfig def bucket_max_seqlen_q(max_seqlen_q: int) -> int: @@ -98,73 +109,328 @@ def inkling_fa4_num_splits( ) -def inkling_fa4_rel_attention( - q: torch.Tensor, - key_cache: torch.Tensor, - value_cache: torch.Tensor, - *, - block_table: torch.Tensor, - cache_seqlens: torch.Tensor, - cu_seqlens_q: torch.Tensor, - max_seqlen_q: int, - softmax_scale: float, - causal: bool, - window_size: tuple[int, int], - rel_extent: int, - rel_logits: torch.Tensor, - num_splits: int = 32, - out: torch.Tensor | None = None, -) -> torch.Tensor: - """Paged varlen FA4 over the bound K/V cache with the Inkling relative bias. +def _num_warps_bucket(num_reqs: int) -> int: + num_warps = min((num_reqs + 30) // 31, 32) + return 1 << (num_warps - 1).bit_length() - ``q`` is ``(num_tokens, num_heads, head_dim)``; ``key_cache`` / ``value_cache`` - are the paged caches ``(num_blocks, block_size, num_kv_heads, head_dim)``; - ``block_table`` is the per-request page table and ``cache_seqlens`` the - per-request KV lengths (``seqused_k``). ``rel_logits`` is - ``(num_tokens, num_heads, rel_extent)``. - Hopper uses standard FA4's score-mod gather. Blackwell uses tml-fa4's - sheared relative-bias layout. - """ - # cute uses (None, None) to mean "no window". - cute_window = (None, None) if window_size == (-1, -1) else window_size +class InklingFA4RelAttentionKernel( + VllmJitKernel["InklingFA4RelAttentionKernel.CompileKey"] +): + @dataclass(frozen=True) + class CompileKey: + is_local: bool + num_heads: int + num_kv_heads: int + head_dim: int + rel_extent: int + dtype: torch.dtype + kv_dtype: torch.dtype + block_size: int + window_size: tuple[int, int] + max_seqlen_q: int + num_splits: int + num_warps_bucket: int | None + large_num_reqs: bool - rel_logits = rel_logits.contiguous() - flash_attn_varlen_func: Callable[..., Any] - if _use_sheared_bias(): - from vllm.third_party.tml_fa4 import ( - flash_attn_varlen_func as tml_flash_attn_varlen_func, + @staticmethod + def kernel( + q: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + *, + block_table: torch.Tensor, + cache_seqlens: torch.Tensor, + cu_seqlens_q: torch.Tensor, + max_seqlen_q: int, + softmax_scale: float, + causal: bool, + window_size: tuple[int, int], + rel_extent: int, + rel_logits: torch.Tensor, + num_splits: int = 32, + out: torch.Tensor | None = None, + ) -> torch.Tensor: + """Paged varlen FA4 over the bound K/V cache with the Inkling relative bias. + + ``q`` is ``(num_tokens, num_heads, head_dim)``; ``key_cache`` / ``value_cache`` + are the paged caches ``(num_blocks, block_size, num_kv_heads, head_dim)``; + ``block_table`` is the per-request page table and ``cache_seqlens`` the + per-request KV lengths (``seqused_k``). ``rel_logits`` is + ``(num_tokens, num_heads, rel_extent)``. + + Hopper uses standard FA4's score-mod gather. Blackwell uses tml-fa4's + sheared relative-bias layout. + """ + # cute uses (None, None) to mean "no window". + cute_window = (None, None) if window_size == (-1, -1) else window_size + + rel_logits = rel_logits.contiguous() + flash_attn_varlen_func: Callable[..., Any] + if _use_sheared_bias(): + from vllm.third_party.tml_fa4 import ( + flash_attn_varlen_func as tml_flash_attn_varlen_func, + ) + + flash_attn_varlen_func = tml_flash_attn_varlen_func + bias_kwargs: dict[str, Any] = {"rel_bias": rel_logits} + else: + from vllm.vllm_flash_attn.cute import ( + flash_attn_varlen_func as cute_flash_attn_varlen_func, + ) + + flash_attn_varlen_func = cute_flash_attn_varlen_func + bias_kwargs = { + "score_mod": _get_score_mod(rel_extent), + "aux_tensors": [rel_logits], + } + + ret = flash_attn_varlen_func( + q=q, + k=key_cache, + v=value_cache, + cu_seqlens_q=cu_seqlens_q, + seqused_k=cache_seqlens, + max_seqlen_q=max_seqlen_q, + page_table=block_table, + softmax_scale=softmax_scale, + causal=causal, + window_size=cute_window, + num_splits=num_splits, + return_lse=False, + out=out, + **bias_kwargs, + ) + if isinstance(ret, tuple): + return ret[0] + return ret + + def dispatch( # type: ignore[override] + self, + *, + is_local: bool, + num_heads: int, + num_kv_heads: int, + head_dim: int, + rel_extent: int, + dtype: torch.dtype, + kv_dtype: torch.dtype, + block_size: int, + window_size: tuple[int, int], + max_kv_len: int, + query_len: int, + num_reqs: int, + ) -> CompileKey: + max_seqlen_q = bucket_max_seqlen_q(query_len) + num_splits = inkling_fa4_num_splits( + is_local=is_local, + batch_size=num_reqs, + max_query_len=max_seqlen_q, + num_heads=num_heads, + num_kv_heads=num_kv_heads, + max_kv_len=max_kv_len, + ) + return self.CompileKey( + is_local=is_local, + num_heads=num_heads, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + rel_extent=rel_extent, + dtype=dtype, + kv_dtype=kv_dtype, + block_size=block_size, + window_size=window_size, + max_seqlen_q=max_seqlen_q, + num_splits=num_splits, + num_warps_bucket=(_num_warps_bucket(num_reqs) if num_splits > 1 else None), + large_num_reqs=num_reqs > 1024, ) - flash_attn_varlen_func = tml_flash_attn_varlen_func - bias_kwargs: dict[str, Any] = {"rel_bias": rel_logits} - else: - from vllm.vllm_flash_attn.cute import ( - flash_attn_varlen_func as cute_flash_attn_varlen_func, + def _is_valid_warmup_dispatch( + self, + *, + query_len: int, + num_reqs: int, + max_num_batched_tokens: int, + ) -> bool: + return query_len + num_reqs <= max_num_batched_tokens + 1 + + def get_warmup_keys(self, vllm_config: VllmConfig) -> list[CompileKey]: + max_num_reqs = vllm_config.scheduler_config.max_num_seqs + max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens + if max_num_reqs <= 0 or max_num_batched_tokens <= 0: + return [] + + hf_config = vllm_config.model_config.hf_config + get_text_config = getattr(hf_config, "get_text_config", None) + config = get_text_config() if callable(get_text_config) else hf_config + + tp_size = get_tensor_model_parallel_world_size() + dtype = vllm_config.model_config.dtype + kv_dtype = kv_cache_dtype_str_to_dtype( + vllm_config.cache_config.cache_dtype, + vllm_config.model_config, + ) + block_size = vllm_config.cache_config.block_size + local_extent = config.sliding_window_size + + global_num_kv_heads = config.num_key_value_heads + local_num_kv_heads = config.swa_num_key_value_heads + assert config.num_attention_heads % tp_size == 0 + assert config.swa_num_attention_heads % tp_size == 0 + if global_num_kv_heads >= tp_size: + assert global_num_kv_heads % tp_size == 0 + else: + assert tp_size % global_num_kv_heads == 0 + if local_num_kv_heads >= tp_size: + assert local_num_kv_heads % tp_size == 0 + else: + assert tp_size % local_num_kv_heads == 0 + + return self._trace_dispatch(self.dispatch)( + zip_inputs( + dict( + is_local=False, + num_heads=config.num_attention_heads // tp_size, + num_kv_heads=max(1, global_num_kv_heads // tp_size), + head_dim=config.head_dim, + rel_extent=config.rel_extent, + dtype=dtype, + kv_dtype=kv_dtype, + block_size=block_size, + window_size=(-1, -1), + max_kv_len=vllm_config.model_config.max_model_len, + ), + dict( + is_local=True, + num_heads=config.swa_num_attention_heads // tp_size, + num_kv_heads=max(1, local_num_kv_heads // tp_size), + head_dim=config.swa_head_dim, + rel_extent=local_extent, + dtype=dtype, + kv_dtype=kv_dtype, + block_size=block_size, + window_size=(local_extent - 1, 0), + max_kv_len=local_extent, + ), + ), + query_len=WarmupIntRange( + 1, + max_num_batched_tokens + 1, + advance=lambda value: bucket_max_seqlen_q(value) + 1, + ), + num_reqs=WarmupIntRange(1, max_num_reqs + 1), + max_num_batched_tokens=max_num_batched_tokens, + _when=self._is_valid_warmup_dispatch, ) - flash_attn_varlen_func = cute_flash_attn_varlen_func - bias_kwargs = { - "score_mod": _get_score_mod(rel_extent), - "aux_tensors": [rel_logits], - } + def compile(self, compile_key: CompileKey) -> None: + from torch._subclasses.fake_tensor import FakeTensorMode - ret = flash_attn_varlen_func( - q=q, - k=key_cache, - v=value_cache, - cu_seqlens_q=cu_seqlens_q, - seqused_k=cache_seqlens, - max_seqlen_q=max_seqlen_q, - page_table=block_table, - softmax_scale=softmax_scale, - causal=causal, - window_size=cute_window, - num_splits=num_splits, - return_lse=False, - out=out, - **bias_kwargs, - ) - if isinstance(ret, tuple): - return ret[0] - return ret + if compile_key.num_splits == 1 or compile_key.num_warps_bucket is None: + num_reqs = 1025 if compile_key.large_num_reqs else 1 + else: + min_num_warps = ( + 1 + if compile_key.num_warps_bucket == 1 + else compile_key.num_warps_bucket // 2 + 1 + ) + min_num_reqs = (min_num_warps - 1) * 31 + 1 + num_reqs = max( + min_num_reqs, + 1025 if compile_key.large_num_reqs else 1, + ) + + with FakeTensorMode(): + device = torch.accelerator.current_accelerator() + total_q = compile_key.max_seqlen_q + num_reqs - 1 + q = torch.empty( + total_q, + compile_key.num_heads, + compile_key.head_dim, + dtype=compile_key.dtype, + device=device, + ) + kv = torch.empty( + 1, + 2, + compile_key.block_size, + compile_key.num_kv_heads, + compile_key.head_dim, + dtype=compile_key.kv_dtype, + device=device, + ) + key_cache, value_cache = kv.unbind(1) + self.kernel( + q, + key_cache, + value_cache, + block_table=torch.empty( + num_reqs, + 1, + dtype=torch.int32, + device=device, + ), + cache_seqlens=torch.empty( + num_reqs, + dtype=torch.int32, + device=device, + ), + cu_seqlens_q=torch.empty( + num_reqs + 1, + dtype=torch.int32, + device=device, + ), + max_seqlen_q=compile_key.max_seqlen_q, + softmax_scale=compile_key.head_dim**-1, + causal=True, + window_size=compile_key.window_size, + rel_extent=compile_key.rel_extent, + rel_logits=torch.empty( + total_q, + compile_key.num_heads, + compile_key.rel_extent, + dtype=compile_key.dtype, + device=device, + ), + num_splits=compile_key.num_splits, + out=torch.empty_like(q), + ) + + def __call__( + self, + q: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + *, + block_table: torch.Tensor, + cache_seqlens: torch.Tensor, + cu_seqlens_q: torch.Tensor, + max_seqlen_q: int, + softmax_scale: float, + causal: bool, + window_size: tuple[int, int], + rel_extent: int, + rel_logits: torch.Tensor, + num_splits: int = 32, + out: torch.Tensor | None = None, + ) -> torch.Tensor: + return self.kernel( + q, + key_cache, + value_cache, + block_table=block_table, + cache_seqlens=cache_seqlens, + cu_seqlens_q=cu_seqlens_q, + max_seqlen_q=max_seqlen_q, + softmax_scale=softmax_scale, + causal=causal, + window_size=window_size, + rel_extent=rel_extent, + rel_logits=rel_logits, + num_splits=num_splits, + out=out, + ) + + +INKLING_FA4_REL_ATTENTION_KERNEL = InklingFA4RelAttentionKernel() diff --git a/vllm/models/inkling/nvidia/ops/fa4_warmup.py b/vllm/models/inkling/nvidia/ops/fa4_warmup.py deleted file mode 100644 index ff8cdb0293..0000000000 --- a/vllm/models/inkling/nvidia/ops/fa4_warmup.py +++ /dev/null @@ -1,154 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Startup compilation of the FA4 kernels used by Inkling.""" - -from __future__ import annotations - -from collections.abc import Iterator -from dataclasses import dataclass -from functools import partial - -import torch - -from vllm.model_executor.warmup.cutedsl_warmup import ( - CuTeDSLCompileUnit, - register_cutedsl_warmup_provider, -) - -from .fa4_rel_attention import ( - bucket_max_seqlen_q, - inkling_fa4_num_splits, - inkling_fa4_rel_attention, -) - - -@dataclass(frozen=True) -class InklingFA4WarmupConfig: - num_heads: int - num_kv_heads: int - head_dim: int - rel_extent: int - window_size: tuple[int, int] - is_local: bool - max_kv_len: int - dtype: torch.dtype - kv_dtype: torch.dtype - block_size: int - max_num_reqs: int - max_num_batched_tokens: int - - -def _num_warps_bucket(num_reqs: int) -> int: - num_warps = min((num_reqs + 30) // 31, 32) - return 1 << (num_warps - 1).bit_length() - - -def _compile(config: InklingFA4WarmupConfig, max_seqlen_q: int, num_reqs: int) -> None: - from torch._subclasses.fake_tensor import FakeTensorMode - - num_splits = inkling_fa4_num_splits( - is_local=config.is_local, - batch_size=num_reqs, - max_query_len=max_seqlen_q, - num_heads=config.num_heads, - num_kv_heads=config.num_kv_heads, - max_kv_len=config.max_kv_len, - ) - with FakeTensorMode(): - device = torch.accelerator.current_accelerator() - total_q = max_seqlen_q + num_reqs - 1 - q = torch.empty( - total_q, - config.num_heads, - config.head_dim, - dtype=config.dtype, - device=device, - ) - kv = torch.empty( - 1, - 2, - config.block_size, - config.num_kv_heads, - config.head_dim, - dtype=config.kv_dtype, - device=device, - ) - key_cache, value_cache = kv.unbind(1) - inkling_fa4_rel_attention( - q, - key_cache, - value_cache, - block_table=torch.empty(num_reqs, 1, dtype=torch.int32, device=device), - cache_seqlens=torch.empty(num_reqs, dtype=torch.int32, device=device), - cu_seqlens_q=torch.empty(num_reqs + 1, dtype=torch.int32, device=device), - max_seqlen_q=max_seqlen_q, - softmax_scale=1.0 / config.head_dim, - causal=True, - window_size=config.window_size, - rel_extent=config.rel_extent, - rel_logits=torch.empty( - total_q, - config.num_heads, - config.rel_extent, - dtype=config.dtype, - device=device, - ), - num_splits=num_splits, - out=torch.empty_like(q), - ) - - -def _iter_compile_units( - config: InklingFA4WarmupConfig, -) -> Iterator[CuTeDSLCompileUnit]: - max_bucket = bucket_max_seqlen_q(config.max_num_batched_tokens) - max_seqlen_q = 1 - while max_seqlen_q <= max_bucket: - min_query_len = max_seqlen_q // 2 + 1 - max_num_reqs = min( - config.max_num_reqs, - config.max_num_batched_tokens - min_query_len + 1, - ) - seen: set[tuple[int, int, int | None, bool]] = set() - for num_reqs in range(1, max_num_reqs + 1): - num_splits = inkling_fa4_num_splits( - is_local=config.is_local, - batch_size=num_reqs, - max_query_len=max_seqlen_q, - num_heads=config.num_heads, - num_kv_heads=config.num_kv_heads, - max_kv_len=config.max_kv_len, - ) - key = ( - max_seqlen_q, - num_splits, - _num_warps_bucket(num_reqs) if num_splits > 1 else None, - num_reqs > 1024, - ) - if key in seen: - continue - seen.add(key) - yield CuTeDSLCompileUnit( - name="inkling_fa4", - key=("inkling_fa4", config, key), - compile=partial(_compile, config, max_seqlen_q, num_reqs), - ) - max_seqlen_q *= 2 - - -class _WarmupProvider: - def __init__(self) -> None: - self.configs: set[InklingFA4WarmupConfig] = set() - - def get_cutedsl_warmup_compile_units(self) -> tuple[CuTeDSLCompileUnit, ...]: - return tuple( - unit for config in self.configs for unit in _iter_compile_units(config) - ) - - -_PROVIDER = _WarmupProvider() - - -def register_fa4_warmup(config: InklingFA4WarmupConfig) -> None: - _PROVIDER.configs.add(config) - register_cutedsl_warmup_provider(_PROVIDER) diff --git a/vllm/v1/attention/backends/mla/prefill/flash_attn.py b/vllm/v1/attention/backends/mla/prefill/flash_attn.py index 97cb57c18a..ce81be07f1 100644 --- a/vllm/v1/attention/backends/mla/prefill/flash_attn.py +++ b/vllm/v1/attention/backends/mla/prefill/flash_attn.py @@ -13,7 +13,7 @@ import vllm.envs as envs from vllm.model_executor.layers.quantization.utils.quant_utils import ( kFp8StaticTensorSym, ) -from vllm.model_executor.warmup.jit_warmup import VllmJitKernel +from vllm.model_executor.warmup.jit_warmup import VllmJitKernel, zip_inputs from vllm.platforms import current_platform from vllm.v1.attention.backends.fa_utils import ( compile_flash_attn_varlen_func_from_specs, @@ -57,12 +57,6 @@ 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.""" @@ -92,6 +86,16 @@ class FA4MLAPrefillKernel(VllmJitKernel["FA4MLAPrefillKernel.CompileKey"]): return_softmax_lse: bool = False num_splits: int = 0 + @staticmethod + def kernel( + *args: Any, + runtime_kernel: Callable[..., Any] | None = None, + **kwargs: Any, + ) -> Any: + kernel = flash_attn_varlen_func if runtime_kernel is None else runtime_kernel + assert kernel is not None + return kernel(*args, **kwargs) + def dispatch( # type: ignore[override] self, *, @@ -99,7 +103,8 @@ class FA4MLAPrefillKernel(VllmJitKernel["FA4MLAPrefillKernel.CompileKey"]): dtype: torch.dtype, num_heads: int, mla_dims: "MLADims", - shape_probe: _FA4MLAPrefillShapeProbe, + max_seqlen_q: int, + max_seqlen_k: int, requires_v_padding: bool, causal: bool, fa_version: int, @@ -109,17 +114,17 @@ class FA4MLAPrefillKernel(VllmJitKernel["FA4MLAPrefillKernel.CompileKey"]): ) -> CompileKey: return self.CompileKey( q_shape=( - batch_size * shape_probe.max_seqlen_q, + batch_size * 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, + batch_size * 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, + batch_size * max_seqlen_k, num_heads, ( mla_dims.qk_nope_head_dim + mla_dims.qk_rope_head_dim @@ -128,8 +133,8 @@ class FA4MLAPrefillKernel(VllmJitKernel["FA4MLAPrefillKernel.CompileKey"]): ), ), q_dtype=dtype, - max_seqlen_q=shape_probe.max_seqlen_q, - max_seqlen_k=shape_probe.max_seqlen_k, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, softmax_scale=(mla_dims.qk_nope_head_dim + mla_dims.qk_rope_head_dim) ** -0.5, causal=causal, @@ -150,6 +155,28 @@ class FA4MLAPrefillKernel(VllmJitKernel["FA4MLAPrefillKernel.CompileKey"]): num_splits=num_splits, ) + def _is_valid_warmup_shape_probe( + self, + *, + max_seqlen_k: int, + num_splits: int, + is_sm90: bool, + qk_head_dim: int, + effective_v_head_dim: int, + ) -> bool: + long_k = FA4_MLA_PREFILL_LONG_K_BLOCKS * FA4_MLA_PREFILL_K_TILE + very_long_k = FA4_MLA_PREFILL_VERY_LONG_K_BLOCKS * FA4_MLA_PREFILL_K_TILE + return ( + max_seqlen_k < long_k + or (num_splits != 1 and max_seqlen_k == long_k) + or ( + num_splits != 1 + and max_seqlen_k == very_long_k + and not is_sm90 + and qk_head_dim != effective_v_head_dim + ) + ) + def get_warmup_keys(self, vllm_config: "VllmConfig") -> list[CompileKey]: from vllm.model_executor.layers.attention.mla_attention import ( get_mla_dims, @@ -185,58 +212,49 @@ class FA4MLAPrefillKernel(VllmJitKernel["FA4MLAPrefillKernel.CompileKey"]): 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, + shape_probes = ( + dict(max_seqlen_q=1, max_seqlen_k=FA4_MLA_PREFILL_K_TILE), + dict( + max_seqlen_q=FA4_MLA_PREFILL_Q_TILE + 1, + max_seqlen_k=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, - ), - ) - ) + dict( + max_seqlen_q=1, + max_seqlen_k=FA4_MLA_PREFILL_LONG_K_BLOCKS * FA4_MLA_PREFILL_K_TILE, + ), + dict( + max_seqlen_q=FA4_MLA_PREFILL_Q_TILE + 1, + max_seqlen_k=FA4_MLA_PREFILL_LONG_K_BLOCKS * FA4_MLA_PREFILL_K_TILE, + ), + dict( + max_seqlen_q=1, + max_seqlen_k=FA4_MLA_PREFILL_VERY_LONG_K_BLOCKS + * FA4_MLA_PREFILL_K_TILE, + ), + dict( + max_seqlen_q=FA4_MLA_PREFILL_Q_TILE + 1, + max_seqlen_k=FA4_MLA_PREFILL_VERY_LONG_K_BLOCKS + * FA4_MLA_PREFILL_K_TILE, + ), + ) return self._trace_dispatch(self.dispatch)( + zip_inputs(*shape_probes), 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, + is_sm90=is_sm90, + qk_head_dim=qk_head_dim, + effective_v_head_dim=effective_v_head_dim, causal=FA4_MLA_PREFILL_CAUSAL_OPTIONS, return_lse=FA4_MLA_PREFILL_LSE_OPTIONS, num_splits=num_splits, fa_version=fa_version, + _when=self._is_valid_warmup_shape_probe, ) - @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 = ( @@ -271,11 +289,13 @@ class FA4MLAPrefillKernel(VllmJitKernel["FA4MLAPrefillKernel.CompileKey"]): 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() + return self.kernel( + q=q, + k=k, + v=v, + runtime_kernel=runtime_kernel, + **kwargs, + ) class FlashAttnPrefillBackend(MLAPrefillBackend): @@ -470,3 +490,6 @@ class FlashAttnPrefillBackend(MLAPrefillBackend): causal=False, # Context is unmasked return_softmax_lse=True, ) + + +FA4_MLA_PREFILL_KERNEL = FA4MLAPrefillKernel()