@tduncan / Nvidia-Optimization-Stack-VLLM / commits / e16b5e518d

[1/N][warmup][DSv4] Migrate generic MLA metadata and indexing kernels (#50175)

Signed-off-by: LopezCastroRoberto <rocastro@redhat.com> Co-authored-by: Codex <codex@openai.com>

Roberto L. Castro committed Sep 1, 2026 at 05:26 UTC e16b5e518db8ed54a9c82dd825d3a3d5502ef13e

17 files changed +1269 -625

docs/contributing/jit_kernel_warmup.md

+10 -8

@@ -17,7 +17,7 @@ Each warmable kernel defines its compile-key mapping and compile-only entry poin

18 Here, a **kernel wrapper** (or just **wrapper**) is an instance of a concrete `VllmJitKernel` subclass.

20 -Expose one wrapper near the kernel's normal runtime entry point. Prefer this shape:

20 +Expose one wrapper near the kernel's normal runtime entry point. Prefer this backend-agnostic shape:

22 ```python

23 class MyKernel(VllmJitKernel["MyKernel.CompileKey"]):

@@ -43,14 +43,14 @@ class MyKernel(VllmJitKernel["MyKernel.CompileKey"]):

43 return self.kernel(...)

46 -MY_KERNEL = MyKernel()

46 +_MY_KERNEL = MyKernel()

47 ```

49 `CompileKey`, `dispatch(...)`, and `get_warmup_keys(...)` are backend-agnostic. Backend-specific behavior belongs in `kernel(...)`, `compile(...)`, and `__call__(...)`.

51 The module-level singleton should be used by warmup and by the runtime call path. This keeps dispatch behavior shared instead of duplicated.

53 -`VllmJitKernel.warmup(...)` compiles every key returned by `get_warmup_keys(...)`; wrappers should not reimplement it.

53 +`VllmJitKernel.warmup(...)` materializes every key returned by `get_warmup_keys(...)`; wrappers should not reimplement it.

55 ### Choose Compile-Key Fields

@@ -70,18 +70,20 @@ def get_warmup_keys(self, vllm_config: VllmConfig) -> list[CompileKey]:

71 Use independent ranges or alternatives for cartesian products, `zip_inputs(...)` for coupled rows, and `_when` for validity constraints. The complete syntax is documented in [Search-Space Reference](#2-search-space-reference).

73 -### Compile Without Launching

73 +### Materialize Without Launching

75 -`compile(compile_key)` means "make this specialization available". Depending on the backend, that may compile from source, call a compile-only API, load an already-built artifact, or compile on cache miss.

75 +The base contract uses `compile(compile_key)` internally to mean "make this specialization available". Depending on the backend, that may compile from source, call a compile-only API, load an already-built artifact, or compile on cache miss. Contributors only implement this hook when their backend helper does not provide it.

77 -`compile(...)` should not launch a real inference workload or allocate real tensors. Each DSL should expose fake tensor/spec descriptors suitable for compilation only.

77 +For Triton, implement `warmup_inputs(...)`. `@kernel_launcher` binds those inputs to `__call__(...)`, whose body returns the runtime grid and derived launch arguments as a `LaunchSpec`. The shared helper then routes that specification to the normal kernel launch or to `kernel.warmup(...)` during startup.

78 +

79 +Materialization should not launch a real inference workload or allocate real tensors. `TritonWarmupTensor(..., strides=None)` represents compact row-major storage; provide explicit representative strides for padded, transposed, or otherwise strided runtime inputs.

81 ### Register the Selected Wrapper

83 Register the wrapper where the runtime implementation is selected:

85 ```python

84 -MY_KERNEL.register_warmup()

86 +_MY_KERNEL.register_warmup()

87 ```

89 Registration records metadata only. It does not compile or launch the kernel. Repeated registrations from equivalent layers are allowed and deduplicated later.

@@ -92,7 +94,7 @@ Registration records metadata only. It does not compile or launch the kernel. Re

94 - Keep specialization mapping in `dispatch(...)` instead of duplicating it in warmup code.

95 - Use fake tensors or backend compile-only descriptors; never perform a dummy runtime launch.

96 - Keep registration metadata-only so model construction remains cheap and side-effect free.

95 -- Compile registered kernels only through `kernel_warmup()`.

97 +- Materialize registered kernels only through `kernel_warmup()`.

98 - Keep runtime execution and startup compilation separate and easy to review.

99 - Use one module-level wrapper instance for registration and runtime calls.

100

tests/model_executor/layers/test_fused_shared_expert.py

@@ -553,6 +553,7 @@ def test_models_fse_init(

553 method="mtp",

554 parallel_drafting=False,

555 enable_adaptive_verification=False,

556 + use_dspark=lambda: False,

557 )

558 mtp = DeepSeekV4MTP(vllm_config=vllm_config)

559 assert model.is_fused_shared_expert_enabled is (fse_enabled and not exclude)

tests/model_executor/test_jit_warmup.py

+29

@@ -590,6 +590,35 @@ def test_registry_passes_vllm_config_to_default_requests() -> None:

590 assert kernel.compiled == [ConfigKernel.CompileKey(value=7)]

591

592

593 +def test_registry_default_request_supports_config_free_kernel() -> None:

594 + class ConfigFreeKernel(VllmJitKernel["ConfigFreeKernel.CompileKey"]):

595 + @dataclass(frozen=True)

596 + class CompileKey:

597 + value: int

598 +

599 + def __init__(self) -> None:

600 + self.compiled: list[ConfigFreeKernel.CompileKey] = []

601 + super().__init__()

602 +

603 + def dispatch(self, *, value: int) -> CompileKey: # type: ignore[override]

604 + return self.CompileKey(value=value)

605 +

606 + def get_warmup_keys(self) -> list[CompileKey]:

607 + return [self.dispatch(value=7)]

608 +

609 + def compile(self, compile_key: CompileKey) -> None:

610 + self.compiled.append(compile_key)

611 +

612 + registry = JitWarmupRegistry(_config())

613 + kernel = ConfigFreeKernel()

614 +

615 + with registry.activate():

616 + kernel.register_warmup()

617 + registry.warmup()

618 +

619 + assert kernel.compiled == [ConfigFreeKernel.CompileKey(value=7)]

620 +

621 +

622 def test_get_ast_full_name_handles_names_attributes_and_other_nodes() -> None:

623 dotted_expr = ast.parse("foo.bar.baz").body[0]

624 call_expr = ast.parse("foo()").body[0]

tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py

+43 -2

@@ -8,10 +8,16 @@ import torch

9 from tests.v1.attention.utils import create_vllm_config

10 from vllm.v1.attention.backend import CommonAttentionMetadata

11 +from vllm.v1.attention.backends.mla.compressor_utils import (

12 + CompressedSlotMappingKernel,

13 +)

14 from vllm.v1.attention.backends.mla.indexer import (

15 BuildPrefillChunkMetadataKernel,

16 DeepseekV32IndexerMetadataBuilder,

17 )

18 +from vllm.v1.attention.backends.mla.sparse_utils import (

19 + ConvertReqIndexToGlobalIndexKernel,

20 +)

21 from vllm.v1.kv_cache_interface import MLAAttentionSpec

22 from vllm.v1.worker.block_table import get_block_table_width

@@ -20,7 +26,7 @@ def test_indexer_warmup_normalizes_zero_compress_ratios():

26 config = SimpleNamespace(

27 scheduler_config=SimpleNamespace(max_num_batched_tokens=8),

28 model_config=SimpleNamespace(

23 - hf_config=SimpleNamespace(compress_ratios=[0, 0, 4, 128, 0])

29 + hf_config=SimpleNamespace(compress_ratios=[0, 0, 4, 128, 0], index_kpool=32)

30 ),

31 parallel_config=SimpleNamespace(

32 decode_context_parallel_size=1,

@@ -30,7 +36,42 @@ def test_indexer_warmup_normalizes_zero_compress_ratios():

37 keys = BuildPrefillChunkMetadataKernel().get_warmup_keys(config)

33 - assert {key.COMPRESS_RATIO for key in keys} == {1, 4, 128}

39 + assert {key.compress_ratio for key in keys} == {1, 4, 32, 128}

40 + assert {(key.query_slice_start, key.query_slice_stop) for key in keys} == {

41 + (query_slice_start, query_slice_stop)

42 + for query_slice_start in (1, 2, 16)

43 + for query_slice_stop in (1, 2, 16)

44 + }

45 +

46 +

47 +def test_compressed_slot_mapping_warmup_includes_index_kpool():

48 + config = SimpleNamespace(

49 + cache_config=SimpleNamespace(block_size=256),

50 + model_config=SimpleNamespace(hf_config=SimpleNamespace(index_kpool=32)),

51 + )

52 +

53 + keys = CompressedSlotMappingKernel().get_warmup_keys(config)

54 + assert {(key.compress_ratio, key.block_size) for key in keys} == {(32, 2)}

55 +

56 +

57 +def test_index_conversion_warmup_uses_physical_block_stride():

58 + config = SimpleNamespace(

59 + cache_config=SimpleNamespace(block_size=64),

60 + model_config=SimpleNamespace(

61 + max_model_len=1024,

62 + hf_config=SimpleNamespace(index_topk=2048),

63 + ),

64 + parallel_config=SimpleNamespace(

65 + decode_context_parallel_size=1,

66 + cp_kv_cache_interleave_size=1,

67 + ),

68 + )

69 +

70 + keys = ConvertReqIndexToGlobalIndexKernel().get_warmup_keys(

71 + config,

72 + block_stride_rows=4096,

73 + )

74 + assert {key.block_stride_rows for key in keys} == {4096}

77 @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")

tests/v1/worker/test_gpu_model_runner_v2_eplb.py

@@ -7,6 +7,7 @@ from typing import Any

8 import torch

10 +from vllm.model_executor.warmup.jit_warmup import JitWarmupRegistry

11 from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT

12 from vllm.v1.worker.gpu import eplb_utils as eplb

13 from vllm.v1.worker.gpu import model_runner as mrv2

@@ -84,6 +85,7 @@ def _make_runner(**overrides: Any) -> Any:

85 post_forward=lambda *_, **__: None,

86 )

87 runner.eplb = eplb.EPLBController(runner.parallel_config, runner.device)

88 + runner.jit_warmup_registry = JitWarmupRegistry(runner.vllm_config)

89 runner.pooling_runner = None

90 runner.execute_model_state = None

91 for key, value in overrides.items():

vllm/model_executor/layers/attention/mla_attention.py

+47

@@ -604,6 +604,33 @@ class MLAAttention(nn.Module, AttentionLayerBase):

604

605 self.use_sparse = use_sparse

606

607 + if vllm_config.kernel_config.enable_jit_warmup:

608 + backend_name = self.attn_backend.get_name()

609 + if backend_name in (

610 + "FLASHMLA_SPARSE",

611 + "FLASHINFER_MLA_SPARSE",

612 + "FLASHINFER_MLA_SPARSE_SM120",

613 + "DEEPSEEK_V32_INDEXER",

614 + ):

615 + from vllm.v1.attention.backends.mla.compressor_utils import (

616 + _COMPRESSED_SLOT_MAPPING_KERNEL,

617 + )

618 + from vllm.v1.attention.backends.mla.indexer import (

619 + _BUILD_PREFILL_CHUNK_METADATA_KERNEL,

620 + _PREPARE_UNIFORM_DECODE_KERNEL,

621 + )

622 +

623 + _COMPRESSED_SLOT_MAPPING_KERNEL.register_warmup()

624 + _PREPARE_UNIFORM_DECODE_KERNEL.register_warmup()

625 + _BUILD_PREFILL_CHUNK_METADATA_KERNEL.register_warmup()

626 +

627 + if backend_name != "DEEPSEEK_V32_INDEXER":

628 + from vllm.v1.attention.backends.mla.sparse_swa import (

629 + _COMPUTE_PREFILL_METADATA_KERNEL,

630 + )

631 +

632 + _COMPUTE_PREFILL_METADATA_KERNEL.register_warmup()

633 +

634 self.dcp_manager: MLADCPManager | None = None

635 if self.impl.dcp_world_size > 1:

636 query_dtype = (

@@ -652,6 +679,26 @@ class MLAAttention(nn.Module, AttentionLayerBase):

679 def bind_kv_cache(self, kv_cache: torch.Tensor) -> None:

680 # [B, H=1, N, C] -> [B, N, C]

681 self.kv_cache = kv_cache.squeeze(1)

682 + if (

683 + self._vllm_config.kernel_config.enable_jit_warmup

684 + and self.attn_backend.get_name()

685 + in (

686 + "FLASHMLA_SPARSE",

687 + "FLASHINFER_MLA_SPARSE",

688 + "FLASHINFER_MLA_SPARSE_SM120",

689 + "DEEPSEEK_V32_INDEXER",

690 + )

691 + ):

692 + from vllm.v1.attention.backends.mla.sparse_utils import (

693 + _CONVERT_REQ_INDEX_TO_GLOBAL_INDEX_KERNEL,

694 + )

695 +

696 + row_width = self.kv_cache.shape[-1]

697 + assert self.kv_cache.stride(0) % row_width == 0

698 + _CONVERT_REQ_INDEX_TO_GLOBAL_INDEX_KERNEL.register_warmup(

699 + self._vllm_config,

700 + block_stride_rows=self.kv_cache.stride(0) // row_width,

701 + )

702

703 @property

704 def chunked_prefill_workspace_size(self) -> int:

vllm/model_executor/warmup/jit_warmup.py

+5 -1

@@ -737,7 +737,11 @@ class JitWarmupRegistry:

737 for kernel, registrations in self._registrations.items():

738 compile_keys: dict[Any, None] = {}

739 for args, kwargs in registrations:

740 - if not args and not kwargs:

740 + if (

741 + not args

742 + and not kwargs

743 + and inspect.signature(kernel.get_warmup_keys).parameters

744 + ):

745 args = (self.vllm_config,)

746 for compile_key in kernel.get_warmup_keys(*args, **kwargs):

747 compile_keys[compile_key] = None

vllm/model_executor/warmup/jit_warmup_triton_helper.py

+90 -10

@@ -2,15 +2,21 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 import ast

4 import inspect

5 -from collections.abc import Callable

5 +from abc import abstractmethod

6 +from collections.abc import Callable, Mapping

7 from dataclasses import dataclass

7 -from typing import Any

8 +from functools import cached_property, wraps

9 +from typing import Any, ClassVar, Generic, TypeVar, cast

11 from vllm.model_executor.warmup.jit_warmup import (

12 + VllmJitKernel,

13 get_ast_full_name,

14 get_function_source_node,

15 )

17 +CompileKeyT = TypeVar("CompileKeyT")

18 +LaunchSpec = tuple[tuple[int, ...], dict[str, Any]]

19 +

21 def triton_scalar_specialization_rep(value: int) -> int:

22 """Return an integer with the same default Triton JIT specialization.

@@ -50,10 +56,16 @@ def triton_scalar_specialization_rep(value: int) -> int:

57 @dataclass(frozen=True)

58 class TritonWarmupTensor:

53 - # Compile-only tensor descriptor for Triton pointer specialization.

59 + """Compile-only tensor metadata used by Triton warmup.

60 +

61 + ``strides=None`` represents compact row-major storage. Pass explicit strides

62 + whenever the runtime tensor can be padded, transposed, or otherwise strided.

63 + """

64 +

65 dtype: Any

66 aligned: bool = True

67 shape: tuple[int, ...] = (1,)

68 + strides: tuple[int, ...] | None = None

70 def data_ptr(self) -> int:

71 return 0 if self.aligned else 1

@@ -61,13 +73,81 @@ class TritonWarmupTensor:

73 def ptr_range(self) -> int:

74 return 0

64 - def stride(self) -> tuple[int, ...]:

65 - strides: list[int] = []

66 - stride = 1

67 - for size in reversed(self.shape):

68 - strides.append(stride)

69 - stride *= size

70 - return tuple(reversed(strides))

76 + def stride(self, dim: int | None = None) -> int | tuple[int, ...]:

77 + if self.strides is None:

78 + strides: list[int] = []

79 + stride = 1

80 + for size in reversed(self.shape):

81 + strides.append(stride)

82 + stride *= size

83 + result = tuple(reversed(strides))

84 + else:

85 + result = self.strides

86 + return result if dim is None else result[dim]

87 +

88 +

89 +class VllmTritonJitKernel(VllmJitKernel[CompileKeyT], Generic[CompileKeyT]):

90 + """Triton owner whose runtime launch specification is reused for warmup."""

91 +

92 + kernel: ClassVar[Any]

93 + _warming = False

94 +

95 + @abstractmethod

96 + def warmup_inputs(self, compile_key: CompileKeyT) -> dict[str, Any]:

97 + """Return runtime-shaped inputs that reproduce one compile key."""

98 + raise NotImplementedError

99 +

100 + def compile(self, compile_key: CompileKeyT) -> None:

101 + inputs = self.warmup_inputs(compile_key)

102 + self._warming = True

103 + try:

104 + cast(Callable[..., None], self)(**inputs)

105 + finally:

106 + self._warming = False

107 +

108 + @cached_property

109 + def _kernel_param_names(self) -> frozenset[str]:

110 + return frozenset(self.kernel.arg_names)

111 +

112 + def launch(

113 + self,

114 + grid: tuple[int, ...],

115 + inputs: Mapping[str, Any],

116 + /,

117 + **kwargs: Any,

118 + ) -> Any:

119 + for name, value in inputs.items():

120 + target = name if name in self._kernel_param_names else f"{name}_ptr"

121 + if target in self._kernel_param_names and target not in kwargs:

122 + kwargs[target] = value

123 + if self._warming:

124 + warmup = getattr(self.kernel, "warmup", None)

125 + assert warmup is not None

126 + return warmup(grid=(1,), **kwargs)

127 + return self.kernel[grid](**kwargs)

128 +

129 +

130 +def kernel_launcher(

131 + call_fn: Callable[..., LaunchSpec],

132 +) -> Callable[..., None]:

133 + """Launch a Triton kernel from a declarative ``__call__`` specification."""

134 + signature = inspect.signature(call_fn)

135 +

136 + @wraps(call_fn)

137 + def wrapper(

138 + self: VllmTritonJitKernel[Any],

139 + *args: Any,

140 + **kwargs: Any,

141 + ) -> None:

142 + grid, launch_kwargs = call_fn(self, *args, **kwargs)

143 + bound = signature.bind(self, *args, **kwargs)

144 + bound.apply_defaults()

145 + inputs = {

146 + name: value for name, value in bound.arguments.items() if name != "self"

147 + }

148 + self.launch(grid, inputs, **launch_kwargs)

149 +

150 + return wrapper

151

152

153 @dataclass(frozen=True)

vllm/model_executor/warmup/kernel_warmup.py

-4

@@ -35,9 +35,6 @@ from vllm.model_executor.warmup.kimi_k3_triton_warmup import (

35 kimi_k3_triton_warmup,

36 )

37 from vllm.model_executor.warmup.qwen_triton_warmup import qwen_triton_warmup

38 -from vllm.model_executor.warmup.sparse_mla_triton_warmup import (

39 - sparse_mla_triton_warmup,

40 -)

38 from vllm.platforms import current_platform

39 from vllm.utils.deep_gemm import is_deep_gemm_supported

40 from vllm.utils.flashinfer import has_flashinfer

@@ -154,7 +151,6 @@ def kernel_warmup(worker: "Worker", *, process_local_only: bool = False):

151 if worker.vllm_config.kernel_config.enable_jit_warmup:

152 kimi_k3_triton_warmup(worker)

153 fa4_cutedsl_warmup(worker)

157 - sparse_mla_triton_warmup(worker)

154

155 if current_platform.has_device_capability(90):

156 _warmup_ll_bf16_router_gemm(worker.get_model())

vllm/model_executor/warmup/sparse_mla_triton_warmup.py deleted

-111

@@ -1,111 +0,0 @@

1 -# SPDX-License-Identifier: Apache-2.0

2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -"""Warm up sparse-MLA Triton metadata kernels."""

4 -

5 -from typing import TYPE_CHECKING

6 -

7 -from vllm.logger import init_logger

8 -

9 -if TYPE_CHECKING:

10 - from vllm.config import VllmConfig

11 - from vllm.v1.worker.gpu_model_runner import GPUModelRunner

12 - from vllm.v1.worker.gpu_worker import Worker

13 -

14 -logger = init_logger(__name__)

15 -

16 -_DEEPSEEK_V4_SPARSE_MLA_BACKENDS = frozenset(

17 - {

18 - "FLASHMLA_SPARSE_DSV4",

19 - "FLASHINFER_MLA_SPARSE_DSV4",

20 - "ROCM_FLASHMLA_SPARSE_DSV4",

21 - "DEEPSEEK_SPARSE_SWA",

22 - }

23 -)

24 -_GENERIC_SPARSE_MLA_BACKENDS = frozenset(

25 - {

26 - "FLASHMLA_SPARSE",

27 - "FLASHINFER_MLA_SPARSE",

28 - "FLASHINFER_MLA_SPARSE_SM120",

29 - }

30 -)

31 -_INDEXER_PREFILL_CHUNK_METADATA_BACKENDS = frozenset({"DEEPSEEK_V32_INDEXER"})

32 -

33 -_INDEXER_PREFILL_CHUNK_METADATA_BACKENDS = frozenset({"DEEPSEEK_V32_INDEXER"})

34 -

35 -

36 -def _attention_backend_name(backend: object) -> str | None:

37 - get_name = getattr(backend, "get_name", None)

38 - if get_name is None:

39 - return None

40 - try:

41 - return get_name()

42 - except NotImplementedError:

43 - return None

44 -

45 -

46 -def _has_attention_backend(

47 - runner: "GPUModelRunner",

48 - backend_names: frozenset[str],

49 -) -> bool:

50 - for groups in getattr(runner, "attn_groups", []) or ():

51 - for group in groups:

52 - name = _attention_backend_name(getattr(group, "backend", None))

53 - if name in backend_names:

54 - return True

55 - return False

56 -

57 -

58 -def _compile_sparse_swa_prefill_metadata_kernel(

59 - vllm_config: "VllmConfig",

60 -) -> None:

61 - from vllm.v1.attention.backends.mla.sparse_swa import (

62 - _COMPUTE_PREFILL_METADATA_KERNEL,

63 - )

64 -

65 - _COMPUTE_PREFILL_METADATA_KERNEL.warmup(vllm_config)

66 -

67 -

68 -def _compile_prefill_chunk_metadata_kernel(

69 - vllm_config: "VllmConfig",

70 -) -> None:

71 - from vllm.v1.attention.backends.mla.indexer import (

72 - _BUILD_PREFILL_CHUNK_METADATA_KERNEL,

73 - )

74 -

75 - _BUILD_PREFILL_CHUNK_METADATA_KERNEL.warmup(vllm_config)

76 -

77 -

78 -def _compile_combine_topk_swa_indices_kernel(

79 - vllm_config: "VllmConfig",

80 -) -> None:

81 - from vllm.models.deepseek_v4.common.ops.cache_utils import (

82 - _COMBINE_TOPK_SWA_INDICES_KERNEL,

83 - )

84 -

85 - _COMBINE_TOPK_SWA_INDICES_KERNEL.warmup(vllm_config)

86 -

87 -

88 -def sparse_mla_triton_warmup(worker: "Worker") -> None:

89 - runner = worker.model_runner

90 - if runner.is_pooling_model:

91 - return

92 -

93 - max_tokens = worker.scheduler_config.max_num_batched_tokens

94 - max_num_prefills = min(worker.scheduler_config.max_num_seqs, max_tokens)

95 - if max_tokens <= 0 or max_num_prefills <= 0:

96 - return

97 -

98 - vllm_config = runner.vllm_config

99 - try:

100 - if _has_attention_backend(runner, _DEEPSEEK_V4_SPARSE_MLA_BACKENDS):

101 - _compile_sparse_swa_prefill_metadata_kernel(vllm_config)

102 - _compile_prefill_chunk_metadata_kernel(vllm_config)

103 - _compile_combine_topk_swa_indices_kernel(vllm_config)

104 - elif _has_attention_backend(runner, _GENERIC_SPARSE_MLA_BACKENDS):

105 - _compile_sparse_swa_prefill_metadata_kernel(vllm_config)

106 - _compile_prefill_chunk_metadata_kernel(vllm_config)

107 - elif _has_attention_backend(runner, _INDEXER_PREFILL_CHUNK_METADATA_BACKENDS):

108 - _compile_prefill_chunk_metadata_kernel(vllm_config)

109 -

110 - except Exception:

111 - logger.warning("Skipping sparse MLA Triton warmup.", exc_info=True)

vllm/models/deepseek_v4/attention.py

+51

@@ -353,6 +353,57 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):

353 k_cache_prefix=self.prefix,

354 )

355

356 + if vllm_config.kernel_config.enable_jit_warmup:

357 + from vllm.v1.attention.backends.mla.sparse_swa import (

358 + _COMPUTE_PREFILL_METADATA_KERNEL,

359 + _COMPUTE_SWA_INDICES_AND_LENS_KERNEL,

360 + )

361 +

362 + _COMPUTE_PREFILL_METADATA_KERNEL.register_warmup()

363 + _COMPUTE_SWA_INDICES_AND_LENS_KERNEL.register_warmup(

364 + window_size=self.window_size,

365 + block_size=self.swa_cache_layer.block_size,

366 + )

367 +

368 + if self.compress_ratio > 1:

369 + from vllm.v1.attention.backends.mla.compressor_utils import (

370 + _COMPRESSED_SLOT_MAPPING_KERNEL,

371 + )

372 +

373 + _COMPRESSED_SLOT_MAPPING_KERNEL.register_warmup()

374 +

375 + if self.indexer is not None:

376 + from vllm.v1.attention.backends.mla.indexer import (

377 + _BUILD_PREFILL_CHUNK_METADATA_KERNEL,

378 + _PREPARE_UNIFORM_DECODE_KERNEL,

379 + )

380 +

381 + _PREPARE_UNIFORM_DECODE_KERNEL.register_warmup()

382 + _BUILD_PREFILL_CHUNK_METADATA_KERNEL.register_warmup()

383 +

384 + spec_config = vllm_config.speculative_config

385 + if spec_config is not None and spec_config.use_dspark():

386 + from vllm.v1.attention.backends.mla.sparse_swa import (

387 + _COMPUTE_DSPARK_NONCAUSAL_SWA_INDICES_KERNEL,

388 + )

389 +

390 + _COMPUTE_DSPARK_NONCAUSAL_SWA_INDICES_KERNEL.register_warmup(

391 + window_size=self.window_size,

392 + num_speculative_tokens=spec_config.num_speculative_tokens,

393 + block_size=self.swa_cache_layer.block_size,

394 + )

395 +

396 + if self.backend_cls.get_name() in (

397 + "FLASHMLA_SPARSE_DSV4",

398 + "ROCM_FLASHMLA_SPARSE_DSV4",

399 + "XPU_V4_MLA_SPARSE",

400 + ):

401 + from vllm.models.deepseek_v4.common.ops.cache_utils import (

402 + _COMBINE_TOPK_SWA_INDICES_KERNEL,

403 + )

404 +

405 + _COMBINE_TOPK_SWA_INDICES_KERNEL.register_warmup()

406 +

407 def forward(

408 self,

409 positions: torch.Tensor,

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

+125 -40

@@ -1,7 +1,18 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +from dataclasses import dataclass

4 +from typing import Any

5 +

6 import torch

8 +from vllm.model_executor.warmup.jit_warmup import zip_inputs

9 +from vllm.model_executor.warmup.jit_warmup_triton_helper import (

10 + LaunchSpec,

11 + TritonWarmupTensor,

12 + VllmTritonJitKernel,

13 + kernel_launcher,

14 + triton_scalar_specialization_rep,

15 +)

16 from vllm.triton_utils import tl, triton

17 from vllm.utils.math_utils import cdiv

@@ -17,49 +28,124 @@ def get_dspark_swa_index_width(

28 return cdiv(width, _DSPARK_SWA_INDEX_ALIGNMENT) * _DSPARK_SWA_INDEX_ALIGNMENT

20 -@triton.jit

21 -def _compressed_slot_mapping_kernel(

22 - # [num_tokens]

23 - slot_mapping_ptr,

24 - # [num_reqs + 1]

25 - query_start_loc_ptr,

26 - # [num_reqs]

27 - seq_lens_ptr,

28 - # [num_reqs, max_num_blocks]

29 - block_table_ptr,

30 - block_table_stride,

31 - block_size,

32 - COMPRESS_RATIO: tl.constexpr,

33 - PAD_ID: tl.constexpr,

34 - TRITON_BLOCK_SIZE: tl.constexpr,

31 +class CompressedSlotMappingKernel(

32 + VllmTritonJitKernel["CompressedSlotMappingKernel.CompileKey"]

33 ):

36 - batch_idx = tl.program_id(0)

34 + TRITON_BLOCK_SIZE = 1024

35 +

36 + @dataclass(frozen=True)

37 + class CompileKey:

38 + compress_ratio: int

39 + triton_block_size: int

40 + block_size: int

41 +

42 + @staticmethod

43 + @triton.jit(do_not_specialize=["block_table_stride"])

44 + def kernel(

45 + # [num_tokens]

46 + slot_mapping_ptr,

47 + # [num_reqs + 1]

48 + query_start_loc_ptr,

49 + # [num_reqs]

50 + seq_lens_ptr,

51 + # [num_reqs, max_num_blocks]

52 + block_table_ptr,

53 + block_table_stride,

54 + block_size,

55 + COMPRESS_RATIO: tl.constexpr,

56 + PAD_ID: tl.constexpr,

57 + TRITON_BLOCK_SIZE: tl.constexpr,

58 + ):

59 + batch_idx = tl.program_id(0)

38 - query_start = tl.load(query_start_loc_ptr + batch_idx)

39 - query_end = tl.load(query_start_loc_ptr + batch_idx + 1)

40 - query_len = query_end - query_start

61 + query_start = tl.load(query_start_loc_ptr + batch_idx)

62 + query_end = tl.load(query_start_loc_ptr + batch_idx + 1)

63 + query_len = query_end - query_start

42 - seq_len = tl.load(seq_lens_ptr + batch_idx)

43 - start_pos = seq_len - query_len

65 + seq_len = tl.load(seq_lens_ptr + batch_idx)

66 + start_pos = seq_len - query_len

45 - for i in range(0, query_len, TRITON_BLOCK_SIZE):

46 - offset = i + tl.arange(0, TRITON_BLOCK_SIZE)

47 - mask = offset < query_len

68 + for i in range(0, query_len, TRITON_BLOCK_SIZE):

69 + offset = i + tl.arange(0, TRITON_BLOCK_SIZE)

70 + mask = offset < query_len

49 - pos = start_pos + i + tl.arange(0, TRITON_BLOCK_SIZE)

50 - is_valid = (pos + 1) % COMPRESS_RATIO == 0

51 - pos_after_compress = pos // COMPRESS_RATIO

72 + pos = start_pos + i + tl.arange(0, TRITON_BLOCK_SIZE)

73 + is_valid = (pos + 1) % COMPRESS_RATIO == 0

74 + pos_after_compress = pos // COMPRESS_RATIO

53 - block_ids = pos_after_compress // block_size

54 - block_numbers = tl.load(

55 - block_table_ptr + batch_idx * block_table_stride + block_ids,

56 - mask=mask & is_valid,

76 + block_ids = pos_after_compress // block_size

77 + block_numbers = tl.load(

78 + block_table_ptr + batch_idx * block_table_stride + block_ids,

79 + mask=mask & is_valid,

80 + )

81 + slot_ids = block_numbers * block_size + pos_after_compress % block_size

82 +

83 + # NOTE

84 + slot_ids = tl.where(is_valid, slot_ids, PAD_ID)

85 + tl.store(slot_mapping_ptr + query_start + offset, slot_ids, mask=mask)

86 +

87 + def dispatch( # type: ignore[override]

88 + self,

89 + *,

90 + compress_ratio: int,

91 + block_size: int,

92 + ) -> CompileKey:

93 + return self.CompileKey(

94 + compress_ratio=compress_ratio,

95 + triton_block_size=self.TRITON_BLOCK_SIZE,

96 + block_size=triton_scalar_specialization_rep(block_size),

97 )

58 - slot_ids = block_numbers * block_size + pos_after_compress % block_size

60 - # NOTE

61 - slot_ids = tl.where(is_valid, slot_ids, PAD_ID)

62 - tl.store(slot_mapping_ptr + query_start + offset, slot_ids, mask=mask)

99 + def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]:

100 + hf_config = vllm_config.model_config.hf_config

101 + configured_ratios = (

102 + *(getattr(hf_config, "compress_ratios", None) or ()),

103 + getattr(hf_config, "index_kpool", 1) or 1,

104 + )

105 + compress_ratios = tuple(

106 + dict.fromkeys(int(ratio) for ratio in configured_ratios if int(ratio) > 1)

107 + )

108 + if not compress_ratios:

109 + return []

110 + return self._trace_dispatch(self.dispatch)(

111 + zip_inputs(

112 + *(

113 + dict(

114 + compress_ratio=ratio,

115 + block_size=vllm_config.cache_config.block_size // ratio,

116 + )

117 + for ratio in compress_ratios

118 + )

119 + )

120 + )

121 +

122 + def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:

123 + int32_ptr = TritonWarmupTensor(torch.int32)

124 + return dict(

125 + slot_mapping=TritonWarmupTensor(torch.int64),

126 + query_start_loc=int32_ptr,

127 + seq_lens=int32_ptr,

128 + block_table=int32_ptr,

129 + block_size=compile_key.block_size,

130 + compress_ratio=compile_key.compress_ratio,

131 + )

132 +

133 + @kernel_launcher

134 + def __call__(

135 + self,

136 + slot_mapping: torch.Tensor,

137 + query_start_loc: torch.Tensor,

138 + seq_lens: torch.Tensor,

139 + block_table: torch.Tensor,

140 + block_size: int,

141 + compress_ratio: int,

142 + ) -> LaunchSpec:

143 + return (block_table.shape[0],), dict(

144 + block_table_stride=block_table.stride(0),

145 + COMPRESS_RATIO=compress_ratio,

146 + PAD_ID=-1,

147 + TRITON_BLOCK_SIZE=self.TRITON_BLOCK_SIZE,

148 + )

149

150

151 def get_compressed_slot_mapping(

@@ -83,16 +169,15 @@ def get_compressed_slot_mapping(

169 (num_tokens,), -1, dtype=torch.int64, device=query_start_loc.device

170 )

171

86 - num_reqs = block_table.shape[0]

87 - _compressed_slot_mapping_kernel[(num_reqs,)](

172 + _COMPRESSED_SLOT_MAPPING_KERNEL(

173 slot_mapping,

174 query_start_loc,

175 seq_lens,

176 block_table,

92 - block_table.stride(0),

177 block_size,

178 compress_ratio,

95 - PAD_ID=-1,

96 - TRITON_BLOCK_SIZE=1024,

179 )

180 return slot_mapping

181 +

182 +

183 +_COMPRESSED_SLOT_MAPPING_KERNEL = CompressedSlotMappingKernel()

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

+163 -101

@@ -1,6 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 from dataclasses import dataclass

4 +from typing import Any

6 import torch

@@ -8,13 +9,13 @@ import vllm.envs as envs

9 from vllm.config import VllmConfig

10 from vllm.distributed import get_dcp_group, get_pcp_group

11 from vllm.logger import init_logger

11 -from vllm.model_executor.warmup.jit_warmup import (

12 - VllmJitKernel,

13 - WarmupIntRange,

14 -)

12 from vllm.model_executor.warmup.jit_warmup_triton_helper import (

13 + LaunchSpec,

14 TritonPointerInputVariant,

15 TritonWarmupTensor,

16 + VllmTritonJitKernel,

17 + kernel_launcher,

18 + triton_scalar_specialization_rep,

19 )

20 from vllm.platforms import current_platform

21 from vllm.triton_utils import tl, triton

@@ -63,41 +64,112 @@ def dsa_indexer_uses_fp4(vllm_config: VllmConfig) -> bool:

64 return use_fp4

66 -@triton.jit

67 -def _prepare_uniform_decode_kernel(

68 - seq_lens_ptr,

69 - decode_seq_lens_ptr,

70 - block_table_ptr,

71 - block_table_stride,

72 - expanded_block_table_ptr,

73 - expanded_bt_stride,

74 - decode_lens_ptr,

75 - max_decode_len,

76 - BLOCK_SIZE: tl.constexpr,

67 +class PrepareUniformDecodeKernel(

68 + VllmTritonJitKernel["PrepareUniformDecodeKernel.CompileKey"]

69 ):

78 - idx = tl.program_id(0)

79 - req_id = idx // max_decode_len

80 - local_idx = idx % max_decode_len

81 -

82 - # Compute number of KVs attended to by this token. Padding requests have

83 - # seq_len == 0, which would otherwise make the first token of each padded

84 - # request negative (e.g. next_n=2 gives 0-2+0+1 = -1). Downstream kernels

85 - # read these as uint32, turning -1 into ~4e9.

86 - seq_len = tl.load(seq_lens_ptr + req_id)

87 - per_token_seq_len = tl.maximum(seq_len - max_decode_len + local_idx + 1, 0)

88 - tl.store(decode_seq_lens_ptr + idx, per_token_seq_len)

89 -

90 - # Copy block table row.

91 - src = block_table_ptr + req_id * block_table_stride

92 - dst = expanded_block_table_ptr + idx * expanded_bt_stride

93 - for i in tl.range(0, expanded_bt_stride, BLOCK_SIZE):

94 - off = i + tl.arange(0, BLOCK_SIZE)

95 - mask = off < expanded_bt_stride

96 - src_block = tl.load(src + off, mask=mask)

97 - tl.store(dst + off, src_block, mask=mask)

98 -

99 - # All reqs now have decode_len = 1.

100 - tl.store(decode_lens_ptr + idx, 1)

70 + BLOCK_SIZE = 1024

71 +

72 + @dataclass(frozen=True)

73 + class CompileKey:

74 + block_size: int

75 + block_table_stride: int

76 + expanded_bt_stride: int

77 + max_decode_len: int

78 +

79 + @staticmethod

80 + @triton.jit

81 + def kernel(

82 + seq_lens_ptr,

83 + decode_seq_lens_ptr,

84 + block_table_ptr,

85 + block_table_stride,

86 + expanded_block_table_ptr,

87 + expanded_bt_stride,

88 + decode_lens_ptr,

89 + max_decode_len,

90 + BLOCK_SIZE: tl.constexpr,

91 + ):

92 + idx = tl.program_id(0)

93 + req_id = idx // max_decode_len

94 + local_idx = idx % max_decode_len

95 +

96 + # Compute number of KVs attended to by this token. Padding requests have

97 + # seq_len == 0, which would otherwise make the first token of each padded

98 + # request negative (e.g. next_n=2 gives 0-2+0+1 = -1). Downstream kernels

99 + # read these as uint32, turning -1 into ~4e9.

100 + seq_len = tl.load(seq_lens_ptr + req_id)

101 + per_token_seq_len = tl.maximum(seq_len - max_decode_len + local_idx + 1, 0)

102 + tl.store(decode_seq_lens_ptr + idx, per_token_seq_len)

103 +

104 + # Copy block table row.

105 + src = block_table_ptr + req_id * block_table_stride

106 + dst = expanded_block_table_ptr + idx * expanded_bt_stride

107 + for i in tl.range(0, expanded_bt_stride, BLOCK_SIZE):

108 + off = i + tl.arange(0, BLOCK_SIZE)

109 + mask = off < expanded_bt_stride

110 + src_block = tl.load(src + off, mask=mask)

111 + tl.store(dst + off, src_block, mask=mask)

112 +

113 + # All reqs now have decode_len = 1.

114 + tl.store(decode_lens_ptr + idx, 1)

115 +

116 + def dispatch( # type: ignore[override]

117 + self,

118 + *,

119 + block_size: int,

120 + block_table_stride: int,

121 + expanded_bt_stride: int,

122 + max_decode_len: int,

123 + ) -> CompileKey:

124 + return self.CompileKey(

125 + block_size=block_size,

126 + block_table_stride=triton_scalar_specialization_rep(block_table_stride),

127 + expanded_bt_stride=triton_scalar_specialization_rep(expanded_bt_stride),

128 + max_decode_len=triton_scalar_specialization_rep(max_decode_len),

129 + )

130 +

131 + def get_warmup_keys(self) -> list[CompileKey]:

132 + return self._trace_dispatch(self.dispatch)(

133 + block_size=self.BLOCK_SIZE,

134 + block_table_stride=(1, 2, 16),

135 + expanded_bt_stride=(1, 2, 16),

136 + max_decode_len=(1, 2, 16),

137 + )

138 +

139 + def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:

140 + int32_ptr = TritonWarmupTensor(torch.int32)

141 + return dict(

142 + seq_lens=int32_ptr,

143 + decode_seq_lens=int32_ptr,

144 + block_table=TritonWarmupTensor(

145 + torch.int32,

146 + shape=(1, compile_key.block_table_stride),

147 + ),

148 + expanded_block_table=TritonWarmupTensor(

149 + torch.int32,

150 + shape=(1, compile_key.expanded_bt_stride),

151 + ),

152 + decode_lens=int32_ptr,

153 + num_decode_tokens=1,

154 + max_decode_len=compile_key.max_decode_len,

155 + )

156 +

157 + @kernel_launcher

158 + def __call__(

159 + self,

160 + seq_lens: torch.Tensor,

161 + decode_seq_lens: torch.Tensor,

162 + block_table: torch.Tensor,

163 + expanded_block_table: torch.Tensor,

164 + decode_lens: torch.Tensor,

165 + num_decode_tokens: int,

166 + max_decode_len: int,

167 + ) -> LaunchSpec:

168 + return (num_decode_tokens,), dict(

169 + block_table_stride=block_table.stride(0),

170 + expanded_bt_stride=expanded_block_table.stride(0),

171 + BLOCK_SIZE=self.BLOCK_SIZE,

172 + )

173

174

175 def split_indexer_prefill_chunks(

@@ -215,14 +287,8 @@ class DeepseekV32IndexerPrefillChunkMetadata:

287 max_local_total_seq_lens: int = 0

288

289

218 -_BUILD_PREFILL_CHUNK_METADATA_INPUT_VARIANTS = (

219 - TritonPointerInputVariant.from_alignment(uncompressed_seq_lens=True),

220 - TritonPointerInputVariant.from_alignment(uncompressed_seq_lens=False),

221 -)

222 -

223 -

290 class BuildPrefillChunkMetadataKernel(

225 - VllmJitKernel["BuildPrefillChunkMetadataKernel.CompileKey"]

291 + VllmTritonJitKernel["BuildPrefillChunkMetadataKernel.CompileKey"]

292 ):

293 BLOCK_SIZE = 1024

294

@@ -230,11 +296,11 @@ class BuildPrefillChunkMetadataKernel(

296 class CompileKey:

297 query_slice_start: int

298 query_slice_stop: int

233 - DCP_RANK: int

234 - DCP_WORLD: int

235 - DCP_INTERLEAVE: int

236 - BLOCK_SIZE: int

237 - COMPRESS_RATIO: int

299 + dcp_rank: int

300 + dcp_world: int

301 + dcp_interleave: int

302 + block_size: int

303 + compress_ratio: int

304 input_variant: TritonPointerInputVariant

305

306 @staticmethod

@@ -327,13 +393,13 @@ class BuildPrefillChunkMetadataKernel(

393 input_variant: TritonPointerInputVariant,

394 ) -> CompileKey:

395 return self.CompileKey(

330 - query_slice_start=query_slice_start,

331 - query_slice_stop=query_slice_stop,

332 - DCP_RANK=DCP_RANK,

333 - DCP_WORLD=DCP_WORLD,

334 - DCP_INTERLEAVE=DCP_INTERLEAVE,

335 - BLOCK_SIZE=BLOCK_SIZE,

336 - COMPRESS_RATIO=COMPRESS_RATIO,

396 + query_slice_start=triton_scalar_specialization_rep(query_slice_start),

397 + query_slice_stop=triton_scalar_specialization_rep(query_slice_stop),

398 + dcp_rank=triton_scalar_specialization_rep(DCP_RANK),

399 + dcp_world=triton_scalar_specialization_rep(DCP_WORLD),

400 + dcp_interleave=triton_scalar_specialization_rep(DCP_INTERLEAVE),

401 + block_size=BLOCK_SIZE,

402 + compress_ratio=COMPRESS_RATIO,

403 input_variant=input_variant,

404 )

405

@@ -345,42 +411,51 @@ class BuildPrefillChunkMetadataKernel(

411 dcp_interleave = parallel_config.cp_kv_cache_interleave_size

412 dcp_rank = get_dcp_group().rank_in_group if dcp_world > 1 else 0

413 compress_ratios = tuple(

348 - max(1, int(ratio))

349 - for ratio in (getattr(hf_config, "compress_ratios", None) or (1,))

414 + dict.fromkeys(

415 + max(1, int(ratio))

416 + for ratio in (

417 + *(getattr(hf_config, "compress_ratios", None) or (1,)),

418 + getattr(hf_config, "index_kpool", 1) or 1,

419 + )

420 + )

421 )

422 return self._trace_dispatch(self.dispatch)(

352 - query_slice_start=WarmupIntRange(0, 2),

423 + # Cover Triton's divisible, exact-one, and generic i32 classes.

424 + query_slice_start=(0, 1, 2),

425 query_slice_stop=(1, 2 * max_tokens - 1, 2 * max_tokens),

426 DCP_RANK=dcp_rank,

427 DCP_WORLD=dcp_world,

428 DCP_INTERLEAVE=dcp_interleave,

429 BLOCK_SIZE=self.BLOCK_SIZE,

430 COMPRESS_RATIO=list(compress_ratios),

359 - input_variant=_BUILD_PREFILL_CHUNK_METADATA_INPUT_VARIANTS,

431 + input_variant=(

432 + TritonPointerInputVariant.from_alignment(uncompressed_seq_lens=True),

433 + TritonPointerInputVariant.from_alignment(uncompressed_seq_lens=False),

434 + ),

435 )

436

362 - def compile(self, compile_key: CompileKey) -> None:

363 - warmup = getattr(self.kernel, "warmup", None)

364 - assert warmup is not None

437 + def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:

438 int32_ptr = TritonWarmupTensor(torch.int32)

366 - warmup(

367 - int32_ptr,

368 - compile_key.input_variant.pointer("uncompressed_seq_lens", torch.int32),

369 - int32_ptr,

370 - int32_ptr,

371 - int32_ptr,

372 - int32_ptr,

373 - int32_ptr,

374 - compile_key.query_slice_start,

375 - compile_key.query_slice_stop,

376 - compile_key.DCP_RANK,

377 - compile_key.DCP_WORLD,

378 - compile_key.DCP_INTERLEAVE,

379 - BLOCK_SIZE=compile_key.BLOCK_SIZE,

380 - COMPRESS_RATIO=compile_key.COMPRESS_RATIO,

381 - grid=(1,),

439 + return dict(

440 + query_start_loc=int32_ptr,

441 + uncompressed_seq_lens=compile_key.input_variant.pointer(

442 + "uncompressed_seq_lens", torch.int32

443 + ),

444 + cu_compressed_seq_lens=int32_ptr,

445 + row_start_cu_compressed_seq_lens=int32_ptr,

446 + token_to_seq=int32_ptr,

447 + cu_compressed_seq_len_ks=int32_ptr,

448 + cu_compressed_seq_len_ke=int32_ptr,

449 + query_slice_start=compile_key.query_slice_start,

450 + query_slice_stop=compile_key.query_slice_stop,

451 + DCP_RANK=compile_key.dcp_rank,

452 + DCP_WORLD=compile_key.dcp_world,

453 + DCP_INTERLEAVE=compile_key.dcp_interleave,

454 + num_reqs=1,

455 + COMPRESS_RATIO=compile_key.compress_ratio,

456 )

457

458 + @kernel_launcher

459 def __call__(

460 self,

461 query_start_loc: torch.Tensor,

@@ -398,28 +473,13 @@ class BuildPrefillChunkMetadataKernel(

473 *,

474 num_reqs: int,

475 COMPRESS_RATIO: int,

401 - ) -> None:

402 - self.kernel[(num_reqs,)](

403 - query_start_loc,

404 - uncompressed_seq_lens,

405 - cu_compressed_seq_lens,

406 - row_start_cu_compressed_seq_lens,

407 - token_to_seq,

408 - cu_compressed_seq_len_ks,

409 - cu_compressed_seq_len_ke,

410 - query_slice_start,

411 - query_slice_stop,

412 - DCP_RANK,

413 - DCP_WORLD,

414 - DCP_INTERLEAVE,

476 + ) -> LaunchSpec:

477 + return (num_reqs,), dict(

478 BLOCK_SIZE=self.BLOCK_SIZE,

479 COMPRESS_RATIO=COMPRESS_RATIO,

480 )

481

482

420 -_BUILD_PREFILL_CHUNK_METADATA_KERNEL = BuildPrefillChunkMetadataKernel()

421 -

422 -

483 @dataclass

484 class DeepseekV32IndexerPrefillMetadata:

485 chunks: list[DeepseekV32IndexerPrefillChunkMetadata]

@@ -694,16 +754,14 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):

754 and num_decodes * max_decode_len == num_decode_tokens

755 ):

756 # Uniform decode lengths with no cudagraph token padding.

697 - _prepare_uniform_decode_kernel[(num_decode_tokens,)](

757 + _PREPARE_UNIFORM_DECODE_KERNEL(

758 seq_lens,

759 self.decode_seq_lens_buffer,

760 block_table,

701 - block_table.stride(0),

761 self.expanded_block_table_buffer,

703 - self.expanded_block_table_buffer.stride(0),

762 self.decode_lens_buffer,

763 + num_decode_tokens,

764 max_decode_len,

706 - BLOCK_SIZE=1024,

765 )

766 self.decode_seq_lens_buffer[num_decode_tokens:] = 0

767 seq_lens = self.decode_seq_lens_buffer[:num_decode_tokens]

@@ -1180,3 +1238,7 @@ def build_prefill_chunk_metadata(

1238 local_total_seq_lens=local_total_seq_lens,

1239 max_local_total_seq_lens=max_local_total_seq_lens,

1240 )

1241 +

1242 +

1243 +_PREPARE_UNIFORM_DECODE_KERNEL = PrepareUniformDecodeKernel()

1244 +_BUILD_PREFILL_CHUNK_METADATA_KERNEL = BuildPrefillChunkMetadataKernel()

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

+304 -155

@@ -1,17 +1,21 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 from dataclasses import dataclass, field

4 -from typing import ClassVar, cast

4 +from typing import Any, ClassVar, cast

6 import torch

8 from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config

9 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase

10 from vllm.model_executor.warmup.jit_warmup import (

11 - VllmJitKernel,

11 WarmupIntRange,

12 )

14 -from vllm.model_executor.warmup.jit_warmup_triton_helper import TritonWarmupTensor

13 +from vllm.model_executor.warmup.jit_warmup_triton_helper import (

14 + LaunchSpec,

15 + TritonWarmupTensor,

16 + VllmTritonJitKernel,

17 + kernel_launcher,

18 +)

19 from vllm.platforms import current_platform

20 from vllm.triton_utils import tl, triton

21 from vllm.utils.math_utils import cdiv, next_power_of_2

@@ -275,11 +279,11 @@ class DeepseekSparseSWAMetadata:

279

280

281 class ComputePrefillMetadataKernel(

278 - VllmJitKernel["ComputePrefillMetadataKernel.CompileKey"]

282 + VllmTritonJitKernel["ComputePrefillMetadataKernel.CompileKey"]

283 ):

284 @dataclass(frozen=True)

285 class CompileKey:

282 - BLOCK_SIZE: int

286 + block_size: int

287

288 @staticmethod

289 @triton.jit(do_not_specialize=["num_prefills", "num_decodes", "window_size"])

@@ -320,7 +324,7 @@ class ComputePrefillMetadataKernel(

324 num_prefills: int,

325 ) -> CompileKey:

326 return self.CompileKey(

323 - BLOCK_SIZE=next_power_of_2(num_prefills),

327 + block_size=next_power_of_2(num_prefills),

328 )

329

330 def get_warmup_keys(self, vllm_config: VllmConfig) -> list[CompileKey]:

@@ -336,21 +340,18 @@ class ComputePrefillMetadataKernel(

340 num_prefills=WarmupIntRange(1, max_prefills + 1),

341 )

342

339 - def compile(self, compile_key: CompileKey) -> None:

340 - warmup = getattr(self.kernel, "warmup", None)

341 - assert warmup is not None

343 + def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:

344 int32_ptr = TritonWarmupTensor(torch.int32)

343 - warmup(

344 - int32_ptr,

345 - int32_ptr,

346 - int32_ptr,

347 - compile_key.BLOCK_SIZE,

348 - 0,

349 - 1,

350 - BLOCK_SIZE=compile_key.BLOCK_SIZE,

351 - grid=(1,),

345 + return dict(

346 + prefill_gather_lens=int32_ptr,

347 + seq_lens=int32_ptr,

348 + query_start_loc=int32_ptr,

349 + num_prefills=compile_key.block_size,

350 + num_decodes=0,

351 + window_size=1,

352 )

353

354 + @kernel_launcher

355 def __call__(

356 self,

357 prefill_gather_lens: torch.Tensor,

@@ -359,22 +360,13 @@ class ComputePrefillMetadataKernel(

360 num_prefills: int,

361 num_decodes: int,

362 window_size: int,

362 - ) -> None:

363 + ) -> LaunchSpec:

364 compile_key = self.dispatch(num_prefills=num_prefills)

364 - self.kernel[(1,)](

365 - prefill_gather_lens,

366 - seq_lens,

367 - query_start_loc,

368 - num_prefills,

369 - num_decodes,

370 - window_size,

371 - BLOCK_SIZE=compile_key.BLOCK_SIZE,

365 + return (1,), dict(

366 + BLOCK_SIZE=compile_key.block_size,

367 )

368

369

375 -_COMPUTE_PREFILL_METADATA_KERNEL = ComputePrefillMetadataKernel()

376 -

377 -

370 class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):

371 """Builds metadata for DeepseekV4 SWA cache.

372

@@ -545,9 +537,8 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):

537 device=self.device,

538 )

539 decode_swa_indices = self.decode_swa_indices_noncausal

548 - _compute_dspark_noncausal_swa_indices_kernel[(num_decode_tokens,)](

540 + _COMPUTE_DSPARK_NONCAUSAL_SWA_INDICES_KERNEL(

541 decode_swa_indices,

550 - decode_swa_indices.stride(0),

542 self.decode_swa_lens,

543 self.window_size,

544 self.noncausal_index_width,

@@ -556,15 +547,13 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):

547 token_to_req_indices,

548 is_valid_token,

549 block_table,

559 - block_table.stride(0),

550 self.block_size,

551 + num_tokens=num_decode_tokens,

552 token_offset=0,

562 - TRITON_BLOCK_SIZE=1024,

553 )

554 else:

565 - _compute_swa_indices_and_lens_kernel[(num_decode_tokens,)](

555 + _COMPUTE_SWA_INDICES_AND_LENS_KERNEL(

556 decode_swa_indices,

567 - decode_swa_indices.stride(0),

557 self.decode_swa_lens,

558 self.window_size,

559 query_start_loc,

@@ -572,10 +561,9 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):

561 token_to_req_indices,

562 is_valid_token,

563 block_table,

575 - block_table.stride(0),

564 self.block_size,

565 + num_tokens=num_decode_tokens,

566 token_offset=0,

578 - TRITON_BLOCK_SIZE=1024,

567 )

568

569 # Prefill SWA indices live in paged coordinates. `token_offset` lets

@@ -584,9 +572,8 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):

572 if num_prefill_tokens > 0:

573 prefill_swa_indices = self.prefill_swa_indices[:num_prefill_tokens]

574 prefill_swa_lens = self.prefill_swa_lens[:num_prefill_tokens]

587 - _compute_swa_indices_and_lens_kernel[(num_prefill_tokens,)](

575 + _COMPUTE_SWA_INDICES_AND_LENS_KERNEL(

576 prefill_swa_indices,

589 - prefill_swa_indices.stride(0),

577 prefill_swa_lens,

578 self.window_size,

579 query_start_loc,

@@ -594,10 +581,9 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):

581 token_to_req_indices,

582 is_valid_token,

583 block_table,

597 - block_table.stride(0),

584 self.block_size,

585 + num_tokens=num_prefill_tokens,

586 token_offset=num_decode_tokens,

600 - TRITON_BLOCK_SIZE=1024,

587 )

588

589 # Pre-compute DeepseekV4 prefill metadata shared across all attention layers.

@@ -668,9 +654,8 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):

