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

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

Browse Source 
Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>
Co-authored-by: Codex <codex@openai.com>
This commit is contained in:
Roberto L. Castro
2026-08-31 23:26:05 -04:00
committed by GitHub
co-authored by Codex
parent 188716ace7
commit e16b5e518d
17 changed files with 1261 additions and 617 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/contributing/jit_kernel_warmup.md
+10 -8
View File
@@ -17,7 +17,7 @@ Each warmable kernel defines its compile-key mapping and compile-only entry poin
Here, a **kernel wrapper** (or just **wrapper**) is an instance of a concrete `VllmJitKernel` subclass.
Expose one wrapper near the kernel's normal runtime entry point. Prefer this shape:
Expose one wrapper near the kernel's normal runtime entry point. Prefer this backend-agnostic shape:
```python
class MyKernel(VllmJitKernel["MyKernel.CompileKey"]):
@@ -43,14 +43,14 @@ class MyKernel(VllmJitKernel["MyKernel.CompileKey"]):
return self.kernel(...)
MY_KERNEL = MyKernel()
_MY_KERNEL = MyKernel()
```
`CompileKey`, `dispatch(...)`, and `get_warmup_keys(...)` are backend-agnostic. Backend-specific behavior belongs in `kernel(...)`, `compile(...)`, and `__call__(...)`.
The module-level singleton should be used by warmup and by the runtime call path. This keeps dispatch behavior shared instead of duplicated.
`VllmJitKernel.warmup(...)` compiles every key returned by `get_warmup_keys(...)`; wrappers should not reimplement it.
`VllmJitKernel.warmup(...)` materializes every key returned by `get_warmup_keys(...)`; wrappers should not reimplement it.
### Choose Compile-Key Fields
@@ -70,18 +70,20 @@ def get_warmup_keys(self, vllm_config: VllmConfig) -> list[CompileKey]:
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).
### Compile Without Launching
### Materialize Without Launching
`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.
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.
`compile(...)` should not launch a real inference workload or allocate real tensors. Each DSL should expose fake tensor/spec descriptors suitable for compilation only.
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.
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.
### Register the Selected Wrapper
Register the wrapper where the runtime implementation is selected:
```python
MY_KERNEL.register_warmup()
_MY_KERNEL.register_warmup()
```
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
- Keep specialization mapping in `dispatch(...)` instead of duplicating it in warmup code.
- Use fake tensors or backend compile-only descriptors; never perform a dummy runtime launch.
- Keep registration metadata-only so model construction remains cheap and side-effect free.
- Compile registered kernels only through `kernel_warmup()`.
- Materialize registered kernels only through `kernel_warmup()`.
- Keep runtime execution and startup compilation separate and easy to review.
- Use one module-level wrapper instance for registration and runtime calls.
tests/model_executor/layers/test_fused_shared_expert.py
+1
View File
@@ -553,6 +553,7 @@ def test_models_fse_init(
method="mtp",
parallel_drafting=False,
enable_adaptive_verification=False,
use_dspark=lambda: False,
)
mtp = DeepSeekV4MTP(vllm_config=vllm_config)
assert model.is_fused_shared_expert_enabled is (fse_enabled and not exclude)
tests/model_executor/test_jit_warmup.py
+29
View File
@@ -590,6 +590,35 @@ def test_registry_passes_vllm_config_to_default_requests() -> None:
assert kernel.compiled == [ConfigKernel.CompileKey(value=7)]
def test_registry_default_request_supports_config_free_kernel() -> None:
class ConfigFreeKernel(VllmJitKernel["ConfigFreeKernel.CompileKey"]):
@dataclass(frozen=True)
class CompileKey:
value: int
def __init__(self) -> None:
self.compiled: list[ConfigFreeKernel.CompileKey] = []
super().__init__()
def dispatch(self, *, value: int) -> CompileKey: # type: ignore[override]
return self.CompileKey(value=value)
def get_warmup_keys(self) -> list[CompileKey]:
return [self.dispatch(value=7)]
def compile(self, compile_key: CompileKey) -> None:
self.compiled.append(compile_key)
registry = JitWarmupRegistry(_config())
kernel = ConfigFreeKernel()
with registry.activate():
kernel.register_warmup()
registry.warmup()
assert kernel.compiled == [ConfigFreeKernel.CompileKey(value=7)]
def test_get_ast_full_name_handles_names_attributes_and_other_nodes() -> None:
dotted_expr = ast.parse("foo.bar.baz").body[0]
call_expr = ast.parse("foo()").body[0]
tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py
+43 -2
View File
@@ -8,10 +8,16 @@ import torch
from tests.v1.attention.utils import create_vllm_config
from vllm.v1.attention.backend import CommonAttentionMetadata
from vllm.v1.attention.backends.mla.compressor_utils import (
CompressedSlotMappingKernel,
)
from vllm.v1.attention.backends.mla.indexer import (
BuildPrefillChunkMetadataKernel,
DeepseekV32IndexerMetadataBuilder,
)
from vllm.v1.attention.backends.mla.sparse_utils import (
ConvertReqIndexToGlobalIndexKernel,
)
from vllm.v1.kv_cache_interface import MLAAttentionSpec
from vllm.v1.worker.block_table import get_block_table_width
@@ -20,7 +26,7 @@ def test_indexer_warmup_normalizes_zero_compress_ratios():
config = SimpleNamespace(
scheduler_config=SimpleNamespace(max_num_batched_tokens=8),
model_config=SimpleNamespace(
hf_config=SimpleNamespace(compress_ratios=[0, 0, 4, 128, 0])
hf_config=SimpleNamespace(compress_ratios=[0, 0, 4, 128, 0], index_kpool=32)
),
parallel_config=SimpleNamespace(
decode_context_parallel_size=1,
@@ -30,7 +36,42 @@ def test_indexer_warmup_normalizes_zero_compress_ratios():
keys = BuildPrefillChunkMetadataKernel().get_warmup_keys(config)
assert {key.COMPRESS_RATIO for key in keys} == {1, 4, 128}
assert {key.compress_ratio for key in keys} == {1, 4, 32, 128}
assert {(key.query_slice_start, key.query_slice_stop) for key in keys} == {
(query_slice_start, query_slice_stop)
for query_slice_start in (1, 2, 16)
for query_slice_stop in (1, 2, 16)
}
def test_compressed_slot_mapping_warmup_includes_index_kpool():
config = SimpleNamespace(
cache_config=SimpleNamespace(block_size=256),
model_config=SimpleNamespace(hf_config=SimpleNamespace(index_kpool=32)),
)
keys = CompressedSlotMappingKernel().get_warmup_keys(config)
assert {(key.compress_ratio, key.block_size) for key in keys} == {(32, 2)}
def test_index_conversion_warmup_uses_physical_block_stride():
config = SimpleNamespace(
cache_config=SimpleNamespace(block_size=64),
model_config=SimpleNamespace(
max_model_len=1024,
hf_config=SimpleNamespace(index_topk=2048),
),
parallel_config=SimpleNamespace(
decode_context_parallel_size=1,
cp_kv_cache_interleave_size=1,
),
)
keys = ConvertReqIndexToGlobalIndexKernel().get_warmup_keys(
config,
block_stride_rows=4096,
)
assert {key.block_stride_rows for key in keys} == {4096}
@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
tests/v1/worker/test_gpu_model_runner_v2_eplb.py
+2
View File
@@ -7,6 +7,7 @@ from typing import Any
import torch
from vllm.model_executor.warmup.jit_warmup import JitWarmupRegistry
from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT
from vllm.v1.worker.gpu import eplb_utils as eplb
from vllm.v1.worker.gpu import model_runner as mrv2
@@ -84,6 +85,7 @@ def _make_runner(**overrides: Any) -> Any:
post_forward=lambda *_, **__: None,
)
runner.eplb = eplb.EPLBController(runner.parallel_config, runner.device)
runner.jit_warmup_registry = JitWarmupRegistry(runner.vllm_config)
runner.pooling_runner = None
runner.execute_model_state = None
for key, value in overrides.items():
vllm/model_executor/layers/attention/mla_attention.py
+47
View File
@@ -604,6 +604,33 @@ class MLAAttention(nn.Module, AttentionLayerBase):
self.use_sparse = use_sparse
if vllm_config.kernel_config.enable_jit_warmup:
backend_name = self.attn_backend.get_name()
if backend_name in (
"FLASHMLA_SPARSE",
"FLASHINFER_MLA_SPARSE",
"FLASHINFER_MLA_SPARSE_SM120",
"DEEPSEEK_V32_INDEXER",
):
from vllm.v1.attention.backends.mla.compressor_utils import (
_COMPRESSED_SLOT_MAPPING_KERNEL,
)
from vllm.v1.attention.backends.mla.indexer import (
_BUILD_PREFILL_CHUNK_METADATA_KERNEL,
_PREPARE_UNIFORM_DECODE_KERNEL,
)
_COMPRESSED_SLOT_MAPPING_KERNEL.register_warmup()
_PREPARE_UNIFORM_DECODE_KERNEL.register_warmup()
_BUILD_PREFILL_CHUNK_METADATA_KERNEL.register_warmup()
if backend_name != "DEEPSEEK_V32_INDEXER":
from vllm.v1.attention.backends.mla.sparse_swa import (
_COMPUTE_PREFILL_METADATA_KERNEL,
)
_COMPUTE_PREFILL_METADATA_KERNEL.register_warmup()
self.dcp_manager: MLADCPManager | None = None
if self.impl.dcp_world_size > 1:
query_dtype = (
@@ -652,6 +679,26 @@ class MLAAttention(nn.Module, AttentionLayerBase):
def bind_kv_cache(self, kv_cache: torch.Tensor) -> None:
# [B, H=1, N, C] -> [B, N, C]
self.kv_cache = kv_cache.squeeze(1)
if (
self._vllm_config.kernel_config.enable_jit_warmup
and self.attn_backend.get_name()
in (
"FLASHMLA_SPARSE",
"FLASHINFER_MLA_SPARSE",
"FLASHINFER_MLA_SPARSE_SM120",
"DEEPSEEK_V32_INDEXER",
)
):
from vllm.v1.attention.backends.mla.sparse_utils import (
_CONVERT_REQ_INDEX_TO_GLOBAL_INDEX_KERNEL,
)
row_width = self.kv_cache.shape[-1]
assert self.kv_cache.stride(0) % row_width == 0
_CONVERT_REQ_INDEX_TO_GLOBAL_INDEX_KERNEL.register_warmup(
self._vllm_config,
block_stride_rows=self.kv_cache.stride(0) // row_width,
)
@property
def chunked_prefill_workspace_size(self) -> int:
vllm/model_executor/warmup/jit_warmup.py
+5 -1
View File
@@ -737,7 +737,11 @@ class JitWarmupRegistry:
for kernel, registrations in self._registrations.items():
compile_keys: dict[Any, None] = {}
for args, kwargs in registrations:
if not args and not kwargs:
if (
not args
and not kwargs
and inspect.signature(kernel.get_warmup_keys).parameters
):
args = (self.vllm_config,)
for compile_key in kernel.get_warmup_keys(*args, **kwargs):
compile_keys[compile_key] = None
vllm/model_executor/warmup/jit_warmup_triton_helper.py
+90 -10
View File
@@ -2,15 +2,21 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import ast
import inspect
from collections.abc import Callable
from abc import abstractmethod
from collections.abc import Callable, Mapping
from dataclasses import dataclass
from typing import Any
from functools import cached_property, wraps
from typing import Any, ClassVar, Generic, TypeVar, cast
from vllm.model_executor.warmup.jit_warmup import (
VllmJitKernel,
get_ast_full_name,
get_function_source_node,
)
CompileKeyT = TypeVar("CompileKeyT")
LaunchSpec = tuple[tuple[int, ...], dict[str, Any]]
def triton_scalar_specialization_rep(value: int) -> int:
"""Return an integer with the same default Triton JIT specialization.
@@ -50,10 +56,16 @@ def triton_scalar_specialization_rep(value: int) -> int:
@dataclass(frozen=True)
class TritonWarmupTensor:
# Compile-only tensor descriptor for Triton pointer specialization.
"""Compile-only tensor metadata used by Triton warmup.
``strides=None`` represents compact row-major storage. Pass explicit strides
whenever the runtime tensor can be padded, transposed, or otherwise strided.
"""
dtype: Any
aligned: bool = True
shape: tuple[int, ...] = (1,)
strides: tuple[int, ...] | None = None
def data_ptr(self) -> int:
return 0 if self.aligned else 1
@@ -61,13 +73,81 @@ class TritonWarmupTensor:
def ptr_range(self) -> int:
return 0
def stride(self) -> tuple[int, ...]:
strides: list[int] = []
stride = 1
for size in reversed(self.shape):
strides.append(stride)
stride *= size
return tuple(reversed(strides))
def stride(self, dim: int | None = None) -> int | tuple[int, ...]:
if self.strides is None:
strides: list[int] = []
stride = 1
for size in reversed(self.shape):
strides.append(stride)
stride *= size
result = tuple(reversed(strides))
else:
result = self.strides
return result if dim is None else result[dim]
class VllmTritonJitKernel(VllmJitKernel[CompileKeyT], Generic[CompileKeyT]):
"""Triton owner whose runtime launch specification is reused for warmup."""
kernel: ClassVar[Any]
_warming = False
@abstractmethod
def warmup_inputs(self, compile_key: CompileKeyT) -> dict[str, Any]:
"""Return runtime-shaped inputs that reproduce one compile key."""
raise NotImplementedError
def compile(self, compile_key: CompileKeyT) -> None:
inputs = self.warmup_inputs(compile_key)
self._warming = True
try:
cast(Callable[..., None], self)(**inputs)
finally:
self._warming = False
@cached_property
def _kernel_param_names(self) -> frozenset[str]:
return frozenset(self.kernel.arg_names)
def launch(
self,
grid: tuple[int, ...],
inputs: Mapping[str, Any],
/,
**kwargs: Any,
) -> Any:
for name, value in inputs.items():
target = name if name in self._kernel_param_names else f"{name}_ptr"
if target in self._kernel_param_names and target not in kwargs:
kwargs[target] = value
if self._warming:
warmup = getattr(self.kernel, "warmup", None)
assert warmup is not None
return warmup(grid=(1,), **kwargs)
return self.kernel[grid](**kwargs)
def kernel_launcher(
call_fn: Callable[..., LaunchSpec],
) -> Callable[..., None]:
"""Launch a Triton kernel from a declarative ``__call__`` specification."""
signature = inspect.signature(call_fn)
@wraps(call_fn)
def wrapper(
self: VllmTritonJitKernel[Any],
*args: Any,
**kwargs: Any,
) -> None:
grid, launch_kwargs = call_fn(self, *args, **kwargs)
bound = signature.bind(self, *args, **kwargs)
bound.apply_defaults()
inputs = {
name: value for name, value in bound.arguments.items() if name != "self"
}
self.launch(grid, inputs, **launch_kwargs)
return wrapper
@dataclass(frozen=True)
vllm/model_executor/warmup/kernel_warmup.py
-4
View File
@@ -35,9 +35,6 @@ from vllm.model_executor.warmup.kimi_k3_triton_warmup import (
kimi_k3_triton_warmup,
)
from vllm.model_executor.warmup.qwen_triton_warmup import qwen_triton_warmup
from vllm.model_executor.warmup.sparse_mla_triton_warmup import (
sparse_mla_triton_warmup,
)
from vllm.platforms import current_platform
from vllm.utils.deep_gemm import is_deep_gemm_supported
from vllm.utils.flashinfer import has_flashinfer
@@ -154,7 +151,6 @@ def kernel_warmup(worker: "Worker", *, process_local_only: bool = False):
if worker.vllm_config.kernel_config.enable_jit_warmup:
kimi_k3_triton_warmup(worker)
fa4_cutedsl_warmup(worker)
sparse_mla_triton_warmup(worker)
if current_platform.has_device_capability(90):
_warmup_ll_bf16_router_gemm(worker.get_model())
vllm/model_executor/warmup/sparse_mla_triton_warmup.py
-111
View File
@@ -1,111 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Warm up sparse-MLA Triton metadata kernels."""
from typing import TYPE_CHECKING
from vllm.logger import init_logger
if TYPE_CHECKING:
from vllm.config import VllmConfig
from vllm.v1.worker.gpu_model_runner import GPUModelRunner
from vllm.v1.worker.gpu_worker import Worker
logger = init_logger(__name__)
_DEEPSEEK_V4_SPARSE_MLA_BACKENDS = frozenset(
{
"FLASHMLA_SPARSE_DSV4",
"FLASHINFER_MLA_SPARSE_DSV4",
"ROCM_FLASHMLA_SPARSE_DSV4",
"DEEPSEEK_SPARSE_SWA",
}
)
_GENERIC_SPARSE_MLA_BACKENDS = frozenset(
{
"FLASHMLA_SPARSE",
"FLASHINFER_MLA_SPARSE",
"FLASHINFER_MLA_SPARSE_SM120",
}
)
_INDEXER_PREFILL_CHUNK_METADATA_BACKENDS = frozenset({"DEEPSEEK_V32_INDEXER"})
_INDEXER_PREFILL_CHUNK_METADATA_BACKENDS = frozenset({"DEEPSEEK_V32_INDEXER"})
def _attention_backend_name(backend: object) -> str | None:
get_name = getattr(backend, "get_name", None)
if get_name is None:
return None
try:
return get_name()
except NotImplementedError:
return None
def _has_attention_backend(
runner: "GPUModelRunner",
backend_names: frozenset[str],
) -> bool:
for groups in getattr(runner, "attn_groups", []) or ():
for group in groups:
name = _attention_backend_name(getattr(group, "backend", None))
if name in backend_names:
return True
return False
def _compile_sparse_swa_prefill_metadata_kernel(
vllm_config: "VllmConfig",
) -> None:
from vllm.v1.attention.backends.mla.sparse_swa import (
_COMPUTE_PREFILL_METADATA_KERNEL,
)
_COMPUTE_PREFILL_METADATA_KERNEL.warmup(vllm_config)
def _compile_prefill_chunk_metadata_kernel(
vllm_config: "VllmConfig",
) -> None:
from vllm.v1.attention.backends.mla.indexer import (
_BUILD_PREFILL_CHUNK_METADATA_KERNEL,
)
_BUILD_PREFILL_CHUNK_METADATA_KERNEL.warmup(vllm_config)
def _compile_combine_topk_swa_indices_kernel(
vllm_config: "VllmConfig",
) -> None:
from vllm.models.deepseek_v4.common.ops.cache_utils import (
_COMBINE_TOPK_SWA_INDICES_KERNEL,
)
_COMBINE_TOPK_SWA_INDICES_KERNEL.warmup(vllm_config)
def sparse_mla_triton_warmup(worker: "Worker") -> None:
runner = worker.model_runner
if runner.is_pooling_model:
return
max_tokens = worker.scheduler_config.max_num_batched_tokens
max_num_prefills = min(worker.scheduler_config.max_num_seqs, max_tokens)
if max_tokens <= 0 or max_num_prefills <= 0:
return
vllm_config = runner.vllm_config
try:
if _has_attention_backend(runner, _DEEPSEEK_V4_SPARSE_MLA_BACKENDS):
_compile_sparse_swa_prefill_metadata_kernel(vllm_config)
_compile_prefill_chunk_metadata_kernel(vllm_config)
_compile_combine_topk_swa_indices_kernel(vllm_config)
elif _has_attention_backend(runner, _GENERIC_SPARSE_MLA_BACKENDS):
_compile_sparse_swa_prefill_metadata_kernel(vllm_config)
_compile_prefill_chunk_metadata_kernel(vllm_config)
elif _has_attention_backend(runner, _INDEXER_PREFILL_CHUNK_METADATA_BACKENDS):
_compile_prefill_chunk_metadata_kernel(vllm_config)
except Exception:
logger.warning("Skipping sparse MLA Triton warmup.", exc_info=True)
vllm/models/deepseek_v4/attention.py
+51
View File
@@ -353,6 +353,57 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
k_cache_prefix=self.prefix,
)
if vllm_config.kernel_config.enable_jit_warmup:
from vllm.v1.attention.backends.mla.sparse_swa import (
_COMPUTE_PREFILL_METADATA_KERNEL,
_COMPUTE_SWA_INDICES_AND_LENS_KERNEL,
)
_COMPUTE_PREFILL_METADATA_KERNEL.register_warmup()
_COMPUTE_SWA_INDICES_AND_LENS_KERNEL.register_warmup(
window_size=self.window_size,
block_size=self.swa_cache_layer.block_size,
)
if self.compress_ratio > 1:
from vllm.v1.attention.backends.mla.compressor_utils import (
_COMPRESSED_SLOT_MAPPING_KERNEL,
)
_COMPRESSED_SLOT_MAPPING_KERNEL.register_warmup()
if self.indexer is not None:
from vllm.v1.attention.backends.mla.indexer import (
_BUILD_PREFILL_CHUNK_METADATA_KERNEL,
_PREPARE_UNIFORM_DECODE_KERNEL,
)
_PREPARE_UNIFORM_DECODE_KERNEL.register_warmup()
_BUILD_PREFILL_CHUNK_METADATA_KERNEL.register_warmup()
spec_config = vllm_config.speculative_config
if spec_config is not None and spec_config.use_dspark():
from vllm.v1.attention.backends.mla.sparse_swa import (
_COMPUTE_DSPARK_NONCAUSAL_SWA_INDICES_KERNEL,
)
_COMPUTE_DSPARK_NONCAUSAL_SWA_INDICES_KERNEL.register_warmup(
window_size=self.window_size,
num_speculative_tokens=spec_config.num_speculative_tokens,
block_size=self.swa_cache_layer.block_size,
)
if self.backend_cls.get_name() in (
"FLASHMLA_SPARSE_DSV4",
"ROCM_FLASHMLA_SPARSE_DSV4",
"XPU_V4_MLA_SPARSE",
):
from vllm.models.deepseek_v4.common.ops.cache_utils import (
_COMBINE_TOPK_SWA_INDICES_KERNEL,
)
_COMBINE_TOPK_SWA_INDICES_KERNEL.register_warmup()
def forward(
self,
positions: torch.Tensor,
vllm/v1/attention/backends/mla/compressor_utils.py
+125 -40
View File
@@ -1,7 +1,18 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from dataclasses import dataclass
from typing import Any
import torch
from vllm.model_executor.warmup.jit_warmup import zip_inputs
from vllm.model_executor.warmup.jit_warmup_triton_helper import (
LaunchSpec,
TritonWarmupTensor,
VllmTritonJitKernel,
kernel_launcher,
triton_scalar_specialization_rep,
)
from vllm.triton_utils import tl, triton
from vllm.utils.math_utils import cdiv
@@ -17,49 +28,124 @@ def get_dspark_swa_index_width(
return cdiv(width, _DSPARK_SWA_INDEX_ALIGNMENT) * _DSPARK_SWA_INDEX_ALIGNMENT
@triton.jit
def _compressed_slot_mapping_kernel(
# [num_tokens]
slot_mapping_ptr,
# [num_reqs + 1]
query_start_loc_ptr,
# [num_reqs]
seq_lens_ptr,
# [num_reqs, max_num_blocks]
block_table_ptr,
block_table_stride,
block_size,
COMPRESS_RATIO: tl.constexpr,
PAD_ID: tl.constexpr,
TRITON_BLOCK_SIZE: tl.constexpr,
class CompressedSlotMappingKernel(
VllmTritonJitKernel["CompressedSlotMappingKernel.CompileKey"]
):
batch_idx = tl.program_id(0)
TRITON_BLOCK_SIZE = 1024
query_start = tl.load(query_start_loc_ptr + batch_idx)
query_end = tl.load(query_start_loc_ptr + batch_idx + 1)
query_len = query_end - query_start
@dataclass(frozen=True)
class CompileKey:
compress_ratio: int
triton_block_size: int
block_size: int
seq_len = tl.load(seq_lens_ptr + batch_idx)
start_pos = seq_len - query_len
@staticmethod
@triton.jit(do_not_specialize=["block_table_stride"])
def kernel(
# [num_tokens]
slot_mapping_ptr,
# [num_reqs + 1]
query_start_loc_ptr,
# [num_reqs]
seq_lens_ptr,
# [num_reqs, max_num_blocks]
block_table_ptr,
block_table_stride,
block_size,
COMPRESS_RATIO: tl.constexpr,
PAD_ID: tl.constexpr,
TRITON_BLOCK_SIZE: tl.constexpr,
):
batch_idx = tl.program_id(0)
for i in range(0, query_len, TRITON_BLOCK_SIZE):
offset = i + tl.arange(0, TRITON_BLOCK_SIZE)
mask = offset < query_len
query_start = tl.load(query_start_loc_ptr + batch_idx)
query_end = tl.load(query_start_loc_ptr + batch_idx + 1)
query_len = query_end - query_start
pos = start_pos + i + tl.arange(0, TRITON_BLOCK_SIZE)
is_valid = (pos + 1) % COMPRESS_RATIO == 0
pos_after_compress = pos // COMPRESS_RATIO
seq_len = tl.load(seq_lens_ptr + batch_idx)
start_pos = seq_len - query_len
block_ids = pos_after_compress // block_size
block_numbers = tl.load(
block_table_ptr + batch_idx * block_table_stride + block_ids,
mask=mask & is_valid,
for i in range(0, query_len, TRITON_BLOCK_SIZE):
offset = i + tl.arange(0, TRITON_BLOCK_SIZE)
mask = offset < query_len
pos = start_pos + i + tl.arange(0, TRITON_BLOCK_SIZE)
is_valid = (pos + 1) % COMPRESS_RATIO == 0
pos_after_compress = pos // COMPRESS_RATIO
block_ids = pos_after_compress // block_size
block_numbers = tl.load(
block_table_ptr + batch_idx * block_table_stride + block_ids,
mask=mask & is_valid,
)
slot_ids = block_numbers * block_size + pos_after_compress % block_size
# NOTE
slot_ids = tl.where(is_valid, slot_ids, PAD_ID)
tl.store(slot_mapping_ptr + query_start + offset, slot_ids, mask=mask)
def dispatch( # type: ignore[override]
self,
*,
compress_ratio: int,
block_size: int,
) -> CompileKey:
return self.CompileKey(
compress_ratio=compress_ratio,
triton_block_size=self.TRITON_BLOCK_SIZE,
block_size=triton_scalar_specialization_rep(block_size),
)
slot_ids = block_numbers * block_size + pos_after_compress % block_size
# NOTE
slot_ids = tl.where(is_valid, slot_ids, PAD_ID)
tl.store(slot_mapping_ptr + query_start + offset, slot_ids, mask=mask)
def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]:
hf_config = vllm_config.model_config.hf_config
configured_ratios = (
*(getattr(hf_config, "compress_ratios", None) or ()),
getattr(hf_config, "index_kpool", 1) or 1,
)
compress_ratios = tuple(
dict.fromkeys(int(ratio) for ratio in configured_ratios if int(ratio) > 1)
)
if not compress_ratios:
return []
return self._trace_dispatch(self.dispatch)(
zip_inputs(
*(
dict(
compress_ratio=ratio,
block_size=vllm_config.cache_config.block_size // ratio,
)
for ratio in compress_ratios
)
)
)
def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
int32_ptr = TritonWarmupTensor(torch.int32)
return dict(
slot_mapping=TritonWarmupTensor(torch.int64),
query_start_loc=int32_ptr,
seq_lens=int32_ptr,
block_table=int32_ptr,
block_size=compile_key.block_size,
compress_ratio=compile_key.compress_ratio,
)
@kernel_launcher
def __call__(
self,
slot_mapping: torch.Tensor,
query_start_loc: torch.Tensor,
seq_lens: torch.Tensor,
block_table: torch.Tensor,
block_size: int,
compress_ratio: int,
) -> LaunchSpec:
return (block_table.shape[0],), dict(
block_table_stride=block_table.stride(0),
COMPRESS_RATIO=compress_ratio,
PAD_ID=-1,
TRITON_BLOCK_SIZE=self.TRITON_BLOCK_SIZE,
)
def get_compressed_slot_mapping(
@@ -83,16 +169,15 @@ def get_compressed_slot_mapping(
(num_tokens,), -1, dtype=torch.int64, device=query_start_loc.device
)
num_reqs = block_table.shape[0]
_compressed_slot_mapping_kernel[(num_reqs,)](
_COMPRESSED_SLOT_MAPPING_KERNEL(
slot_mapping,
query_start_loc,
seq_lens,
block_table,
block_table.stride(0),
block_size,
compress_ratio,
PAD_ID=-1,
TRITON_BLOCK_SIZE=1024,
)
return slot_mapping
_COMPRESSED_SLOT_MAPPING_KERNEL = CompressedSlotMappingKernel()
vllm/v1/attention/backends/mla/indexer.py
+160 -98
View File
@@ -1,6 +1,7 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from dataclasses import dataclass
from typing import Any
import torch
@@ -8,13 +9,13 @@ import vllm.envs as envs
from vllm.config import VllmConfig
from vllm.distributed import get_dcp_group, get_pcp_group
from vllm.logger import init_logger
from vllm.model_executor.warmup.jit_warmup import (
VllmJitKernel,
WarmupIntRange,
)
from vllm.model_executor.warmup.jit_warmup_triton_helper import (
LaunchSpec,
TritonPointerInputVariant,
TritonWarmupTensor,
VllmTritonJitKernel,
kernel_launcher,
triton_scalar_specialization_rep,
)
from vllm.platforms import current_platform
from vllm.triton_utils import tl, triton
@@ -63,41 +64,112 @@ def dsa_indexer_uses_fp4(vllm_config: VllmConfig) -> bool:
return use_fp4
@triton.jit
def _prepare_uniform_decode_kernel(
seq_lens_ptr,
decode_seq_lens_ptr,
block_table_ptr,
block_table_stride,
expanded_block_table_ptr,
expanded_bt_stride,
decode_lens_ptr,
max_decode_len,
BLOCK_SIZE: tl.constexpr,
class PrepareUniformDecodeKernel(
VllmTritonJitKernel["PrepareUniformDecodeKernel.CompileKey"]
):
idx = tl.program_id(0)
req_id = idx // max_decode_len
local_idx = idx % max_decode_len
BLOCK_SIZE = 1024
# Compute number of KVs attended to by this token. Padding requests have
# seq_len == 0, which would otherwise make the first token of each padded
# request negative (e.g. next_n=2 gives 0-2+0+1 = -1). Downstream kernels
# read these as uint32, turning -1 into ~4e9.
seq_len = tl.load(seq_lens_ptr + req_id)
per_token_seq_len = tl.maximum(seq_len - max_decode_len + local_idx + 1, 0)
tl.store(decode_seq_lens_ptr + idx, per_token_seq_len)
@dataclass(frozen=True)
class CompileKey:
block_size: int
block_table_stride: int
expanded_bt_stride: int
max_decode_len: int
# Copy block table row.
src = block_table_ptr + req_id * block_table_stride
dst = expanded_block_table_ptr + idx * expanded_bt_stride
for i in tl.range(0, expanded_bt_stride, BLOCK_SIZE):
off = i + tl.arange(0, BLOCK_SIZE)
mask = off < expanded_bt_stride
src_block = tl.load(src + off, mask=mask)
tl.store(dst + off, src_block, mask=mask)
@staticmethod
@triton.jit
def kernel(
seq_lens_ptr,
decode_seq_lens_ptr,
block_table_ptr,
block_table_stride,
expanded_block_table_ptr,
expanded_bt_stride,
decode_lens_ptr,
max_decode_len,
BLOCK_SIZE: tl.constexpr,
):
idx = tl.program_id(0)
req_id = idx // max_decode_len
local_idx = idx % max_decode_len
# All reqs now have decode_len = 1.
tl.store(decode_lens_ptr + idx, 1)
# Compute number of KVs attended to by this token. Padding requests have
# seq_len == 0, which would otherwise make the first token of each padded
# request negative (e.g. next_n=2 gives 0-2+0+1 = -1). Downstream kernels
# read these as uint32, turning -1 into ~4e9.
seq_len = tl.load(seq_lens_ptr + req_id)
per_token_seq_len = tl.maximum(seq_len - max_decode_len + local_idx + 1, 0)
tl.store(decode_seq_lens_ptr + idx, per_token_seq_len)
# Copy block table row.
src = block_table_ptr + req_id * block_table_stride
dst = expanded_block_table_ptr + idx * expanded_bt_stride
for i in tl.range(0, expanded_bt_stride, BLOCK_SIZE):
off = i + tl.arange(0, BLOCK_SIZE)
mask = off < expanded_bt_stride
src_block = tl.load(src + off, mask=mask)
tl.store(dst + off, src_block, mask=mask)
# All reqs now have decode_len = 1.
tl.store(decode_lens_ptr + idx, 1)
def dispatch( # type: ignore[override]
self,
*,
block_size: int,
block_table_stride: int,
expanded_bt_stride: int,
max_decode_len: int,
) -> CompileKey:
return self.CompileKey(
block_size=block_size,
block_table_stride=triton_scalar_specialization_rep(block_table_stride),
expanded_bt_stride=triton_scalar_specialization_rep(expanded_bt_stride),
max_decode_len=triton_scalar_specialization_rep(max_decode_len),
)
def get_warmup_keys(self) -> list[CompileKey]:
return self._trace_dispatch(self.dispatch)(
block_size=self.BLOCK_SIZE,
block_table_stride=(1, 2, 16),
expanded_bt_stride=(1, 2, 16),
max_decode_len=(1, 2, 16),
)
def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
int32_ptr = TritonWarmupTensor(torch.int32)
return dict(
seq_lens=int32_ptr,
decode_seq_lens=int32_ptr,
block_table=TritonWarmupTensor(
torch.int32,
shape=(1, compile_key.block_table_stride),
),
expanded_block_table=TritonWarmupTensor(
torch.int32,
shape=(1, compile_key.expanded_bt_stride),
),
decode_lens=int32_ptr,
num_decode_tokens=1,
max_decode_len=compile_key.max_decode_len,
)
@kernel_launcher
def __call__(
self,
seq_lens: torch.Tensor,
decode_seq_lens: torch.Tensor,
block_table: torch.Tensor,
expanded_block_table: torch.Tensor,
decode_lens: torch.Tensor,
num_decode_tokens: int,
max_decode_len: int,
) -> LaunchSpec:
return (num_decode_tokens,), dict(
block_table_stride=block_table.stride(0),
expanded_bt_stride=expanded_block_table.stride(0),
BLOCK_SIZE=self.BLOCK_SIZE,
)
def split_indexer_prefill_chunks(
@@ -215,14 +287,8 @@ class DeepseekV32IndexerPrefillChunkMetadata:
max_local_total_seq_lens: int = 0
_BUILD_PREFILL_CHUNK_METADATA_INPUT_VARIANTS = (
TritonPointerInputVariant.from_alignment(uncompressed_seq_lens=True),
TritonPointerInputVariant.from_alignment(uncompressed_seq_lens=False),
)
class BuildPrefillChunkMetadataKernel(
VllmJitKernel["BuildPrefillChunkMetadataKernel.CompileKey"]
VllmTritonJitKernel["BuildPrefillChunkMetadataKernel.CompileKey"]
):
BLOCK_SIZE = 1024
@@ -230,11 +296,11 @@ class BuildPrefillChunkMetadataKernel(
class CompileKey:
query_slice_start: int
query_slice_stop: int
DCP_RANK: int
DCP_WORLD: int
DCP_INTERLEAVE: int
BLOCK_SIZE: int
COMPRESS_RATIO: int
dcp_rank: int
dcp_world: int
dcp_interleave: int
block_size: int
compress_ratio: int
input_variant: TritonPointerInputVariant
@staticmethod
@@ -327,13 +393,13 @@ class BuildPrefillChunkMetadataKernel(
input_variant: TritonPointerInputVariant,
) -> CompileKey:
return self.CompileKey(
query_slice_start=query_slice_start,
query_slice_stop=query_slice_stop,
DCP_RANK=DCP_RANK,
DCP_WORLD=DCP_WORLD,
DCP_INTERLEAVE=DCP_INTERLEAVE,
BLOCK_SIZE=BLOCK_SIZE,
COMPRESS_RATIO=COMPRESS_RATIO,
query_slice_start=triton_scalar_specialization_rep(query_slice_start),
query_slice_stop=triton_scalar_specialization_rep(query_slice_stop),
dcp_rank=triton_scalar_specialization_rep(DCP_RANK),
dcp_world=triton_scalar_specialization_rep(DCP_WORLD),
dcp_interleave=triton_scalar_specialization_rep(DCP_INTERLEAVE),
block_size=BLOCK_SIZE,
compress_ratio=COMPRESS_RATIO,
input_variant=input_variant,
)
@@ -345,42 +411,51 @@ class BuildPrefillChunkMetadataKernel(
dcp_interleave = parallel_config.cp_kv_cache_interleave_size
dcp_rank = get_dcp_group().rank_in_group if dcp_world > 1 else 0
compress_ratios = tuple(
max(1, int(ratio))
for ratio in (getattr(hf_config, "compress_ratios", None) or (1,))
dict.fromkeys(
max(1, int(ratio))
for ratio in (
*(getattr(hf_config, "compress_ratios", None) or (1,)),
getattr(hf_config, "index_kpool", 1) or 1,
)
)
)
return self._trace_dispatch(self.dispatch)(
query_slice_start=WarmupIntRange(0, 2),
# Cover Triton's divisible, exact-one, and generic i32 classes.
query_slice_start=(0, 1, 2),
query_slice_stop=(1, 2 * max_tokens - 1, 2 * max_tokens),
DCP_RANK=dcp_rank,
DCP_WORLD=dcp_world,
DCP_INTERLEAVE=dcp_interleave,
BLOCK_SIZE=self.BLOCK_SIZE,
COMPRESS_RATIO=list(compress_ratios),
input_variant=_BUILD_PREFILL_CHUNK_METADATA_INPUT_VARIANTS,
input_variant=(
TritonPointerInputVariant.from_alignment(uncompressed_seq_lens=True),
TritonPointerInputVariant.from_alignment(uncompressed_seq_lens=False),
),
)
def compile(self, compile_key: CompileKey) -> None:
warmup = getattr(self.kernel, "warmup", None)
assert warmup is not None
def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
int32_ptr = TritonWarmupTensor(torch.int32)
warmup(
int32_ptr,
compile_key.input_variant.pointer("uncompressed_seq_lens", torch.int32),
int32_ptr,
int32_ptr,
int32_ptr,
int32_ptr,
int32_ptr,
compile_key.query_slice_start,
compile_key.query_slice_stop,
compile_key.DCP_RANK,
compile_key.DCP_WORLD,
compile_key.DCP_INTERLEAVE,
BLOCK_SIZE=compile_key.BLOCK_SIZE,
COMPRESS_RATIO=compile_key.COMPRESS_RATIO,
grid=(1,),
return dict(
query_start_loc=int32_ptr,
uncompressed_seq_lens=compile_key.input_variant.pointer(
"uncompressed_seq_lens", torch.int32
),
cu_compressed_seq_lens=int32_ptr,
row_start_cu_compressed_seq_lens=int32_ptr,
token_to_seq=int32_ptr,
cu_compressed_seq_len_ks=int32_ptr,
cu_compressed_seq_len_ke=int32_ptr,
query_slice_start=compile_key.query_slice_start,
query_slice_stop=compile_key.query_slice_stop,
DCP_RANK=compile_key.dcp_rank,
DCP_WORLD=compile_key.dcp_world,
DCP_INTERLEAVE=compile_key.dcp_interleave,
num_reqs=1,
COMPRESS_RATIO=compile_key.compress_ratio,
)
@kernel_launcher
def __call__(
self,
query_start_loc: torch.Tensor,
@@ -398,28 +473,13 @@ class BuildPrefillChunkMetadataKernel(
*,
num_reqs: int,
COMPRESS_RATIO: int,
) -> None:
self.kernel[(num_reqs,)](
query_start_loc,
uncompressed_seq_lens,
cu_compressed_seq_lens,
row_start_cu_compressed_seq_lens,
token_to_seq,
cu_compressed_seq_len_ks,
cu_compressed_seq_len_ke,
query_slice_start,
query_slice_stop,
DCP_RANK,
DCP_WORLD,
DCP_INTERLEAVE,
) -> LaunchSpec:
return (num_reqs,), dict(
BLOCK_SIZE=self.BLOCK_SIZE,
COMPRESS_RATIO=COMPRESS_RATIO,
)
_BUILD_PREFILL_CHUNK_METADATA_KERNEL = BuildPrefillChunkMetadataKernel()
@dataclass
class DeepseekV32IndexerPrefillMetadata:
chunks: list[DeepseekV32IndexerPrefillChunkMetadata]
@@ -694,16 +754,14 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
and num_decodes * max_decode_len == num_decode_tokens
):
# Uniform decode lengths with no cudagraph token padding.
_prepare_uniform_decode_kernel[(num_decode_tokens,)](
_PREPARE_UNIFORM_DECODE_KERNEL(
seq_lens,
self.decode_seq_lens_buffer,
block_table,
block_table.stride(0),
self.expanded_block_table_buffer,
self.expanded_block_table_buffer.stride(0),
self.decode_lens_buffer,
num_decode_tokens,
max_decode_len,
BLOCK_SIZE=1024,
)
self.decode_seq_lens_buffer[num_decode_tokens:] = 0
seq_lens = self.decode_seq_lens_buffer[:num_decode_tokens]
@@ -1180,3 +1238,7 @@ def build_prefill_chunk_metadata(
local_total_seq_lens=local_total_seq_lens,
max_local_total_seq_lens=max_local_total_seq_lens,
)
_PREPARE_UNIFORM_DECODE_KERNEL = PrepareUniformDecodeKernel()
_BUILD_PREFILL_CHUNK_METADATA_KERNEL = BuildPrefillChunkMetadataKernel()
vllm/v1/attention/backends/mla/sparse_swa.py
+309 -160
View File
@@ -1,17 +1,21 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from dataclasses import dataclass, field
from typing import ClassVar, cast
from typing import Any, ClassVar, cast
import torch
from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
from vllm.model_executor.warmup.jit_warmup import (
VllmJitKernel,
WarmupIntRange,
)
from vllm.model_executor.warmup.jit_warmup_triton_helper import TritonWarmupTensor
from vllm.model_executor.warmup.jit_warmup_triton_helper import (
LaunchSpec,
TritonWarmupTensor,
VllmTritonJitKernel,
kernel_launcher,
)
from vllm.platforms import current_platform
from vllm.triton_utils import tl, triton
from vllm.utils.math_utils import cdiv, next_power_of_2
@@ -275,11 +279,11 @@ class DeepseekSparseSWAMetadata:
class ComputePrefillMetadataKernel(
VllmJitKernel["ComputePrefillMetadataKernel.CompileKey"]
VllmTritonJitKernel["ComputePrefillMetadataKernel.CompileKey"]
):
@dataclass(frozen=True)
class CompileKey:
BLOCK_SIZE: int
block_size: int
@staticmethod
@triton.jit(do_not_specialize=["num_prefills", "num_decodes", "window_size"])
@@ -320,7 +324,7 @@ class ComputePrefillMetadataKernel(
num_prefills: int,
) -> CompileKey:
return self.CompileKey(
BLOCK_SIZE=next_power_of_2(num_prefills),
block_size=next_power_of_2(num_prefills),
)
def get_warmup_keys(self, vllm_config: VllmConfig) -> list[CompileKey]:
@@ -336,21 +340,18 @@ class ComputePrefillMetadataKernel(
num_prefills=WarmupIntRange(1, max_prefills + 1),
)
def compile(self, compile_key: CompileKey) -> None:
warmup = getattr(self.kernel, "warmup", None)
assert warmup is not None
def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
int32_ptr = TritonWarmupTensor(torch.int32)
warmup(
int32_ptr,
int32_ptr,
int32_ptr,
compile_key.BLOCK_SIZE,
0,
1,
BLOCK_SIZE=compile_key.BLOCK_SIZE,
grid=(1,),
return dict(
prefill_gather_lens=int32_ptr,
seq_lens=int32_ptr,
query_start_loc=int32_ptr,
num_prefills=compile_key.block_size,
num_decodes=0,
window_size=1,
)
@kernel_launcher
def __call__(
self,
prefill_gather_lens: torch.Tensor,
@@ -359,22 +360,13 @@ class ComputePrefillMetadataKernel(
num_prefills: int,
num_decodes: int,
window_size: int,
) -> None:
) -> LaunchSpec:
compile_key = self.dispatch(num_prefills=num_prefills)
self.kernel[(1,)](
prefill_gather_lens,
seq_lens,
query_start_loc,
num_prefills,
num_decodes,
window_size,
BLOCK_SIZE=compile_key.BLOCK_SIZE,
return (1,), dict(
BLOCK_SIZE=compile_key.block_size,
)
_COMPUTE_PREFILL_METADATA_KERNEL = ComputePrefillMetadataKernel()
class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
"""Builds metadata for DeepseekV4 SWA cache.
@@ -545,9 +537,8 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
device=self.device,
)
decode_swa_indices = self.decode_swa_indices_noncausal
_compute_dspark_noncausal_swa_indices_kernel[(num_decode_tokens,)](
_COMPUTE_DSPARK_NONCAUSAL_SWA_INDICES_KERNEL(
decode_swa_indices,
decode_swa_indices.stride(0),
self.decode_swa_lens,
self.window_size,
self.noncausal_index_width,
@@ -556,15 +547,13 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
token_to_req_indices,
is_valid_token,
block_table,
block_table.stride(0),
self.block_size,
num_tokens=num_decode_tokens,
token_offset=0,
TRITON_BLOCK_SIZE=1024,
)
else:
_compute_swa_indices_and_lens_kernel[(num_decode_tokens,)](
_COMPUTE_SWA_INDICES_AND_LENS_KERNEL(
decode_swa_indices,
decode_swa_indices.stride(0),
self.decode_swa_lens,
self.window_size,
query_start_loc,
@@ -572,10 +561,9 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
token_to_req_indices,
is_valid_token,
block_table,
block_table.stride(0),
self.block_size,
num_tokens=num_decode_tokens,
token_offset=0,
TRITON_BLOCK_SIZE=1024,
)
# Prefill SWA indices live in paged coordinates. `token_offset` lets
@@ -584,9 +572,8 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
if num_prefill_tokens > 0:
prefill_swa_indices = self.prefill_swa_indices[:num_prefill_tokens]
prefill_swa_lens = self.prefill_swa_lens[:num_prefill_tokens]
_compute_swa_indices_and_lens_kernel[(num_prefill_tokens,)](
_COMPUTE_SWA_INDICES_AND_LENS_KERNEL(
prefill_swa_indices,
prefill_swa_indices.stride(0),
prefill_swa_lens,
self.window_size,
query_start_loc,
@@ -594,10 +581,9 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
token_to_req_indices,
is_valid_token,
block_table,
block_table.stride(0),
self.block_size,
num_tokens=num_prefill_tokens,
token_offset=num_decode_tokens,
TRITON_BLOCK_SIZE=1024,
)
# Pre-compute DeepseekV4 prefill metadata shared across all attention layers.
@@ -668,9 +654,8 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
assert metadata.decode_swa_indices is not None
assert metadata.decode_swa_lens is not None
_compute_swa_indices_and_lens_kernel[(metadata.num_decode_tokens,)](
_COMPUTE_SWA_INDICES_AND_LENS_KERNEL(
metadata.decode_swa_indices,
metadata.decode_swa_indices.stride(0),
metadata.decode_swa_lens,
metadata.decode_swa_indices.shape[-1],
metadata.query_start_loc,
@@ -678,10 +663,9 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
metadata.token_to_req_indices,
metadata.is_valid_token,
metadata.block_table,
metadata.block_table.stride(0),
self.block_size,
num_tokens=metadata.num_decode_tokens,
token_offset=0,
TRITON_BLOCK_SIZE=1024,
)
tile_sched = self.build_tile_scheduler(metadata.num_decode_tokens)
metadata.tile_sched_swaonly = tile_sched[_LAYER_TYPE_SWAONLY]
@@ -771,142 +755,307 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
return result
@triton.jit(do_not_specialize=["token_offset"])
def _compute_swa_indices_and_lens_kernel(
swa_indices_ptr,
swa_indices_stride,
swa_lens_ptr,
window_size,
query_start_loc_ptr,
seq_lens_ptr,
token_to_req_indices_ptr,
is_valid_token_ptr,
block_table_ptr,
block_table_stride,
block_size,
token_offset,
TRITON_BLOCK_SIZE: tl.constexpr,
# TODO(ben): unify this kernel to reduce duplication
class ComputeSWAIndicesAndLensKernel(
VllmTritonJitKernel["ComputeSWAIndicesAndLensKernel.CompileKey"]
):
pid = tl.program_id(0)
token_idx = pid + token_offset
is_valid = tl.load(is_valid_token_ptr + token_idx)
if not is_valid:
tl.store(swa_lens_ptr + pid, 0)
# Clear the row so a padded token cannot gather through stale indices.
@dataclass(frozen=True)
class CompileKey:
window_size: int
block_size: int
triton_block_size: int
@staticmethod
@triton.jit(
do_not_specialize=[
"swa_indices_stride",
"block_table_stride",
"token_offset",
]
)
def kernel(
swa_indices_ptr,
swa_indices_stride,
swa_lens_ptr,
window_size,
query_start_loc_ptr,
seq_lens_ptr,
token_to_req_indices_ptr,
is_valid_token_ptr,
block_table_ptr,
block_table_stride,
block_size,
token_offset,
TRITON_BLOCK_SIZE: tl.constexpr,
):
pid = tl.program_id(0)
token_idx = pid + token_offset
is_valid = tl.load(is_valid_token_ptr + token_idx)
if not is_valid:
tl.store(swa_lens_ptr + pid, 0)
# Clear the row so a padded token cannot gather through stale indices.
for i in range(0, window_size, TRITON_BLOCK_SIZE):
offset = i + tl.arange(0, TRITON_BLOCK_SIZE)
tl.store(
swa_indices_ptr + pid * swa_indices_stride + offset,
-1,
mask=offset < window_size,
)
return
req_idx = tl.load(token_to_req_indices_ptr + token_idx)
query_start = tl.load(query_start_loc_ptr + req_idx)
query_end = tl.load(query_start_loc_ptr + req_idx + 1)
query_len = query_end - query_start
seq_len = tl.load(seq_lens_ptr + req_idx)
prefix_len = seq_len - query_len
pos = prefix_len + token_idx - query_start
start_pos = tl.maximum(pos - window_size + 1, 0)
end_pos = pos + 1
swa_len = end_pos - start_pos
tl.store(swa_lens_ptr + pid, swa_len)
for i in range(0, window_size, TRITON_BLOCK_SIZE):
offset = i + tl.arange(0, TRITON_BLOCK_SIZE)
pos_offset = start_pos + offset
block_indices = pos_offset // block_size
block_numbers = tl.load(
block_table_ptr + req_idx * block_table_stride + block_indices,
mask=pos_offset < end_pos,
)
block_offsets = pos_offset % block_size
slot_ids = block_numbers * block_size + block_offsets
slot_ids = tl.where(offset < swa_len, slot_ids, -1)
tl.store(
swa_indices_ptr + pid * swa_indices_stride + offset,
-1,
slot_ids,
mask=offset < window_size,
)
return
req_idx = tl.load(token_to_req_indices_ptr + token_idx)
query_start = tl.load(query_start_loc_ptr + req_idx)
query_end = tl.load(query_start_loc_ptr + req_idx + 1)
query_len = query_end - query_start
seq_len = tl.load(seq_lens_ptr + req_idx)
prefix_len = seq_len - query_len
pos = prefix_len + token_idx - query_start
start_pos = tl.maximum(pos - window_size + 1, 0)
end_pos = pos + 1
swa_len = end_pos - start_pos
tl.store(swa_lens_ptr + pid, swa_len)
for i in range(0, window_size, TRITON_BLOCK_SIZE):
offset = i + tl.arange(0, TRITON_BLOCK_SIZE)
pos_offset = start_pos + offset
block_indices = pos_offset // block_size
block_numbers = tl.load(
block_table_ptr + req_idx * block_table_stride + block_indices,
mask=pos_offset < end_pos,
def dispatch( # type: ignore[override]
self,
**compile_key_fields: int,
) -> CompileKey:
return self.CompileKey(
**compile_key_fields,
triton_block_size=1024,
)
block_offsets = pos_offset % block_size
slot_ids = block_numbers * block_size + block_offsets
slot_ids = tl.where(offset < swa_len, slot_ids, -1)
tl.store(
swa_indices_ptr + pid * swa_indices_stride + offset,
slot_ids,
mask=offset < window_size,
def get_warmup_keys(
self,
*,
window_size: int,
block_size: int,
) -> list[CompileKey]:
return self._trace_dispatch(self.dispatch)(
window_size=window_size,
block_size=block_size,
)
def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
int32_ptr = TritonWarmupTensor(torch.int32)
return dict(
swa_indices=TritonWarmupTensor(torch.int32, shape=(1, 1), strides=(1, 1)),
swa_lens=int32_ptr,
window_size=compile_key.window_size,
query_start_loc=int32_ptr,
seq_lens=int32_ptr,
token_to_req_indices=int32_ptr,
is_valid_token=TritonWarmupTensor(torch.bool),
block_table=TritonWarmupTensor(torch.int32, shape=(1, 1), strides=(1, 1)),
block_size=compile_key.block_size,
num_tokens=1,
token_offset=0,
)
@kernel_launcher
def __call__(
self,
swa_indices: torch.Tensor,
swa_lens: torch.Tensor,
window_size: int,
query_start_loc: torch.Tensor,
seq_lens: torch.Tensor,
token_to_req_indices: torch.Tensor,
is_valid_token: torch.Tensor,
block_table: torch.Tensor,
block_size: int,
*,
num_tokens: int,
token_offset: int,
) -> LaunchSpec:
return (num_tokens,), dict(
swa_indices_stride=swa_indices.stride(0),
block_table_stride=block_table.stride(0),
TRITON_BLOCK_SIZE=1024,
)
# TODO(ben): unify this kernel to reduce duplication
@triton.jit(do_not_specialize=["token_offset"])
def _compute_dspark_noncausal_swa_indices_kernel(
swa_indices_ptr,
swa_indices_stride,
swa_lens_ptr,
window_size,
index_width,
query_start_loc_ptr,
seq_lens_ptr,
token_to_req_indices_ptr,
is_valid_token_ptr,
block_table_ptr,
block_table_stride,
block_size,
token_offset,
TRITON_BLOCK_SIZE: tl.constexpr,
class ComputeDSparkNoncausalSWAIndicesKernel(
VllmTritonJitKernel["ComputeDSparkNoncausalSWAIndicesKernel.CompileKey"]
):
"""Non-causal per-token indices for the DSpark draft block.
@dataclass(frozen=True)
class CompileKey:
window_size: int
index_width: int
block_size: int
triton_block_size: int
@staticmethod
@triton.jit(
do_not_specialize=[
"swa_indices_stride",
"block_table_stride",
"token_offset",
]
)
def kernel(
swa_indices_ptr,
swa_indices_stride,
swa_lens_ptr,
window_size,
index_width,
query_start_loc_ptr,
seq_lens_ptr,
token_to_req_indices_ptr,
is_valid_token_ptr,
block_table_ptr,
block_table_stride,
block_size,
token_offset,
TRITON_BLOCK_SIZE: tl.constexpr,
):
"""Non-causal per-token indices for the DSpark draft block.
Here, we populate the topk indices with the trailing window of context tokens,
plus all query tokens (including future ones).
"""
pid = tl.program_id(0)
token_idx = pid + token_offset
is_valid = tl.load(is_valid_token_ptr + token_idx)
if not is_valid:
tl.store(swa_lens_ptr + pid, 0)
# Clear the row so a padded token cannot gather through stale indices.
for i in range(0, index_width, TRITON_BLOCK_SIZE):
offset = i + tl.arange(0, TRITON_BLOCK_SIZE)
tl.store(
swa_indices_ptr + pid * swa_indices_stride + offset,
-1,
mask=offset < index_width,
)
return
req_idx = tl.load(token_to_req_indices_ptr + token_idx)
query_start = tl.load(query_start_loc_ptr + req_idx)
query_end = tl.load(query_start_loc_ptr + req_idx + 1)
query_len = query_end - query_start
seq_len = tl.load(seq_lens_ptr + req_idx)
prefix_len = seq_len - query_len
# Block-anchored window (shared by every token in the block) + full block.
start_pos = tl.maximum(prefix_len - window_size, 0)
end_pos = seq_len
swa_len = end_pos - start_pos
tl.store(swa_lens_ptr + pid, swa_len)
Here, we populate the topk indices with the trailing window of context tokens,
plus all query tokens (including future ones).
"""
pid = tl.program_id(0)
token_idx = pid + token_offset
is_valid = tl.load(is_valid_token_ptr + token_idx)
if not is_valid:
tl.store(swa_lens_ptr + pid, 0)
# Clear the row so a padded token cannot gather through stale indices.
for i in range(0, index_width, TRITON_BLOCK_SIZE):
offset = i + tl.arange(0, TRITON_BLOCK_SIZE)
pos_offset = start_pos + offset
block_indices = pos_offset // block_size
block_numbers = tl.load(
block_table_ptr + req_idx * block_table_stride + block_indices,
mask=pos_offset < end_pos,
)
block_offsets = pos_offset % block_size
slot_ids = block_numbers * block_size + block_offsets
slot_ids = tl.where(offset < swa_len, slot_ids, -1)
tl.store(
swa_indices_ptr + pid * swa_indices_stride + offset,
-1,
slot_ids,
mask=offset < index_width,
)
return
req_idx = tl.load(token_to_req_indices_ptr + token_idx)
query_start = tl.load(query_start_loc_ptr + req_idx)
query_end = tl.load(query_start_loc_ptr + req_idx + 1)
query_len = query_end - query_start
seq_len = tl.load(seq_lens_ptr + req_idx)
prefix_len = seq_len - query_len
# Block-anchored window (shared by every token in the block) + full block.
start_pos = tl.maximum(prefix_len - window_size, 0)
end_pos = seq_len
swa_len = end_pos - start_pos
tl.store(swa_lens_ptr + pid, swa_len)
for i in range(0, index_width, TRITON_BLOCK_SIZE):
offset = i + tl.arange(0, TRITON_BLOCK_SIZE)
pos_offset = start_pos + offset
block_indices = pos_offset // block_size
block_numbers = tl.load(
block_table_ptr + req_idx * block_table_stride + block_indices,
mask=pos_offset < end_pos,
def dispatch( # type: ignore[override]
self,
*,
window_size: int,
num_speculative_tokens: int,
block_size: int,
) -> CompileKey:
return self.CompileKey(
window_size=window_size,
index_width=get_dspark_swa_index_width(window_size, num_speculative_tokens),
block_size=block_size,
triton_block_size=1024,
)
block_offsets = pos_offset % block_size
slot_ids = block_numbers * block_size + block_offsets
slot_ids = tl.where(offset < swa_len, slot_ids, -1)
tl.store(
swa_indices_ptr + pid * swa_indices_stride + offset,
slot_ids,
mask=offset < index_width,
def get_warmup_keys(
self,
*,
window_size: int,
num_speculative_tokens: int,
block_size: int,
) -> list[CompileKey]:
return self._trace_dispatch(self.dispatch)(
window_size=window_size,
num_speculative_tokens=num_speculative_tokens,
block_size=block_size,
)
def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
int32_ptr = TritonWarmupTensor(torch.int32)
return dict(
swa_indices=TritonWarmupTensor(torch.int32, shape=(1, 1), strides=(1, 1)),
swa_lens=int32_ptr,
window_size=compile_key.window_size,
index_width=compile_key.index_width,
query_start_loc=int32_ptr,
seq_lens=int32_ptr,
token_to_req_indices=int32_ptr,
is_valid_token=TritonWarmupTensor(torch.bool),
block_table=TritonWarmupTensor(torch.int32, shape=(1, 1), strides=(1, 1)),
block_size=compile_key.block_size,
num_tokens=1,
token_offset=0,
)
@kernel_launcher
def __call__(
self,
swa_indices: torch.Tensor,
swa_lens: torch.Tensor,
window_size: int,
index_width: int,
query_start_loc: torch.Tensor,
seq_lens: torch.Tensor,
token_to_req_indices: torch.Tensor,
is_valid_token: torch.Tensor,
block_table: torch.Tensor,
block_size: int,
*,
num_tokens: int,
token_offset: int,
) -> LaunchSpec:
return (num_tokens,), dict(
swa_indices_stride=swa_indices.stride(0),
block_table_stride=block_table.stride(0),
TRITON_BLOCK_SIZE=1024,
)
_COMPUTE_PREFILL_METADATA_KERNEL = ComputePrefillMetadataKernel()
_COMPUTE_SWA_INDICES_AND_LENS_KERNEL = ComputeSWAIndicesAndLensKernel()
_COMPUTE_DSPARK_NONCAUSAL_SWA_INDICES_KERNEL = ComputeDSparkNoncausalSWAIndicesKernel()
vllm/v1/attention/backends/mla/sparse_utils.py
+360 -163
View File
@@ -2,9 +2,26 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Utility functions for sparse MLA backends."""
from dataclasses import dataclass
from typing import Any
import torch
from vllm.config import VllmConfig
from vllm.distributed import get_dcp_group
from vllm.model_executor.warmup.jit_warmup import (
zip_inputs,
)
from vllm.model_executor.warmup.jit_warmup_triton_helper import (
LaunchSpec,
TritonWarmupTensor,
VllmTritonJitKernel,
kernel_launcher,
triton_scalar_specialization_rep,
)
from vllm.triton_utils import tl, triton
from vllm.utils.math_utils import cdiv
from vllm.v1.worker.block_table import get_block_table_width
def flat_kv_row_view(
@@ -30,126 +47,330 @@ def flat_kv_row_view(
# Kernel with prefill workspace support and valid count tracking
@triton.jit
def _convert_req_index_to_global_index_kernel(
req_id_ptr, # int32 [num_tokens]
block_table_ptr, # int32 [num_requests, max_num_blocks_per_req]
token_indices_ptr, # int32 [num_tokens, NUM_TOPK_TOKENS]
out_ptr, # int32 [num_tokens, NUM_TOPK_TOKENS]
valid_count_ptr, # int32 [num_tokens] - output valid count per row
prefill_request_id_ptr, # int32 [num_tokens], -1 for decode, >=0 for prefill
workspace_starts_ptr, # int32 [num_prefill_reqs+1] or nullptr
# shapes (compile-time where possible)
max_num_blocks_per_req: tl.constexpr,
BLOCK_SIZE: tl.constexpr,
BLOCK_STRIDE_ROWS: tl.constexpr,
BLOCK_N: tl.constexpr, # tile width along columns
HAS_PREFILL: tl.constexpr,
COUNT_VALID: tl.constexpr, # whether to count valid indices
# BLOCK_N == NUM_TOPK_TOKENS: one program owns the row, so the valid count
# is an in-register reduction and needs no atomic.
SINGLE_TILE: tl.constexpr,
# When set, scatter valid slots to a contiguous prefix [0, valid_count) using
# valid_count_ptr as an atomic slot allocator (DCP filtering leaves interior
# -1 gaps; the trtllm-gen sparse kernel reads the first valid_count entries).
# Requires COUNT_VALID and an out buffer pre-filled with -1. Order within the
# prefix is unspecified (only the selected set matters).
COMPACT_TO_FRONT: tl.constexpr,
# DCP de-interleave: with DCP_SIZE == 1 these are an exact no-op
DCP_SIZE: tl.constexpr,
DCP_RANK: tl.constexpr,
DCP_INTERLEAVE: tl.constexpr,
# strides (in elements)
bt_stride0,
bt_stride1,
ti_stride0,
ti_stride1,
out_stride0,
out_stride1,
class ConvertReqIndexToGlobalIndexKernel(
VllmTritonJitKernel["ConvertReqIndexToGlobalIndexKernel.CompileKey"]
):
# program_id(0) -> token_id (row)
# program_id(1) -> tile index along columns
token_id = tl.program_id(0)
tile_id = tl.program_id(1)
@dataclass(frozen=True)
class CompileKey:
block_size: int
block_stride_rows: int
block_n: int
has_prefill_workspace: bool
count_valid: bool
single_tile: bool
compact_to_front: bool
dcp_size: int
dcp_rank: int
dcp_interleave: int
num_warps: int
num_topk_tokens: int
block_table_stride: int
# Each program covers BLOCK_N consecutive columns
indice_id = tile_id * BLOCK_N + tl.arange(0, BLOCK_N)
@staticmethod
@triton.jit(do_not_specialize=["max_num_blocks_per_req"])
def kernel(
req_id_ptr, # int32 [num_tokens]
block_table_ptr, # int32 [num_requests, max_num_blocks_per_req]
token_indices_ptr, # int32 [num_tokens, NUM_TOPK_TOKENS]
out_ptr, # int32 [num_tokens, NUM_TOPK_TOKENS]
valid_count_ptr, # int32 [num_tokens] - output valid count per row
prefill_request_id_ptr, # int32 [num_tokens], -1 for decode, >=0 for prefill
workspace_starts_ptr, # int32 [num_prefill_reqs+1] or nullptr
# shapes (compile-time where possible)
max_num_blocks_per_req,
BLOCK_SIZE: tl.constexpr,
BLOCK_STRIDE_ROWS: tl.constexpr,
BLOCK_N: tl.constexpr, # tile width along columns
HAS_PREFILL: tl.constexpr,
COUNT_VALID: tl.constexpr, # whether to count valid indices
# BLOCK_N == NUM_TOPK_TOKENS: one program owns the row, so the valid count
# is an in-register reduction and needs no atomic.
SINGLE_TILE: tl.constexpr,
# When set, scatter valid slots to a contiguous prefix [0, valid_count) using
# valid_count_ptr as an atomic slot allocator (DCP filtering leaves interior
# -1 gaps; the trtllm-gen sparse kernel reads the first valid_count entries).
# Requires COUNT_VALID and an out buffer pre-filled with -1. Order within the
# prefix is unspecified (only the selected set matters).
COMPACT_TO_FRONT: tl.constexpr,
# DCP de-interleave: with DCP_SIZE == 1 these are an exact no-op
DCP_SIZE: tl.constexpr,
DCP_RANK: tl.constexpr,
DCP_INTERLEAVE: tl.constexpr,
# strides (in elements)
bt_stride0,
bt_stride1,
ti_stride0,
ti_stride1,
out_stride0,
out_stride1,
):
# program_id(0) -> token_id (row)
# program_id(1) -> tile index along columns
token_id = tl.program_id(0)
tile_id = tl.program_id(1)
# Load request id for this token (no mask: grid is exact)
req = tl.load(req_id_ptr + token_id)
# Each program covers BLOCK_N consecutive columns
indice_id = tile_id * BLOCK_N + tl.arange(0, BLOCK_N)
# Load token indices for this tile
ti_ptr = token_indices_ptr + token_id * ti_stride0 + indice_id * ti_stride1
tok = tl.load(ti_ptr) # int32
# Load request id for this token (no mask: grid is exact)
req = tl.load(req_id_ptr + token_id)
# Only token == -1 should propagate as -1
is_invalid_tok = tok < 0
is_prefill = False
if HAS_PREFILL:
prefill_req_id = tl.load(prefill_request_id_ptr + token_id)
is_prefill = prefill_req_id >= 0
# Load token indices for this tile
ti_ptr = token_indices_ptr + token_id * ti_stride0 + indice_id * ti_stride1
tok = tl.load(ti_ptr) # int32
# DCP de-interleave the global token id into this rank's local slot.
# Tokens are interleaved in groups of DCP_INTERLEAVE across ranks. With
# DCP_SIZE == 1 (and any interleave) owning_rank == 0 == DCP_RANK (never
# remote) and local_idx == tok, so this reduces to the non-DCP path; with
# DCP_INTERLEAVE == 1 it reduces to plain round-robin (tok % / // DCP_SIZE).
owning_rank = (tok // DCP_INTERLEAVE) % DCP_SIZE
is_remote = owning_rank != DCP_RANK
local_idx = (
tok // (DCP_SIZE * DCP_INTERLEAVE)
) * DCP_INTERLEAVE + tok % DCP_INTERLEAVE
# Only token == -1 should propagate as -1
is_invalid_tok = tok < 0
is_prefill = False
if HAS_PREFILL:
prefill_req_id = tl.load(prefill_request_id_ptr + token_id)
is_prefill = prefill_req_id >= 0
# Compute block id and in-block offset
block_id = local_idx // BLOCK_SIZE
inblock_off = local_idx % BLOCK_SIZE
# DCP de-interleave the global token id into this rank's local slot.
# Tokens are interleaved in groups of DCP_INTERLEAVE across ranks. With
# DCP_SIZE == 1 (and any interleave) owning_rank == 0 == DCP_RANK (never
# remote) and local_idx == tok, so this reduces to the non-DCP path; with
# DCP_INTERLEAVE == 1 it reduces to plain round-robin (tok % / // DCP_SIZE).
owning_rank = (tok // DCP_INTERLEAVE) % DCP_SIZE
is_remote = owning_rank != DCP_RANK
local_idx = (
tok // (DCP_SIZE * DCP_INTERLEAVE)
) * DCP_INTERLEAVE + tok % DCP_INTERLEAVE
# Guard block_table access
valid_block = (block_id < max_num_blocks_per_req) & (block_id >= 0)
bt_ptr = block_table_ptr + req * bt_stride0 + block_id * bt_stride1
is_invalid_tok |= ~valid_block | is_remote
base = tl.load(bt_ptr, mask=valid_block & ~is_prefill & ~is_remote, other=0)
out_val = base * BLOCK_STRIDE_ROWS + inblock_off
# Compute block id and in-block offset
block_id = local_idx // BLOCK_SIZE
inblock_off = local_idx % BLOCK_SIZE
# Override with prefill output if prefill is enabled
if HAS_PREFILL:
workspace_start = tl.load(
workspace_starts_ptr + prefill_req_id, mask=is_prefill, other=0
)
prefill_out = workspace_start + tok
out_val = tl.where(is_prefill, prefill_out, out_val)
out_val = tl.where(is_invalid_tok, -1, out_val)
# Guard block_table access
valid_block = (block_id < max_num_blocks_per_req) & (block_id >= 0)
bt_ptr = block_table_ptr + req * bt_stride0 + block_id * bt_stride1
is_invalid_tok |= ~valid_block | is_remote
base = tl.load(bt_ptr, mask=valid_block & ~is_prefill & ~is_remote, other=0)
out_val = base * BLOCK_STRIDE_ROWS + inblock_off
if COMPACT_TO_FRONT:
# Scatter valid slots to a contiguous prefix. A per-tile exclusive prefix
# sum gives each valid lane a distinct local offset; one atomic add of the
# tile's valid count reserves a contiguous base across racing tiles. The
# out buffer is pre-filled with -1, so unwritten tail slots stay -1.
# With no racing tiles the base is 0 and the allocator becomes a store.
is_valid = (~is_invalid_tok).to(tl.int32)
local_offset = tl.cumsum(is_valid) - is_valid
tile_valid_count = tl.sum(is_valid)
if SINGLE_TILE:
base = 0
tl.store(valid_count_ptr + token_id, tile_valid_count)
else:
base = tl.atomic_add(valid_count_ptr + token_id, tile_valid_count)
dest = base + local_offset
out_ptr_dest = out_ptr + token_id * out_stride0 + dest * out_stride1
tl.store(out_ptr_dest, out_val, mask=is_valid == 1)
else:
# Store results in place (input column == output column).
out_ptr_ij = out_ptr + token_id * out_stride0 + indice_id * out_stride1
tl.store(out_ptr_ij, out_val)
# Override with prefill output if prefill is enabled
if HAS_PREFILL:
workspace_start = tl.load(
workspace_starts_ptr + prefill_req_id, mask=is_prefill, other=0
)
prefill_out = workspace_start + tok
out_val = tl.where(is_prefill, prefill_out, out_val)
out_val = tl.where(is_invalid_tok, -1, out_val)
# Accumulate the tile's valid count into the row total; a single tile's
# reduction *is* the total.
if COUNT_VALID:
tile_valid_count = tl.sum((~is_invalid_tok).to(tl.int32))
if COMPACT_TO_FRONT:
# Scatter valid slots to a contiguous prefix. A per-tile exclusive prefix
# sum gives each valid lane a distinct local offset; one atomic add of the
# tile's valid count reserves a contiguous base across racing tiles. The
# out buffer is pre-filled with -1, so unwritten tail slots stay -1.
# With no racing tiles the base is 0 and the allocator becomes a store.
is_valid = (~is_invalid_tok).to(tl.int32)
local_offset = tl.cumsum(is_valid) - is_valid
tile_valid_count = tl.sum(is_valid)
if SINGLE_TILE:
base = 0
tl.store(valid_count_ptr + token_id, tile_valid_count)
else:
tl.atomic_add(valid_count_ptr + token_id, tile_valid_count)
base = tl.atomic_add(valid_count_ptr + token_id, tile_valid_count)
dest = base + local_offset
out_ptr_dest = out_ptr + token_id * out_stride0 + dest * out_stride1
tl.store(out_ptr_dest, out_val, mask=is_valid == 1)
else:
# Store results in place (input column == output column).
out_ptr_ij = out_ptr + token_id * out_stride0 + indice_id * out_stride1
tl.store(out_ptr_ij, out_val)
# Accumulate the tile's valid count into the row total; a single tile's
# reduction *is* the total.
if COUNT_VALID:
tile_valid_count = tl.sum((~is_invalid_tok).to(tl.int32))
if SINGLE_TILE:
tl.store(valid_count_ptr + token_id, tile_valid_count)
else:
tl.atomic_add(valid_count_ptr + token_id, tile_valid_count)
def dispatch( # type: ignore[override]
self,
*,
BLOCK_SIZE: int,
BLOCK_STRIDE_ROWS: int,
BLOCK_N: int,
NUM_TOPK_TOKENS: int,
HAS_PREFILL_WORKSPACE: bool,
COUNT_VALID: bool,
COMPACT_TO_FRONT: bool,
DCP_SIZE: int,
DCP_RANK: int,
DCP_INTERLEAVE: int,
max_num_blocks_per_req: int,
) -> CompileKey:
tiling = _remap_tiling(NUM_TOPK_TOKENS, BLOCK_N, COUNT_VALID)
return self.CompileKey(
block_size=BLOCK_SIZE,
block_stride_rows=BLOCK_STRIDE_ROWS,
block_n=tiling[1],
has_prefill_workspace=HAS_PREFILL_WORKSPACE,
count_valid=COUNT_VALID,
single_tile=tiling[0],
compact_to_front=COMPACT_TO_FRONT,
dcp_size=DCP_SIZE,
dcp_rank=DCP_RANK,
dcp_interleave=DCP_INTERLEAVE,
num_warps=tiling[3],
num_topk_tokens=NUM_TOPK_TOKENS,
block_table_stride=triton_scalar_specialization_rep(max_num_blocks_per_req),
)
def get_warmup_keys(
self,
vllm_config: VllmConfig,
*,
block_stride_rows: int,
) -> list[CompileKey]:
block_size = vllm_config.cache_config.block_size
dcp_size = vllm_config.parallel_config.decode_context_parallel_size
dcp_interleave = vllm_config.parallel_config.cp_kv_cache_interleave_size
dcp_rank = get_dcp_group().rank_in_group if dcp_size > 1 else 0
num_topk_tokens = vllm_config.model_config.hf_config.index_topk
max_num_blocks = cdiv(
vllm_config.model_config.max_model_len,
block_size * dcp_size,
)
max_num_blocks_per_req = get_block_table_width(max_num_blocks, block_size)
return self._trace_dispatch(self.dispatch)(
zip_inputs(
dict(
HAS_PREFILL_WORKSPACE=False,
COUNT_VALID=False,
COMPACT_TO_FRONT=False,
DCP_SIZE=1,
DCP_RANK=0,
DCP_INTERLEAVE=1,
),
dict(
HAS_PREFILL_WORKSPACE=False,
COUNT_VALID=True,
COMPACT_TO_FRONT=False,
DCP_SIZE=1,
DCP_RANK=0,
DCP_INTERLEAVE=1,
),
dict(
HAS_PREFILL_WORKSPACE=True,
COUNT_VALID=True,
COMPACT_TO_FRONT=False,
DCP_SIZE=1,
DCP_RANK=0,
DCP_INTERLEAVE=1,
),
dict(
HAS_PREFILL_WORKSPACE=False,
COUNT_VALID=True,
COMPACT_TO_FRONT=False,
DCP_SIZE=dcp_size,
DCP_RANK=dcp_rank,
DCP_INTERLEAVE=dcp_interleave,
),
dict(
HAS_PREFILL_WORKSPACE=False,
COUNT_VALID=True,
COMPACT_TO_FRONT=dcp_size > 1,
DCP_SIZE=dcp_size,
DCP_RANK=dcp_rank,
DCP_INTERLEAVE=dcp_interleave,
),
),
BLOCK_SIZE=block_size,
BLOCK_STRIDE_ROWS=block_stride_rows,
BLOCK_N=128,
NUM_TOPK_TOKENS=num_topk_tokens,
max_num_blocks_per_req=max_num_blocks_per_req,
)
def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
int32_ptr = TritonWarmupTensor(torch.int32)
token_shape = (1, compile_key.num_topk_tokens)
token_strides = (16, 1)
return dict(
req_id=int32_ptr,
block_table=TritonWarmupTensor(
torch.int32,
shape=(1, compile_key.block_table_stride),
),
token_indices=TritonWarmupTensor(
torch.int32,
shape=token_shape,
strides=token_strides,
),
out=TritonWarmupTensor(
torch.int32,
shape=token_shape,
strides=token_strides,
),
valid_counts=int32_ptr if compile_key.count_valid else None,
prefill_workspace_request_ids=(
int32_ptr if compile_key.has_prefill_workspace else None
),
prefill_workspace_starts=(
int32_ptr if compile_key.has_prefill_workspace else None
),
max_num_blocks_per_req=1,
BLOCK_SIZE=compile_key.block_size,
BLOCK_STRIDE_ROWS=compile_key.block_stride_rows,
BLOCK_N=compile_key.block_n,
NUM_TOPK_TOKENS=compile_key.num_topk_tokens,
HAS_PREFILL_WORKSPACE=compile_key.has_prefill_workspace,
COUNT_VALID=compile_key.count_valid,
COMPACT_TO_FRONT=compile_key.compact_to_front,
DCP_SIZE=compile_key.dcp_size,
DCP_RANK=compile_key.dcp_rank,
DCP_INTERLEAVE=compile_key.dcp_interleave,
)
@kernel_launcher
def __call__(
self,
req_id: torch.Tensor,
block_table: torch.Tensor,
token_indices: torch.Tensor,
out: torch.Tensor,
valid_counts: torch.Tensor | None,
prefill_workspace_request_ids: torch.Tensor | None,
prefill_workspace_starts: torch.Tensor | None,
*,
max_num_blocks_per_req: int,
BLOCK_SIZE: int,
BLOCK_STRIDE_ROWS: int,
BLOCK_N: int,
NUM_TOPK_TOKENS: int,
HAS_PREFILL_WORKSPACE: bool,
COUNT_VALID: bool,
COMPACT_TO_FRONT: bool,
DCP_SIZE: int,
DCP_RANK: int,
DCP_INTERLEAVE: int,
) -> LaunchSpec:
single_tile, block_n, tiles_per_row, num_warps = _remap_tiling(
NUM_TOPK_TOKENS, BLOCK_N, COUNT_VALID
)
# Exact 2D grid: tokens × column tiles
return (req_id.shape[0], tiles_per_row), dict(
valid_count_ptr=valid_counts,
prefill_request_id_ptr=prefill_workspace_request_ids,
workspace_starts_ptr=prefill_workspace_starts,
BLOCK_N=block_n,
HAS_PREFILL=HAS_PREFILL_WORKSPACE,
SINGLE_TILE=single_tile,
# Strides in elements
bt_stride0=block_table.stride(0),
bt_stride1=block_table.stride(1),
ti_stride0=token_indices.stride(0),
ti_stride1=token_indices.stride(1),
out_stride0=out.stride(0),
out_stride1=out.stride(1),
num_warps=num_warps,
)
def _remap_tiling(
@@ -231,9 +452,7 @@ def triton_convert_req_index_to_global_index(
num_tokens = req_id.shape[0]
max_num_blocks_per_req = block_table.shape[1]
single_tile, block_n, tiles_per_row, num_warps = _remap_tiling(
NUM_TOPK_TOKENS, BLOCK_N, return_valid_counts
)
single_tile, _, _, _ = _remap_tiling(NUM_TOPK_TOKENS, BLOCK_N, return_valid_counts)
# Ensure contiguous tensors on the same device
req_id_c = req_id.contiguous()
@@ -247,11 +466,6 @@ def triton_convert_req_index_to_global_index(
alloc = torch.empty if single_tile else torch.zeros
valid_counts = alloc(num_tokens, dtype=torch.int32, device=token_indices.device)
# Strides in elements
bt_stride0, bt_stride1 = block_table_c.stride()
ti_stride0, ti_stride1 = token_indices_c.stride()
out_stride0, out_stride1 = out.stride()
# Prepare prefill pointers
if HAS_PREFILL_WORKSPACE:
assert prefill_workspace_request_ids is not None # for mypy
@@ -259,10 +473,7 @@ def triton_convert_req_index_to_global_index(
assert prefill_workspace_request_ids.is_contiguous()
assert prefill_workspace_starts.is_contiguous()
# Exact 2D grid: tokens × column tiles
grid = (num_tokens, tiles_per_row)
_convert_req_index_to_global_index_kernel[grid](
_CONVERT_REQ_INDEX_TO_GLOBAL_INDEX_KERNEL(
req_id_c,
block_table_c,
token_indices_c,
@@ -271,26 +482,20 @@ def triton_convert_req_index_to_global_index(
prefill_workspace_request_ids,
prefill_workspace_starts,
# shapes / constexprs
max_num_blocks_per_req,
BLOCK_SIZE,
BLOCK_STRIDE_ROWS if BLOCK_STRIDE_ROWS is not None else BLOCK_SIZE,
block_n,
HAS_PREFILL_WORKSPACE,
return_valid_counts,
single_tile,
False, # COMPACT_TO_FRONT: keep input column == output column
max_num_blocks_per_req=max_num_blocks_per_req,
BLOCK_SIZE=BLOCK_SIZE,
BLOCK_STRIDE_ROWS=(
BLOCK_STRIDE_ROWS if BLOCK_STRIDE_ROWS is not None else BLOCK_SIZE
),
BLOCK_N=BLOCK_N,
NUM_TOPK_TOKENS=NUM_TOPK_TOKENS,
HAS_PREFILL_WORKSPACE=HAS_PREFILL_WORKSPACE,
COUNT_VALID=return_valid_counts,
COMPACT_TO_FRONT=False,
# DCP disabled (no-op de-interleave)
1,
0,
1,
# strides
bt_stride0,
bt_stride1,
ti_stride0,
ti_stride1,
out_stride0,
out_stride1,
num_warps=num_warps,
DCP_SIZE=1,
DCP_RANK=0,
DCP_INTERLEAVE=1,
)
if return_valid_counts:
@@ -362,9 +567,7 @@ def triton_filter_and_convert_dcp_index(
count_valid = return_valid_counts or compact_valid_to_front
# The compaction builds on the counting, so it shares the tiling.
single_tile, block_n, tiles_per_row, num_warps = _remap_tiling(
NUM_TOPK_TOKENS, BLOCK_N, count_valid
)
single_tile, _, _, _ = _remap_tiling(NUM_TOPK_TOKENS, BLOCK_N, count_valid)
if compact_valid_to_front:
out = torch.full_like(token_indices_c, -1)
@@ -377,11 +580,7 @@ def triton_filter_and_convert_dcp_index(
alloc = torch.empty if single_tile else torch.zeros
valid_counts = alloc(num_tokens, dtype=torch.int32, device=token_indices.device)
bt_stride0, bt_stride1 = block_table_c.stride()
ti_stride0, ti_stride1 = token_indices_c.stride()
out_stride0, out_stride1 = out.stride()
_convert_req_index_to_global_index_kernel[(num_tokens, tiles_per_row)](
_CONVERT_REQ_INDEX_TO_GLOBAL_INDEX_KERNEL(
req_id_c,
block_table_c,
token_indices_c,
@@ -390,27 +589,25 @@ def triton_filter_and_convert_dcp_index(
# No prefill workspace on the DCP decode path.
None,
None,
max_num_blocks_per_req,
BLOCK_SIZE,
BLOCK_STRIDE_ROWS if BLOCK_STRIDE_ROWS is not None else BLOCK_SIZE,
block_n,
False, # HAS_PREFILL
count_valid,
single_tile,
compact_valid_to_front,
dcp_size,
dcp_rank,
cp_kv_cache_interleave_size,
bt_stride0,
bt_stride1,
ti_stride0,
ti_stride1,
out_stride0,
out_stride1,
num_warps=num_warps,
max_num_blocks_per_req=max_num_blocks_per_req,
BLOCK_SIZE=BLOCK_SIZE,
BLOCK_STRIDE_ROWS=(
BLOCK_STRIDE_ROWS if BLOCK_STRIDE_ROWS is not None else BLOCK_SIZE
),
BLOCK_N=BLOCK_N,
NUM_TOPK_TOKENS=NUM_TOPK_TOKENS,
HAS_PREFILL_WORKSPACE=False,
COUNT_VALID=count_valid,
COMPACT_TO_FRONT=compact_valid_to_front,
DCP_SIZE=dcp_size,
DCP_RANK=dcp_rank,
DCP_INTERLEAVE=cp_kv_cache_interleave_size,
)
if return_valid_counts:
assert valid_counts is not None
return out, valid_counts
return out
_CONVERT_REQ_INDEX_TO_GLOBAL_INDEX_KERNEL = ConvertReqIndexToGlobalIndexKernel()
vllm/v1/worker/gpu/model_runner.py
+17 -12
View File
@@ -382,9 +382,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):
model_loader = get_model_loader(self.vllm_config.load_config)
logger.info_once("Loading model from scratch...")
self.model = model_loader.load_model(
vllm_config=self.vllm_config, model_config=self.vllm_config.model_config
)
# Capture warmup providers selected while constructing the model.
with self.jit_warmup_registry.activate():
self.model = model_loader.load_model(
vllm_config=self.vllm_config,
model_config=self.vllm_config.model_config,
)
if self.lora_config:
self.model = self.load_lora_model(
self.model, self.vllm_config, self.device
@@ -668,15 +671,17 @@ class GPUModelRunner(LoRAModelRunnerMixin):
self.speculator.init_cudagraph_manager(cudagraph_mode)
self.kv_caches: list[torch.Tensor] = []
kv_caches_dict = init_kv_cache(
self.kv_caches,
self.compilation_config.static_forward_context,
self.kv_cache_config,
self.device,
self.kernel_block_sizes,
self.vllm_config,
kv_cache_allocation_context=kv_cache_allocation_context,
)
# Capture warmup providers that depend on allocated KV-cache strides.
with self.jit_warmup_registry.activate():
kv_caches_dict = init_kv_cache(
self.kv_caches,
self.compilation_config.static_forward_context,
self.kv_cache_config,
self.device,
self.kernel_block_sizes,
self.vllm_config,
kv_cache_allocation_context=kv_cache_allocation_context,
)
if is_profiling:
self.kv_connector = NO_OP_KV_CONNECTOR
else:
vllm/v1/worker/gpu_model_runner.py
+12 -8
View File
@@ -5393,9 +5393,11 @@ class GPUModelRunner(
if load_dummy_weights:
self.load_config.load_format = "dummy"
model_loader = get_model_loader(self.load_config)
self.model = model_loader.load_model(
vllm_config=self.vllm_config, model_config=self.model_config
)
# Capture warmup providers selected while constructing the model.
with self.jit_warmup_registry.activate():
self.model = model_loader.load_model(
vllm_config=self.vllm_config, model_config=self.model_config
)
if self.lora_config:
self.model = self.load_lora_model(
self.model, self.vllm_config, self.device
@@ -7500,11 +7502,13 @@ class GPUModelRunner(
# Reinitialize need to after initialize_attn_backend
self.may_reinitialize_input_batch(kv_cache_config, kernel_block_sizes)
kv_caches = self.initialize_kv_cache_tensors(
kv_cache_config,
kernel_block_sizes,
kv_cache_allocation_context=kv_cache_allocation_context,
)
# Capture warmup providers that depend on allocated KV-cache strides.
with self.jit_warmup_registry.activate():
kv_caches = self.initialize_kv_cache_tensors(
kv_cache_config,
kernel_block_sizes,
kv_cache_allocation_context=kv_cache_allocation_context,
)
if (
self.speculative_config