654 assert metadata.decode_swa_indices is not None

655 assert metadata.decode_swa_lens is not None

656

671 - _compute_swa_indices_and_lens_kernel[(metadata.num_decode_tokens,)](

657 + _COMPUTE_SWA_INDICES_AND_LENS_KERNEL(

658 metadata.decode_swa_indices,

673 - metadata.decode_swa_indices.stride(0),

659 metadata.decode_swa_lens,

660 metadata.decode_swa_indices.shape[-1],

661 metadata.query_start_loc,

@@ -678,10 +663,9 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):

663 metadata.token_to_req_indices,

664 metadata.is_valid_token,

665 metadata.block_table,

681 - metadata.block_table.stride(0),

666 self.block_size,

667 + num_tokens=metadata.num_decode_tokens,

668 token_offset=0,

684 - TRITON_BLOCK_SIZE=1024,

669 )

670 tile_sched = self.build_tile_scheduler(metadata.num_decode_tokens)

671 metadata.tile_sched_swaonly = tile_sched[_LAYER_TYPE_SWAONLY]

@@ -771,142 +755,307 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):

755 return result

756

757

774 -@triton.jit(do_not_specialize=["token_offset"])

775 -def _compute_swa_indices_and_lens_kernel(

776 - swa_indices_ptr,

777 - swa_indices_stride,

778 - swa_lens_ptr,

779 - window_size,

780 - query_start_loc_ptr,

781 - seq_lens_ptr,

782 - token_to_req_indices_ptr,

783 - is_valid_token_ptr,

784 - block_table_ptr,

785 - block_table_stride,

786 - block_size,

787 - token_offset,

788 - TRITON_BLOCK_SIZE: tl.constexpr,

758 +# TODO(ben): unify this kernel to reduce duplication

759 +

760 +

761 +class ComputeSWAIndicesAndLensKernel(

762 + VllmTritonJitKernel["ComputeSWAIndicesAndLensKernel.CompileKey"]

763 ):

790 - pid = tl.program_id(0)

791 - token_idx = pid + token_offset

792 - is_valid = tl.load(is_valid_token_ptr + token_idx)

793 - if not is_valid:

794 - tl.store(swa_lens_ptr + pid, 0)

795 - # Clear the row so a padded token cannot gather through stale indices.

764 + @dataclass(frozen=True)

765 + class CompileKey:

766 + window_size: int

767 + block_size: int

768 + triton_block_size: int

769 +

770 + @staticmethod

771 + @triton.jit(

772 + do_not_specialize=[

773 + "swa_indices_stride",

774 + "block_table_stride",

775 + "token_offset",

776 + ]

777 + )

778 + def kernel(

779 + swa_indices_ptr,

780 + swa_indices_stride,

781 + swa_lens_ptr,

782 + window_size,

783 + query_start_loc_ptr,

784 + seq_lens_ptr,

785 + token_to_req_indices_ptr,

786 + is_valid_token_ptr,

787 + block_table_ptr,

788 + block_table_stride,

789 + block_size,

790 + token_offset,

791 + TRITON_BLOCK_SIZE: tl.constexpr,

792 + ):

793 + pid = tl.program_id(0)

794 + token_idx = pid + token_offset

795 + is_valid = tl.load(is_valid_token_ptr + token_idx)

796 + if not is_valid:

797 + tl.store(swa_lens_ptr + pid, 0)

798 + # Clear the row so a padded token cannot gather through stale indices.

799 + for i in range(0, window_size, TRITON_BLOCK_SIZE):

800 + offset = i + tl.arange(0, TRITON_BLOCK_SIZE)

801 + tl.store(

802 + swa_indices_ptr + pid * swa_indices_stride + offset,

803 + -1,

804 + mask=offset < window_size,

805 + )

806 + return

807 +

808 + req_idx = tl.load(token_to_req_indices_ptr + token_idx)

809 +

810 + query_start = tl.load(query_start_loc_ptr + req_idx)

811 + query_end = tl.load(query_start_loc_ptr + req_idx + 1)

812 + query_len = query_end - query_start

813 +

814 + seq_len = tl.load(seq_lens_ptr + req_idx)

815 + prefix_len = seq_len - query_len

816 +

817 + pos = prefix_len + token_idx - query_start

818 + start_pos = tl.maximum(pos - window_size + 1, 0)

819 + end_pos = pos + 1

820 +

821 + swa_len = end_pos - start_pos

822 + tl.store(swa_lens_ptr + pid, swa_len)

823 +

824 for i in range(0, window_size, TRITON_BLOCK_SIZE):

825 offset = i + tl.arange(0, TRITON_BLOCK_SIZE)

826 +

827 + pos_offset = start_pos + offset

828 + block_indices = pos_offset // block_size

829 + block_numbers = tl.load(

830 + block_table_ptr + req_idx * block_table_stride + block_indices,

831 + mask=pos_offset < end_pos,

832 + )

833 + block_offsets = pos_offset % block_size

834 + slot_ids = block_numbers * block_size + block_offsets

835 +

836 + slot_ids = tl.where(offset < swa_len, slot_ids, -1)

837 tl.store(

838 swa_indices_ptr + pid * swa_indices_stride + offset,

800 - -1,

839 + slot_ids,

840 mask=offset < window_size,

841 )

803 - return

842

805 - req_idx = tl.load(token_to_req_indices_ptr + token_idx)

843 + def dispatch( # type: ignore[override]

844 + self,

845 + **compile_key_fields: int,

846 + ) -> CompileKey:

847 + return self.CompileKey(

848 + **compile_key_fields,

849 + triton_block_size=1024,

850 + )

851 +

852 + def get_warmup_keys(

853 + self,

854 + *,

855 + window_size: int,

856 + block_size: int,

857 + ) -> list[CompileKey]:

858 + return self._trace_dispatch(self.dispatch)(

859 + window_size=window_size,

860 + block_size=block_size,

861 + )

862 +

863 + def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:

864 + int32_ptr = TritonWarmupTensor(torch.int32)

865 + return dict(

866 + swa_indices=TritonWarmupTensor(torch.int32, shape=(1, 1), strides=(1, 1)),

867 + swa_lens=int32_ptr,

868 + window_size=compile_key.window_size,

869 + query_start_loc=int32_ptr,

870 + seq_lens=int32_ptr,

871 + token_to_req_indices=int32_ptr,

872 + is_valid_token=TritonWarmupTensor(torch.bool),

873 + block_table=TritonWarmupTensor(torch.int32, shape=(1, 1), strides=(1, 1)),

874 + block_size=compile_key.block_size,

875 + num_tokens=1,

876 + token_offset=0,

877 + )

878 +

879 + @kernel_launcher

880 + def __call__(

881 + self,

882 + swa_indices: torch.Tensor,

883 + swa_lens: torch.Tensor,

884 + window_size: int,

885 + query_start_loc: torch.Tensor,

886 + seq_lens: torch.Tensor,

887 + token_to_req_indices: torch.Tensor,

888 + is_valid_token: torch.Tensor,

889 + block_table: torch.Tensor,

890 + block_size: int,

891 + *,

892 + num_tokens: int,

893 + token_offset: int,

894 + ) -> LaunchSpec:

895 + return (num_tokens,), dict(

896 + swa_indices_stride=swa_indices.stride(0),

897 + block_table_stride=block_table.stride(0),

898 + TRITON_BLOCK_SIZE=1024,

899 + )

900

807 - query_start = tl.load(query_start_loc_ptr + req_idx)

808 - query_end = tl.load(query_start_loc_ptr + req_idx + 1)

809 - query_len = query_end - query_start

901

811 - seq_len = tl.load(seq_lens_ptr + req_idx)

812 - prefix_len = seq_len - query_len

902 +class ComputeDSparkNoncausalSWAIndicesKernel(

903 + VllmTritonJitKernel["ComputeDSparkNoncausalSWAIndicesKernel.CompileKey"]

904 +):

905 + @dataclass(frozen=True)

906 + class CompileKey:

907 + window_size: int

908 + index_width: int

909 + block_size: int

910 + triton_block_size: int

911

814 - pos = prefix_len + token_idx - query_start

815 - start_pos = tl.maximum(pos - window_size + 1, 0)

816 - end_pos = pos + 1

912 + @staticmethod

913 + @triton.jit(

914 + do_not_specialize=[

915 + "swa_indices_stride",

916 + "block_table_stride",

917 + "token_offset",

918 + ]

919 + )

920 + def kernel(

921 + swa_indices_ptr,

922 + swa_indices_stride,

923 + swa_lens_ptr,

924 + window_size,

925 + index_width,

926 + query_start_loc_ptr,

927 + seq_lens_ptr,

928 + token_to_req_indices_ptr,

929 + is_valid_token_ptr,

930 + block_table_ptr,

931 + block_table_stride,

932 + block_size,

933 + token_offset,

934 + TRITON_BLOCK_SIZE: tl.constexpr,

935 + ):

936 + """Non-causal per-token indices for the DSpark draft block.

937

818 - swa_len = end_pos - start_pos

819 - tl.store(swa_lens_ptr + pid, swa_len)

938 + Here, we populate the topk indices with the trailing window of context tokens,

939 + plus all query tokens (including future ones).

940 + """

941 + pid = tl.program_id(0)

942 + token_idx = pid + token_offset

943 + is_valid = tl.load(is_valid_token_ptr + token_idx)

944 + if not is_valid:

945 + tl.store(swa_lens_ptr + pid, 0)

946 + # Clear the row so a padded token cannot gather through stale indices.

947 + for i in range(0, index_width, TRITON_BLOCK_SIZE):

948 + offset = i + tl.arange(0, TRITON_BLOCK_SIZE)

949 + tl.store(

950 + swa_indices_ptr + pid * swa_indices_stride + offset,

951 + -1,

952 + mask=offset < index_width,

953 + )

954 + return

955

821 - for i in range(0, window_size, TRITON_BLOCK_SIZE):

822 - offset = i + tl.arange(0, TRITON_BLOCK_SIZE)

956 + req_idx = tl.load(token_to_req_indices_ptr + token_idx)

957

824 - pos_offset = start_pos + offset

825 - block_indices = pos_offset // block_size

826 - block_numbers = tl.load(

827 - block_table_ptr + req_idx * block_table_stride + block_indices,

828 - mask=pos_offset < end_pos,

829 - )

830 - block_offsets = pos_offset % block_size

831 - slot_ids = block_numbers * block_size + block_offsets

958 + query_start = tl.load(query_start_loc_ptr + req_idx)

959 + query_end = tl.load(query_start_loc_ptr + req_idx + 1)

960 + query_len = query_end - query_start

961

833 - slot_ids = tl.where(offset < swa_len, slot_ids, -1)

834 - tl.store(

835 - swa_indices_ptr + pid * swa_indices_stride + offset,

836 - slot_ids,

837 - mask=offset < window_size,

838 - )

962 + seq_len = tl.load(seq_lens_ptr + req_idx)

963 + prefix_len = seq_len - query_len

964

965 + # Block-anchored window (shared by every token in the block) + full block.

966 + start_pos = tl.maximum(prefix_len - window_size, 0)

967 + end_pos = seq_len

968

841 -# TODO(ben): unify this kernel to reduce duplication

842 -@triton.jit(do_not_specialize=["token_offset"])

843 -def _compute_dspark_noncausal_swa_indices_kernel(

844 - swa_indices_ptr,

845 - swa_indices_stride,

846 - swa_lens_ptr,

847 - window_size,

848 - index_width,

849 - query_start_loc_ptr,

850 - seq_lens_ptr,

851 - token_to_req_indices_ptr,

852 - is_valid_token_ptr,

853 - block_table_ptr,

854 - block_table_stride,

855 - block_size,

856 - token_offset,

857 - TRITON_BLOCK_SIZE: tl.constexpr,

858 -):

859 - """Non-causal per-token indices for the DSpark draft block.

969 + swa_len = end_pos - start_pos

970 + tl.store(swa_lens_ptr + pid, swa_len)

971

861 - Here, we populate the topk indices with the trailing window of context tokens,

862 - plus all query tokens (including future ones).

863 - """

864 - pid = tl.program_id(0)

865 - token_idx = pid + token_offset

866 - is_valid = tl.load(is_valid_token_ptr + token_idx)

867 - if not is_valid:

868 - tl.store(swa_lens_ptr + pid, 0)

869 - # Clear the row so a padded token cannot gather through stale indices.

972 for i in range(0, index_width, TRITON_BLOCK_SIZE):

973 offset = i + tl.arange(0, TRITON_BLOCK_SIZE)

974 +

975 + pos_offset = start_pos + offset

976 + block_indices = pos_offset // block_size

977 + block_numbers = tl.load(

978 + block_table_ptr + req_idx * block_table_stride + block_indices,

979 + mask=pos_offset < end_pos,

980 + )

981 + block_offsets = pos_offset % block_size

982 + slot_ids = block_numbers * block_size + block_offsets

983 +

984 + slot_ids = tl.where(offset < swa_len, slot_ids, -1)

985 tl.store(

986 swa_indices_ptr + pid * swa_indices_stride + offset,

874 - -1,

987 + slot_ids,

988 mask=offset < index_width,

989 )

877 - return

878 -

879 - req_idx = tl.load(token_to_req_indices_ptr + token_idx)

990

881 - query_start = tl.load(query_start_loc_ptr + req_idx)

882 - query_end = tl.load(query_start_loc_ptr + req_idx + 1)

883 - query_len = query_end - query_start

884 -

885 - seq_len = tl.load(seq_lens_ptr + req_idx)

886 - prefix_len = seq_len - query_len

887 -

888 - # Block-anchored window (shared by every token in the block) + full block.

889 - start_pos = tl.maximum(prefix_len - window_size, 0)

890 - end_pos = seq_len

891 -

892 - swa_len = end_pos - start_pos

893 - tl.store(swa_lens_ptr + pid, swa_len)

991 + def dispatch( # type: ignore[override]

992 + self,

993 + *,

994 + window_size: int,

995 + num_speculative_tokens: int,

996 + block_size: int,

997 + ) -> CompileKey:

998 + return self.CompileKey(

999 + window_size=window_size,

1000 + index_width=get_dspark_swa_index_width(window_size, num_speculative_tokens),

1001 + block_size=block_size,

1002 + triton_block_size=1024,

1003 + )

1004

895 - for i in range(0, index_width, TRITON_BLOCK_SIZE):

896 - offset = i + tl.arange(0, TRITON_BLOCK_SIZE)

1005 + def get_warmup_keys(

1006 + self,

1007 + *,

1008 + window_size: int,

1009 + num_speculative_tokens: int,

1010 + block_size: int,

1011 + ) -> list[CompileKey]:

1012 + return self._trace_dispatch(self.dispatch)(

1013 + window_size=window_size,

1014 + num_speculative_tokens=num_speculative_tokens,

1015 + block_size=block_size,

1016 + )

1017

898 - pos_offset = start_pos + offset

899 - block_indices = pos_offset // block_size

900 - block_numbers = tl.load(

901 - block_table_ptr + req_idx * block_table_stride + block_indices,

902 - mask=pos_offset < end_pos,

1018 + def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:

1019 + int32_ptr = TritonWarmupTensor(torch.int32)

1020 + return dict(

1021 + swa_indices=TritonWarmupTensor(torch.int32, shape=(1, 1), strides=(1, 1)),

1022 + swa_lens=int32_ptr,

1023 + window_size=compile_key.window_size,

1024 + index_width=compile_key.index_width,

1025 + query_start_loc=int32_ptr,

1026 + seq_lens=int32_ptr,

1027 + token_to_req_indices=int32_ptr,

1028 + is_valid_token=TritonWarmupTensor(torch.bool),

1029 + block_table=TritonWarmupTensor(torch.int32, shape=(1, 1), strides=(1, 1)),

1030 + block_size=compile_key.block_size,

1031 + num_tokens=1,

1032 + token_offset=0,

1033 )

904 - block_offsets = pos_offset % block_size

905 - slot_ids = block_numbers * block_size + block_offsets

1034

907 - slot_ids = tl.where(offset < swa_len, slot_ids, -1)

908 - tl.store(

909 - swa_indices_ptr + pid * swa_indices_stride + offset,

910 - slot_ids,

911 - mask=offset < index_width,

1035 + @kernel_launcher

1036 + def __call__(

1037 + self,

1038 + swa_indices: torch.Tensor,

1039 + swa_lens: torch.Tensor,

1040 + window_size: int,

1041 + index_width: int,

1042 + query_start_loc: torch.Tensor,

1043 + seq_lens: torch.Tensor,

1044 + token_to_req_indices: torch.Tensor,

1045 + is_valid_token: torch.Tensor,

1046 + block_table: torch.Tensor,

1047 + block_size: int,

1048 + *,

1049 + num_tokens: int,

1050 + token_offset: int,

1051 + ) -> LaunchSpec:

1052 + return (num_tokens,), dict(

1053 + swa_indices_stride=swa_indices.stride(0),

1054 + block_table_stride=block_table.stride(0),

1055 + TRITON_BLOCK_SIZE=1024,

1056 )

1057 +

1058 +

1059 +_COMPUTE_PREFILL_METADATA_KERNEL = ComputePrefillMetadataKernel()

1060 +_COMPUTE_SWA_INDICES_AND_LENS_KERNEL = ComputeSWAIndicesAndLensKernel()

1061 +_COMPUTE_DSPARK_NONCAUSAL_SWA_INDICES_KERNEL = ComputeDSparkNoncausalSWAIndicesKernel()

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

+370 -173

@@ -2,9 +2,26 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 """Utility functions for sparse MLA backends."""

5 +from dataclasses import dataclass

6 +from typing import Any

7 +

8 import torch

10 +from vllm.config import VllmConfig

11 +from vllm.distributed import get_dcp_group

12 +from vllm.model_executor.warmup.jit_warmup import (

13 + zip_inputs,

14 +)

15 +from vllm.model_executor.warmup.jit_warmup_triton_helper import (

16 + LaunchSpec,

17 + TritonWarmupTensor,

18 + VllmTritonJitKernel,

19 + kernel_launcher,

20 + triton_scalar_specialization_rep,

21 +)

22 from vllm.triton_utils import tl, triton

23 +from vllm.utils.math_utils import cdiv

24 +from vllm.v1.worker.block_table import get_block_table_width

27 def flat_kv_row_view(

@@ -30,126 +47,330 @@ def flat_kv_row_view(

49 # Kernel with prefill workspace support and valid count tracking

33 -@triton.jit

34 -def _convert_req_index_to_global_index_kernel(

35 - req_id_ptr, # int32 [num_tokens]

36 - block_table_ptr, # int32 [num_requests, max_num_blocks_per_req]

37 - token_indices_ptr, # int32 [num_tokens, NUM_TOPK_TOKENS]

38 - out_ptr, # int32 [num_tokens, NUM_TOPK_TOKENS]

39 - valid_count_ptr, # int32 [num_tokens] - output valid count per row

40 - prefill_request_id_ptr, # int32 [num_tokens], -1 for decode, >=0 for prefill

41 - workspace_starts_ptr, # int32 [num_prefill_reqs+1] or nullptr

42 - # shapes (compile-time where possible)

43 - max_num_blocks_per_req: tl.constexpr,

44 - BLOCK_SIZE: tl.constexpr,

45 - BLOCK_STRIDE_ROWS: tl.constexpr,

46 - BLOCK_N: tl.constexpr, # tile width along columns

47 - HAS_PREFILL: tl.constexpr,

48 - COUNT_VALID: tl.constexpr, # whether to count valid indices

49 - # BLOCK_N == NUM_TOPK_TOKENS: one program owns the row, so the valid count

50 - # is an in-register reduction and needs no atomic.

51 - SINGLE_TILE: tl.constexpr,

52 - # When set, scatter valid slots to a contiguous prefix [0, valid_count) using

53 - # valid_count_ptr as an atomic slot allocator (DCP filtering leaves interior

54 - # -1 gaps; the trtllm-gen sparse kernel reads the first valid_count entries).

55 - # Requires COUNT_VALID and an out buffer pre-filled with -1. Order within the

56 - # prefix is unspecified (only the selected set matters).

57 - COMPACT_TO_FRONT: tl.constexpr,

58 - # DCP de-interleave: with DCP_SIZE == 1 these are an exact no-op

59 - DCP_SIZE: tl.constexpr,

60 - DCP_RANK: tl.constexpr,

61 - DCP_INTERLEAVE: tl.constexpr,

62 - # strides (in elements)

63 - bt_stride0,

64 - bt_stride1,

65 - ti_stride0,

66 - ti_stride1,

67 - out_stride0,

68 - out_stride1,

50 +

51 +

52 +class ConvertReqIndexToGlobalIndexKernel(

53 + VllmTritonJitKernel["ConvertReqIndexToGlobalIndexKernel.CompileKey"]

54 ):

70 - # program_id(0) -> token_id (row)

71 - # program_id(1) -> tile index along columns

72 - token_id = tl.program_id(0)

73 - tile_id = tl.program_id(1)

74 -

75 - # Each program covers BLOCK_N consecutive columns

76 - indice_id = tile_id * BLOCK_N + tl.arange(0, BLOCK_N)

77 -

78 - # Load request id for this token (no mask: grid is exact)

79 - req = tl.load(req_id_ptr + token_id)

80 -

81 - # Load token indices for this tile

82 - ti_ptr = token_indices_ptr + token_id * ti_stride0 + indice_id * ti_stride1

83 - tok = tl.load(ti_ptr) # int32

84 -

85 - # Only token == -1 should propagate as -1

86 - is_invalid_tok = tok < 0

87 - is_prefill = False

88 - if HAS_PREFILL:

89 - prefill_req_id = tl.load(prefill_request_id_ptr + token_id)

90 - is_prefill = prefill_req_id >= 0

91 -

92 - # DCP de-interleave the global token id into this rank's local slot.

93 - # Tokens are interleaved in groups of DCP_INTERLEAVE across ranks. With

94 - # DCP_SIZE == 1 (and any interleave) owning_rank == 0 == DCP_RANK (never

95 - # remote) and local_idx == tok, so this reduces to the non-DCP path; with

96 - # DCP_INTERLEAVE == 1 it reduces to plain round-robin (tok % / // DCP_SIZE).

97 - owning_rank = (tok // DCP_INTERLEAVE) % DCP_SIZE

98 - is_remote = owning_rank != DCP_RANK

99 - local_idx = (

100 - tok // (DCP_SIZE * DCP_INTERLEAVE)

101 - ) * DCP_INTERLEAVE + tok % DCP_INTERLEAVE

102 -

103 - # Compute block id and in-block offset

104 - block_id = local_idx // BLOCK_SIZE

105 - inblock_off = local_idx % BLOCK_SIZE

106 -

107 - # Guard block_table access

108 - valid_block = (block_id < max_num_blocks_per_req) & (block_id >= 0)

109 - bt_ptr = block_table_ptr + req * bt_stride0 + block_id * bt_stride1

110 - is_invalid_tok |= ~valid_block | is_remote

111 - base = tl.load(bt_ptr, mask=valid_block & ~is_prefill & ~is_remote, other=0)

112 - out_val = base * BLOCK_STRIDE_ROWS + inblock_off

113 -

114 - # Override with prefill output if prefill is enabled

115 - if HAS_PREFILL:

116 - workspace_start = tl.load(

117 - workspace_starts_ptr + prefill_req_id, mask=is_prefill, other=0

118 - )

119 - prefill_out = workspace_start + tok

120 - out_val = tl.where(is_prefill, prefill_out, out_val)

121 - out_val = tl.where(is_invalid_tok, -1, out_val)

122 -

123 - if COMPACT_TO_FRONT:

124 - # Scatter valid slots to a contiguous prefix. A per-tile exclusive prefix

125 - # sum gives each valid lane a distinct local offset; one atomic add of the

126 - # tile's valid count reserves a contiguous base across racing tiles. The

127 - # out buffer is pre-filled with -1, so unwritten tail slots stay -1.

128 - # With no racing tiles the base is 0 and the allocator becomes a store.

129 - is_valid = (~is_invalid_tok).to(tl.int32)

130 - local_offset = tl.cumsum(is_valid) - is_valid

131 - tile_valid_count = tl.sum(is_valid)

132 - if SINGLE_TILE:

133 - base = 0

134 - tl.store(valid_count_ptr + token_id, tile_valid_count)

135 - else:

136 - base = tl.atomic_add(valid_count_ptr + token_id, tile_valid_count)

137 - dest = base + local_offset

138 - out_ptr_dest = out_ptr + token_id * out_stride0 + dest * out_stride1

139 - tl.store(out_ptr_dest, out_val, mask=is_valid == 1)

140 - else:

141 - # Store results in place (input column == output column).

142 - out_ptr_ij = out_ptr + token_id * out_stride0 + indice_id * out_stride1

143 - tl.store(out_ptr_ij, out_val)

144 -

145 - # Accumulate the tile's valid count into the row total; a single tile's

146 - # reduction *is* the total.

147 - if COUNT_VALID:

148 - tile_valid_count = tl.sum((~is_invalid_tok).to(tl.int32))

55 + @dataclass(frozen=True)

56 + class CompileKey:

57 + block_size: int

58 + block_stride_rows: int

59 + block_n: int

60 + has_prefill_workspace: bool

61 + count_valid: bool

62 + single_tile: bool

63 + compact_to_front: bool

64 + dcp_size: int

65 + dcp_rank: int

66 + dcp_interleave: int

67 + num_warps: int

68 + num_topk_tokens: int

69 + block_table_stride: int

70 +

71 + @staticmethod

72 + @triton.jit(do_not_specialize=["max_num_blocks_per_req"])

73 + def kernel(

74 + req_id_ptr, # int32 [num_tokens]

75 + block_table_ptr, # int32 [num_requests, max_num_blocks_per_req]

76 + token_indices_ptr, # int32 [num_tokens, NUM_TOPK_TOKENS]

77 + out_ptr, # int32 [num_tokens, NUM_TOPK_TOKENS]

78 + valid_count_ptr, # int32 [num_tokens] - output valid count per row

79 + prefill_request_id_ptr, # int32 [num_tokens], -1 for decode, >=0 for prefill

80 + workspace_starts_ptr, # int32 [num_prefill_reqs+1] or nullptr

81 + # shapes (compile-time where possible)

82 + max_num_blocks_per_req,

83 + BLOCK_SIZE: tl.constexpr,

84 + BLOCK_STRIDE_ROWS: tl.constexpr,

85 + BLOCK_N: tl.constexpr, # tile width along columns

86 + HAS_PREFILL: tl.constexpr,

87 + COUNT_VALID: tl.constexpr, # whether to count valid indices

88 + # BLOCK_N == NUM_TOPK_TOKENS: one program owns the row, so the valid count

89 + # is an in-register reduction and needs no atomic.

90 + SINGLE_TILE: tl.constexpr,

91 + # When set, scatter valid slots to a contiguous prefix [0, valid_count) using

92 + # valid_count_ptr as an atomic slot allocator (DCP filtering leaves interior

93 + # -1 gaps; the trtllm-gen sparse kernel reads the first valid_count entries).

94 + # Requires COUNT_VALID and an out buffer pre-filled with -1. Order within the

95 + # prefix is unspecified (only the selected set matters).

96 + COMPACT_TO_FRONT: tl.constexpr,

97 + # DCP de-interleave: with DCP_SIZE == 1 these are an exact no-op

98 + DCP_SIZE: tl.constexpr,

99 + DCP_RANK: tl.constexpr,

100 + DCP_INTERLEAVE: tl.constexpr,

101 + # strides (in elements)

102 + bt_stride0,

103 + bt_stride1,

104 + ti_stride0,

105 + ti_stride1,

106 + out_stride0,

107 + out_stride1,

108 + ):

109 + # program_id(0) -> token_id (row)

110 + # program_id(1) -> tile index along columns

111 + token_id = tl.program_id(0)

112 + tile_id = tl.program_id(1)

113 +

114 + # Each program covers BLOCK_N consecutive columns

115 + indice_id = tile_id * BLOCK_N + tl.arange(0, BLOCK_N)

116 +

117 + # Load request id for this token (no mask: grid is exact)

118 + req = tl.load(req_id_ptr + token_id)

119 +

120 + # Load token indices for this tile

121 + ti_ptr = token_indices_ptr + token_id * ti_stride0 + indice_id * ti_stride1

122 + tok = tl.load(ti_ptr) # int32

123 +

124 + # Only token == -1 should propagate as -1

125 + is_invalid_tok = tok < 0

126 + is_prefill = False

127 + if HAS_PREFILL:

128 + prefill_req_id = tl.load(prefill_request_id_ptr + token_id)

129 + is_prefill = prefill_req_id >= 0

130 +

131 + # DCP de-interleave the global token id into this rank's local slot.

132 + # Tokens are interleaved in groups of DCP_INTERLEAVE across ranks. With

133 + # DCP_SIZE == 1 (and any interleave) owning_rank == 0 == DCP_RANK (never

134 + # remote) and local_idx == tok, so this reduces to the non-DCP path; with

135 + # DCP_INTERLEAVE == 1 it reduces to plain round-robin (tok % / // DCP_SIZE).

136 + owning_rank = (tok // DCP_INTERLEAVE) % DCP_SIZE

137 + is_remote = owning_rank != DCP_RANK

138 + local_idx = (

139 + tok // (DCP_SIZE * DCP_INTERLEAVE)

140 + ) * DCP_INTERLEAVE + tok % DCP_INTERLEAVE

141 +

142 + # Compute block id and in-block offset

143 + block_id = local_idx // BLOCK_SIZE

144 + inblock_off = local_idx % BLOCK_SIZE

145 +

146 + # Guard block_table access

147 + valid_block = (block_id < max_num_blocks_per_req) & (block_id >= 0)

148 + bt_ptr = block_table_ptr + req * bt_stride0 + block_id * bt_stride1

149 + is_invalid_tok |= ~valid_block | is_remote

150 + base = tl.load(bt_ptr, mask=valid_block & ~is_prefill & ~is_remote, other=0)

151 + out_val = base * BLOCK_STRIDE_ROWS + inblock_off

152 +

153 + # Override with prefill output if prefill is enabled

154 + if HAS_PREFILL:

155 + workspace_start = tl.load(

156 + workspace_starts_ptr + prefill_req_id, mask=is_prefill, other=0

157 + )

158 + prefill_out = workspace_start + tok

159 + out_val = tl.where(is_prefill, prefill_out, out_val)

160 + out_val = tl.where(is_invalid_tok, -1, out_val)

161 +

162 + if COMPACT_TO_FRONT:

163 + # Scatter valid slots to a contiguous prefix. A per-tile exclusive prefix

164 + # sum gives each valid lane a distinct local offset; one atomic add of the

165 + # tile's valid count reserves a contiguous base across racing tiles. The

166 + # out buffer is pre-filled with -1, so unwritten tail slots stay -1.

167 + # With no racing tiles the base is 0 and the allocator becomes a store.

168 + is_valid = (~is_invalid_tok).to(tl.int32)

169 + local_offset = tl.cumsum(is_valid) - is_valid

170 + tile_valid_count = tl.sum(is_valid)

171 if SINGLE_TILE:

172 + base = 0

173 tl.store(valid_count_ptr + token_id, tile_valid_count)

174 else:

152 - tl.atomic_add(valid_count_ptr + token_id, tile_valid_count)

175 + base = tl.atomic_add(valid_count_ptr + token_id, tile_valid_count)

176 + dest = base + local_offset

177 + out_ptr_dest = out_ptr + token_id * out_stride0 + dest * out_stride1

178 + tl.store(out_ptr_dest, out_val, mask=is_valid == 1)

179 + else:

180 + # Store results in place (input column == output column).

181 + out_ptr_ij = out_ptr + token_id * out_stride0 + indice_id * out_stride1

182 + tl.store(out_ptr_ij, out_val)

183 +

184 + # Accumulate the tile's valid count into the row total; a single tile's

185 + # reduction *is* the total.

186 + if COUNT_VALID:

187 + tile_valid_count = tl.sum((~is_invalid_tok).to(tl.int32))

188 + if SINGLE_TILE:

189 + tl.store(valid_count_ptr + token_id, tile_valid_count)

190 + else:

191 + tl.atomic_add(valid_count_ptr + token_id, tile_valid_count)

192 +

193 + def dispatch( # type: ignore[override]

194 + self,

195 + *,

196 + BLOCK_SIZE: int,

197 + BLOCK_STRIDE_ROWS: int,

198 + BLOCK_N: int,

199 + NUM_TOPK_TOKENS: int,

200 + HAS_PREFILL_WORKSPACE: bool,

201 + COUNT_VALID: bool,

202 + COMPACT_TO_FRONT: bool,

203 + DCP_SIZE: int,

204 + DCP_RANK: int,

205 + DCP_INTERLEAVE: int,

206 + max_num_blocks_per_req: int,

207 + ) -> CompileKey:

208 + tiling = _remap_tiling(NUM_TOPK_TOKENS, BLOCK_N, COUNT_VALID)

209 + return self.CompileKey(

210 + block_size=BLOCK_SIZE,

211 + block_stride_rows=BLOCK_STRIDE_ROWS,

212 + block_n=tiling[1],

213 + has_prefill_workspace=HAS_PREFILL_WORKSPACE,

214 + count_valid=COUNT_VALID,

215 + single_tile=tiling[0],

216 + compact_to_front=COMPACT_TO_FRONT,

217 + dcp_size=DCP_SIZE,

218 + dcp_rank=DCP_RANK,

219 + dcp_interleave=DCP_INTERLEAVE,

220 + num_warps=tiling[3],

221 + num_topk_tokens=NUM_TOPK_TOKENS,

222 + block_table_stride=triton_scalar_specialization_rep(max_num_blocks_per_req),

223 + )

224 +

225 + def get_warmup_keys(

226 + self,

227 + vllm_config: VllmConfig,

228 + *,

229 + block_stride_rows: int,

230 + ) -> list[CompileKey]:

231 + block_size = vllm_config.cache_config.block_size

232 + dcp_size = vllm_config.parallel_config.decode_context_parallel_size

233 + dcp_interleave = vllm_config.parallel_config.cp_kv_cache_interleave_size

234 + dcp_rank = get_dcp_group().rank_in_group if dcp_size > 1 else 0

235 + num_topk_tokens = vllm_config.model_config.hf_config.index_topk

236 + max_num_blocks = cdiv(

237 + vllm_config.model_config.max_model_len,

238 + block_size * dcp_size,

239 + )

240 + max_num_blocks_per_req = get_block_table_width(max_num_blocks, block_size)

241 + return self._trace_dispatch(self.dispatch)(

242 + zip_inputs(

243 + dict(

244 + HAS_PREFILL_WORKSPACE=False,

245 + COUNT_VALID=False,

246 + COMPACT_TO_FRONT=False,

247 + DCP_SIZE=1,

248 + DCP_RANK=0,

249 + DCP_INTERLEAVE=1,

250 + ),

251 + dict(

252 + HAS_PREFILL_WORKSPACE=False,

253 + COUNT_VALID=True,

254 + COMPACT_TO_FRONT=False,

255 + DCP_SIZE=1,

256 + DCP_RANK=0,

257 + DCP_INTERLEAVE=1,

258 + ),

259 + dict(

260 + HAS_PREFILL_WORKSPACE=True,

261 + COUNT_VALID=True,

262 + COMPACT_TO_FRONT=False,

263 + DCP_SIZE=1,

264 + DCP_RANK=0,

265 + DCP_INTERLEAVE=1,

266 + ),

267 + dict(

268 + HAS_PREFILL_WORKSPACE=False,

269 + COUNT_VALID=True,

270 + COMPACT_TO_FRONT=False,

271 + DCP_SIZE=dcp_size,

272 + DCP_RANK=dcp_rank,

273 + DCP_INTERLEAVE=dcp_interleave,

274 + ),

275 + dict(

276 + HAS_PREFILL_WORKSPACE=False,

277 + COUNT_VALID=True,

278 + COMPACT_TO_FRONT=dcp_size > 1,

279 + DCP_SIZE=dcp_size,

280 + DCP_RANK=dcp_rank,

281 + DCP_INTERLEAVE=dcp_interleave,

282 + ),

283 + ),

284 + BLOCK_SIZE=block_size,

285 + BLOCK_STRIDE_ROWS=block_stride_rows,

286 + BLOCK_N=128,

287 + NUM_TOPK_TOKENS=num_topk_tokens,

288 + max_num_blocks_per_req=max_num_blocks_per_req,

289 + )

290 +

291 + def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:

292 + int32_ptr = TritonWarmupTensor(torch.int32)

293 + token_shape = (1, compile_key.num_topk_tokens)

294 + token_strides = (16, 1)

295 + return dict(

296 + req_id=int32_ptr,

297 + block_table=TritonWarmupTensor(

298 + torch.int32,

299 + shape=(1, compile_key.block_table_stride),

300 + ),

301 + token_indices=TritonWarmupTensor(

302 + torch.int32,

303 + shape=token_shape,

304 + strides=token_strides,

305 + ),

306 + out=TritonWarmupTensor(

307 + torch.int32,

308 + shape=token_shape,

309 + strides=token_strides,

310 + ),

311 + valid_counts=int32_ptr if compile_key.count_valid else None,

312 + prefill_workspace_request_ids=(

313 + int32_ptr if compile_key.has_prefill_workspace else None

314 + ),

315 + prefill_workspace_starts=(

316 + int32_ptr if compile_key.has_prefill_workspace else None

317 + ),

318 + max_num_blocks_per_req=1,

319 + BLOCK_SIZE=compile_key.block_size,

320 + BLOCK_STRIDE_ROWS=compile_key.block_stride_rows,

321 + BLOCK_N=compile_key.block_n,

322 + NUM_TOPK_TOKENS=compile_key.num_topk_tokens,

323 + HAS_PREFILL_WORKSPACE=compile_key.has_prefill_workspace,

324 + COUNT_VALID=compile_key.count_valid,

325 + COMPACT_TO_FRONT=compile_key.compact_to_front,

326 + DCP_SIZE=compile_key.dcp_size,

327 + DCP_RANK=compile_key.dcp_rank,

328 + DCP_INTERLEAVE=compile_key.dcp_interleave,

329 + )

330 +

331 + @kernel_launcher

332 + def __call__(

333 + self,

334 + req_id: torch.Tensor,

335 + block_table: torch.Tensor,

336 + token_indices: torch.Tensor,

337 + out: torch.Tensor,

338 + valid_counts: torch.Tensor | None,

339 + prefill_workspace_request_ids: torch.Tensor | None,

340 + prefill_workspace_starts: torch.Tensor | None,

341 + *,

342 + max_num_blocks_per_req: int,

343 + BLOCK_SIZE: int,

344 + BLOCK_STRIDE_ROWS: int,

345 + BLOCK_N: int,

346 + NUM_TOPK_TOKENS: int,

347 + HAS_PREFILL_WORKSPACE: bool,

348 + COUNT_VALID: bool,

349 + COMPACT_TO_FRONT: bool,

350 + DCP_SIZE: int,

351 + DCP_RANK: int,

352 + DCP_INTERLEAVE: int,

353 + ) -> LaunchSpec:

354 + single_tile, block_n, tiles_per_row, num_warps = _remap_tiling(

355 + NUM_TOPK_TOKENS, BLOCK_N, COUNT_VALID

356 + )

357 + # Exact 2D grid: tokens × column tiles

358 + return (req_id.shape[0], tiles_per_row), dict(

359 + valid_count_ptr=valid_counts,

360 + prefill_request_id_ptr=prefill_workspace_request_ids,

361 + workspace_starts_ptr=prefill_workspace_starts,

362 + BLOCK_N=block_n,

363 + HAS_PREFILL=HAS_PREFILL_WORKSPACE,

364 + SINGLE_TILE=single_tile,

365 + # Strides in elements

366 + bt_stride0=block_table.stride(0),

367 + bt_stride1=block_table.stride(1),

368 + ti_stride0=token_indices.stride(0),

369 + ti_stride1=token_indices.stride(1),

370 + out_stride0=out.stride(0),

371 + out_stride1=out.stride(1),

372 + num_warps=num_warps,

373 + )

374

375

376 def _remap_tiling(

@@ -231,9 +452,7 @@ def triton_convert_req_index_to_global_index(

452 num_tokens = req_id.shape[0]

453 max_num_blocks_per_req = block_table.shape[1]

454

234 - single_tile, block_n, tiles_per_row, num_warps = _remap_tiling(

235 - NUM_TOPK_TOKENS, BLOCK_N, return_valid_counts

236 - )

455 + single_tile, _, _, _ = _remap_tiling(NUM_TOPK_TOKENS, BLOCK_N, return_valid_counts)

456

457 # Ensure contiguous tensors on the same device

458 req_id_c = req_id.contiguous()

@@ -247,11 +466,6 @@ def triton_convert_req_index_to_global_index(

466 alloc = torch.empty if single_tile else torch.zeros

467 valid_counts = alloc(num_tokens, dtype=torch.int32, device=token_indices.device)

468

250 - # Strides in elements

251 - bt_stride0, bt_stride1 = block_table_c.stride()

252 - ti_stride0, ti_stride1 = token_indices_c.stride()

253 - out_stride0, out_stride1 = out.stride()

254 -

469 # Prepare prefill pointers

470 if HAS_PREFILL_WORKSPACE:

471 assert prefill_workspace_request_ids is not None # for mypy

@@ -259,10 +473,7 @@ def triton_convert_req_index_to_global_index(

473 assert prefill_workspace_request_ids.is_contiguous()

474 assert prefill_workspace_starts.is_contiguous()

475

262 - # Exact 2D grid: tokens × column tiles

263 - grid = (num_tokens, tiles_per_row)

264 -

265 - _convert_req_index_to_global_index_kernel[grid](

476 + _CONVERT_REQ_INDEX_TO_GLOBAL_INDEX_KERNEL(

477 req_id_c,

478 block_table_c,

479 token_indices_c,

@@ -271,26 +482,20 @@ def triton_convert_req_index_to_global_index(

482 prefill_workspace_request_ids,

483 prefill_workspace_starts,

484 # shapes / constexprs

274 - max_num_blocks_per_req,

275 - BLOCK_SIZE,

276 - BLOCK_STRIDE_ROWS if BLOCK_STRIDE_ROWS is not None else BLOCK_SIZE,

277 - block_n,

278 - HAS_PREFILL_WORKSPACE,

279 - return_valid_counts,

280 - single_tile,

281 - False, # COMPACT_TO_FRONT: keep input column == output column

485 + max_num_blocks_per_req=max_num_blocks_per_req,

486 + BLOCK_SIZE=BLOCK_SIZE,

487 + BLOCK_STRIDE_ROWS=(

488 + BLOCK_STRIDE_ROWS if BLOCK_STRIDE_ROWS is not None else BLOCK_SIZE

489 + ),

490 + BLOCK_N=BLOCK_N,

491 + NUM_TOPK_TOKENS=NUM_TOPK_TOKENS,

492 + HAS_PREFILL_WORKSPACE=HAS_PREFILL_WORKSPACE,

493 + COUNT_VALID=return_valid_counts,

494 + COMPACT_TO_FRONT=False,

495 # DCP disabled (no-op de-interleave)

283 - 1,

284 - 0,

285 - 1,

286 - # strides

287 - bt_stride0,

288 - bt_stride1,

289 - ti_stride0,

290 - ti_stride1,

291 - out_stride0,

292 - out_stride1,

293 - num_warps=num_warps,

496 + DCP_SIZE=1,

497 + DCP_RANK=0,

498 + DCP_INTERLEAVE=1,

499 )

500

501 if return_valid_counts:

@@ -362,9 +567,7 @@ def triton_filter_and_convert_dcp_index(

567 count_valid = return_valid_counts or compact_valid_to_front

568

569 # The compaction builds on the counting, so it shares the tiling.

365 - single_tile, block_n, tiles_per_row, num_warps = _remap_tiling(

366 - NUM_TOPK_TOKENS, BLOCK_N, count_valid

367 - )

570 + single_tile, _, _, _ = _remap_tiling(NUM_TOPK_TOKENS, BLOCK_N, count_valid)

571

572 if compact_valid_to_front:

573 out = torch.full_like(token_indices_c, -1)

@@ -377,11 +580,7 @@ def triton_filter_and_convert_dcp_index(

580 alloc = torch.empty if single_tile else torch.zeros

581 valid_counts = alloc(num_tokens, dtype=torch.int32, device=token_indices.device)

582

380 - bt_stride0, bt_stride1 = block_table_c.stride()

381 - ti_stride0, ti_stride1 = token_indices_c.stride()

382 - out_stride0, out_stride1 = out.stride()

383 -

384 - _convert_req_index_to_global_index_kernel[(num_tokens, tiles_per_row)](

583 + _CONVERT_REQ_INDEX_TO_GLOBAL_INDEX_KERNEL(

584 req_id_c,

585 block_table_c,

586 token_indices_c,

@@ -390,27 +589,25 @@ def triton_filter_and_convert_dcp_index(

589 # No prefill workspace on the DCP decode path.

590 None,

591 None,

393 - max_num_blocks_per_req,

394 - BLOCK_SIZE,

395 - BLOCK_STRIDE_ROWS if BLOCK_STRIDE_ROWS is not None else BLOCK_SIZE,

396 - block_n,

397 - False, # HAS_PREFILL

398 - count_valid,

399 - single_tile,

400 - compact_valid_to_front,

401 - dcp_size,

402 - dcp_rank,

403 - cp_kv_cache_interleave_size,

404 - bt_stride0,

405 - bt_stride1,

406 - ti_stride0,

407 - ti_stride1,

408 - out_stride0,

409 - out_stride1,

410 - num_warps=num_warps,

592 + max_num_blocks_per_req=max_num_blocks_per_req,

593 + BLOCK_SIZE=BLOCK_SIZE,

594 + BLOCK_STRIDE_ROWS=(

595 + BLOCK_STRIDE_ROWS if BLOCK_STRIDE_ROWS is not None else BLOCK_SIZE

596 + ),

597 + BLOCK_N=BLOCK_N,

598 + NUM_TOPK_TOKENS=NUM_TOPK_TOKENS,

599 + HAS_PREFILL_WORKSPACE=False,

600 + COUNT_VALID=count_valid,

601 + COMPACT_TO_FRONT=compact_valid_to_front,

602 + DCP_SIZE=dcp_size,

603 + DCP_RANK=dcp_rank,

604 + DCP_INTERLEAVE=cp_kv_cache_interleave_size,

605 )

606

607 if return_valid_counts:

608 assert valid_counts is not None

609 return out, valid_counts

610 return out

611 +

612 +

613 +_CONVERT_REQ_INDEX_TO_GLOBAL_INDEX_KERNEL = ConvertReqIndexToGlobalIndexKernel()

vllm/v1/worker/gpu/model_runner.py

+17 -12

@@ -382,9 +382,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):

382 model_loader = get_model_loader(self.vllm_config.load_config)

383 logger.info_once("Loading model from scratch...")

384

385 - self.model = model_loader.load_model(

386 - vllm_config=self.vllm_config, model_config=self.vllm_config.model_config

387 - )

385 + # Capture warmup providers selected while constructing the model.

386 + with self.jit_warmup_registry.activate():

387 + self.model = model_loader.load_model(

388 + vllm_config=self.vllm_config,

389 + model_config=self.vllm_config.model_config,

390 + )

391 if self.lora_config:

392 self.model = self.load_lora_model(

393 self.model, self.vllm_config, self.device

@@ -668,15 +671,17 @@ class GPUModelRunner(LoRAModelRunnerMixin):

671 self.speculator.init_cudagraph_manager(cudagraph_mode)

672

673 self.kv_caches: list[torch.Tensor] = []

671 - kv_caches_dict = init_kv_cache(

672 - self.kv_caches,

673 - self.compilation_config.static_forward_context,

674 - self.kv_cache_config,

675 - self.device,

676 - self.kernel_block_sizes,

677 - self.vllm_config,

678 - kv_cache_allocation_context=kv_cache_allocation_context,

679 - )

674 + # Capture warmup providers that depend on allocated KV-cache strides.

675 + with self.jit_warmup_registry.activate():

676 + kv_caches_dict = init_kv_cache(

677 + self.kv_caches,

678 + self.compilation_config.static_forward_context,

679 + self.kv_cache_config,

680 + self.device,

681 + self.kernel_block_sizes,

682 + self.vllm_config,

683 + kv_cache_allocation_context=kv_cache_allocation_context,

684 + )

685 if is_profiling:

686 self.kv_connector = NO_OP_KV_CONNECTOR

687 else:

vllm/v1/worker/gpu_model_runner.py

+12 -8

@@ -5393,9 +5393,11 @@ class GPUModelRunner(

5393 if load_dummy_weights:

5394 self.load_config.load_format = "dummy"

5395 model_loader = get_model_loader(self.load_config)

5396 - self.model = model_loader.load_model(

5397 - vllm_config=self.vllm_config, model_config=self.model_config

5398 - )

5396 + # Capture warmup providers selected while constructing the model.

5397 + with self.jit_warmup_registry.activate():

5398 + self.model = model_loader.load_model(

5399 + vllm_config=self.vllm_config, model_config=self.model_config

5400 + )

5401 if self.lora_config:

5402 self.model = self.load_lora_model(

5403 self.model, self.vllm_config, self.device

@@ -7500,11 +7502,13 @@ class GPUModelRunner(

7502

7503 # Reinitialize need to after initialize_attn_backend

7504 self.may_reinitialize_input_batch(kv_cache_config, kernel_block_sizes)

7503 - kv_caches = self.initialize_kv_cache_tensors(

7504 - kv_cache_config,

7505 - kernel_block_sizes,

7506 - kv_cache_allocation_context=kv_cache_allocation_context,

7507 - )

7505 + # Capture warmup providers that depend on allocated KV-cache strides.

7506 + with self.jit_warmup_registry.activate():

7507 + kv_caches = self.initialize_kv_cache_tensors(

7508 + kv_cache_config,

7509 + kernel_block_sizes,

7510 + kv_cache_allocation_context=kv_cache_allocation_context,

7511 + )

7512

7513 if (

7514 self.speculative_config