[4/N][warmup][DSv4] Migrate common attention kernels (#50176)
Signed-off-by: LopezCastroRoberto <rocastro@redhat.com> Co-authored-by: Codex <codex@openai.com>
Roberto L. Castro committed
Sep 8, 2026 at 16:25 UTC
6ddbab03defe20d536a8c9e6cb9d74d67aa62416
19 files changed
+2978
-1272
tests/kernels/core/test_fused_q_kv_rmsnorm.py
+3
-1
@@ -13,7 +13,9 @@ from __future__ import annotations
13
import pytest
14
import torch
15
16
-from vllm.models.common.ops import fused_q_kv_rmsnorm
16
+from vllm.models.common.ops.fused_qk_rmsnorm import (
17
+ fused_q_kv_rmsnorm,
18
+)
19
from vllm.platforms import current_platform
20
21
pytestmark = pytest.mark.skipif(
tests/model_executor/test_jit_warmup.py
+34
@@ -431,6 +431,40 @@ def test_dispatch_can_forward_compile_key_fields() -> None:
431
kernel.compile_key({"tokens": 5, "mode": 2, "extra": 1})
432
433
434
+def test_forwarded_compile_key_fields_exclude_predicate_only_inputs() -> None:
435
+ class ForwardingKernel(VllmJitKernel["ForwardingKernel.CompileKey"]):
436
+ @dataclass(frozen=True)
437
+ class CompileKey:
438
+ mode: int
439
+ block_size: int
440
+
441
+ def dispatch( # type: ignore[override]
442
+ self,
443
+ *,
444
+ tokens: int,
445
+ **compile_key_fields: int,
446
+ ) -> CompileKey:
447
+ return self.CompileKey(
448
+ **compile_key_fields,
449
+ block_size=_round_up(tokens, multiple=4),
450
+ )
451
+
452
+ def get_warmup_keys(self) -> list[CompileKey]:
453
+ return self._trace_dispatch(self.dispatch)(
454
+ tokens=(1, 5),
455
+ mode=2,
456
+ max_tokens=4,
457
+ _when=lambda *, tokens, max_tokens: tokens <= max_tokens,
458
+ )
459
+
460
+ def compile(self, compile_key: CompileKey) -> None:
461
+ pass
462
+
463
+ assert ForwardingKernel().get_warmup_keys() == [
464
+ ForwardingKernel.CompileKey(mode=2, block_size=4)
465
+ ]
466
+
467
+
468
def test_dispatch_supports_tuple_and_mapping_subscriptions() -> None:
469
class SubscriptKernel(VllmJitKernel["SubscriptKernel.CompileKey"]):
470
@dataclass(frozen=True)
vllm/model_executor/warmup/jit_warmup.py
+18
-3
@@ -14,8 +14,8 @@ from abc import ABC, abstractmethod
14
from collections.abc import Callable, Iterable, Iterator, Mapping
15
from contextlib import contextmanager
16
from contextvars import ContextVar
17
-from dataclasses import dataclass
18
-from typing import Any, Generic, TypeVar
17
+from dataclasses import dataclass, fields
18
+from typing import Any, Generic, TypeVar, cast
19
20
__all__ = [
21
"JitWarmupRegistry",
@@ -597,6 +597,16 @@ class VllmJitKernel(Generic[CompileKeyT], ABC):
597
predicate_trace = (
598
_trace_warmup_predicate(_when) if _when is not None else None
599
)
600
+ predicate_only_names: frozenset[str] = frozenset()
601
+ if predicate_trace is not None:
602
+ compile_key_fields = frozenset(
603
+ field.name for field in fields(cast(Any, self.CompileKey))
604
+ )
605
+ predicate_only_names = (
606
+ predicate_trace.input_names
607
+ - compile_key_dispatch_trace.input_names
608
+ - compile_key_fields
609
+ )
610
# Unmatched **kwargs fields also belong to the expansion space.
611
available_names = set(kwargs).union(
612
*(group.rows[0] for group in input_groups)
@@ -638,7 +648,12 @@ class VllmJitKernel(Generic[CompileKeyT], ABC):
648
):
649
continue
650
compile_key = compile_key_dispatch_trace.compile_key(
641
- self.CompileKey, dispatch_values
651
+ self.CompileKey,
652
+ {
653
+ name: value
654
+ for name, value in dispatch_values.items()
655
+ if name not in predicate_only_names
656
+ },
657
)
658
compile_keys[compile_key] = None
659
return list(compile_keys)
vllm/model_executor/warmup/jit_warmup_triton_helper.py
+41
-9
@@ -15,7 +15,13 @@ from vllm.model_executor.warmup.jit_warmup import (
15
)
16
17
CompileKeyT = TypeVar("CompileKeyT")
18
-LaunchSpec = tuple[tuple[int, ...], dict[str, Any]]
18
+# ``(grid, launch_kwargs)`` or ``(grid, launch_kwargs, outputs)`` for
19
+# self-allocating kernels. ``grid=None`` skips the launch (e.g. empty batch)
20
+# but still returns the declared outputs.
21
+LaunchSpec = (
22
+ tuple[tuple[int, ...] | None, dict[str, Any]]
23
+ | tuple[tuple[int, ...] | None, dict[str, Any], Any]
24
+)
25
26
27
def triton_scalar_specialization_rep(value: int) -> int:
@@ -91,6 +97,7 @@ class VllmTritonJitKernel(VllmJitKernel[CompileKeyT], Generic[CompileKeyT]):
97
98
kernel: ClassVar[Any]
99
_warming = False
100
+ _warming_compile_key: CompileKeyT | None = None
101
102
@abstractmethod
103
def warmup_inputs(self, compile_key: CompileKeyT) -> dict[str, Any]:
@@ -100,10 +107,12 @@ class VllmTritonJitKernel(VllmJitKernel[CompileKeyT], Generic[CompileKeyT]):
107
def compile(self, compile_key: CompileKeyT) -> None:
108
inputs = self.warmup_inputs(compile_key)
109
self._warming = True
110
+ self._warming_compile_key = compile_key
111
try:
112
cast(Callable[..., None], self)(**inputs)
113
finally:
114
self._warming = False
115
+ self._warming_compile_key = None
116
117
@cached_property
118
def _kernel_arg_names(self) -> tuple[str, ...]:
@@ -119,11 +128,14 @@ class VllmTritonJitKernel(VllmJitKernel[CompileKeyT], Generic[CompileKeyT]):
128
129
def launch(
130
self,
122
- grid: tuple[int, ...],
131
+ grid: tuple[int, ...] | None,
132
inputs: Mapping[str, Any],
133
/,
134
+ *,
135
+ kernel: Any = None,
136
**kwargs: Any,
137
) -> Any:
138
+ kernel = kernel if kernel is not None else self.kernel
139
runtime_launcher = kwargs.pop("_runtime_launcher", None)
140
runtime_launcher_arg_count = kwargs.pop("_runtime_launcher_arg_count", 0)
141
for name, value in inputs.items():
@@ -131,22 +143,36 @@ class VllmTritonJitKernel(VllmJitKernel[CompileKeyT], Generic[CompileKeyT]):
143
if target in self._kernel_arg_names and target not in kwargs:
144
kwargs[target] = value
145
if self._warming:
134
- warmup = getattr(self.kernel, "warmup", None)
146
+ if (
147
+ self._warming_compile_key is not None
148
+ and "launch_pdl" in kwargs
149
+ and hasattr(self._warming_compile_key, "launch_pdl")
150
+ ):
151
+ kwargs["launch_pdl"] = self._warming_compile_key.launch_pdl
152
+ warmup = getattr(kernel, "warmup", None)
153
assert warmup is not None
154
return warmup(grid=(1,), **kwargs)
155
+ if grid is None:
156
+ return None
157
if runtime_launcher is not None:
158
regular_args = [
159
kwargs.pop(name)
160
for name in self._kernel_arg_names[:runtime_launcher_arg_count]
161
]
142
- return runtime_launcher(self.kernel, grid, *regular_args, **kwargs)
143
- return self.kernel[grid](**kwargs)
162
+ return runtime_launcher(kernel, grid, *regular_args, **kwargs)
163
+ return kernel[grid](**kwargs)
164
165
166
def kernel_launcher(
167
call_fn: Callable[..., LaunchSpec],
148
-) -> Callable[..., None]:
149
- """Launch a Triton kernel from a declarative ``__call__`` specification."""
168
+) -> Callable[..., Any]:
169
+ """Launch a Triton kernel from a declarative ``__call__`` specification.
170
+
171
+ ``call_fn`` returns either ``(grid, launch_kwargs)`` or, when it allocates
172
+ its own outputs, ``(grid, launch_kwargs, outputs)``. The declared outputs are
173
+ returned to the caller. A ``grid`` of ``None`` skips the launch (e.g. an
174
+ empty-token batch) while still returning the outputs.
175
+ """
176
signature = inspect.signature(call_fn)
177
178
@wraps(call_fn)
@@ -154,14 +180,20 @@ def kernel_launcher(
180
self: VllmTritonJitKernel[Any],
181
*args: Any,
182
**kwargs: Any,
157
- ) -> None:
158
- grid, launch_kwargs = call_fn(self, *args, **kwargs)
183
+ ) -> Any:
184
+ spec = call_fn(self, *args, **kwargs)
185
+ if len(spec) == 3:
186
+ grid, launch_kwargs, outputs = spec
187
+ else:
188
+ grid, launch_kwargs = spec
189
+ outputs = None
190
bound = signature.bind(self, *args, **kwargs)
191
bound.apply_defaults()
192
inputs = {
193
name: value for name, value in bound.arguments.items() if name != "self"
194
}
195
self.launch(grid, inputs, **launch_kwargs)
196
+ return outputs
197
198
return wrapper
199
vllm/models/common/ops/__init__.py
+2
-1
@@ -2,8 +2,9 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""Ops shared across model implementations."""
4
5
-from .fused_qk_rmsnorm import fused_q_kv_rmsnorm
5
+from .fused_qk_rmsnorm import _FUSED_Q_KV_RMSNORM_KERNEL, fused_q_kv_rmsnorm
6
7
__all__ = [
8
+ "_FUSED_Q_KV_RMSNORM_KERNEL",
9
"fused_q_kv_rmsnorm",
10
]
vllm/models/common/ops/fused_qk_rmsnorm.py
+196
-90
@@ -1,67 +1,194 @@
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
7
8
+from vllm.model_executor.warmup.jit_warmup_triton_helper import (
9
+ LaunchSpec,
10
+ TritonWarmupTensor,
11
+ VllmTritonJitKernel,
12
+ kernel_launcher,
13
+)
14
from vllm.platforms import current_platform
15
from vllm.triton_utils import tl, triton
16
+from vllm.utils.math_utils import next_power_of_2
17
+
18
+
19
+class FusedQKVRMSNormKernel(VllmTritonJitKernel["FusedQKVRMSNormKernel.CompileKey"]):
20
+ @dataclass(frozen=True)
21
+ class CompileKey:
22
+ dtype: torch.dtype
23
+ q_size: int
24
+ kv_size: int
25
+ block_size: int
26
+ q_in_stride: int
27
+ q_out_stride: int
28
+ kv_in_stride: int
29
+ kv_out_stride: int
30
+ eps: float
31
+ launch_pdl: bool
32
+
33
+ @staticmethod
34
+ @triton.jit
35
+ def kernel(
36
+ q_ptr,
37
+ q_out_ptr,
38
+ q_weight_ptr,
39
+ q_in_stride,
40
+ q_out_stride,
41
+ kv_ptr,
42
+ kv_out_ptr,
43
+ kv_weight_ptr,
44
+ kv_in_stride,
45
+ kv_out_stride,
46
+ eps,
47
+ Q_SIZE: tl.constexpr,
48
+ KV_SIZE: tl.constexpr,
49
+ BLOCK_SIZE: tl.constexpr,
50
+ launch_pdl: tl.constexpr,
51
+ ):
52
+ # num_tokens goes on grid-x (max 2**31 - 1); task goes on grid-y.
53
+ # CUDA's grid-y/z are capped at 65535, so putting num_tokens there crashes
54
+ # the launch at max-num-batched-tokens >= 65536 with "invalid argument".
55
+ # int64: q_in_stride can be ~24K (128 heads × 192) and overflows int32
56
+ # past num_tokens ~87K under large chunked prefill.
57
+ token_idx = tl.program_id(0).to(tl.int64)
58
+ pid_task = tl.program_id(1)
59
+
60
+ if pid_task == 0:
61
+ SIZE = Q_SIZE
62
+ row_in = q_ptr + token_idx * q_in_stride
63
+ weight_ptr = q_weight_ptr
64
+ row_out = q_out_ptr + token_idx * q_out_stride
65
+ else:
66
+ SIZE = KV_SIZE
67
+ row_in = kv_ptr + token_idx * kv_in_stride
68
+ weight_ptr = kv_weight_ptr
69
+ row_out = kv_out_ptr + token_idx * kv_out_stride
70
+
71
+ # RMSNorm in fp32 throughout — matches csrc/layernorm_kernels.cu's
72
+ # `(scalar_t)(x * s_variance * w)` and DeepseekV4's compressor kernel, which
73
+ # keep x, rrms, and w all in fp32 and perform a single cast at store.
74
+ block = tl.arange(0, BLOCK_SIZE)
75
+ mask = block < SIZE
76
+ # The weight load does not depend on the producer's output, so issue it
77
+ # before the PDL wait: gamma streams in while the producer finishes,
78
+ # leaving a single dependent global round trip (x) after the wait.
79
+ w = tl.load(weight_ptr + block, mask=mask, other=0.0).to(tl.float32)
80
+
81
+ if launch_pdl:
82
+ tl.extra.cuda.gdc_wait()
83
+ tl.extra.cuda.gdc_launch_dependents()
84
85
+ x = tl.load(row_in + block, mask=mask, other=0.0).to(tl.float32)
86
+ variance = tl.sum(x * x, axis=0) / SIZE
87
+ rrms = tl.rsqrt(variance + eps)
88
+ y = x * rrms * w
89
+ tl.store(row_out + block, y.to(row_out.dtype.element_ty), mask=mask)
90
9
-@triton.jit
10
-def _fused_q_kv_rmsnorm_kernel(
11
- q_ptr,
12
- q_out_ptr,
13
- q_weight_ptr,
14
- q_in_stride,
15
- q_out_stride,
16
- kv_ptr,
17
- kv_out_ptr,
18
- kv_weight_ptr,
19
- kv_in_stride,
20
- kv_out_stride,
21
- eps,
22
- Q_SIZE: tl.constexpr,
23
- KV_SIZE: tl.constexpr,
24
- BLOCK_SIZE: tl.constexpr,
25
- launch_pdl: tl.constexpr,
26
-):
27
- # num_tokens goes on grid-x (max 2**31 - 1); task goes on grid-y.
28
- # CUDA's grid-y/z are capped at 65535, so putting num_tokens there crashes
29
- # the launch at max-num-batched-tokens >= 65536 with "invalid argument".
30
- # int64: q_in_stride can be ~24K (128 heads × 192) and overflows int32
31
- # past num_tokens ~87K under large chunked prefill.
32
- token_idx = tl.program_id(0).to(tl.int64)
33
- pid_task = tl.program_id(1)
34
-
35
- if pid_task == 0:
36
- SIZE = Q_SIZE
37
- row_in = q_ptr + token_idx * q_in_stride
38
- weight_ptr = q_weight_ptr
39
- row_out = q_out_ptr + token_idx * q_out_stride
40
- else:
41
- SIZE = KV_SIZE
42
- row_in = kv_ptr + token_idx * kv_in_stride
43
- weight_ptr = kv_weight_ptr
44
- row_out = kv_out_ptr + token_idx * kv_out_stride
45
-
46
- # RMSNorm in fp32 throughout — matches csrc/layernorm_kernels.cu's
47
- # `(scalar_t)(x * s_variance * w)` and DeepseekV4's compressor kernel, which
48
- # keep x, rrms, and w all in fp32 and perform a single cast at store.
49
- block = tl.arange(0, BLOCK_SIZE)
50
- mask = block < SIZE
51
- # The weight load does not depend on the producer's output, so issue it
52
- # before the PDL wait: gamma streams in while the producer finishes,
53
- # leaving a single dependent global round trip (x) after the wait.
54
- w = tl.load(weight_ptr + block, mask=mask, other=0.0).to(tl.float32)
55
-
56
- if launch_pdl:
57
- tl.extra.cuda.gdc_wait()
58
- tl.extra.cuda.gdc_launch_dependents()
59
-
60
- x = tl.load(row_in + block, mask=mask, other=0.0).to(tl.float32)
61
- variance = tl.sum(x * x, axis=0) / SIZE
62
- rrms = tl.rsqrt(variance + eps)
63
- y = x * rrms * w
64
- tl.store(row_out + block, y.to(row_out.dtype.element_ty), mask=mask)
91
+ def dispatch( # type: ignore[override]
92
+ self,
93
+ *,
94
+ dtype: torch.dtype,
95
+ q_size: int,
96
+ kv_size: int,
97
+ eps: float,
98
+ launch_pdl: bool,
99
+ **compile_key_fields: int,
100
+ ) -> CompileKey:
101
+ max_size = q_size if q_size >= kv_size else kv_size
102
+ return self.CompileKey(
103
+ **compile_key_fields,
104
+ dtype=dtype,
105
+ q_size=q_size,
106
+ kv_size=kv_size,
107
+ block_size=next_power_of_2(max_size),
108
+ eps=eps,
109
+ launch_pdl=launch_pdl,
110
+ )
111
+
112
+ def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]:
113
+ hf_config = vllm_config.model_config.hf_config
114
+ q_size = int(getattr(hf_config, "q_lora_rank", 0) or 0)
115
+ kv_size = int(getattr(hf_config, "head_dim", 0) or 0)
116
+ if q_size <= 0 or kv_size <= 0:
117
+ return []
118
+
119
+ input_stride = q_size + kv_size
120
+ return self._trace_dispatch(self.dispatch)(
121
+ dtype=vllm_config.model_config.dtype,
122
+ q_size=q_size,
123
+ kv_size=kv_size,
124
+ q_in_stride=input_stride,
125
+ q_out_stride=(input_stride, q_size),
126
+ kv_in_stride=input_stride,
127
+ kv_out_stride=(input_stride, kv_size),
128
+ eps=float(hf_config.rms_norm_eps),
129
+ launch_pdl=(False, True),
130
+ )
131
+
132
+ def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
133
+ return dict(
134
+ q=TritonWarmupTensor(
135
+ compile_key.dtype,
136
+ shape=(1, compile_key.q_size),
137
+ strides=(compile_key.q_in_stride, 1),
138
+ ),
139
+ kv=TritonWarmupTensor(
140
+ compile_key.dtype,
141
+ shape=(1, compile_key.kv_size),
142
+ strides=(compile_key.kv_in_stride, 1),
143
+ ),
144
+ q_out=TritonWarmupTensor(
145
+ compile_key.dtype,
146
+ shape=(1, compile_key.q_size),
147
+ strides=(compile_key.q_out_stride, 1),
148
+ ),
149
+ kv_out=TritonWarmupTensor(
150
+ compile_key.dtype,
151
+ shape=(1, compile_key.kv_size),
152
+ strides=(compile_key.kv_out_stride, 1),
153
+ ),
154
+ q_weight=TritonWarmupTensor(
155
+ compile_key.dtype,
156
+ shape=(compile_key.q_size,),
157
+ ),
158
+ kv_weight=TritonWarmupTensor(
159
+ compile_key.dtype,
160
+ shape=(compile_key.kv_size,),
161
+ ),
162
+ eps=compile_key.eps,
163
+ )
164
+
165
+ @kernel_launcher
166
+ def __call__(
167
+ self,
168
+ q: torch.Tensor,
169
+ kv: torch.Tensor,
170
+ q_out: torch.Tensor,
171
+ kv_out: torch.Tensor,
172
+ q_weight: torch.Tensor,
173
+ kv_weight: torch.Tensor,
174
+ eps: float,
175
+ ) -> LaunchSpec:
176
+ q_size = q.shape[1]
177
+ kv_size = kv.shape[1]
178
+ block_size = next_power_of_2(max(q_size, kv_size))
179
+ return (q.shape[0], 2), dict(
180
+ q_in_stride=q.stride(0),
181
+ q_out_stride=q_out.stride(0),
182
+ kv_in_stride=kv.stride(0),
183
+ kv_out_stride=kv_out.stride(0),
184
+ Q_SIZE=q_size,
185
+ KV_SIZE=kv_size,
186
+ BLOCK_SIZE=block_size,
187
+ launch_pdl=current_platform.is_arch_support_pdl(),
188
+ # One vectorized 128-bit access per thread for K3/DSv4-sized rows,
189
+ # instead of a serial per-thread chain at the default 4 warps.
190
+ num_warps=8 if block_size >= 2048 else 4,
191
+ )
192
193
194
def fused_q_kv_rmsnorm(
@@ -78,40 +205,19 @@ def fused_q_kv_rmsnorm(
205
assert qr.stride(-1) == 1 and kv.stride(-1) == 1
206
assert q_weight.is_contiguous() and kv_weight.is_contiguous()
207
81
- q_size = qr.shape[1]
82
- kv_size = kv.shape[1]
83
- num_tokens = qr.shape[0]
84
- # Allocate packed outputs rather than empty_like: qr/kv are typically
85
- # column slices of a fused qkv_a projection, and for num_tokens == 1
86
- # empty_like preserves the parent's row stride (torch keeps strides of
87
- # size-1 dims), e.g. (2112, 1) for a [1, 1536] q slice. Downstream
88
- # dispatchers that require packed row-major inputs (e.g. the Kimi-K3
89
- # low-latency GEMM's _is_packed_row_major check) then reject the tensor
90
- # and silently fall back to cuBLAS on every decode step.
208
qr_out = torch.empty(qr.shape, dtype=qr.dtype, device=qr.device)
209
kv_out = torch.empty(kv.shape, dtype=kv.dtype, device=kv.device)
93
- if num_tokens == 0:
94
- return qr_out, kv_out
95
-
96
- block_size = triton.next_power_of_2(max(q_size, kv_size))
97
- _fused_q_kv_rmsnorm_kernel[(num_tokens, 2)](
98
- qr,
99
- qr_out,
100
- q_weight,
101
- qr.stride(0),
102
- qr_out.stride(0),
103
- kv,
104
- kv_out,
105
- kv_weight,
106
- kv.stride(0),
107
- kv_out.stride(0),
108
- eps,
109
- Q_SIZE=q_size,
110
- KV_SIZE=kv_size,
111
- BLOCK_SIZE=block_size,
112
- launch_pdl=current_platform.is_arch_support_pdl(),
113
- # One vectorized 128-bit access per thread for K3/DSv4-sized rows,
114
- # instead of a serial per-thread chain at the default 4 warps.
115
- num_warps=8 if block_size >= 2048 else 4,
116
- )
210
+ if qr.shape[0] > 0:
211
+ _FUSED_Q_KV_RMSNORM_KERNEL(
212
+ qr,
213
+ kv,
214
+ qr_out,
215
+ kv_out,
216
+ q_weight,
217
+ kv_weight,
218
+ eps,
219
+ )
220
return qr_out, kv_out
221
+
222
+
223
+_FUSED_Q_KV_RMSNORM_KERNEL = FusedQKVRMSNormKernel()
vllm/models/deepseek_v4/amd/mtp.py
+9
-5
@@ -39,9 +39,9 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader
39
from vllm.model_executor.models.deepseek_mtp import SharedHead
40
from vllm.model_executor.models.deepseek_v2 import get_spec_layer_idx_from_weight_name
41
from vllm.model_executor.models.utils import maybe_prefix
42
-from vllm.models.deepseek_v4.common.ops import (
43
- fused_mtp_input_rmsnorm,
44
- mtp_shared_head_rmsnorm,
42
+from vllm.models.deepseek_v4.common.ops.fused_mtp_input_rmsnorm import (
43
+ _FUSED_MTP_INPUT_RMSNORM_KERNEL,
44
+ _MTP_SHARED_HEAD_RMSNORM_KERNEL,
45
)
46
from vllm.platforms import current_platform
47
from vllm.sequence import IntermediateTensors
@@ -127,6 +127,10 @@ class DeepSeekV4MultiTokenPredictorLayer(nn.Module):
127
128
self.hc_head_op = HCHeadOp()
129
130
+ if vllm_config.kernel_config.enable_jit_warmup:
131
+ _FUSED_MTP_INPUT_RMSNORM_KERNEL.register_warmup()
132
+ _MTP_SHARED_HEAD_RMSNORM_KERNEL.register_warmup()
133
+
134
def forward(
135
self,
136
input_ids: torch.Tensor,
@@ -143,7 +147,7 @@ class DeepSeekV4MultiTokenPredictorLayer(nn.Module):
147
-1, self.hc_mult, self.config.hidden_size
148
)
149
# Fused: mask inputs at position 0 (not needed by MTP), enorm, hnorm.
146
- inputs_embeds, previous_hidden_states = fused_mtp_input_rmsnorm(
150
+ inputs_embeds, previous_hidden_states = _FUSED_MTP_INPUT_RMSNORM_KERNEL(
151
inputs_embeds,
152
positions,
153
previous_hidden_states,
@@ -256,7 +260,7 @@ class DeepSeekV4MultiTokenPredictor(nn.Module):
260
mtp_layer.rms_norm_eps,
261
mtp_layer.hc_eps,
262
)
259
- hidden_states = mtp_shared_head_rmsnorm(
263
+ hidden_states = _MTP_SHARED_HEAD_RMSNORM_KERNEL(
264
hidden_states,
265
mtp_layer.shared_head.norm.weight.data,
266
mtp_layer.shared_head.norm.variance_epsilon,
vllm/models/deepseek_v4/attention.py
+58
-1
@@ -22,7 +22,10 @@ from vllm.model_executor.layers.linear import (
22
RowParallelLinear,
23
)
24
from vllm.model_executor.layers.sparse_attn_indexer import SparseAttnIndexer
25
-from vllm.models.common.ops import fused_q_kv_rmsnorm
25
+from vllm.models.common.ops.fused_qk_rmsnorm import (
26
+ _FUSED_Q_KV_RMSNORM_KERNEL,
27
+ fused_q_kv_rmsnorm,
28
+)
29
from vllm.models.deepseek_v4.common.ops import (
30
fused_indexer_q_rope_quant,
31
)
@@ -417,6 +420,44 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
420
421
_COMBINE_TOPK_SWA_INDICES_KERNEL.register_warmup()
422
423
+ from vllm.platforms import current_platform
424
+ from vllm.utils.import_utils import has_cutedsl
425
+
426
+ _FUSED_Q_KV_RMSNORM_KERNEL.register_warmup()
427
+
428
+ backend_name = self.backend_cls.get_name()
429
+ if current_platform.is_cuda():
430
+ from vllm.models.deepseek_v4.common.ops.fused_inv_rope_fp8_quant import ( # noqa: E501
431
+ _FUSED_INV_ROPE_FP8_QUANT_KERNEL,
432
+ )
433
+
434
+ _FUSED_INV_ROPE_FP8_QUANT_KERNEL.register_warmup()
435
+
436
+ if self.compress_ratio == 128:
437
+ from vllm.models.deepseek_v4.sparse_mla import (
438
+ _BUILD_C128A_TOPK_METADATA_KERNEL,
439
+ )
440
+
441
+ _BUILD_C128A_TOPK_METADATA_KERNEL.register_warmup()
442
+
443
+ if backend_name == "FLASHMLA_SPARSE_DSV4":
444
+ from vllm.models.deepseek_v4.common.ops.cache_utils import (
445
+ _COMPUTE_GLOBAL_TOPK_INDICES_AND_LENS_KERNEL,
446
+ _DEQUANTIZE_AND_GATHER_K_CACHE_KERNEL,
447
+ )
448
+
449
+ _COMPUTE_GLOBAL_TOPK_INDICES_AND_LENS_KERNEL.register_warmup()
450
+ if not has_cutedsl():
451
+ _DEQUANTIZE_AND_GATHER_K_CACHE_KERNEL.register_warmup()
452
+ elif backend_name == "FLASHINFER_MLA_SPARSE_DSV4":
453
+ from vllm.models.deepseek_v4.common.ops.cache_utils import (
454
+ _BUILD_FLASHINFER_MIXED_SPARSE_INDICES_KERNEL,
455
+ _COMPUTE_GLOBAL_TOPK_INDICES_AND_LENS_KERNEL,
456
+ )
457
+
458
+ _COMPUTE_GLOBAL_TOPK_INDICES_AND_LENS_KERNEL.register_warmup()
459
+ _BUILD_FLASHINFER_MIXED_SPARSE_INDICES_KERNEL.register_warmup()
460
+
461
def forward(
462
self,
463
positions: torch.Tensor,
@@ -974,6 +1015,22 @@ class DeepseekV4Indexer(nn.Module):
1015
torch.cuda.Event(),
1016
]
1017
1018
+ if vllm_config.kernel_config.enable_jit_warmup:
1019
+ from vllm.platforms import current_platform
1020
+ from vllm.utils.import_utils import has_cutedsl
1021
+
1022
+ if not has_cutedsl() and not current_platform.is_xpu():
1023
+ from vllm.models.deepseek_v4.common.ops.fused_indexer_q import (
1024
+ _FUSED_INDEXER_Q_ROPE_MXFP4_TRITON_KERNEL,
1025
+ _FUSED_INDEXER_Q_ROPE_QUANT_TRITON_KERNEL,
1026
+ )
1027
+
1028
+ (
1029
+ _FUSED_INDEXER_Q_ROPE_MXFP4_TRITON_KERNEL
1030
+ if self.use_fp4_kv
1031
+ else _FUSED_INDEXER_Q_ROPE_QUANT_TRITON_KERNEL
1032
+ ).register_warmup()
1033
+
1034
def forward(
1035
self,
1036
hidden_states: torch.Tensor,
vllm/models/deepseek_v4/common/ops/__init__.py
-5
@@ -10,8 +10,6 @@ from .cache_utils import (
10
)
11
from .fused_indexer_q import MXFP4_BLOCK_SIZE, fused_indexer_q_rope_quant
12
from .fused_inv_rope_fp8_quant import fused_inv_rope_fp8_quant
13
-from .fused_mtp_input_rmsnorm import fused_mtp_input_rmsnorm, mtp_shared_head_rmsnorm
14
-from .save_partial_states import save_partial_states
13
14
__all__ = [
15
"MXFP4_BLOCK_SIZE",
@@ -21,8 +19,5 @@ __all__ = [
19
"dequantize_and_gather_k_cache",
20
"fused_indexer_q_rope_quant",
21
"fused_inv_rope_fp8_quant",
24
- "fused_mtp_input_rmsnorm",
25
- "mtp_shared_head_rmsnorm",
22
"quantize_and_insert_k_cache",
27
- "save_partial_states",
23
]
vllm/models/deepseek_v4/common/ops/cache_utils.py
+903
-454
@@ -22,10 +22,17 @@ import torch
22
from vllm.model_executor.layers.quantization.utils.quant_utils import (
23
get_fp8_min_max,
24
)
25
-from vllm.model_executor.warmup.jit_warmup import VllmJitKernel, zip_inputs
25
+from vllm.model_executor.warmup.jit_warmup import (
26
+ WarmupIntRange,
27
+ zip_inputs,
28
+)
29
from vllm.model_executor.warmup.jit_warmup_triton_helper import (
30
+ LaunchSpec,
31
TritonPointerInputVariant,
32
TritonWarmupTensor,
33
+ VllmTritonJitKernel,
34
+ kernel_launcher,
35
+ triton_scalar_specialization_rep,
36
)
37
from vllm.platforms import current_platform
38
from vllm.triton_utils import tl, triton
@@ -225,166 +232,268 @@ def quantize_and_insert_k_cache(
232
)
233
234
228
-@triton.jit
229
-def _dequantize_and_gather_k_kernel(
230
- out_ptr,
231
- out_stride0,
232
- out_stride1,
233
- k_cache_ptr,
234
- seq_lens_ptr,
235
- block_table_ptr,
236
- offset,
237
- gather_lens_ptr,
238
- # Constants
239
- max_blocks_per_seq: tl.constexpr,
240
- fp8_dim: tl.constexpr, # 448
241
- bf16_dim: tl.constexpr, # 64
242
- scale_dim: tl.constexpr, # 8
243
- quant_block: tl.constexpr, # 64 (quantization block size)
244
- cache_block_size: tl.constexpr, # 64 or 128 (paged cache block size)
245
- token_data_size: tl.constexpr, # 576 bytes per token data
246
- block_stride: tl.constexpr, # total bytes per block (padded) int32
247
- output_dim: tl.constexpr, # 512
248
- fp8_max: tl.constexpr,
249
- n_quant_blocks: tl.constexpr, # 7 real blocks
250
- use_fnuz: tl.constexpr = False,
235
+class DequantizeAndGatherKCacheKernel(
236
+ VllmTritonJitKernel["DequantizeAndGatherKCacheKernel.CompileKey"]
237
):
252
- batch_idx = tl.program_id(0)
253
- worker_id = tl.program_id(1)
254
- num_workers = tl.num_programs(1)
238
+ NUM_WORKERS = 128
239
256
- seq_len = tl.load(seq_lens_ptr + batch_idx)
257
- if gather_lens_ptr is not None: # noqa: SIM108
258
- gather_len = tl.load(gather_lens_ptr + batch_idx)
259
- else:
260
- # Gather all tokens
261
- gather_len = seq_len
262
- start_pos = seq_len - gather_len
263
-
264
- for i in range(worker_id, gather_len, num_workers):
265
- # Calculate the actual token index in the sequence
266
- pos = start_pos + i
267
-
268
- # Calculate which block and position within block
269
- block_in_seq = pos // cache_block_size
270
- pos_in_block = pos % cache_block_size
271
-
272
- # Get physical block index from block table
273
- block_table_row_ptr = block_table_ptr + batch_idx * max_blocks_per_seq
274
- physical_block_idx = tl.load(block_table_row_ptr + block_in_seq) # int32
275
-
276
- # int64: physical_block_idx * block_stride can exceed 2^31 with many
277
- # KV-cache blocks (e.g. >= 57K at block_stride ~37K).
278
- cache_block_ptr = k_cache_ptr + physical_block_idx.to(tl.int64) * block_stride
279
-
280
- # Token data pointer
281
- token_data_ptr = cache_block_ptr + pos_in_block * token_data_size
282
-
283
- # Scale pointer: after all token data
284
- token_scale_ptr = (
285
- cache_block_ptr
286
- + cache_block_size * token_data_size
287
- + pos_in_block * scale_dim
288
- )
240
+ @dataclass(frozen=True)
241
+ class CompileKey:
242
+ max_blocks_per_seq: int
243
+ cache_block_size: int
244
+ block_stride: int
245
+ use_fnuz: bool
246
+ has_gather_lens: bool
247
+ offset: int
248
290
- # Token data layout: [0:448] fp8, [448:576] bf16
291
- token_fp8_ptr = token_data_ptr
292
- token_bf16_ptr = token_data_ptr + fp8_dim
249
+ @staticmethod
250
+ @triton.jit
251
+ def kernel(
252
+ out_ptr,
253
+ out_stride0,
254
+ out_stride1,
255
+ k_cache_ptr,
256
+ seq_lens_ptr,
257
+ block_table_ptr,
258
+ offset,
259
+ gather_lens_ptr,
260
+ # Constants
261
+ max_blocks_per_seq: tl.constexpr,
262
+ fp8_dim: tl.constexpr, # 448
263
+ bf16_dim: tl.constexpr, # 64
264
+ scale_dim: tl.constexpr, # 8
265
+ quant_block: tl.constexpr, # 64 (quantization block size)
266
+ cache_block_size: tl.constexpr, # 64 or 128 (paged cache block size)
267
+ token_data_size: tl.constexpr, # 576 bytes per token data
268
+ block_stride: tl.constexpr, # total bytes per block (padded) int32
269
+ output_dim: tl.constexpr, # 512
270
+ fp8_max: tl.constexpr,
271
+ n_quant_blocks: tl.constexpr, # 7 real blocks
272
+ use_fnuz: tl.constexpr = False,
273
+ ):
274
+ batch_idx = tl.program_id(0)
275
+ worker_id = tl.program_id(1)
276
+ num_workers = tl.num_programs(1)
277
294
- # Output pointer for this token (flattened)
295
- output_row_ptr = out_ptr + batch_idx * out_stride0 + (offset + i) * out_stride1
278
+ seq_len = tl.load(seq_lens_ptr + batch_idx)
279
+ if gather_lens_ptr is not None: # noqa: SIM108
280
+ gather_len = tl.load(gather_lens_ptr + batch_idx)
281
+ else:
282
+ # Gather all tokens
283
+ gather_len = seq_len
284
+ start_pos = seq_len - gather_len
285
+
286
+ for i in range(worker_id, gather_len, num_workers):
287
+ # Calculate the actual token index in the sequence
288
+ pos = start_pos + i
289
+
290
+ # Calculate which block and position within block
291
+ block_in_seq = pos // cache_block_size
292
+ pos_in_block = pos % cache_block_size
293
+
294
+ # Get physical block index from block table
295
+ block_table_row_ptr = block_table_ptr + batch_idx * max_blocks_per_seq
296
+ physical_block_idx = tl.load(block_table_row_ptr + block_in_seq) # int32
297
+
298
+ # int64: physical_block_idx * block_stride can exceed 2^31 with many
299
+ # KV-cache blocks (e.g. >= 57K at block_stride ~37K).
300
+ cache_block_ptr = (
301
+ k_cache_ptr + physical_block_idx.to(tl.int64) * block_stride
302
+ )
303
297
- # ========== Dequantize FP8 portion using UE8M0 ==========
298
- for qblock_idx in tl.static_range(n_quant_blocks):
299
- qblock_start = qblock_idx * quant_block
304
+ # Token data pointer
305
+ token_data_ptr = cache_block_ptr + pos_in_block * token_data_size
306
301
- if qblock_start < fp8_dim:
302
- offsets = qblock_start + tl.arange(0, quant_block)
303
- mask = offsets < fp8_dim
307
+ # Scale pointer: after all token data
308
+ token_scale_ptr = (
309
+ cache_block_ptr
310
+ + cache_block_size * token_data_size
311
+ + pos_in_block * scale_dim
312
+ )
313
305
- # Load quantized fp8 values (stored as uint8)
306
- x_uint8 = tl.load(token_fp8_ptr + offsets, mask=mask, other=0)
314
+ # Token data layout: [0:448] fp8, [448:576] bf16
315
+ token_fp8_ptr = token_data_ptr
316
+ token_bf16_ptr = token_data_ptr + fp8_dim
317
308
- # Bitcast uint8 back to fp8 (FNUZ on gfx942, OCP elsewhere).
309
- if use_fnuz:
310
- x_fp8 = x_uint8.to(tl.float8e4b8, bitcast=True)
311
- else:
312
- x_fp8 = x_uint8.to(tl.float8e4nv, bitcast=True)
318
+ # Output pointer for this token (flattened)
319
+ output_row_ptr = (
320
+ out_ptr + batch_idx * out_stride0 + (offset + i) * out_stride1
321
+ )
322
314
- # Convert fp8 to float32 for computation
315
- x_float = x_fp8.to(tl.float32)
323
+ # ========== Dequantize FP8 portion using UE8M0 ==========
324
+ for qblock_idx in tl.static_range(n_quant_blocks):
325
+ qblock_start = qblock_idx * quant_block
326
317
- # Load and decode UE8M0 scale
318
- # UE8M0: scale = 2^(stored_value - 127)
319
- encoded_scale = tl.load(token_scale_ptr + qblock_idx)
320
- exponent = encoded_scale.to(tl.float32) - 127.0
321
- scale = tl.exp2(exponent)
327
+ if qblock_start < fp8_dim:
328
+ offsets = qblock_start + tl.arange(0, quant_block)
329
+ mask = offsets < fp8_dim
330
323
- # Dequantize: bf16_value = fp8_value * scale
324
- x_dequant = x_float * scale
331
+ # Load quantized fp8 values (stored as uint8)
332
+ x_uint8 = tl.load(token_fp8_ptr + offsets, mask=mask, other=0)
333
326
- # Store as bf16
327
- tl.store(output_row_ptr + offsets, x_dequant.to(tl.bfloat16), mask=mask)
334
+ # Bitcast uint8 back to fp8 (FNUZ on gfx942, OCP elsewhere).
335
+ if use_fnuz:
336
+ x_fp8 = x_uint8.to(tl.float8e4b8, bitcast=True)
337
+ else:
338
+ x_fp8 = x_uint8.to(tl.float8e4nv, bitcast=True)
339
329
- # ========== Copy BF16 portion directly ==========
330
- bf16_output_offset = fp8_dim # After 448 elements in output
340
+ # Convert fp8 to float32 for computation
341
+ x_float = x_fp8.to(tl.float32)
342
332
- # Read bf16 from cache
333
- bf16_cache_ptr = token_bf16_ptr.to(tl.pointer_type(tl.bfloat16))
343
+ # Load and decode UE8M0 scale
344
+ # UE8M0: scale = 2^(stored_value - 127)
345
+ encoded_scale = tl.load(token_scale_ptr + qblock_idx)
346
+ exponent = encoded_scale.to(tl.float32) - 127.0
347
+ scale = tl.exp2(exponent)
348
335
- # Process in chunks of 16
336
- for j in tl.static_range(bf16_dim // 16):
337
- chunk_offsets = j * 16 + tl.arange(0, 16)
338
- bf16_vals = tl.load(bf16_cache_ptr + chunk_offsets)
339
- tl.store(output_row_ptr + bf16_output_offset + chunk_offsets, bf16_vals)
349
+ # Dequantize: bf16_value = fp8_value * scale
350
+ x_dequant = x_float * scale
351
352
+ # Store as bf16
353
+ tl.store(
354
+ output_row_ptr + offsets, x_dequant.to(tl.bfloat16), mask=mask
355
+ )
356
342
-def dequantize_and_gather_k_cache_triton(
343
- # [num_reqs, max_num_tokens, head_size]
344
- out: torch.Tensor,
345
- # [num_blocks, block_size, head_bytes]
346
- k_cache: torch.Tensor,
347
- # [num_reqs]
348
- seq_lens: torch.Tensor,
349
- # [num_reqs]
350
- gather_lens: torch.Tensor | None,
351
- # [num_reqs, max_blocks_per_seq]
352
- block_table: torch.Tensor,
353
- block_size: int,
354
- offset: int,
355
- use_fnuz: bool = False,
356
-) -> None:
357
- TOKEN_FP8_DIM = 448
358
- TOKEN_BF16_DIM = 64
359
- TOKEN_SCALE_DIM = 8
360
- QUANT_BLOCK_SIZE = 64
361
- FP8_MAX = 448.0
362
- TOKEN_DATA_SIZE = TOKEN_FP8_DIM + TOKEN_BF16_DIM * 2
357
+ # ========== Copy BF16 portion directly ==========
358
+ bf16_output_offset = fp8_dim # After 448 elements in output
359
364
- num_reqs = seq_lens.shape[0]
365
- NUM_WORKERS = 128
366
- _dequantize_and_gather_k_kernel[(num_reqs, NUM_WORKERS)](
367
- out,
368
- out.stride(0),
369
- out.stride(1),
370
- k_cache,
371
- seq_lens,
372
- block_table,
373
- offset,
374
- gather_lens,
375
- max_blocks_per_seq=block_table.shape[-1],
376
- fp8_dim=TOKEN_FP8_DIM,
377
- bf16_dim=TOKEN_BF16_DIM,
378
- scale_dim=TOKEN_SCALE_DIM,
379
- quant_block=QUANT_BLOCK_SIZE,
380
- cache_block_size=block_size,
381
- token_data_size=TOKEN_DATA_SIZE,
382
- block_stride=k_cache.stride(0),
383
- output_dim=512,
384
- fp8_max=FP8_MAX,
385
- n_quant_blocks=7,
386
- use_fnuz=use_fnuz,
387
- )
360
+ # Read bf16 from cache
361
+ bf16_cache_ptr = token_bf16_ptr.to(tl.pointer_type(tl.bfloat16))
362
+
363
+ # Process in chunks of 16
364
+ for j in tl.static_range(bf16_dim // 16):
365
+ chunk_offsets = j * 16 + tl.arange(0, 16)
366
+ bf16_vals = tl.load(bf16_cache_ptr + chunk_offsets)
367
+ tl.store(output_row_ptr + bf16_output_offset + chunk_offsets, bf16_vals)
368
+
369
+ def dispatch( # type: ignore[override]
370
+ self,
371
+ *,
372
+ max_model_len: int,
373
+ block_table_block_size: int,
374
+ cache_block_size: int,
375
+ offset: int,
376
+ **compile_key_fields: bool,
377
+ ) -> CompileKey:
378
+ token_stride = 576
379
+ scale_stride = 8
380
+ unpadded = cache_block_size * (token_stride + scale_stride)
381
+ block_stride = ((unpadded + token_stride - 1) // token_stride) * token_stride
382
+ return self.CompileKey(
383
+ **compile_key_fields,
384
+ max_blocks_per_seq=(max_model_len + block_table_block_size - 1)
385
+ // block_table_block_size,
386
+ cache_block_size=cache_block_size,
387
+ block_stride=block_stride,
388
+ offset=triton_scalar_specialization_rep(offset),
389
+ )
390
+
391
+ def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]:
392
+ max_model_len = vllm_config.model_config.max_model_len
393
+ block_size = vllm_config.cache_config.block_size
394
+ if max_model_len <= 0 or block_size <= 0:
395
+ return []
396
+
397
+ compress_ratios = frozenset(
398
+ max(1, int(compress_ratio))
399
+ for compress_ratio in vllm_config.model_config.hf_config.compress_ratios
400
+ )
401
+ return self._trace_dispatch(self.dispatch)(
402
+ zip_inputs(
403
+ dict(
404
+ cache_block_size=block_size,
405
+ has_gather_lens=True,
406
+ use_fnuz=current_platform.is_fp8_fnuz(),
407
+ offset=1,
408
+ enabled=True,
409
+ ),
410
+ dict(
411
+ cache_block_size=block_size,
412
+ has_gather_lens=True,
413
+ use_fnuz=current_platform.is_fp8_fnuz(),
414
+ offset=2,
415
+ enabled=True,
416
+ ),
417
+ dict(
418
+ cache_block_size=block_size,
419
+ has_gather_lens=True,
420
+ use_fnuz=current_platform.is_fp8_fnuz(),
421
+ offset=16,
422
+ enabled=True,
423
+ ),
424
+ dict(
425
+ cache_block_size=max(block_size // 4, 1),
426
+ has_gather_lens=False,
427
+ use_fnuz=False,
428
+ offset=16,
429
+ enabled=4 in compress_ratios,
430
+ ),
431
+ dict(
432
+ cache_block_size=max(block_size // 128, 1),
433
+ has_gather_lens=False,
434
+ use_fnuz=False,
435
+ offset=16,
436
+ enabled=128 in compress_ratios,
437
+ ),
438
+ ),
439
+ max_model_len=max_model_len,
440
+ block_table_block_size=block_size,
441
+ _when=lambda *, enabled: enabled,
442
+ )
443
+
444
+ def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
445
+ int32_ptr = TritonWarmupTensor(torch.int32)
446
+ return dict(
447
+ out=TritonWarmupTensor(
448
+ torch.bfloat16,
449
+ shape=(1, 1, 512),
450
+ strides=(512, 512, 1),
451
+ ),
452
+ k_cache=TritonWarmupTensor(
453
+ torch.uint8,
454
+ shape=(1, 1),
455
+ strides=(compile_key.block_stride, 1),
456
+ ),
457
+ seq_lens=int32_ptr,
458
+ gather_lens=(int32_ptr if compile_key.has_gather_lens else None),
459
+ block_table=TritonWarmupTensor(
460
+ torch.int32,
461
+ shape=(1, compile_key.max_blocks_per_seq),
462
+ ),
463
+ block_size=compile_key.cache_block_size,
464
+ offset=compile_key.offset,
465
+ use_fnuz=compile_key.use_fnuz,
466
+ )
467
+
468
+ @kernel_launcher
469
+ def __call__(
470
+ self,
471
+ out: torch.Tensor,
472
+ k_cache: torch.Tensor,
473
+ seq_lens: torch.Tensor,
474
+ gather_lens: torch.Tensor | None,
475
+ block_table: torch.Tensor,
476
+ block_size: int,
477
+ offset: int,
478
+ *,
479
+ use_fnuz: bool = False,
480
+ ) -> LaunchSpec:
481
+ num_reqs = seq_lens.shape[0]
482
+ return (num_reqs, self.NUM_WORKERS), dict(
483
+ out_stride0=out.stride(0),
484
+ out_stride1=out.stride(1),
485
+ max_blocks_per_seq=block_table.shape[-1],
486
+ fp8_dim=448,
487
+ bf16_dim=64,
488
+ scale_dim=8,
489
+ quant_block=64,
490
+ cache_block_size=block_size,
491
+ token_data_size=576,
492
+ block_stride=k_cache.stride(0),
493
+ output_dim=512,
494
+ fp8_max=448.0,
495
+ n_quant_blocks=7,
496
+ )
497
498
499
def dequantize_and_gather_k_cache(
@@ -420,7 +529,7 @@ def dequantize_and_gather_k_cache(
529
)
530
return
531
423
- dequantize_and_gather_k_cache_triton(
532
+ _DEQUANTIZE_AND_GATHER_K_CACHE_KERNEL(
533
out,
534
k_cache,
535
seq_lens,
@@ -449,72 +558,181 @@ def compute_global_topk_indices_and_lens(
558
num_tokens = topk_indices.shape[0]
559
global_topk_indices = torch.empty_like(topk_indices)
560
topk_lens = torch.empty(num_tokens, dtype=torch.int32, device=topk_indices.device)
452
- _compute_global_topk_indices_and_lens_kernel[(num_tokens,)](
561
+ _COMPUTE_GLOBAL_TOPK_INDICES_AND_LENS_KERNEL(
562
global_topk_indices,
454
- global_topk_indices.stride(0),
563
topk_lens,
564
topk_indices,
457
- topk_indices.stride(0),
458
- topk_indices.shape[-1],
565
token_to_req_indices,
566
block_table,
461
- block_table.stride(0),
567
block_size,
568
is_valid_token,
464
- TRITON_BLOCK_SIZE=1024,
569
)
570
return global_topk_indices, topk_lens
571
572
469
-@triton.jit
470
-def _compute_global_topk_indices_and_lens_kernel(
471
- global_topk_indices_ptr,
472
- global_topk_indices_stride: tl.constexpr,
473
- topk_lens_ptr,
474
- topk_indices_ptr,
475
- topk_indices_stride: tl.constexpr,
476
- topk: tl.constexpr,
477
- token_to_req_indices_ptr,
478
- block_table_ptr,
479
- block_table_stride: tl.constexpr,
480
- block_size: tl.constexpr,
481
- is_valid_token_ptr,
482
- TRITON_BLOCK_SIZE: tl.constexpr,
573
+class ComputeGlobalTopkIndicesAndLensKernel(
574
+ VllmTritonJitKernel["ComputeGlobalTopkIndicesAndLensKernel.CompileKey"]
575
):
484
- token_idx = tl.program_id(0)
485
- is_valid_token = tl.load(is_valid_token_ptr + token_idx)
486
- req_idx = tl.load(token_to_req_indices_ptr + token_idx)
487
-
488
- count = tl.zeros((), dtype=tl.int32)
489
- for i in range(0, topk, TRITON_BLOCK_SIZE):
490
- offset = i + tl.arange(0, TRITON_BLOCK_SIZE)
491
- mask = offset < topk
492
-
493
- local_idx = tl.load(
494
- topk_indices_ptr + token_idx * topk_indices_stride + offset,
495
- mask=mask,
496
- other=-1,
576
+ TRITON_BLOCK_SIZE = 1024
577
+
578
+ @dataclass(frozen=True)
579
+ class CompileKey:
580
+ global_topk_indices_stride: int
581
+ topk_indices_stride: int
582
+ topk: int
583
+ block_table_stride: int
584
+ block_size: int
585
+
586
+ @staticmethod
587
+ @triton.jit
588
+ def kernel(
589
+ global_topk_indices_ptr,
590
+ global_topk_indices_stride: tl.constexpr,
591
+ topk_lens_ptr,
592
+ topk_indices_ptr,
593
+ topk_indices_stride: tl.constexpr,
594
+ topk: tl.constexpr,
595
+ token_to_req_indices_ptr,
596
+ block_table_ptr,
597
+ block_table_stride: tl.constexpr,
598
+ block_size: tl.constexpr,
599
+ is_valid_token_ptr,
600
+ TRITON_BLOCK_SIZE: tl.constexpr,
601
+ ):
602
+ token_idx = tl.program_id(0)
603
+ is_valid_token = tl.load(is_valid_token_ptr + token_idx)
604
+ req_idx = tl.load(token_to_req_indices_ptr + token_idx)
605
+
606
+ count = tl.zeros((), dtype=tl.int32)
607
+ for i in range(0, topk, TRITON_BLOCK_SIZE):
608
+ offset = i + tl.arange(0, TRITON_BLOCK_SIZE)
609
+ mask = offset < topk
610
+
611
+ local_idx = tl.load(
612
+ topk_indices_ptr + token_idx * topk_indices_stride + offset,
613
+ mask=mask,
614
+ other=-1,
615
+ )
616
+ is_valid = local_idx >= 0
617
+
618
+ block_indices = local_idx // block_size
619
+ block_numbers = tl.load(
620
+ block_table_ptr + req_idx * block_table_stride + block_indices,
621
+ mask=mask & is_valid,
622
+ )
623
+ block_offsets = local_idx % block_size
624
+
625
+ slot_ids = block_numbers * block_size + block_offsets
626
+ slot_ids = tl.where(is_valid, slot_ids, -1)
627
+ tl.store(
628
+ global_topk_indices_ptr
629
+ + token_idx * global_topk_indices_stride
630
+ + offset,
631
+ slot_ids,
632
+ mask=mask,
633
+ )
634
+ count += tl.sum(is_valid.to(tl.int32), axis=0)
635
+
636
+ # Zero out length for padding tokens.
637
+ tl.store(topk_lens_ptr + token_idx, tl.where(is_valid_token, count, 0))
638
+
639
+ def dispatch( # type: ignore[override]
640
+ self,
641
+ *,
642
+ topk_width: int,
643
+ block_size: int,
644
+ block_table_block_size: int,
645
+ max_model_len: int,
646
+ ) -> CompileKey:
647
+ return self.CompileKey(
648
+ global_topk_indices_stride=topk_width,
649
+ topk_indices_stride=topk_width,
650
+ topk=topk_width,
651
+ block_table_stride=(max_model_len + block_table_block_size - 1)
652
+ // block_table_block_size,
653
+ block_size=block_size,
654
)
498
- is_valid = local_idx >= 0
655
500
- block_indices = local_idx // block_size
501
- block_numbers = tl.load(
502
- block_table_ptr + req_idx * block_table_stride + block_indices,
503
- mask=mask & is_valid,
656
+ def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]:
657
+ index_topk = vllm_config.model_config.hf_config.index_topk
658
+ compress_ratios = vllm_config.model_config.hf_config.compress_ratios
659
+ cache_block_size = vllm_config.cache_config.block_size
660
+ if index_topk <= 0 or cache_block_size <= 0:
661
+ return []
662
+
663
+ max_c128a_topk = (
664
+ ((vllm_config.model_config.max_model_len + 127) // 128 + 127) // 128 * 128
665
+ )
666
+ active_c128a_topk_widths = WarmupIntRange(
667
+ 128,
668
+ max_c128a_topk + 1,
669
+ advance=lambda width: (
670
+ min(width * 2, max_c128a_topk)
671
+ if width < max_c128a_topk
672
+ else max_c128a_topk + 1
673
+ ),
674
)
505
- block_offsets = local_idx % block_size
506
-
507
- slot_ids = block_numbers * block_size + block_offsets
508
- slot_ids = tl.where(is_valid, slot_ids, -1)
509
- tl.store(
510
- global_topk_indices_ptr + token_idx * global_topk_indices_stride + offset,
511
- slot_ids,
512
- mask=mask,
675
+ trace_dispatch = self._trace_dispatch(self.dispatch)
676
+ c4_keys = trace_dispatch(
677
+ topk_width=index_topk if 4 in compress_ratios else (),
678
+ block_size=max(1, cache_block_size // 4),
679
+ block_table_block_size=cache_block_size,
680
+ max_model_len=vllm_config.model_config.max_model_len,
681
)
514
- count += tl.sum(is_valid.to(tl.int32), axis=0)
682
+ c128a_keys = trace_dispatch(
683
+ topk_width=active_c128a_topk_widths if 128 in compress_ratios else (),
684
+ block_size=max(1, cache_block_size // 128),
685
+ block_table_block_size=cache_block_size,
686
+ max_model_len=vllm_config.model_config.max_model_len,
687
+ )
688
+ return list(dict.fromkeys((*c4_keys, *c128a_keys)))
689
516
- # Zero out length for padding tokens.
517
- tl.store(topk_lens_ptr + token_idx, tl.where(is_valid_token, count, 0))
690
+ def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
691
+ int32_ptr = TritonWarmupTensor(torch.int32)
692
+ return dict(
693
+ global_topk_indices=TritonWarmupTensor(
694
+ torch.int32,
695
+ shape=(1, compile_key.topk),
696
+ strides=(compile_key.global_topk_indices_stride, 1),
697
+ ),
698
+ topk_lens=int32_ptr,
699
+ topk_indices=TritonWarmupTensor(
700
+ torch.int32,
701
+ shape=(1, compile_key.topk),
702
+ strides=(compile_key.topk_indices_stride, 1),
703
+ ),
704
+ token_to_req_indices=int32_ptr,
705
+ block_table=TritonWarmupTensor(
706
+ torch.int32,
707
+ shape=(1, compile_key.block_table_stride),
708
+ ),
709
+ block_size=compile_key.block_size,
710
+ is_valid_token=TritonWarmupTensor(torch.bool),
711
+ )
712
+
713
+ @kernel_launcher
714
+ def __call__(
715
+ self,
716
+ global_topk_indices: torch.Tensor,
717
+ topk_lens: torch.Tensor,
718
+ topk_indices: torch.Tensor,
719
+ token_to_req_indices: torch.Tensor,
720
+ block_table: torch.Tensor,
721
+ block_size: int,
722
+ is_valid_token: torch.Tensor,
723
+ ) -> LaunchSpec:
724
+ num_tokens = topk_indices.shape[0]
725
+ return (num_tokens,), dict(
726
+ global_topk_indices_stride=global_topk_indices.stride(0),
727
+ topk_indices_stride=topk_indices.stride(0),
728
+ topk=topk_indices.shape[-1],
729
+ block_table_stride=block_table.stride(0),
730
+ TRITON_BLOCK_SIZE=self.TRITON_BLOCK_SIZE,
731
+ )
732
+
733
+
734
+_DEQUANTIZE_AND_GATHER_K_CACHE_KERNEL = DequantizeAndGatherKCacheKernel()
735
+_COMPUTE_GLOBAL_TOPK_INDICES_AND_LENS_KERNEL = ComputeGlobalTopkIndicesAndLensKernel()
736
737
738
# FlashMLA sparse prefill asserts `params.topk % B_TOPK == 0` (see
@@ -582,43 +800,6 @@ def combine_topk_swa_indices(
800
return combined_indices, combined_lens
801
802
585
-_COMBINE_TOPK_SWA_NUM_WORKERS = 256
586
-
587
-
588
-# Representative pointer alignment variants for Triton pointer specialization.
589
-_COMBINE_TOPK_SWA_POINTER_INPUTS = zip_inputs(
590
- dict(
591
- topk_indices=True,
592
- query_start_loc=True,
593
- seq_lens=True,
594
- gather_lens=True,
595
- ),
596
- dict(
597
- topk_indices=True,
598
- query_start_loc=False,
599
- seq_lens=False,
600
- gather_lens=True,
601
- ),
602
- dict(
603
- topk_indices=False,
604
- query_start_loc=False,
605
- seq_lens=False,
606
- gather_lens=False,
607
- ),
608
-)
609
-
610
-
611
-_DSV4_COMBINE_TOPK_SWA_WARMUP_INPUTS = zip_inputs(
612
- # DSv4-Flash / SWA-only and C4A.
613
- dict(compress_ratio=1, topk=0, topk_width=512),
614
- dict(compress_ratio=4, topk=512, topk_width=512),
615
- # DSv4-Pro C4A.
616
- dict(compress_ratio=4, topk=1024, topk_width=1024),
617
- # DSv4-Pro C128A.
618
- dict(compress_ratio=128, topk=8192, topk_width=8192),
619
-)
620
-
621
-
803
def _hf_config_int(vllm_config: Any, name: str, default: int) -> int:
804
model_config = getattr(vllm_config, "model_config", None)
805
hf_config = getattr(model_config, "hf_config", None)
@@ -631,8 +812,10 @@ def _scheduler_config_int(vllm_config: Any, name: str, default: int) -> int:
812
813
814
class CombineTopkSwaIndicesKernel(
634
- VllmJitKernel["CombineTopkSwaIndicesKernel.CompileKey"]
815
+ VllmTritonJitKernel["CombineTopkSwaIndicesKernel.CompileKey"]
816
):
817
+ NUM_WORKERS = 256
818
+
819
@dataclass(frozen=True)
820
class CompileKey:
821
TOP_K: int
@@ -746,11 +929,14 @@ class CombineTopkSwaIndicesKernel(
929
query_start_loc: bool,
930
seq_lens: bool,
931
gather_lens: bool,
749
- topk: int,
932
compress_ratio: int,
933
+ index_topk: int,
934
+ active_topk_width: int,
935
WINDOW_SIZE: int,
936
image_width: int = 0,
937
) -> CompileKey:
938
+ topk_width = active_topk_width if compress_ratio == 128 else index_topk
939
+ topk = 0 if compress_ratio == 1 else topk_width
940
padded_topk = next_power_of_2(topk_width)
941
input_variant = TritonPointerInputVariant.from_alignment(
942
topk_indices=topk_indices,
@@ -772,6 +958,23 @@ class CombineTopkSwaIndicesKernel(
958
return []
959
960
window_size = _hf_config_int(vllm_config, "sliding_window", 128)
961
+ hf_config = vllm_config.model_config.hf_config
962
+ index_topk = int(getattr(hf_config, "index_topk", 0) or 0)
963
+ compress_ratios = tuple(
964
+ sorted({max(1, int(ratio)) for ratio in hf_config.compress_ratios})
965
+ )
966
+ max_c128a_topk = (
967
+ ((vllm_config.model_config.max_model_len + 127) // 128 + 127) // 128 * 128
968
+ )
969
+ active_topk_widths = WarmupIntRange(
970
+ 128,
971
+ max_c128a_topk + 1,
972
+ advance=lambda width: (
973
+ min(width * 2, max_c128a_topk)
974
+ if width < max_c128a_topk
975
+ else max_c128a_topk + 1
976
+ ),
977
+ )
978
image_width = (
979
_hf_config_int(vllm_config, "vision_max_n_token", 0)
980
if _hf_config_int(vllm_config, "vision_n_layers", 0) > 0
@@ -781,38 +984,59 @@ class CombineTopkSwaIndicesKernel(
984
# and the in-image bidirectional variant.
985
image_widths = [0, image_width] if image_width > 0 else [0]
986
return self._trace_dispatch(self.dispatch)(
784
- _DSV4_COMBINE_TOPK_SWA_WARMUP_INPUTS,
785
- _COMBINE_TOPK_SWA_POINTER_INPUTS,
987
+ zip_inputs(
988
+ dict(
989
+ topk_indices=True,
990
+ query_start_loc=True,
991
+ seq_lens=True,
992
+ gather_lens=True,
993
+ ),
994
+ dict(
995
+ topk_indices=True,
996
+ query_start_loc=False,
997
+ seq_lens=False,
998
+ gather_lens=True,
999
+ ),
1000
+ dict(
1001
+ topk_indices=False,
1002
+ query_start_loc=False,
1003
+ seq_lens=False,
1004
+ gather_lens=False,
1005
+ ),
1006
+ ),
1007
+ compress_ratio=compress_ratios,
1008
+ index_topk=index_topk,
1009
+ active_topk_width=active_topk_widths,
1010
WINDOW_SIZE=window_size,
1011
image_width=image_widths,
1012
)
1013
790
- def compile(self, compile_key: CompileKey) -> None:
791
- warmup = getattr(self.kernel, "warmup", None)
792
- assert warmup is not None
1014
+ def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
1015
int32_ptr = TritonWarmupTensor(torch.int32)
794
- input_variant = compile_key.input_variant
795
- warmup(
796
- int32_ptr,
797
- 1, # do not specialize combined_indices_stride
798
- int32_ptr,
799
- input_variant.pointer("topk_indices", torch.int32),
800
- 1, # do not specialize topk_indices_stride
801
- input_variant.pointer("query_start_loc", torch.int32),
802
- input_variant.pointer("seq_lens", torch.int32),
803
- input_variant.pointer("gather_lens", torch.int32),
804
- int32_ptr,
805
- int32_ptr,
806
- 1, # do not specialize M
807
- 1, # do not specialize N
1016
+ variant = compile_key.input_variant
1017
+ has_image = compile_key.IMAGE_WIDTH > 0
1018
+ return dict(
1019
+ combined_indices=int32_ptr,
1020
+ combined_lens=int32_ptr,
1021
+ topk_indices=variant.pointer(
1022
+ "topk_indices", torch.int32, shape=(1, compile_key.PADDED_TOP_K)
1023
+ ),
1024
+ query_start_loc=variant.pointer("query_start_loc", torch.int32),
1025
+ seq_lens=variant.pointer("seq_lens", torch.int32),
1026
+ gather_lens=variant.pointer("gather_lens", torch.int32),
1027
+ M=1,
1028
+ N=1,
1029
TOP_K=compile_key.TOP_K,
1030
COMPRESS_RATIO=compile_key.COMPRESS_RATIO,
1031
WINDOW_SIZE=compile_key.WINDOW_SIZE,
811
- IMAGE_WIDTH=compile_key.IMAGE_WIDTH,
812
- PADDED_TOP_K=compile_key.PADDED_TOP_K,
813
- grid=(1, _COMBINE_TOPK_SWA_NUM_WORKERS),
1032
+ # Reproduce the runtime substitution: no-image batches launch with
1033
+ # ``topk_indices`` in the visibility pointer slots (see __call__).
1034
+ left_visible=(int32_ptr if has_image else None),
1035
+ right_visible=(int32_ptr if has_image else None),
1036
+ max_image_tokens=compile_key.IMAGE_WIDTH,
1037
)
1038
1039
+ @kernel_launcher
1040
def __call__(
1041
self,
1042
combined_indices: torch.Tensor,
@@ -830,23 +1054,15 @@ class CombineTopkSwaIndicesKernel(
1054
left_visible: torch.Tensor | None = None,
1055
right_visible: torch.Tensor | None = None,
1056
max_image_tokens: int = 0,
833
- ) -> None:
1057
+ ) -> LaunchSpec:
1058
num_reqs = seq_lens.shape[0]
1059
has_image = left_visible is not None
1060
image_width = max_image_tokens if has_image else 0
837
- self.kernel[(num_reqs, _COMBINE_TOPK_SWA_NUM_WORKERS)](
838
- combined_indices,
839
- combined_indices.stride(0),
840
- combined_lens,
841
- topk_indices,
842
- topk_indices.stride(0),
843
- query_start_loc,
844
- seq_lens,
845
- gather_lens,
846
- left_visible if has_image else topk_indices,
847
- right_visible if has_image else topk_indices,
848
- M,
849
- N,
1061
+ return (num_reqs, self.NUM_WORKERS), dict(
1062
+ combined_indices_stride=combined_indices.stride(0),
1063
+ topk_indices_stride=topk_indices.stride(0),
1064
+ left_visible_ptr=left_visible if has_image else topk_indices,
1065
+ right_visible_ptr=right_visible if has_image else topk_indices,
1066
TOP_K=TOP_K,
1067
COMPRESS_RATIO=COMPRESS_RATIO,
1068
WINDOW_SIZE=WINDOW_SIZE,
@@ -974,44 +1190,38 @@ def build_flashinfer_mixed_sparse_indices(
1190
if compressed_block_span is None
1191
else compressed_block_span
1192
)
977
- _build_flashinfer_mixed_sparse_indices_kernel[(num_tokens,)](
1193
+ _BUILD_FLASHINFER_MIXED_SPARSE_INDICES_KERNEL(
1194
sparse_indices,
979
- sparse_indices.stride(0),
1195
sparse_topk_lens,
1196
decode_swa_indices,
982
- decode_swa_indices.stride(0),
1197
decode_compressed_indices,
984
- decode_compressed_indices.stride(0),
1198
decode_compressed_topk_lens,
1199
decode_is_valid_token,
1200
prefill_topk_indices,
988
- prefill_topk_indices.stride(0),
1201
prefill_left_visible if has_image else token_to_req_indices,
1202
prefill_right_visible if has_image else token_to_req_indices,
1203
query_start_loc,
1204
seq_lens,
1205
token_to_req_indices,
1206
swa_block_table,
995
- swa_block_table.stride(0),
1207
swa_block_size,
1208
swa_span,
1209
compressed_block_table,
999
- compressed_block_table.stride(0),
1210
compressed_block_size,
1211
compressed_span,
1002
- NUM_DECODE_TOKENS=num_decode_tokens,
1003
- WINDOW_SIZE=window_size,
1004
- SWA_INDEX_WIDTH=swa_index_width,
1005
- IMAGE_WIDTH=image_width,
1006
- COMPRESS_RATIO=compress_ratio,
1007
- TOP_K=topk,
1008
- PADDED_TOP_K=padded_topk,
1009
- PREFILL_TOPK_STRIDE=prefill_topk_indices.shape[-1],
1010
- DECODE_COMPRESSED_TOPK=decode_compressed_topk,
1011
- DECODE_COMPRESSED_INDICES_ARE_LOCAL=decode_compressed_indices_are_local,
1012
- HAS_DECODE_COMPRESSED_LENS=has_decode_compressed_lens,
1013
- WINDOW_BLOCK_SIZE=window_block_size,
1014
- TOPK_BLOCK_SIZE=topk_block_size,
1212
+ num_decode_tokens=num_decode_tokens,
1213
+ window_size=window_size,
1214
+ swa_index_width=swa_index_width,
1215
+ image_width=image_width,
1216
+ compress_ratio=compress_ratio,
1217
+ topk=topk,
1218
+ padded_topk=padded_topk,
1219
+ prefill_topk_stride=prefill_topk_indices.shape[-1],
1220
+ decode_compressed_topk=decode_compressed_topk,
1221
+ decode_compressed_indices_are_local=decode_compressed_indices_are_local,
1222
+ has_decode_compressed_lens=has_decode_compressed_lens,
1223
+ window_block_size=window_block_size,
1224
+ topk_block_size=topk_block_size,
1225
num_warps=num_warps,
1226
)
1227
return sparse_indices, sparse_topk_lens
@@ -1029,205 +1239,444 @@ def _remap_flashinfer_index(values, block_size, block_span):
1239
return tl.where(is_valid, values, -1)
1240
1241
1032
-@triton.jit(
1033
- do_not_specialize=[
1034
- "sparse_indices_stride",
1035
- "decode_swa_stride",
1036
- "decode_compressed_stride",
1037
- "prefill_topk_stride",
1038
- "swa_block_table_stride",
1039
- "swa_block_size",
1040
- "swa_block_span",
1041
- "compressed_block_table_stride",
1042
- "compressed_block_size",
1043
- "compressed_block_span",
1044
- "NUM_DECODE_TOKENS",
1045
- "PREFILL_TOPK_STRIDE",
1046
- ]
1047
-)
1048
-def _build_flashinfer_mixed_sparse_indices_kernel(
1049
- sparse_indices_ptr,
1050
- sparse_indices_stride,
1051
- sparse_topk_lens_ptr,
1052
- decode_swa_indices_ptr,
1053
- decode_swa_stride,
1054
- decode_compressed_indices_ptr,
1055
- decode_compressed_stride,
1056
- decode_compressed_topk_lens_ptr,
1057
- decode_is_valid_token_ptr,
1058
- prefill_topk_indices_ptr,
1059
- prefill_topk_stride,
1060
- left_visible_ptr,
1061
- right_visible_ptr,
1062
- query_start_loc_ptr,
1063
- seq_lens_ptr,
1064
- token_to_req_indices_ptr,
1065
- swa_block_table_ptr,
1066
- swa_block_table_stride,
1067
- swa_block_size,
1068
- swa_block_span,
1069
- compressed_block_table_ptr,
1070
- compressed_block_table_stride,
1071
- compressed_block_size,
1072
- compressed_block_span,
1073
- NUM_DECODE_TOKENS,
1074
- WINDOW_SIZE: tl.constexpr,
1075
- SWA_INDEX_WIDTH: tl.constexpr,
1076
- IMAGE_WIDTH: tl.constexpr,
1077
- COMPRESS_RATIO: tl.constexpr,
1078
- TOP_K: tl.constexpr,
1079
- PADDED_TOP_K: tl.constexpr,
1080
- PREFILL_TOPK_STRIDE,
1081
- DECODE_COMPRESSED_TOPK: tl.constexpr,
1082
- DECODE_COMPRESSED_INDICES_ARE_LOCAL: tl.constexpr,
1083
- HAS_DECODE_COMPRESSED_LENS: tl.constexpr,
1084
- WINDOW_BLOCK_SIZE: tl.constexpr,
1085
- TOPK_BLOCK_SIZE: tl.constexpr,
1242
+class BuildFlashinferMixedSparseIndicesKernel(
1243
+ VllmTritonJitKernel["BuildFlashinferMixedSparseIndicesKernel.CompileKey"]
1244
):
1087
- token_idx = tl.program_id(0)
1088
- # Total SWA column count; > SWA_INDEX_WIDTH only for in-image
1089
- # bidirectional visibility (vision variant), where the extra columns are
1090
- # -1-padded for decode rows.
1091
- SWA_TOTAL_WIDTH: tl.constexpr = SWA_INDEX_WIDTH + IMAGE_WIDTH
1245
+ @dataclass(frozen=True)
1246
+ class CompileKey:
1247
+ window_size: int
1248
+ swa_index_width: int
1249
+ image_width: int
1250
+ compress_ratio: int
1251
+ top_k: int
1252
+ padded_top_k: int
1253
+ decode_compressed_topk: int
1254
+ decode_compressed_indices_are_local: bool
1255
+ has_decode_compressed_lens: bool
1256
+ window_block_size: int
1257
+ topk_block_size: int
1258
+
1259
+ @staticmethod
1260
+ @triton.jit(
1261
+ do_not_specialize=[
1262
+ "sparse_indices_stride",
1263
+ "decode_swa_stride",
1264
+ "decode_compressed_stride",
1265
+ "prefill_topk_stride",
1266
+ "swa_block_table_stride",
1267
+ "swa_block_size",
1268
+ "swa_block_span",
1269
+ "compressed_block_table_stride",
1270
+ "compressed_block_size",
1271
+ "compressed_block_span",
1272
+ "NUM_DECODE_TOKENS",
1273
+ "PREFILL_TOPK_STRIDE",
1274
+ ]
1275
+ )
1276
+ def kernel(
1277
+ sparse_indices_ptr,
1278
+ sparse_indices_stride,
1279
+ sparse_topk_lens_ptr,
1280
+ decode_swa_indices_ptr,
1281
+ decode_swa_stride,
1282
+ decode_compressed_indices_ptr,
1283
+ decode_compressed_stride,
1284
+ decode_compressed_topk_lens_ptr,
1285
+ decode_is_valid_token_ptr,
1286
+ prefill_topk_indices_ptr,
1287
+ prefill_topk_stride,
1288
+ left_visible_ptr,
1289
+ right_visible_ptr,
1290
+ query_start_loc_ptr,
1291
+ seq_lens_ptr,
1292
+ token_to_req_indices_ptr,
1293
+ swa_block_table_ptr,
1294
+ swa_block_table_stride,
1295
+ swa_block_size,
1296
+ swa_block_span,
1297
+ compressed_block_table_ptr,
1298
+ compressed_block_table_stride,
1299
+ compressed_block_size,
1300
+ compressed_block_span,
1301
+ NUM_DECODE_TOKENS,
1302
+ WINDOW_SIZE: tl.constexpr,
1303
+ SWA_INDEX_WIDTH: tl.constexpr,
1304
+ IMAGE_WIDTH: tl.constexpr,
1305
+ COMPRESS_RATIO: tl.constexpr,
1306
+ TOP_K: tl.constexpr,
1307
+ PADDED_TOP_K: tl.constexpr,
1308
+ PREFILL_TOPK_STRIDE,
1309
+ DECODE_COMPRESSED_TOPK: tl.constexpr,
1310
+ DECODE_COMPRESSED_INDICES_ARE_LOCAL: tl.constexpr,
1311
+ HAS_DECODE_COMPRESSED_LENS: tl.constexpr,
1312
+ WINDOW_BLOCK_SIZE: tl.constexpr,
1313
+ TOPK_BLOCK_SIZE: tl.constexpr,
1314
+ ):
1315
+ token_idx = tl.program_id(0)
1316
+ # Total SWA column count; > SWA_INDEX_WIDTH only for in-image
1317
+ # bidirectional visibility (vision variant), where the extra columns are
1318
+ # -1-padded for decode rows.
1319
+ SWA_TOTAL_WIDTH: tl.constexpr = SWA_INDEX_WIDTH + IMAGE_WIDTH
1320
+
1321
+ if token_idx < NUM_DECODE_TOKENS:
1322
+ for i in range(0, SWA_TOTAL_WIDTH, WINDOW_BLOCK_SIZE):
1323
+ offset = i + tl.arange(0, WINDOW_BLOCK_SIZE)
1324
+ mask = offset < SWA_TOTAL_WIDTH
1325
+ values = tl.load(
1326
+ decode_swa_indices_ptr + token_idx * decode_swa_stride + offset,
1327
+ mask=offset < SWA_INDEX_WIDTH,
1328
+ other=-1,
1329
+ )
1330
+ values = _remap_flashinfer_index(values, swa_block_size, swa_block_span)
1331
+ tl.store(
1332
+ sparse_indices_ptr + token_idx * sparse_indices_stride + offset,
1333
+ values,
1334
+ mask=mask,
1335
+ )
1336
+
1337
+ compressed_len = tl.zeros((), dtype=tl.int32)
1338
+ for i in range(0, PADDED_TOP_K, TOPK_BLOCK_SIZE):
1339
+ offset = i + tl.arange(0, TOPK_BLOCK_SIZE)
1340
+ mask = offset < PADDED_TOP_K
1341
+ values = tl.load(
1342
+ decode_compressed_indices_ptr
1343
+ + token_idx * decode_compressed_stride
1344
+ + offset,
1345
+ mask=offset < DECODE_COMPRESSED_TOPK,
1346
+ other=-1,
1347
+ )
1348
+ if DECODE_COMPRESSED_INDICES_ARE_LOCAL:
1349
+ token_valid = tl.load(decode_is_valid_token_ptr + token_idx)
1350
+ is_valid = values >= 0
1351
+ req_idx = tl.load(token_to_req_indices_ptr + token_idx)
1352
+ block_indices = values // compressed_block_size
1353
+ block_numbers = tl.load(
1354
+ compressed_block_table_ptr
1355
+ + req_idx * compressed_block_table_stride
1356
+ + block_indices,
1357
+ mask=mask & is_valid,
1358
+ other=-1,
1359
+ )
1360
+ block_offsets = values % compressed_block_size
1361
+ values = block_numbers * compressed_block_size + block_offsets
1362
+ values = tl.where(is_valid, values, -1)
1363
+ compressed_len += tl.sum(
1364
+ (is_valid & token_valid).to(tl.int32), axis=0
1365
+ )
1366
+ values = _remap_flashinfer_index(
1367
+ values, compressed_block_size, compressed_block_span
1368
+ )
1369
+ tl.store(
1370
+ sparse_indices_ptr
1371
+ + token_idx * sparse_indices_stride
1372
+ + SWA_TOTAL_WIDTH
1373
+ + offset,
1374
+ values,
1375
+ mask=mask,
1376
+ )
1377
+
1378
+ if DECODE_COMPRESSED_TOPK == 0:
1379
+ compressed_len = tl.zeros((), dtype=tl.int32)
1380
+ elif not DECODE_COMPRESSED_INDICES_ARE_LOCAL:
1381
+ if HAS_DECODE_COMPRESSED_LENS:
1382
+ compressed_len = tl.load(
1383
+ decode_compressed_topk_lens_ptr + token_idx
1384
+ )
1385
+ else:
1386
+ compressed_len = tl.full((), DECODE_COMPRESSED_TOPK, dtype=tl.int32)
1387
+
1388
+ tl.store(sparse_topk_lens_ptr + token_idx, SWA_TOTAL_WIDTH + compressed_len)
1389
+ return
1390
+
1391
+ prefill_idx = token_idx - NUM_DECODE_TOKENS
1392
+ req_idx = tl.load(token_to_req_indices_ptr + token_idx)
1393
+ query_start = tl.load(query_start_loc_ptr + req_idx)
1394
+ query_end = tl.load(query_start_loc_ptr + req_idx + 1)
1395
+ query_len = query_end - query_start
1396
+ seq_len = tl.load(seq_lens_ptr + req_idx)
1397
+ start_pos = seq_len - query_len
1398
+ token_idx_in_query = token_idx - query_start
1399
+ pos = start_pos + token_idx_in_query
1400
+ if IMAGE_WIDTH > 0:
1401
+ # In-image bidirectional visibility: window starts up to
1402
+ # max(left - (window - 1), 0) earlier and extends `right` past pos.
1403
+ left = tl.load(left_visible_ptr + token_idx)
1404
+ right = tl.load(right_visible_ptr + token_idx)
1405
+ else:
1406
+ left = 0
1407
+ right = 0
1408
+ left_add = tl.maximum(left - (WINDOW_SIZE - 1), 0)
1409
+ swa_start_pos = tl.maximum(pos - (WINDOW_SIZE - 1) - left_add, 0)
1410
+ swa_len = pos + right - swa_start_pos + 1
1411
+ topk_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOP_K)
1412
1093
- if token_idx < NUM_DECODE_TOKENS:
1413
for i in range(0, SWA_TOTAL_WIDTH, WINDOW_BLOCK_SIZE):
1414
offset = i + tl.arange(0, WINDOW_BLOCK_SIZE)
1415
mask = offset < SWA_TOTAL_WIDTH
1097
- values = tl.load(
1098
- decode_swa_indices_ptr + token_idx * decode_swa_stride + offset,
1099
- mask=offset < SWA_INDEX_WIDTH,
1416
+ pos_offset = swa_start_pos + offset
1417
+ block_indices = pos_offset // swa_block_size
1418
+ block_numbers = tl.load(
1419
+ swa_block_table_ptr + req_idx * swa_block_table_stride + block_indices,
1420
+ mask=mask & (offset < swa_len),
1421
other=-1,
1422
)
1102
- values = _remap_flashinfer_index(values, swa_block_size, swa_block_span)
1423
+ block_offsets = pos_offset % swa_block_size
1424
+ slot_ids = block_numbers * swa_block_size + block_offsets
1425
+ slot_ids = tl.where(offset < swa_len, slot_ids, -1)
1426
+ slot_ids = _remap_flashinfer_index(slot_ids, swa_block_size, swa_block_span)
1427
tl.store(
1428
sparse_indices_ptr + token_idx * sparse_indices_stride + offset,
1105
- values,
1429
+ slot_ids,
1430
mask=mask,
1431
)
1432
1109
- compressed_len = tl.zeros((), dtype=tl.int32)
1433
for i in range(0, PADDED_TOP_K, TOPK_BLOCK_SIZE):
1434
offset = i + tl.arange(0, TOPK_BLOCK_SIZE)
1435
mask = offset < PADDED_TOP_K
1113
- values = tl.load(
1114
- decode_compressed_indices_ptr
1115
- + token_idx * decode_compressed_stride
1116
- + offset,
1117
- mask=offset < DECODE_COMPRESSED_TOPK,
1436
+ local_idx = tl.load(
1437
+ prefill_topk_indices_ptr + prefill_idx * prefill_topk_stride + offset,
1438
+ mask=(offset < PREFILL_TOPK_STRIDE) & (offset < topk_len),
1439
other=-1,
1440
)
1120
- if DECODE_COMPRESSED_INDICES_ARE_LOCAL:
1121
- token_valid = tl.load(decode_is_valid_token_ptr + token_idx)
1122
- is_valid = values >= 0
1123
- req_idx = tl.load(token_to_req_indices_ptr + token_idx)
1124
- block_indices = values // compressed_block_size
1125
- block_numbers = tl.load(
1126
- compressed_block_table_ptr
1127
- + req_idx * compressed_block_table_stride
1128
- + block_indices,
1129
- mask=mask & is_valid,
1130
- other=-1,
1131
- )
1132
- block_offsets = values % compressed_block_size
1133
- values = block_numbers * compressed_block_size + block_offsets
1134
- values = tl.where(is_valid, values, -1)
1135
- compressed_len += tl.sum((is_valid & token_valid).to(tl.int32), axis=0)
1136
- values = _remap_flashinfer_index(
1137
- values, compressed_block_size, compressed_block_span
1441
+ is_valid = local_idx >= 0
1442
+ block_indices = local_idx // compressed_block_size
1443
+ block_numbers = tl.load(
1444
+ compressed_block_table_ptr
1445
+ + req_idx * compressed_block_table_stride
1446
+ + block_indices,
1447
+ mask=mask & is_valid,
1448
+ other=-1,
1449
+ )
1450
+ block_offsets = local_idx % compressed_block_size
1451
+ slot_ids = block_numbers * compressed_block_size + block_offsets
1452
+ slot_ids = tl.where((offset < topk_len) & is_valid, slot_ids, -1)
1453
+ slot_ids = _remap_flashinfer_index(
1454
+ slot_ids, compressed_block_size, compressed_block_span
1455
)
1456
tl.store(
1457
sparse_indices_ptr
1458
+ token_idx * sparse_indices_stride
1459
+ SWA_TOTAL_WIDTH
1460
+ offset,
1144
- values,
1461
+ slot_ids,
1462
mask=mask,
1463
)
1464
1148
- if DECODE_COMPRESSED_TOPK == 0:
1149
- compressed_len = tl.zeros((), dtype=tl.int32)
1150
- elif not DECODE_COMPRESSED_INDICES_ARE_LOCAL:
1151
- if HAS_DECODE_COMPRESSED_LENS:
1152
- compressed_len = tl.load(decode_compressed_topk_lens_ptr + token_idx)
1153
- else:
1154
- compressed_len = tl.full((), DECODE_COMPRESSED_TOPK, dtype=tl.int32)
1465
+ tl.store(sparse_topk_lens_ptr + token_idx, SWA_TOTAL_WIDTH + topk_len)
1466
1156
- tl.store(sparse_topk_lens_ptr + token_idx, SWA_TOTAL_WIDTH + compressed_len)
1157
- return
1467
+ def dispatch( # type: ignore[override]
1468
+ self,
1469
+ *,
1470
+ swa_index_width: int,
1471
+ image_width: int,
1472
+ compress_ratio: int,
1473
+ index_topk: int,
1474
+ active_topk_width: int,
1475
+ has_prefill: bool,
1476
+ **compile_key_fields: Any,
1477
+ ) -> CompileKey:
1478
+ decode_compressed_topk = (
1479
+ active_topk_width
1480
+ if compress_ratio == 128
1481
+ else index_topk
1482
+ if compress_ratio == 4
1483
+ else 0
1484
+ )
1485
+ topk_width = decode_compressed_topk if has_prefill else 0
1486
+ padded_top_k = max(topk_width, decode_compressed_topk)
1487
+ padded_top_k = (padded_top_k + 3) // 4 * 4
1488
+ return self.CompileKey(
1489
+ **compile_key_fields,
1490
+ compress_ratio=compress_ratio,
1491
+ top_k=topk_width,
1492
+ decode_compressed_indices_are_local=compress_ratio == 4,
1493
+ has_decode_compressed_lens=compress_ratio == 128,
1494
+ swa_index_width=swa_index_width,
1495
+ image_width=image_width,
1496
+ padded_top_k=padded_top_k,
1497
+ decode_compressed_topk=decode_compressed_topk,
1498
+ window_block_size=next_power_of_2(max(swa_index_width, 1)),
1499
+ topk_block_size=next_power_of_2(max(padded_top_k, 1)),
1500
+ )
1501
1159
- prefill_idx = token_idx - NUM_DECODE_TOKENS
1160
- req_idx = tl.load(token_to_req_indices_ptr + token_idx)
1161
- query_start = tl.load(query_start_loc_ptr + req_idx)
1162
- query_end = tl.load(query_start_loc_ptr + req_idx + 1)
1163
- query_len = query_end - query_start
1164
- seq_len = tl.load(seq_lens_ptr + req_idx)
1165
- start_pos = seq_len - query_len
1166
- token_idx_in_query = token_idx - query_start
1167
- pos = start_pos + token_idx_in_query
1168
- if IMAGE_WIDTH > 0:
1169
- # In-image bidirectional visibility: window starts up to
1170
- # max(left - (window - 1), 0) earlier and extends `right` past pos.
1171
- left = tl.load(left_visible_ptr + token_idx)
1172
- right = tl.load(right_visible_ptr + token_idx)
1173
- else:
1174
- left = 0
1175
- right = 0
1176
- left_add = tl.maximum(left - (WINDOW_SIZE - 1), 0)
1177
- swa_start_pos = tl.maximum(pos - (WINDOW_SIZE - 1) - left_add, 0)
1178
- swa_len = pos + right - swa_start_pos + 1
1179
- topk_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOP_K)
1180
-
1181
- for i in range(0, SWA_TOTAL_WIDTH, WINDOW_BLOCK_SIZE):
1182
- offset = i + tl.arange(0, WINDOW_BLOCK_SIZE)
1183
- mask = offset < SWA_TOTAL_WIDTH
1184
- pos_offset = swa_start_pos + offset
1185
- block_indices = pos_offset // swa_block_size
1186
- block_numbers = tl.load(
1187
- swa_block_table_ptr + req_idx * swa_block_table_stride + block_indices,
1188
- mask=mask & (offset < swa_len),
1189
- other=-1,
1502
+ def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]:
1503
+ hf_config = vllm_config.model_config.hf_config
1504
+ window_size = int(getattr(hf_config, "sliding_window", 128) or 128)
1505
+ speculative_config = vllm_config.speculative_config
1506
+ if speculative_config is not None and speculative_config.use_dspark():
1507
+ from vllm.v1.attention.backends.mla.compressor_utils import (
1508
+ get_dspark_swa_index_width,
1509
+ )
1510
+
1511
+ swa_index_width: int | tuple[int, int] = (
1512
+ window_size,
1513
+ get_dspark_swa_index_width(
1514
+ window_size,
1515
+ speculative_config.num_speculative_tokens,
1516
+ ),
1517
+ )
1518
+ else:
1519
+ swa_index_width = window_size
1520
+ max_image_tokens = (
1521
+ int(getattr(hf_config, "vision_max_n_token", 0) or 0)
1522
+ if int(getattr(hf_config, "vision_n_layers", 0) or 0) > 0
1523
+ else 0
1524
)
1191
- block_offsets = pos_offset % swa_block_size
1192
- slot_ids = block_numbers * swa_block_size + block_offsets
1193
- slot_ids = tl.where(offset < swa_len, slot_ids, -1)
1194
- slot_ids = _remap_flashinfer_index(slot_ids, swa_block_size, swa_block_span)
1195
- tl.store(
1196
- sparse_indices_ptr + token_idx * sparse_indices_stride + offset,
1197
- slot_ids,
1198
- mask=mask,
1525
+ image_widths = (0, max_image_tokens) if max_image_tokens > 0 else 0
1526
+ index_topk = int(getattr(hf_config, "index_topk", 0) or 0)
1527
+ compress_ratios = tuple(
1528
+ sorted({max(1, int(ratio)) for ratio in hf_config.compress_ratios})
1529
+ )
1530
+ max_c128a_topk = (
1531
+ ((vllm_config.model_config.max_model_len + 127) // 128 + 127) // 128 * 128
1532
+ )
1533
+ active_topk_widths = WarmupIntRange(
1534
+ 128,
1535
+ max_c128a_topk + 1,
1536
+ advance=lambda width: (
1537
+ min(width * 2, max_c128a_topk)
1538
+ if width < max_c128a_topk
1539
+ else max_c128a_topk + 1
1540
+ ),
1541
+ )
1542
+ return self._trace_dispatch(self.dispatch)(
1543
+ compress_ratio=compress_ratios,
1544
+ index_topk=index_topk,
1545
+ active_topk_width=active_topk_widths,
1546
+ has_prefill=(False, True),
1547
+ window_size=window_size,
1548
+ swa_index_width=swa_index_width,
1549
+ image_width=image_widths,
1550
)
1551
1201
- for i in range(0, PADDED_TOP_K, TOPK_BLOCK_SIZE):
1202
- offset = i + tl.arange(0, TOPK_BLOCK_SIZE)
1203
- mask = offset < PADDED_TOP_K
1204
- local_idx = tl.load(
1205
- prefill_topk_indices_ptr + prefill_idx * prefill_topk_stride + offset,
1206
- mask=(offset < PREFILL_TOPK_STRIDE) & (offset < topk_len),
1207
- other=-1,
1552
+ def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
1553
+ int32_ptr = TritonWarmupTensor(torch.int32)
1554
+ swa_block_size = 64
1555
+ compressed_block_size = (
1556
+ swa_block_size
1557
+ if compile_key.compress_ratio == 1
1558
+ else 256 // compile_key.compress_ratio
1559
)
1209
- is_valid = local_idx >= 0
1210
- block_indices = local_idx // compressed_block_size
1211
- block_numbers = tl.load(
1212
- compressed_block_table_ptr
1213
- + req_idx * compressed_block_table_stride
1214
- + block_indices,
1215
- mask=mask & is_valid,
1216
- other=-1,
1560
+ num_warps = (
1561
+ 4
1562
+ if max(compile_key.window_block_size, compile_key.topk_block_size) >= 256
1563
+ else 1
1564
)
1218
- block_offsets = local_idx % compressed_block_size
1219
- slot_ids = block_numbers * compressed_block_size + block_offsets
1220
- slot_ids = tl.where((offset < topk_len) & is_valid, slot_ids, -1)
1221
- slot_ids = _remap_flashinfer_index(
1222
- slot_ids, compressed_block_size, compressed_block_span
1565
+ return dict(
1566
+ sparse_indices=TritonWarmupTensor(
1567
+ torch.int32,
1568
+ shape=(
1569
+ 1,
1570
+ compile_key.swa_index_width
1571
+ + compile_key.image_width
1572
+ + compile_key.padded_top_k,
1573
+ ),
1574
+ ),
1575
+ sparse_topk_lens=int32_ptr,
1576
+ decode_swa_indices=TritonWarmupTensor(
1577
+ torch.int32,
1578
+ shape=(1, compile_key.swa_index_width),
1579
+ ),
1580
+ decode_compressed_indices=int32_ptr,
1581
+ decode_compressed_topk_lens=int32_ptr,
1582
+ decode_is_valid_token=TritonWarmupTensor(
1583
+ torch.bool
1584
+ if compile_key.decode_compressed_indices_are_local
1585
+ else torch.int32
1586
+ ),
1587
+ prefill_topk_indices=int32_ptr,
1588
+ prefill_left_visible=int32_ptr,
1589
+ prefill_right_visible=int32_ptr,
1590
+ query_start_loc=int32_ptr,
1591
+ seq_lens=int32_ptr,
1592
+ token_to_req_indices=int32_ptr,
1593
+ swa_block_table=TritonWarmupTensor(torch.int32, shape=(1, 1)),
1594
+ swa_block_size=swa_block_size,
1595
+ swa_block_span=1,
1596
+ compressed_block_table=TritonWarmupTensor(torch.int32, shape=(1, 1)),
1597
+ compressed_block_size=compressed_block_size,
1598
+ compressed_block_span=1,
1599
+ num_decode_tokens=1,
1600
+ window_size=compile_key.window_size,
1601
+ swa_index_width=compile_key.swa_index_width,
1602
+ image_width=compile_key.image_width,
1603
+ compress_ratio=compile_key.compress_ratio,
1604
+ topk=compile_key.top_k,
1605
+ padded_topk=compile_key.padded_top_k,
1606
+ prefill_topk_stride=1,
1607
+ decode_compressed_topk=compile_key.decode_compressed_topk,
1608
+ decode_compressed_indices_are_local=(
1609
+ compile_key.decode_compressed_indices_are_local
1610
+ ),
1611
+ has_decode_compressed_lens=compile_key.has_decode_compressed_lens,
1612
+ window_block_size=compile_key.window_block_size,
1613
+ topk_block_size=compile_key.topk_block_size,
1614
+ num_warps=num_warps,
1615
)
1224
- tl.store(
1225
- sparse_indices_ptr
1226
- + token_idx * sparse_indices_stride
1227
- + SWA_TOTAL_WIDTH
1228
- + offset,
1229
- slot_ids,
1230
- mask=mask,
1616
+
1617
+ @kernel_launcher
1618
+ def __call__(
1619
+ self,
1620
+ sparse_indices: torch.Tensor,
1621
+ sparse_topk_lens: torch.Tensor,
1622
+ decode_swa_indices: torch.Tensor,
1623
+ decode_compressed_indices: torch.Tensor,
1624
+ decode_compressed_topk_lens: torch.Tensor,
1625
+ decode_is_valid_token: torch.Tensor,
1626
+ prefill_topk_indices: torch.Tensor,
1627
+ prefill_left_visible: torch.Tensor,
1628
+ prefill_right_visible: torch.Tensor,
1629
+ query_start_loc: torch.Tensor,
1630
+ seq_lens: torch.Tensor,
1631
+ token_to_req_indices: torch.Tensor,
1632
+ swa_block_table: torch.Tensor,
1633
+ swa_block_size: int,
1634
+ swa_block_span: int,
1635
+ compressed_block_table: torch.Tensor,
1636
+ compressed_block_size: int,
1637
+ compressed_block_span: int,
1638
+ *,
1639
+ num_decode_tokens: int,
1640
+ window_size: int,
1641
+ swa_index_width: int,
1642
+ image_width: int,
1643
+ compress_ratio: int,
1644
+ topk: int,
1645
+ padded_topk: int,
1646
+ prefill_topk_stride: int,
1647
+ decode_compressed_topk: int,
1648
+ decode_compressed_indices_are_local: bool,
1649
+ has_decode_compressed_lens: bool,
1650
+ window_block_size: int,
1651
+ topk_block_size: int,
1652
+ num_warps: int,
1653
+ ) -> LaunchSpec:
1654
+ return (sparse_indices.shape[0],), dict(
1655
+ sparse_indices_stride=sparse_indices.stride(0),
1656
+ decode_swa_stride=decode_swa_indices.stride(0),
1657
+ decode_compressed_stride=decode_compressed_indices.stride(0),
1658
+ prefill_topk_stride=prefill_topk_indices.stride(0),
1659
+ left_visible_ptr=prefill_left_visible,
1660
+ right_visible_ptr=prefill_right_visible,
1661
+ swa_block_table_stride=swa_block_table.stride(0),
1662
+ compressed_block_table_stride=compressed_block_table.stride(0),
1663
+ NUM_DECODE_TOKENS=num_decode_tokens,
1664
+ WINDOW_SIZE=window_size,
1665
+ SWA_INDEX_WIDTH=swa_index_width,
1666
+ IMAGE_WIDTH=image_width,
1667
+ COMPRESS_RATIO=compress_ratio,
1668
+ TOP_K=topk,
1669
+ PADDED_TOP_K=padded_topk,
1670
+ PREFILL_TOPK_STRIDE=prefill_topk_stride,
1671
+ DECODE_COMPRESSED_TOPK=decode_compressed_topk,
1672
+ DECODE_COMPRESSED_INDICES_ARE_LOCAL=decode_compressed_indices_are_local,
1673
+ HAS_DECODE_COMPRESSED_LENS=has_decode_compressed_lens,
1674
+ WINDOW_BLOCK_SIZE=window_block_size,
1675
+ TOPK_BLOCK_SIZE=topk_block_size,
1676
+ num_warps=num_warps,
1677
)
1678
1233
- tl.store(sparse_topk_lens_ptr + token_idx, SWA_TOTAL_WIDTH + topk_len)
1679
+
1680
+_BUILD_FLASHINFER_MIXED_SPARSE_INDICES_KERNEL = (
1681
+ BuildFlashinferMixedSparseIndicesKernel()
1682
+)
vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py
+281
-9
@@ -19,13 +19,22 @@ even/odd halves, producing (N_QUANT_BLOCKS, MXFP4_BLOCK/2) packed nibbles
19
and N_QUANT_BLOCKS ue8m0 bytes.
20
"""
21
22
+from dataclasses import dataclass
23
from functools import lru_cache
24
+from types import SimpleNamespace
25
from typing import Any
26
27
import torch
28
29
+from vllm.model_executor.warmup.jit_warmup_triton_helper import (
30
+ LaunchSpec,
31
+ TritonWarmupTensor,
32
+ VllmTritonJitKernel,
33
+ kernel_launcher,
34
+)
35
from vllm.platforms import current_platform
36
from vllm.triton_utils import tl, triton
37
+from vllm.utils.math_utils import round_up
38
39
if current_platform.is_rocm():
40
from vllm.platforms.rocm import _ON_GFX950
@@ -68,14 +77,32 @@ def compress_norm_rope_store_triton(
77
kernel = _fused_kv_compress_norm_rope_insert_sparse_attn
78
num_warps = 4
79
kernel_kwargs = {"SANITIZE_CACHE_NANS": _ON_GFX950}
71
- elif use_fp4_cache:
72
- kernel = _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn
73
- num_warps = 1
74
- kernel_kwargs = {}
80
else:
76
- kernel = _fused_kv_compress_norm_rope_insert_indexer_attn
77
- num_warps = 1
78
- kernel_kwargs = {}
81
+ _FUSED_KV_COMPRESS_NORM_ROPE_INSERT_INDEXER_TRITON_KERNEL(
82
+ state_cache=state_cache,
83
+ num_actual=num_actual,
84
+ token_to_req_indices=token_to_req_indices,
85
+ positions=positions,
86
+ slot_mapping=slot_mapping,
87
+ block_table=block_table,
88
+ block_size=block_size,
89
+ state_width=state_width,
90
+ cos_sin_cache=cos_sin_cache,
91
+ kv_cache=kv_cache,
92
+ k_cache_metadata=k_cache_metadata,
93
+ pdl_kwargs=pdl_kwargs,
94
+ head_dim=head_dim,
95
+ rope_head_dim=rope_head_dim,
96
+ compress_ratio=compress_ratio,
97
+ overlap=overlap,
98
+ use_fp4_cache=use_fp4_cache,
99
+ rms_norm_weight=rms_norm_weight,
100
+ rms_norm_eps=rms_norm_eps,
101
+ quant_block=quant_block,
102
+ token_stride=token_stride,
103
+ scale_dim=scale_dim,
104
+ )
105
+ return
106
107
kernel[(num_actual,)](
108
# state cache
@@ -673,7 +700,15 @@ def compress_norm_rope_store_two_stage_triton(
700
# =============================================================================
701
# Indexer path (head=128, all FP8, single quant block)
702
# =============================================================================
676
-@triton.jit
703
+@triton.jit(
704
+ do_not_specialize=[
705
+ "state_cache_stride0",
706
+ "state_cache_stride1",
707
+ "block_table_stride",
708
+ "rms_norm_eps",
709
+ "cos_sin_stride",
710
+ ]
711
+)
712
def _fused_kv_compress_norm_rope_insert_indexer_attn(
713
# ── state cache (compressor internal state) ──
714
state_cache_ptr,
@@ -850,7 +885,15 @@ def _fused_kv_compress_norm_rope_insert_indexer_attn(
885
# =============================================================================
886
# Indexer path (head=128, MXFP4: 2 nibbles/byte + ue8m0 per 32-elem block)
887
# =============================================================================
853
-@triton.jit
888
+@triton.jit(
889
+ do_not_specialize=[
890
+ "state_cache_stride0",
891
+ "state_cache_stride1",
892
+ "block_table_stride",
893
+ "rms_norm_eps",
894
+ "cos_sin_stride",
895
+ ]
896
+)
897
def _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn(
898
# ── state cache (compressor internal state) ──
899
state_cache_ptr,
@@ -1038,3 +1081,232 @@ def _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn(
1081
1082
tl.store(val_ptr + tl.arange(0, TOKEN_STRIDE), packed_flat)
1083
tl.store(scale_ptr + tl.arange(0, SCALE_DIM), ue8m0)
1084
+
1085
+
1086
+class FusedKVCompressNormRopeInsertIndexerTritonKernel(
1087
+ VllmTritonJitKernel["FusedKVCompressNormRopeInsertIndexerTritonKernel.CompileKey"]
1088
+):
1089
+ @dataclass(frozen=True)
1090
+ class CompileKey:
1091
+ dtype: torch.dtype
1092
+ use_fp4_cache: bool
1093
+ head_size: int
1094
+ triton_block_size: int
1095
+ state_width: int
1096
+ compress_ratio: int
1097
+ overlap: bool
1098
+ rope_head_dim: int
1099
+ fp8_max: float
1100
+ quant_block: int
1101
+ token_stride: int
1102
+ scale_dim: int
1103
+ block_size: int
1104
+ kv_cache_block_size: int
1105
+ kv_block_stride: int
1106
+
1107
+ # Canonical kernel for the launcher's arg_names mapping; the launched
1108
+ # variant (fp8/mxfp4, same arg_names) is picked per call in __call__.
1109
+ kernel = staticmethod(_fused_kv_compress_norm_rope_insert_indexer_attn)
1110
+
1111
+ def dispatch( # type: ignore[override]
1112
+ self,
1113
+ *,
1114
+ dtype: torch.dtype,
1115
+ use_fp4_cache: bool,
1116
+ head_dim: int,
1117
+ rope_head_dim: int,
1118
+ compress_ratio: int,
1119
+ cache_block_size: int,
1120
+ cache_alignment: int,
1121
+ runtime_state_width: int | None = None,
1122
+ runtime_quant_block: int | None = None,
1123
+ runtime_token_stride: int | None = None,
1124
+ runtime_scale_dim: int | None = None,
1125
+ runtime_kv_block_stride: int | None = None,
1126
+ ) -> CompileKey:
1127
+ overlap = compress_ratio == 4
1128
+ default_quant_block = 32 if use_fp4_cache else 128
1129
+ quant_block = (
1130
+ runtime_quant_block
1131
+ if runtime_quant_block is not None
1132
+ else default_quant_block
1133
+ )
1134
+ default_token_stride = head_dim // 2 if use_fp4_cache else head_dim
1135
+ token_stride = (
1136
+ runtime_token_stride
1137
+ if runtime_token_stride is not None
1138
+ else default_token_stride
1139
+ )
1140
+ default_scale_dim = head_dim // quant_block if use_fp4_cache else 4
1141
+ scale_dim = (
1142
+ runtime_scale_dim if runtime_scale_dim is not None else default_scale_dim
1143
+ )
1144
+ raw_kv_cache_block_size = cache_block_size // compress_ratio
1145
+ kv_cache_block_size = (
1146
+ raw_kv_cache_block_size if raw_kv_cache_block_size >= 1 else 1
1147
+ )
1148
+ state_block_size = (
1149
+ 4
1150
+ if compress_ratio == 4
1151
+ else 8
1152
+ if compress_ratio == 128
1153
+ else cache_block_size
1154
+ )
1155
+ default_kv_block_stride = round_up(
1156
+ kv_cache_block_size * (token_stride + scale_dim),
1157
+ cache_alignment,
1158
+ )
1159
+ kv_block_stride = (
1160
+ runtime_kv_block_stride
1161
+ if runtime_kv_block_stride is not None
1162
+ else default_kv_block_stride
1163
+ )
1164
+ state_width = (
1165
+ runtime_state_width
1166
+ if runtime_state_width is not None
1167
+ else head_dim * (1 + overlap)
1168
+ )
1169
+ return self.CompileKey(
1170
+ dtype=dtype,
1171
+ use_fp4_cache=use_fp4_cache,
1172
+ head_size=head_dim,
1173
+ triton_block_size=triton.next_power_of_2(head_dim),
1174
+ state_width=state_width,
1175
+ compress_ratio=compress_ratio,
1176
+ overlap=overlap,
1177
+ rope_head_dim=rope_head_dim,
1178
+ fp8_max=448.0,
1179
+ quant_block=quant_block,
1180
+ token_stride=token_stride,
1181
+ scale_dim=scale_dim,
1182
+ block_size=state_block_size,
1183
+ kv_cache_block_size=kv_cache_block_size,
1184
+ kv_block_stride=kv_block_stride,
1185
+ )
1186
+
1187
+ def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]:
1188
+ hf_config = vllm_config.model_config.hf_config
1189
+ if hf_config is None:
1190
+ return []
1191
+
1192
+ cache_block_size = vllm_config.cache_config.block_size
1193
+ from vllm.v1.attention.backends.mla.indexer import (
1194
+ dsa_indexer_uses_fp4,
1195
+ )
1196
+
1197
+ use_fp4_cache = dsa_indexer_uses_fp4(vllm_config)
1198
+ head_dim = int(getattr(hf_config, "index_head_dim", 0) or 0)
1199
+ rope_head_dim = int(getattr(hf_config, "qk_rope_head_dim", 0) or 0)
1200
+ if head_dim <= 0 or rope_head_dim <= 0 or cache_block_size <= 0:
1201
+ return []
1202
+
1203
+ cache_dtype = vllm_config.cache_config.cache_dtype
1204
+ cache_alignment = 576 if cache_dtype == "fp8_ds_mla" else 512
1205
+ return self._trace_dispatch(self.dispatch)(
1206
+ dtype=vllm_config.model_config.dtype,
1207
+ use_fp4_cache=use_fp4_cache,
1208
+ head_dim=head_dim,
1209
+ rope_head_dim=rope_head_dim,
1210
+ compress_ratio=4,
1211
+ cache_block_size=cache_block_size,
1212
+ cache_alignment=cache_alignment,
1213
+ )
1214
+
1215
+ def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
1216
+ fp32_ptr = TritonWarmupTensor(torch.float32, shape=(1, 1))
1217
+ int32_ptr = TritonWarmupTensor(torch.int32)
1218
+ int64_ptr = TritonWarmupTensor(torch.int64)
1219
+ # kv_cache supplies kv_cache_block_size (shape[1]) and KV_BLOCK_STRIDE
1220
+ # (stride(0)); reproduce both so warmup and launch specialize identically.
1221
+ kv_cache = TritonWarmupTensor(
1222
+ torch.uint8,
1223
+ shape=(1, compile_key.kv_cache_block_size, 1),
1224
+ strides=(compile_key.kv_block_stride, 1, 1),
1225
+ )
1226
+ return dict(
1227
+ state_cache=fp32_ptr,
1228
+ num_actual=1,
1229
+ token_to_req_indices=int32_ptr,
1230
+ positions=int64_ptr,
1231
+ slot_mapping=int64_ptr,
1232
+ block_table=int32_ptr,
1233
+ block_size=compile_key.block_size,
1234
+ state_width=compile_key.state_width,
1235
+ cos_sin_cache=fp32_ptr,
1236
+ kv_cache=kv_cache,
1237
+ k_cache_metadata=SimpleNamespace(slot_mapping=int64_ptr),
1238
+ pdl_kwargs={},
1239
+ head_dim=compile_key.head_size,
1240
+ rope_head_dim=compile_key.rope_head_dim,
1241
+ compress_ratio=compile_key.compress_ratio,
1242
+ overlap=compile_key.overlap,
1243
+ use_fp4_cache=compile_key.use_fp4_cache,
1244
+ rms_norm_weight=TritonWarmupTensor(
1245
+ compile_key.dtype,
1246
+ shape=(compile_key.head_size,),
1247
+ ),
1248
+ rms_norm_eps=1e-6,
1249
+ quant_block=compile_key.quant_block,
1250
+ token_stride=compile_key.token_stride,
1251
+ scale_dim=compile_key.scale_dim,
1252
+ )
1253
+
1254
+ @kernel_launcher
1255
+ def __call__(
1256
+ self,
1257
+ *,
1258
+ state_cache: torch.Tensor,
1259
+ num_actual: int,
1260
+ token_to_req_indices: torch.Tensor,
1261
+ positions: torch.Tensor,
1262
+ slot_mapping: torch.Tensor,
1263
+ block_table: torch.Tensor,
1264
+ block_size: int,
1265
+ state_width: int,
1266
+ cos_sin_cache: torch.Tensor,
1267
+ kv_cache: torch.Tensor,
1268
+ k_cache_metadata: Any,
1269
+ pdl_kwargs: dict,
1270
+ head_dim: int,
1271
+ rope_head_dim: int,
1272
+ compress_ratio: int,
1273
+ overlap: bool,
1274
+ use_fp4_cache: bool,
1275
+ rms_norm_weight: torch.Tensor,
1276
+ rms_norm_eps: float,
1277
+ quant_block: int,
1278
+ token_stride: int,
1279
+ scale_dim: int,
1280
+ ) -> LaunchSpec:
1281
+ return (num_actual,), dict(
1282
+ kernel=(
1283
+ _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn
1284
+ if use_fp4_cache
1285
+ else _fused_kv_compress_norm_rope_insert_indexer_attn
1286
+ ),
1287
+ state_cache_stride0=state_cache.stride(0),
1288
+ state_cache_stride1=state_cache.stride(1),
1289
+ block_table_stride=block_table.stride(0),
1290
+ cos_sin_stride=cos_sin_cache.stride(0),
1291
+ k_cache_ptr=kv_cache,
1292
+ kv_slot_mapping_ptr=k_cache_metadata.slot_mapping,
1293
+ kv_cache_block_size=kv_cache.shape[1],
1294
+ HEAD_SIZE=head_dim,
1295
+ TRITON_BLOCK_SIZE=triton.next_power_of_2(head_dim),
1296
+ STATE_WIDTH=state_width,
1297
+ COMPRESS_RATIO=compress_ratio,
1298
+ OVERLAP=overlap,
1299
+ ROPE_HEAD_DIM=rope_head_dim,
1300
+ FP8_MAX=448.0,
1301
+ QUANT_BLOCK=quant_block,
1302
+ TOKEN_STRIDE=token_stride,
1303
+ SCALE_DIM=scale_dim,
1304
+ KV_BLOCK_STRIDE=kv_cache.stride(0),
1305
+ num_warps=1,
1306
+ **pdl_kwargs,
1307
+ )
1308
+
1309
+
1310
+_FUSED_KV_COMPRESS_NORM_ROPE_INSERT_INDEXER_TRITON_KERNEL = (
1311
+ FusedKVCompressNormRopeInsertIndexerTritonKernel()
1312
+)
vllm/models/deepseek_v4/common/ops/fused_indexer_q.py
+464
-245
@@ -1,8 +1,17 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+from dataclasses import dataclass
5
+from typing import Any
6
+
7
import torch
8
9
+from vllm.model_executor.warmup.jit_warmup_triton_helper import (
10
+ LaunchSpec,
11
+ TritonWarmupTensor,
12
+ VllmTritonJitKernel,
13
+ kernel_launcher,
14
+)
15
from vllm.platforms import current_platform
16
from vllm.triton_utils import tl, triton
17
from vllm.utils.import_utils import has_cutedsl
@@ -67,229 +76,459 @@ def _quantize_mxfp4_pair(x_lo, x_hi):
76
return packed, ue8m0
77
78
70
-@triton.jit
71
-def _fused_indexer_q_rope_quant_kernel(
72
- pos_ptr,
73
- # Index Q RoPE
74
- index_q_ptr,
75
- index_q_stride0,
76
- index_q_stride1,
77
- index_q_cos_sin_ptr,
78
- index_q_cos_sin_stride,
79
- INDEX_Q_HALF_ROT_DIM: tl.constexpr,
80
- # Index Q Quantize
81
- index_q_fp8_ptr,
82
- index_q_fp8_stride0,
83
- index_q_fp8_stride1,
84
- INDEX_Q_HEAD_DIM: tl.constexpr,
85
- # Index weights
86
- index_weights_ptr,
87
- index_weights_stride,
88
- index_weights_softmax_scale,
89
- index_weights_head_scale,
90
- index_weights_out_ptr,
91
- index_weights_out_stride,
92
- FP8_MAX: tl.constexpr = 448.0,
93
- USE_FNUZ: tl.constexpr = False,
94
- USE_EXPLICIT_FMA: tl.constexpr = False,
79
+class FusedIndexerQRopeQuantTritonKernel(
80
+ VllmTritonJitKernel["FusedIndexerQRopeQuantTritonKernel.CompileKey"]
81
):
96
- # Layout matches the unfused reference (DeepseekV4ScalingRotaryEmbedding
97
- # + per_token_group_quant_fp8): GPT-J interleaved RoPE applied to the
98
- # LAST rope_dim dims of each head; the leading [0, NOPE_DIM) is passed
99
- # through unchanged.
100
- INDEX_Q_ROT_DIM: tl.constexpr = 2 * INDEX_Q_HALF_ROT_DIM
101
- INDEX_Q_NOPE_DIM: tl.constexpr = INDEX_Q_HEAD_DIM - INDEX_Q_ROT_DIM
102
- tl.static_assert(INDEX_Q_NOPE_DIM >= 0)
103
-
104
- tok_idx = tl.program_id(0)
105
- head_idx = tl.program_id(1)
106
-
107
- pos = tl.load(pos_ptr + tok_idx)
108
- cos, sin = _get_cos_sin(
82
+ @dataclass(frozen=True)
83
+ class CompileKey:
84
+ dtype: torch.dtype
85
+ num_heads: int
86
+ index_q_half_rot_dim: int
87
+ index_q_head_dim: int
88
+ fp8_max: float
89
+ use_fnuz: bool
90
+ use_explicit_fma: bool
91
+
92
+ @staticmethod
93
+ @triton.jit
94
+ def kernel(
95
+ pos_ptr,
96
+ # Index Q RoPE
97
+ index_q_ptr,
98
+ index_q_stride0,
99
+ index_q_stride1,
100
index_q_cos_sin_ptr,
101
index_q_cos_sin_stride,
111
- pos,
112
- INDEX_Q_HALF_ROT_DIM,
113
- )
114
- half_offset = tl.arange(0, INDEX_Q_HALF_ROT_DIM)
115
- base_ptr = index_q_ptr + tok_idx * index_q_stride0 + head_idx * index_q_stride1
116
-
117
- # Interleaved (GPT-J) RoPE on dims [NOPE_DIM, HEAD_DIM):
118
- # even = q[NOPE_DIM + 2*i], odd = q[NOPE_DIM + 2*i + 1]
119
- rot_base = base_ptr + INDEX_Q_NOPE_DIM
120
- x_even = tl.load(rot_base + half_offset * 2).to(tl.float32)
121
- x_odd = tl.load(rot_base + half_offset * 2 + 1).to(tl.float32)
122
- if USE_EXPLICIT_FMA:
123
- # Match HIP rotary_embedding contraction before bf16 materialization.
124
- r_even = tl.fma(x_even, cos, -(x_odd * sin))
125
- r_odd = tl.fma(x_odd, cos, x_even * sin)
126
- else:
127
- r_even = x_even * cos - x_odd * sin
128
- r_odd = x_odd * cos + x_even * sin
129
-
130
- # Match reference numerics: fp32 → bf16 → fp32 before the ue8m0 absmax.
131
- # Same pattern as the K-side compressor kernel (fused_compress_quant_cache.py).
132
- r_even = r_even.to(tl.bfloat16).to(tl.float32)
133
- r_odd = r_odd.to(tl.bfloat16).to(tl.float32)
134
-
135
- amax = tl.maximum(tl.max(tl.abs(r_even)), tl.max(tl.abs(r_odd)))
136
- if INDEX_Q_NOPE_DIM > 0:
137
- nope_offset = tl.arange(0, INDEX_Q_NOPE_DIM)
138
- x_nope = tl.load(base_ptr + nope_offset).to(tl.float32)
139
- amax = tl.maximum(amax, tl.max(tl.abs(x_nope)))
140
- index_q_scale = tl.div_rn(tl.maximum(amax, 1e-4), FP8_MAX)
141
- index_q_scale = tl.math.exp2(tl.math.ceil(tl.math.log2(index_q_scale)))
142
-
143
- # Store quantized values to index_q_fp8. FNUZ (e4m3fnuz) on gfx942, OCP
144
- # (e4m3fn) elsewhere -- matches the K cache.
145
- fp8_dtype = tl.float8e4b8 if USE_FNUZ else tl.float8e4nv
146
- fp8_base_ptr = (
147
- index_q_fp8_ptr + tok_idx * index_q_fp8_stride0 + head_idx * index_q_fp8_stride1
148
- )
149
- if INDEX_Q_NOPE_DIM > 0:
102
+ INDEX_Q_HALF_ROT_DIM: tl.constexpr,
103
+ # Index Q Quantize
104
+ index_q_fp8_ptr,
105
+ index_q_fp8_stride0,
106
+ index_q_fp8_stride1,
107
+ INDEX_Q_HEAD_DIM: tl.constexpr,
108
+ # Index weights
109
+ index_weights_ptr,
110
+ index_weights_stride,
111
+ index_weights_softmax_scale,
112
+ index_weights_head_scale,
113
+ index_weights_out_ptr,
114
+ index_weights_out_stride,
115
+ FP8_MAX: tl.constexpr = 448.0,
116
+ USE_FNUZ: tl.constexpr = False,
117
+ USE_EXPLICIT_FMA: tl.constexpr = False,
118
+ ):
119
+ # Layout matches the unfused reference (DeepseekV4ScalingRotaryEmbedding
120
+ # + per_token_group_quant_fp8): GPT-J interleaved RoPE applied to the
121
+ # LAST rope_dim dims of each head; the leading [0, NOPE_DIM) is passed
122
+ # through unchanged.
123
+ INDEX_Q_ROT_DIM: tl.constexpr = 2 * INDEX_Q_HALF_ROT_DIM
124
+ INDEX_Q_NOPE_DIM: tl.constexpr = INDEX_Q_HEAD_DIM - INDEX_Q_ROT_DIM
125
+ tl.static_assert(INDEX_Q_NOPE_DIM >= 0)
126
+
127
+ tok_idx = tl.program_id(0)
128
+ head_idx = tl.program_id(1)
129
+
130
+ pos = tl.load(pos_ptr + tok_idx)
131
+ cos, sin = _get_cos_sin(
132
+ index_q_cos_sin_ptr,
133
+ index_q_cos_sin_stride,
134
+ pos,
135
+ INDEX_Q_HALF_ROT_DIM,
136
+ )
137
+ half_offset = tl.arange(0, INDEX_Q_HALF_ROT_DIM)
138
+ base_ptr = index_q_ptr + tok_idx * index_q_stride0 + head_idx * index_q_stride1
139
+
140
+ # Interleaved (GPT-J) RoPE on dims [NOPE_DIM, HEAD_DIM):
141
+ # even = q[NOPE_DIM + 2*i], odd = q[NOPE_DIM + 2*i + 1]
142
+ rot_base = base_ptr + INDEX_Q_NOPE_DIM
143
+ x_even = tl.load(rot_base + half_offset * 2).to(tl.float32)
144
+ x_odd = tl.load(rot_base + half_offset * 2 + 1).to(tl.float32)
145
+ if USE_EXPLICIT_FMA:
146
+ # Match HIP rotary_embedding contraction before bf16 materialization.
147
+ r_even = tl.fma(x_even, cos, -(x_odd * sin))
148
+ r_odd = tl.fma(x_odd, cos, x_even * sin)
149
+ else:
150
+ r_even = x_even * cos - x_odd * sin
151
+ r_odd = x_odd * cos + x_even * sin
152
+
153
+ # Match reference numerics: fp32 → bf16 → fp32 before the ue8m0 absmax.
154
+ # Same pattern as the K-side compressor kernel (fused_compress_quant_cache.py).
155
+ r_even = r_even.to(tl.bfloat16).to(tl.float32)
156
+ r_odd = r_odd.to(tl.bfloat16).to(tl.float32)
157
+
158
+ amax = tl.maximum(tl.max(tl.abs(r_even)), tl.max(tl.abs(r_odd)))
159
+ if INDEX_Q_NOPE_DIM > 0:
160
+ nope_offset = tl.arange(0, INDEX_Q_NOPE_DIM)
161
+ x_nope = tl.load(base_ptr + nope_offset).to(tl.float32)
162
+ amax = tl.maximum(amax, tl.max(tl.abs(x_nope)))
163
+ index_q_scale = tl.div_rn(tl.maximum(amax, 1e-4), FP8_MAX)
164
+ index_q_scale = tl.math.exp2(tl.math.ceil(tl.math.log2(index_q_scale)))
165
+
166
+ # Store quantized values to index_q_fp8. FNUZ (e4m3fnuz) on gfx942, OCP
167
+ # (e4m3fn) elsewhere -- matches the K cache.
168
+ fp8_dtype = tl.float8e4b8 if USE_FNUZ else tl.float8e4nv
169
+ fp8_base_ptr = (
170
+ index_q_fp8_ptr
171
+ + tok_idx * index_q_fp8_stride0
172
+ + head_idx * index_q_fp8_stride1
173
+ )
174
+ if INDEX_Q_NOPE_DIM > 0:
175
+ tl.store(
176
+ fp8_base_ptr + nope_offset,
177
+ tl.div_rn(x_nope, index_q_scale).to(fp8_dtype),
178
+ )
179
+ fp8_rot_base = fp8_base_ptr + INDEX_Q_NOPE_DIM
180
tl.store(
151
- fp8_base_ptr + nope_offset,
152
- tl.div_rn(x_nope, index_q_scale).to(fp8_dtype),
181
+ fp8_rot_base + half_offset * 2,
182
+ tl.div_rn(r_even, index_q_scale).to(fp8_dtype),
183
+ )
184
+ tl.store(
185
+ fp8_rot_base + half_offset * 2 + 1,
186
+ tl.div_rn(r_odd, index_q_scale).to(fp8_dtype),
187
)
154
- fp8_rot_base = fp8_base_ptr + INDEX_Q_NOPE_DIM
155
- tl.store(
156
- fp8_rot_base + half_offset * 2,
157
- tl.div_rn(r_even, index_q_scale).to(fp8_dtype),
158
- )
159
- tl.store(
160
- fp8_rot_base + half_offset * 2 + 1,
161
- tl.div_rn(r_odd, index_q_scale).to(fp8_dtype),
162
- )
163
-
164
- # FP8 weight-fold contract:
165
- # index_weights_out = index_weights * q_scale * softmax_scale * head_scale
166
- # The per-token-per-head q_scale (fp32) IS folded into the output weights
167
- # here because FP8 Q is stored WITHOUT a companion scale tensor — the
168
- # downstream fp8_fp4_mqa_logits/fp8_fp4_paged_mqa_logits kernels use `weights` to
169
- # apply per-token Q scale inline. See the MXFP4 kernel below for the
170
- # contrasting convention (scales live with the Q values, weights are NOT
171
- # q-scaled).
172
- index_weights = tl.load(
173
- index_weights_ptr + tok_idx * index_weights_stride + head_idx
174
- )
175
- index_weights = index_weights.to(tl.float32)
176
- index_weights *= index_q_scale
177
- index_weights *= index_weights_softmax_scale
178
- index_weights *= index_weights_head_scale
179
- tl.store(
180
- index_weights_out_ptr + tok_idx * index_weights_out_stride + head_idx,
181
- index_weights,
182
- )
188
189
+ # FP8 weight-fold contract:
190
+ # index_weights_out = index_weights * q_scale * softmax_scale * head_scale
191
+ # The per-token-per-head q_scale (fp32) IS folded into the output weights
192
+ # here because FP8 Q is stored WITHOUT a companion scale tensor — the
193
+ # downstream fp8_fp4_mqa_logits/fp8_fp4_paged_mqa_logits kernels use
194
+ # `weights` to apply per-token Q scale inline. See the MXFP4 kernel below for
195
+ # contrasting convention (scales live with the Q values, weights are NOT
196
+ # q-scaled).
197
+ index_weights = tl.load(
198
+ index_weights_ptr + tok_idx * index_weights_stride + head_idx
199
+ )
200
+ index_weights = index_weights.to(tl.float32)
201
+ index_weights *= index_q_scale
202
+ index_weights *= index_weights_softmax_scale
203
+ index_weights *= index_weights_head_scale
204
+ tl.store(
205
+ index_weights_out_ptr + tok_idx * index_weights_out_stride + head_idx,
206
+ index_weights,
207
+ )
208
185
-@triton.jit
186
-def _fused_indexer_q_rope_mxfp4_kernel(
187
- pos_ptr,
188
- # Index Q RoPE input (fp/bf16)
189
- index_q_ptr,
190
- index_q_stride0,
191
- index_q_stride1,
192
- index_q_cos_sin_ptr,
193
- index_q_cos_sin_stride,
194
- INDEX_Q_HALF_ROT_DIM: tl.constexpr,
195
- # MXFP4 Q outputs
196
- index_q_mxfp4_ptr, # uint8, (T, H, HEAD_DIM // 2)
197
- index_q_mxfp4_stride0,
198
- index_q_mxfp4_stride1,
199
- index_q_scale_ptr, # uint8 ue8m0, (T, H, HEAD_DIM // BLOCK)
200
- index_q_scale_stride0,
201
- index_q_scale_stride1,
202
- INDEX_Q_HEAD_DIM: tl.constexpr,
203
- MXFP4_BLOCK: tl.constexpr,
204
- # Weights (NO per-token q_scale fold for MXFP4; per-block scales stay
205
- # with the Q values in the output scale tensor).
206
- index_weights_ptr,
207
- index_weights_stride,
208
- index_weights_softmax_scale,
209
- index_weights_head_scale,
210
- index_weights_out_ptr,
211
- index_weights_out_stride,
209
+ def dispatch( # type: ignore[override]
210
+ self,
211
+ *,
212
+ dtype: torch.dtype,
213
+ num_heads: int,
214
+ head_dim: int,
215
+ rope_dim: int,
216
+ use_fnuz: bool,
217
+ ) -> CompileKey:
218
+ return self.CompileKey(
219
+ dtype=dtype,
220
+ num_heads=num_heads,
221
+ index_q_half_rot_dim=rope_dim // 2,
222
+ index_q_head_dim=head_dim,
223
+ fp8_max=224.0 if use_fnuz else 448.0,
224
+ use_fnuz=use_fnuz,
225
+ use_explicit_fma=current_platform.is_rocm(),
226
+ )
227
+
228
+ def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]:
229
+ hf_config = vllm_config.model_config.hf_config
230
+ num_heads = int(getattr(hf_config, "index_n_heads", 0) or 0)
231
+ head_dim = int(getattr(hf_config, "index_head_dim", 0) or 0)
232
+ rope_dim = int(getattr(hf_config, "qk_rope_head_dim", 0) or 0)
233
+ if num_heads <= 0 or head_dim <= 0 or rope_dim <= 0:
234
+ return []
235
+
236
+ return self._trace_dispatch(self.dispatch)(
237
+ dtype=vllm_config.model_config.dtype,
238
+ num_heads=num_heads,
239
+ head_dim=head_dim,
240
+ rope_dim=rope_dim,
241
+ use_fnuz=current_platform.fp8_dtype() == torch.float8_e4m3fnuz,
242
+ )
243
+
244
+ def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
245
+ q_stride0 = compile_key.num_heads * compile_key.index_q_head_dim
246
+ return dict(
247
+ positions=TritonWarmupTensor(torch.int64),
248
+ index_q=TritonWarmupTensor(
249
+ compile_key.dtype,
250
+ shape=(1, compile_key.num_heads, compile_key.index_q_head_dim),
251
+ strides=(q_stride0, compile_key.index_q_head_dim, 1),
252
+ ),
253
+ index_q_cos_sin_cache=TritonWarmupTensor(
254
+ torch.float32,
255
+ shape=(1, compile_key.index_q_half_rot_dim * 2),
256
+ ),
257
+ index_weights=TritonWarmupTensor(
258
+ compile_key.dtype,
259
+ shape=(1, compile_key.num_heads),
260
+ ),
261
+ index_weights_softmax_scale=1.0,
262
+ index_weights_head_scale=1.0,
263
+ index_q_fp8=TritonWarmupTensor(
264
+ current_platform.fp8_dtype(),
265
+ shape=(1, compile_key.num_heads, compile_key.index_q_head_dim),
266
+ strides=(q_stride0, compile_key.index_q_head_dim, 1),
267
+ ),
268
+ index_weights_out=TritonWarmupTensor(
269
+ torch.float32,
270
+ shape=(1, compile_key.num_heads),
271
+ ),
272
+ fp8_max=compile_key.fp8_max,
273
+ use_fnuz=compile_key.use_fnuz,
274
+ )
275
+
276
+ @kernel_launcher
277
+ def __call__(
278
+ self,
279
+ positions: torch.Tensor,
280
+ index_q: torch.Tensor,
281
+ index_q_cos_sin_cache: torch.Tensor,
282
+ index_weights: torch.Tensor,
283
+ index_weights_softmax_scale: float,
284
+ index_weights_head_scale: float,
285
+ index_q_fp8: torch.Tensor,
286
+ index_weights_out: torch.Tensor,
287
+ *,
288
+ fp8_max: float,
289
+ use_fnuz: bool,
290
+ ) -> LaunchSpec:
291
+ num_tokens = positions.shape[0]
292
+ num_index_q_heads = index_q.shape[1]
293
+ return (num_tokens, num_index_q_heads), dict(
294
+ pos_ptr=positions,
295
+ index_q_stride0=index_q.stride(0),
296
+ index_q_stride1=index_q.stride(1),
297
+ index_q_cos_sin_ptr=index_q_cos_sin_cache,
298
+ index_q_cos_sin_stride=index_q_cos_sin_cache.stride(0),
299
+ INDEX_Q_HALF_ROT_DIM=index_q_cos_sin_cache.shape[-1] // 2,
300
+ index_q_fp8_stride0=index_q_fp8.stride(0),
301
+ index_q_fp8_stride1=index_q_fp8.stride(1),
302
+ INDEX_Q_HEAD_DIM=index_q.shape[2],
303
+ index_weights_stride=index_weights.stride(0),
304
+ index_weights_out_stride=index_weights_out.stride(0),
305
+ FP8_MAX=fp8_max,
306
+ USE_FNUZ=use_fnuz,
307
+ USE_EXPLICIT_FMA=current_platform.is_rocm(),
308
+ num_warps=1,
309
+ )
310
+
311
+
312
+class FusedIndexerQRopeMxFp4TritonKernel(
313
+ VllmTritonJitKernel["FusedIndexerQRopeMxFp4TritonKernel.CompileKey"]
314
):
213
- INDEX_Q_ROT_DIM: tl.constexpr = 2 * INDEX_Q_HALF_ROT_DIM
214
- INDEX_Q_NOPE_DIM: tl.constexpr = INDEX_Q_HEAD_DIM - INDEX_Q_ROT_DIM
215
- NUM_NOPE_BLOCKS: tl.constexpr = INDEX_Q_NOPE_DIM // MXFP4_BLOCK
216
- NUM_ROPE_BLOCKS: tl.constexpr = INDEX_Q_ROT_DIM // MXFP4_BLOCK
217
- HALF_BLOCK: tl.constexpr = MXFP4_BLOCK // 2
218
- tl.static_assert(INDEX_Q_NOPE_DIM >= 0)
219
- tl.static_assert(INDEX_Q_NOPE_DIM % MXFP4_BLOCK == 0)
220
- tl.static_assert(INDEX_Q_ROT_DIM % MXFP4_BLOCK == 0)
221
- tl.static_assert(MXFP4_BLOCK % 2 == 0)
222
-
223
- tok_idx = tl.program_id(0)
224
- head_idx = tl.program_id(1)
225
-
226
- pos = tl.load(pos_ptr + tok_idx)
227
-
228
- q_base = index_q_ptr + tok_idx * index_q_stride0 + head_idx * index_q_stride1
229
- out_base = (
230
- index_q_mxfp4_ptr
231
- + tok_idx * index_q_mxfp4_stride0
232
- + head_idx * index_q_mxfp4_stride1
233
- )
234
- scale_base = (
235
- index_q_scale_ptr
236
- + tok_idx * index_q_scale_stride0
237
- + head_idx * index_q_scale_stride1
238
- )
239
-
240
- half_off = tl.arange(0, HALF_BLOCK)
241
-
242
- # ---- NoPE blocks: direct load, pair as (even-index, odd-index) values ----
243
- for b in tl.static_range(NUM_NOPE_BLOCKS):
244
- base = b * MXFP4_BLOCK
245
- x_lo = tl.load(q_base + base + half_off * 2).to(tl.float32)
246
- x_hi = tl.load(q_base + base + half_off * 2 + 1).to(tl.float32)
247
- packed, ue8m0 = _quantize_mxfp4_pair(x_lo, x_hi)
248
- tl.store(out_base + base // 2 + half_off, packed)
249
- tl.store(scale_base + b, ue8m0)
250
-
251
- # ---- RoPE blocks: apply GPT-J interleaved RoPE to the block's 16 pairs,
252
- # then quantize. Each block covers HALF_BLOCK (=16) cos/sin pairs. ----
253
- rot_q_base = q_base + INDEX_Q_NOPE_DIM
254
- for b in tl.static_range(NUM_ROPE_BLOCKS):
255
- pair_off = b * HALF_BLOCK + half_off # indices in [0, HALF_ROT_DIM)
256
- cos_b = tl.load(
257
- index_q_cos_sin_ptr + pos * index_q_cos_sin_stride + pair_off
258
- ).to(tl.float32)
259
- sin_b = tl.load(
260
- index_q_cos_sin_ptr
261
- + pos * index_q_cos_sin_stride
262
- + pair_off
263
- + INDEX_Q_HALF_ROT_DIM
315
+ @dataclass(frozen=True)
316
+ class CompileKey:
317
+ dtype: torch.dtype
318
+ num_heads: int
319
+ index_q_half_rot_dim: int
320
+ index_q_head_dim: int
321
+ mxfp4_block: int
322
+
323
+ @staticmethod
324
+ @triton.jit
325
+ def kernel(
326
+ pos_ptr,
327
+ # Index Q RoPE input (fp/bf16)
328
+ index_q_ptr,
329
+ index_q_stride0,
330
+ index_q_stride1,
331
+ index_q_cos_sin_ptr,
332
+ index_q_cos_sin_stride,
333
+ INDEX_Q_HALF_ROT_DIM: tl.constexpr,
334
+ # MXFP4 Q outputs
335
+ index_q_mxfp4_ptr, # uint8, (T, H, HEAD_DIM // 2)
336
+ index_q_mxfp4_stride0,
337
+ index_q_mxfp4_stride1,
338
+ index_q_scale_ptr, # uint8 ue8m0, (T, H, HEAD_DIM // BLOCK)
339
+ index_q_scale_stride0,
340
+ index_q_scale_stride1,
341
+ INDEX_Q_HEAD_DIM: tl.constexpr,
342
+ MXFP4_BLOCK: tl.constexpr,
343
+ # Weights (NO per-token q_scale fold for MXFP4; per-block scales stay
344
+ # with the Q values in the output scale tensor).
345
+ index_weights_ptr,
346
+ index_weights_stride,
347
+ index_weights_softmax_scale,
348
+ index_weights_head_scale,
349
+ index_weights_out_ptr,
350
+ index_weights_out_stride,
351
+ ):
352
+ INDEX_Q_ROT_DIM: tl.constexpr = 2 * INDEX_Q_HALF_ROT_DIM
353
+ INDEX_Q_NOPE_DIM: tl.constexpr = INDEX_Q_HEAD_DIM - INDEX_Q_ROT_DIM
354
+ NUM_NOPE_BLOCKS: tl.constexpr = INDEX_Q_NOPE_DIM // MXFP4_BLOCK
355
+ NUM_ROPE_BLOCKS: tl.constexpr = INDEX_Q_ROT_DIM // MXFP4_BLOCK
356
+ HALF_BLOCK: tl.constexpr = MXFP4_BLOCK // 2
357
+ tl.static_assert(INDEX_Q_NOPE_DIM >= 0)
358
+ tl.static_assert(INDEX_Q_NOPE_DIM % MXFP4_BLOCK == 0)
359
+ tl.static_assert(INDEX_Q_ROT_DIM % MXFP4_BLOCK == 0)
360
+ tl.static_assert(MXFP4_BLOCK % 2 == 0)
361
+
362
+ tok_idx = tl.program_id(0)
363
+ head_idx = tl.program_id(1)
364
+
365
+ pos = tl.load(pos_ptr + tok_idx)
366
+
367
+ q_base = index_q_ptr + tok_idx * index_q_stride0 + head_idx * index_q_stride1
368
+ out_base = (
369
+ index_q_mxfp4_ptr
370
+ + tok_idx * index_q_mxfp4_stride0
371
+ + head_idx * index_q_mxfp4_stride1
372
+ )
373
+ scale_base = (
374
+ index_q_scale_ptr
375
+ + tok_idx * index_q_scale_stride0
376
+ + head_idx * index_q_scale_stride1
377
+ )
378
+
379
+ half_off = tl.arange(0, HALF_BLOCK)
380
+
381
+ # ---- NoPE blocks: direct load, pair as (even-index, odd-index) values ----
382
+ for b in tl.static_range(NUM_NOPE_BLOCKS):
383
+ base = b * MXFP4_BLOCK
384
+ x_lo = tl.load(q_base + base + half_off * 2).to(tl.float32)
385
+ x_hi = tl.load(q_base + base + half_off * 2 + 1).to(tl.float32)
386
+ packed, ue8m0 = _quantize_mxfp4_pair(x_lo, x_hi)
387
+ tl.store(out_base + base // 2 + half_off, packed)
388
+ tl.store(scale_base + b, ue8m0)
389
+
390
+ # ---- RoPE blocks: apply GPT-J interleaved RoPE to the block's 16 pairs,
391
+ # then quantize. Each block covers HALF_BLOCK (=16) cos/sin pairs. ----
392
+ rot_q_base = q_base + INDEX_Q_NOPE_DIM
393
+ for b in tl.static_range(NUM_ROPE_BLOCKS):
394
+ pair_off = b * HALF_BLOCK + half_off # indices in [0, HALF_ROT_DIM)
395
+ cos_b = tl.load(
396
+ index_q_cos_sin_ptr + pos * index_q_cos_sin_stride + pair_off
397
+ ).to(tl.float32)
398
+ sin_b = tl.load(
399
+ index_q_cos_sin_ptr
400
+ + pos * index_q_cos_sin_stride
401
+ + pair_off
402
+ + INDEX_Q_HALF_ROT_DIM
403
+ ).to(tl.float32)
404
+ x_even = tl.load(rot_q_base + pair_off * 2).to(tl.float32)
405
+ x_odd = tl.load(rot_q_base + pair_off * 2 + 1).to(tl.float32)
406
+ r_even = x_even * cos_b - x_odd * sin_b
407
+ r_odd = x_odd * cos_b + x_even * sin_b
408
+ # bf16 roundtrip for parity with the FP8 kernel / reference numerics.
409
+ r_even = r_even.to(tl.bfloat16).to(tl.float32)
410
+ r_odd = r_odd.to(tl.bfloat16).to(tl.float32)
411
+ packed, ue8m0 = _quantize_mxfp4_pair(r_even, r_odd)
412
+ rope_byte_off = (INDEX_Q_NOPE_DIM + b * MXFP4_BLOCK) // 2
413
+ tl.store(out_base + rope_byte_off + half_off, packed)
414
+ tl.store(scale_base + NUM_NOPE_BLOCKS + b, ue8m0)
415
+
416
+ # MXFP4 weight-fold contract:
417
+ # index_weights_out = index_weights * softmax_scale * head_scale
418
+ # NOTE: q_scale is NOT folded here (contrast with the FP8 kernel above).
419
+ # MXFP4 Q emits a separate ue8m0 scale tensor of shape
420
+ # (T, H, HEAD_DIM // MXFP4_BLOCK) alongside the packed values, so each
421
+ # per-block scale is applied by the downstream MXFP4 logits kernel when
422
+ # dequantizing Q — there is no per-token scalar to fold into `weights`.
423
+ index_weights = tl.load(
424
+ index_weights_ptr + tok_idx * index_weights_stride + head_idx
425
).to(tl.float32)
265
- x_even = tl.load(rot_q_base + pair_off * 2).to(tl.float32)
266
- x_odd = tl.load(rot_q_base + pair_off * 2 + 1).to(tl.float32)
267
- r_even = x_even * cos_b - x_odd * sin_b
268
- r_odd = x_odd * cos_b + x_even * sin_b
269
- # bf16 roundtrip for parity with the FP8 kernel / reference numerics.
270
- r_even = r_even.to(tl.bfloat16).to(tl.float32)
271
- r_odd = r_odd.to(tl.bfloat16).to(tl.float32)
272
- packed, ue8m0 = _quantize_mxfp4_pair(r_even, r_odd)
273
- rope_byte_off = (INDEX_Q_NOPE_DIM + b * MXFP4_BLOCK) // 2
274
- tl.store(out_base + rope_byte_off + half_off, packed)
275
- tl.store(scale_base + NUM_NOPE_BLOCKS + b, ue8m0)
276
-
277
- # MXFP4 weight-fold contract:
278
- # index_weights_out = index_weights * softmax_scale * head_scale
279
- # NOTE: q_scale is NOT folded here (contrast with the FP8 kernel above).
280
- # MXFP4 Q emits a separate ue8m0 scale tensor of shape
281
- # (T, H, HEAD_DIM // MXFP4_BLOCK) alongside the packed values, so each
282
- # per-block scale is applied by the downstream MXFP4 logits kernel when
283
- # dequantizing Q — there is no per-token scalar to fold into `weights`.
284
- index_weights = tl.load(
285
- index_weights_ptr + tok_idx * index_weights_stride + head_idx
286
- ).to(tl.float32)
287
- index_weights *= index_weights_softmax_scale
288
- index_weights *= index_weights_head_scale
289
- tl.store(
290
- index_weights_out_ptr + tok_idx * index_weights_out_stride + head_idx,
291
- index_weights,
292
- )
426
+ index_weights *= index_weights_softmax_scale
427
+ index_weights *= index_weights_head_scale
428
+ tl.store(
429
+ index_weights_out_ptr + tok_idx * index_weights_out_stride + head_idx,
430
+ index_weights,
431
+ )
432
+
433
+ def dispatch( # type: ignore[override]
434
+ self,
435
+ *,
436
+ dtype: torch.dtype,
437
+ num_heads: int,
438
+ head_dim: int,
439
+ rope_dim: int,
440
+ ) -> CompileKey:
441
+ return self.CompileKey(
442
+ dtype=dtype,
443
+ num_heads=num_heads,
444
+ index_q_half_rot_dim=rope_dim // 2,
445
+ index_q_head_dim=head_dim,
446
+ mxfp4_block=MXFP4_BLOCK_SIZE,
447
+ )
448
+
449
+ def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]:
450
+ hf_config = vllm_config.model_config.hf_config
451
+ num_heads = int(getattr(hf_config, "index_n_heads", 0) or 0)
452
+ head_dim = int(getattr(hf_config, "index_head_dim", 0) or 0)
453
+ rope_dim = int(getattr(hf_config, "qk_rope_head_dim", 0) or 0)
454
+ if num_heads <= 0 or head_dim <= 0 or rope_dim <= 0:
455
+ return []
456
+
457
+ return self._trace_dispatch(self.dispatch)(
458
+ dtype=vllm_config.model_config.dtype,
459
+ num_heads=num_heads,
460
+ head_dim=head_dim,
461
+ rope_dim=rope_dim,
462
+ )
463
+
464
+ def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
465
+ q_stride0 = compile_key.num_heads * compile_key.index_q_head_dim
466
+ packed_head_dim = compile_key.index_q_head_dim // 2
467
+ scale_head_dim = compile_key.index_q_head_dim // compile_key.mxfp4_block
468
+ return dict(
469
+ positions=TritonWarmupTensor(torch.int64),
470
+ index_q=TritonWarmupTensor(
471
+ compile_key.dtype,
472
+ shape=(1, compile_key.num_heads, compile_key.index_q_head_dim),
473
+ strides=(q_stride0, compile_key.index_q_head_dim, 1),
474
+ ),
475
+ index_q_cos_sin_cache=TritonWarmupTensor(
476
+ torch.float32,
477
+ shape=(1, compile_key.index_q_half_rot_dim * 2),
478
+ ),
479
+ index_weights=TritonWarmupTensor(
480
+ compile_key.dtype,
481
+ shape=(1, compile_key.num_heads),
482
+ ),
483
+ index_weights_softmax_scale=1.0,
484
+ index_weights_head_scale=1.0,
485
+ index_q_packed=TritonWarmupTensor(
486
+ torch.uint8,
487
+ shape=(1, compile_key.num_heads, packed_head_dim),
488
+ ),
489
+ index_q_scale=TritonWarmupTensor(
490
+ torch.uint8,
491
+ shape=(1, compile_key.num_heads, scale_head_dim),
492
+ ),
493
+ index_weights_out=TritonWarmupTensor(
494
+ torch.float32,
495
+ shape=(1, compile_key.num_heads),
496
+ ),
497
+ )
498
+
499
+ @kernel_launcher
500
+ def __call__(
501
+ self,
502
+ positions: torch.Tensor,
503
+ index_q: torch.Tensor,
504
+ index_q_cos_sin_cache: torch.Tensor,
505
+ index_weights: torch.Tensor,
506
+ index_weights_softmax_scale: float,
507
+ index_weights_head_scale: float,
508
+ index_q_packed: torch.Tensor,
509
+ index_q_scale: torch.Tensor,
510
+ index_weights_out: torch.Tensor,
511
+ ) -> LaunchSpec:
512
+ num_tokens = positions.shape[0]
513
+ num_index_q_heads = index_q.shape[1]
514
+ return (num_tokens, num_index_q_heads), dict(
515
+ pos_ptr=positions,
516
+ index_q_stride0=index_q.stride(0),
517
+ index_q_stride1=index_q.stride(1),
518
+ index_q_cos_sin_ptr=index_q_cos_sin_cache,
519
+ index_q_cos_sin_stride=index_q_cos_sin_cache.stride(0),
520
+ INDEX_Q_HALF_ROT_DIM=index_q_cos_sin_cache.shape[-1] // 2,
521
+ index_q_mxfp4_ptr=index_q_packed,
522
+ index_q_mxfp4_stride0=index_q_packed.stride(0),
523
+ index_q_mxfp4_stride1=index_q_packed.stride(1),
524
+ index_q_scale_stride0=index_q_scale.stride(0),
525
+ index_q_scale_stride1=index_q_scale.stride(1),
526
+ INDEX_Q_HEAD_DIM=index_q.shape[2],
527
+ MXFP4_BLOCK=MXFP4_BLOCK_SIZE,
528
+ index_weights_stride=index_weights.stride(0),
529
+ index_weights_out_stride=index_weights_out.stride(0),
530
+ num_warps=1,
531
+ )
532
533
534
def fused_indexer_q_rope_quant(
@@ -386,29 +625,16 @@ def fused_indexer_q_rope_quant(
625
index_weights_out,
626
)
627
else:
389
- _fused_indexer_q_rope_mxfp4_kernel[(num_tokens, num_index_q_heads)](
628
+ _FUSED_INDEXER_Q_ROPE_MXFP4_TRITON_KERNEL(
629
positions,
630
index_q,
392
- index_q.stride(0),
393
- index_q.stride(1),
631
index_q_cos_sin_cache,
395
- index_q_cos_sin_cache.stride(0),
396
- index_q_cos_sin_cache.shape[-1] // 2,
397
- index_q_packed,
398
- index_q_packed.stride(0),
399
- index_q_packed.stride(1),
400
- index_q_scale,
401
- index_q_scale.stride(0),
402
- index_q_scale.stride(1),
403
- index_q_head_dim,
404
- MXFP4_BLOCK_SIZE,
632
index_weights,
406
- index_weights.stride(0),
633
index_weights_softmax_scale,
634
index_weights_head_scale,
635
+ index_q_packed,
636
+ index_q_scale,
637
index_weights_out,
410
- index_weights_out.stride(0),
411
- num_warps=1, # TODO: Tune this
638
)
639
640
# Values stay uint8 (2 E2M1 nibbles per byte). Scales are 4 ue8m0
@@ -453,27 +679,20 @@ def fused_indexer_q_rope_quant(
679
index_weights_out,
680
)
681
else:
456
- _fused_indexer_q_rope_quant_kernel[(num_tokens, num_index_q_heads)](
682
+ _FUSED_INDEXER_Q_ROPE_QUANT_TRITON_KERNEL(
683
positions,
684
index_q,
459
- index_q.stride(0),
460
- index_q.stride(1),
685
index_q_cos_sin_cache,
462
- index_q_cos_sin_cache.stride(0),
463
- index_q_cos_sin_cache.shape[-1] // 2,
464
- index_q_fp8,
465
- index_q_fp8.stride(0),
466
- index_q_fp8.stride(1),
467
- index_q_head_dim,
686
index_weights,
469
- index_weights.stride(0),
687
index_weights_softmax_scale,
688
index_weights_head_scale,
689
+ index_q_fp8,
690
index_weights_out,
473
- index_weights_out.stride(0),
474
- FP8_MAX=fp8_max,
475
- USE_FNUZ=use_fnuz,
476
- USE_EXPLICIT_FMA=current_platform.is_rocm(),
477
- num_warps=1, # TODO: Tune this
691
+ fp8_max=fp8_max,
692
+ use_fnuz=use_fnuz,
693
)
694
return index_q_fp8, index_weights_out
695
+
696
+
697
+_FUSED_INDEXER_Q_ROPE_QUANT_TRITON_KERNEL = FusedIndexerQRopeQuantTritonKernel()
698
+_FUSED_INDEXER_Q_ROPE_MXFP4_TRITON_KERNEL = FusedIndexerQRopeMxFp4TritonKernel()
vllm/models/deepseek_v4/common/ops/fused_inv_rope_fp8_quant.py
+304
-129
@@ -7,145 +7,327 @@ Output scale format is pre-transformed (MN-major TMA-aligned; FP32 on SM90,
7
INT32-packed UE8M0 on SM100) so fp8_einsum skips transform_sf_into_required_layout.
8
"""
9
10
+from dataclasses import dataclass
11
+from typing import Any
12
+
13
import torch
14
15
+from vllm.model_executor.warmup.jit_warmup_triton_helper import (
16
+ LaunchSpec,
17
+ TritonWarmupTensor,
18
+ VllmTritonJitKernel,
19
+ kernel_launcher,
20
+)
21
from vllm.platforms import current_platform
22
from vllm.triton_utils import tl, triton
23
from vllm.utils.torch_utils import direct_register_custom_op
24
25
17
-@triton.jit(do_not_specialize=["num_tokens"])
18
-def _fused_inv_rope_fp8_quant_per_head(
19
- o_ptr,
20
- positions_ptr,
21
- cos_sin_cache_ptr,
22
- fp8_ptr,
23
- scale_ptr,
24
- num_tokens,
25
- heads_per_group: tl.constexpr,
26
- o_stride_token,
27
- o_stride_head,
28
- cache_stride_pos,
29
- fp8_stride_group,
30
- fp8_stride_token,
31
- scale_stride_group,
32
- scale_stride_k,
33
- fp8_max: tl.constexpr,
34
- eps: tl.constexpr,
35
- QUANT_GROUP_SIZE: tl.constexpr,
36
- CHUNKS_PER_HEAD: tl.constexpr,
37
- ROPE_START: tl.constexpr,
38
- HALF_ROPE: tl.constexpr,
39
- TMA_ALIGNED_SCALES: tl.constexpr,
40
- USE_GDC: tl.constexpr,
41
- launch_pdl: tl.constexpr, # triton metadata
26
+class FusedInvRopeFP8QuantKernel(
27
+ VllmTritonJitKernel["FusedInvRopeFP8QuantKernel.CompileKey"]
28
):
43
- # Cast every stride to int64 — without this, Python-int strides are
44
- # inferred as int32 and `pid_token(int64) × stride(int32)` can lower to
45
- # int32 arithmetic, wrapping past 2³¹ for large prefill batches → IMA.
46
- pid_token = tl.program_id(0).to(tl.int64)
47
- pid_gh = tl.program_id(1).to(tl.int64)
48
- o_stride_token = o_stride_token.to(tl.int64)
49
- o_stride_head = o_stride_head.to(tl.int64)
50
- cache_stride_pos = cache_stride_pos.to(tl.int64)
51
- fp8_stride_group = fp8_stride_group.to(tl.int64)
52
- fp8_stride_token = fp8_stride_token.to(tl.int64)
53
- scale_stride_group = scale_stride_group.to(tl.int64)
54
- scale_stride_k = scale_stride_k.to(tl.int64)
55
-
56
- g = pid_gh // heads_per_group
57
- head_in_group = pid_gh % heads_per_group
58
- global_head = pid_gh
59
- qb_start = head_in_group * CHUNKS_PER_HEAD
60
- if USE_GDC:
61
- tl.extra.cuda.gdc_launch_dependents()
62
- tl.extra.cuda.gdc_wait()
63
- # Padding rows in the TMA-aligned scale buffer: fill with zero and skip quant.
64
- if pid_token >= num_tokens:
29
+ @dataclass(frozen=True)
30
+ class CompileKey:
31
+ heads_per_group: int
32
+ fp8_max: float
33
+ quant_group_size: int
34
+ chunks_per_head: int
35
+ rope_start: int
36
+ half_rope: int
37
+ tma_aligned_scales: bool
38
+ launch_pdl: bool
39
+
40
+ @staticmethod
41
+ # scale_stride_k = align(num_tokens, 4) is only 4-aligned, so its Triton int
42
+ # class varies with batch size; not specialized so one warmup key covers all
43
+ # (int64-cast, scalar addressing only, so perf-neutral).
44
+ @triton.jit(do_not_specialize=["num_tokens", "scale_stride_k"])
45
+ def kernel(
46
+ o_ptr,
47
+ positions_ptr,
48
+ cos_sin_cache_ptr,
49
+ fp8_ptr,
50
+ scale_ptr,
51
+ num_tokens,
52
+ heads_per_group: tl.constexpr,
53
+ o_stride_token,
54
+ o_stride_head,
55
+ cache_stride_pos,
56
+ fp8_stride_group,
57
+ fp8_stride_token,
58
+ scale_stride_group,
59
+ scale_stride_k,
60
+ fp8_max: tl.constexpr,
61
+ eps: tl.constexpr,
62
+ QUANT_GROUP_SIZE: tl.constexpr,
63
+ CHUNKS_PER_HEAD: tl.constexpr,
64
+ ROPE_START: tl.constexpr,
65
+ HALF_ROPE: tl.constexpr,
66
+ TMA_ALIGNED_SCALES: tl.constexpr,
67
+ launch_pdl: tl.constexpr,
68
+ ):
69
+ # Cast every stride to int64 — without this, Python-int strides are
70
+ # inferred as int32 and `pid_token(int64) × stride(int32)` can lower to
71
+ # int32 arithmetic, wrapping past 2³¹ for large prefill batches → IMA.
72
+ pid_token = tl.program_id(0).to(tl.int64)
73
+ pid_gh = tl.program_id(1).to(tl.int64)
74
+ o_stride_token = o_stride_token.to(tl.int64)
75
+ o_stride_head = o_stride_head.to(tl.int64)
76
+ cache_stride_pos = cache_stride_pos.to(tl.int64)
77
+ fp8_stride_group = fp8_stride_group.to(tl.int64)
78
+ fp8_stride_token = fp8_stride_token.to(tl.int64)
79
+ scale_stride_group = scale_stride_group.to(tl.int64)
80
+ scale_stride_k = scale_stride_k.to(tl.int64)
81
+
82
+ g = pid_gh // heads_per_group
83
+ head_in_group = pid_gh % heads_per_group
84
+ global_head = pid_gh
85
+ qb_start = head_in_group * CHUNKS_PER_HEAD
86
+ if launch_pdl:
87
+ tl.extra.cuda.gdc_launch_dependents()
88
+ tl.extra.cuda.gdc_wait()
89
+ # Padding rows in the TMA-aligned scale buffer: fill with zero and skip quant.
90
+ if pid_token >= num_tokens:
91
+ if TMA_ALIGNED_SCALES:
92
+ scale_addr = (
93
+ scale_ptr
94
+ + g * scale_stride_group
95
+ + pid_token
96
+ + head_in_group * scale_stride_k
97
+ )
98
+ tl.store(scale_addr, tl.zeros((), dtype=tl.int32))
99
+ else:
100
+ block_offsets = tl.arange(0, CHUNKS_PER_HEAD)
101
+ qb_indices = qb_start + block_offsets
102
+ scale_addrs = (
103
+ scale_ptr
104
+ + g * scale_stride_group
105
+ + pid_token
106
+ + qb_indices * scale_stride_k
107
+ )
108
+ tl.store(scale_addrs, tl.zeros((CHUNKS_PER_HEAD,), dtype=tl.float32))
109
+ return
110
+
111
+ input_base = o_ptr + pid_token * o_stride_token + global_head * o_stride_head
112
+
113
+ HEAD_DIM: tl.constexpr = CHUNKS_PER_HEAD * QUANT_GROUP_SIZE
114
+ offsets = tl.arange(0, HEAD_DIM)
115
+ x = tl.load(input_base + offsets).to(tl.float32)
116
+
117
+ rope_abs_start: tl.constexpr = (
118
+ CHUNKS_PER_HEAD - 1
119
+ ) * QUANT_GROUP_SIZE + ROPE_START
120
+ pos = tl.load(positions_ptr + pid_token)
121
+ cache_base = cos_sin_cache_ptr + pos * cache_stride_pos
122
+ is_rope = offsets >= rope_abs_start
123
+ rope_local = offsets - rope_abs_start
124
+
125
+ x_partner = tl.load(input_base + (offsets ^ 1), mask=is_rope, other=0.0).to(
126
+ tl.float32
127
+ )
128
+ cs_idx = tl.maximum(rope_local >> 1, 0)
129
+ cos_v = tl.load(cache_base + cs_idx, mask=is_rope, other=1.0)
130
+ sin_v = tl.load(cache_base + HALF_ROPE + cs_idx, mask=is_rope, other=0.0)
131
+ x_add = x * cos_v + x_partner * sin_v
132
+ x_sub = x * cos_v - x_partner * sin_v
133
+ is_even = (rope_local & 1) == 0
134
+ rotated = tl.where(is_even, x_add, x_sub)
135
+ x = tl.where(is_rope, rotated, x)
136
+
137
+ x_2d = tl.reshape(tl.abs(x), (CHUNKS_PER_HEAD, QUANT_GROUP_SIZE))
138
+ block_absmax = tl.maximum(tl.max(x_2d, axis=1), eps)
139
+ scale_raw = block_absmax * (1.0 / fp8_max)
140
+ scales = tl.math.exp2(tl.ceil(tl.log2(scale_raw)))
141
+
142
+ scales_exp = tl.reshape(
143
+ tl.broadcast_to(
144
+ tl.reshape(scales, (CHUNKS_PER_HEAD, 1)),
145
+ (CHUNKS_PER_HEAD, QUANT_GROUP_SIZE),
146
+ ),
147
+ (HEAD_DIM,),
148
+ )
149
+ x_quant = tl.clamp(x / scales_exp, -fp8_max, fp8_max).to(tl.float8e4nv)
150
+
151
+ fp8_base = (
152
+ fp8_ptr
153
+ + g * fp8_stride_group
154
+ + pid_token * fp8_stride_token
155
+ + qb_start * QUANT_GROUP_SIZE
156
+ )
157
+ tl.store(fp8_base + offsets, x_quant)
158
+
159
+ block_offsets = tl.arange(0, CHUNKS_PER_HEAD)
160
+ qb_indices = qb_start + block_offsets
161
if TMA_ALIGNED_SCALES:
162
+ scale_bits = scales.to(tl.int32, bitcast=True)
163
+ ue8m0_bytes = (scale_bits >> 23) & 0xFF
164
+ packed_val = tl.sum(ue8m0_bytes << (block_offsets * 8))
165
scale_addr = (
166
scale_ptr
167
+ g * scale_stride_group
168
+ pid_token
169
+ head_in_group * scale_stride_k
170
)
72
- tl.store(scale_addr, tl.zeros((), dtype=tl.int32))
171
+ tl.store(scale_addr, packed_val)
172
else:
74
- block_offsets = tl.arange(0, CHUNKS_PER_HEAD)
75
- qb_indices = qb_start + block_offsets
173
scale_addrs = (
174
scale_ptr
175
+ g * scale_stride_group
176
+ pid_token
177
+ qb_indices * scale_stride_k
178
)
82
- tl.store(scale_addrs, tl.zeros((CHUNKS_PER_HEAD,), dtype=tl.float32))
83
- return
179
+ tl.store(scale_addrs, scales)
180
85
- input_base = o_ptr + pid_token * o_stride_token + global_head * o_stride_head
181
+ def dispatch( # type: ignore[override]
182
+ self,
183
+ *,
184
+ heads_per_group: int,
185
+ head_dim: int,
186
+ nope_dim: int,
187
+ rope_dim: int,
188
+ quant_group_size: int,
189
+ runtime_fp8_max: float | None = None,
190
+ runtime_chunks_per_head: int | None = None,
191
+ runtime_rope_start: int | None = None,
192
+ runtime_half_rope: int | None = None,
193
+ **compile_key_fields: bool,
194
+ ) -> CompileKey:
195
+ fp8_max = (
196
+ runtime_fp8_max
197
+ if runtime_fp8_max is not None
198
+ else torch.finfo(torch.float8_e4m3fn).max
199
+ )
200
+ chunks_per_head = (
201
+ runtime_chunks_per_head
202
+ if runtime_chunks_per_head is not None
203
+ else head_dim // quant_group_size
204
+ )
205
+ rope_start = (
206
+ runtime_rope_start
207
+ if runtime_rope_start is not None
208
+ else nope_dim % quant_group_size
209
+ )
210
+ half_rope = (
211
+ runtime_half_rope if runtime_half_rope is not None else rope_dim // 2
212
+ )
213
+ return self.CompileKey(
214
+ **compile_key_fields,
215
+ heads_per_group=heads_per_group,
216
+ fp8_max=fp8_max,
217
+ quant_group_size=quant_group_size,
218
+ chunks_per_head=chunks_per_head,
219
+ rope_start=rope_start,
220
+ half_rope=half_rope,
221
+ )
222
87
- HEAD_DIM: tl.constexpr = CHUNKS_PER_HEAD * QUANT_GROUP_SIZE
88
- offsets = tl.arange(0, HEAD_DIM)
89
- x = tl.load(input_base + offsets).to(tl.float32)
223
+ def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]:
224
+ hf_config = vllm_config.model_config.hf_config
225
+ if hf_config is None:
226
+ return []
227
91
- rope_abs_start: tl.constexpr = (CHUNKS_PER_HEAD - 1) * QUANT_GROUP_SIZE + ROPE_START
92
- pos = tl.load(positions_ptr + pid_token)
93
- cache_base = cos_sin_cache_ptr + pos * cache_stride_pos
94
- is_rope = offsets >= rope_abs_start
95
- rope_local = offsets - rope_abs_start
228
+ tp_size = vllm_config.parallel_config.tensor_parallel_size
229
+ num_heads = int(getattr(hf_config, "num_attention_heads", 0) or 0)
230
+ num_groups = int(getattr(hf_config, "o_groups", 0) or 0)
231
+ if num_heads <= 0 or num_groups <= 0:
232
+ return []
233
97
- x_partner = tl.load(input_base + (offsets ^ 1), mask=is_rope, other=0.0).to(
98
- tl.float32
99
- )
100
- cs_idx = tl.maximum(rope_local >> 1, 0)
101
- cos_v = tl.load(cache_base + cs_idx, mask=is_rope, other=1.0)
102
- sin_v = tl.load(cache_base + HALF_ROPE + cs_idx, mask=is_rope, other=0.0)
103
- x_add = x * cos_v + x_partner * sin_v
104
- x_sub = x * cos_v - x_partner * sin_v
105
- is_even = (rope_local & 1) == 0
106
- rotated = tl.where(is_even, x_add, x_sub)
107
- x = tl.where(is_rope, rotated, x)
108
-
109
- x_2d = tl.reshape(tl.abs(x), (CHUNKS_PER_HEAD, QUANT_GROUP_SIZE))
110
- block_absmax = tl.maximum(tl.max(x_2d, axis=1), eps)
111
- scale_raw = block_absmax * (1.0 / fp8_max)
112
- scales = tl.math.exp2(tl.ceil(tl.log2(scale_raw)))
113
-
114
- scales_exp = tl.reshape(
115
- tl.broadcast_to(
116
- tl.reshape(scales, (CHUNKS_PER_HEAD, 1)),
117
- (CHUNKS_PER_HEAD, QUANT_GROUP_SIZE),
118
- ),
119
- (HEAD_DIM,),
120
- )
121
- x_quant = tl.clamp(x / scales_exp, -fp8_max, fp8_max).to(tl.float8e4nv)
234
+ local_heads = num_heads // tp_size
235
+ local_groups = num_groups // tp_size
236
+ if local_groups <= 0:
237
+ return []
238
123
- fp8_base = (
124
- fp8_ptr
125
- + g * fp8_stride_group
126
- + pid_token * fp8_stride_token
127
- + qb_start * QUANT_GROUP_SIZE
128
- )
129
- tl.store(fp8_base + offsets, x_quant)
130
-
131
- block_offsets = tl.arange(0, CHUNKS_PER_HEAD)
132
- qb_indices = qb_start + block_offsets
133
- if TMA_ALIGNED_SCALES:
134
- scale_bits = scales.to(tl.int32, bitcast=True)
135
- ue8m0_bytes = (scale_bits >> 23) & 0xFF
136
- packed_val = tl.sum(ue8m0_bytes << (block_offsets * 8))
137
- scale_addr = (
138
- scale_ptr
139
- + g * scale_stride_group
140
- + pid_token
141
- + head_in_group * scale_stride_k
239
+ capability = current_platform.get_device_capability()
240
+ if capability is None:
241
+ return []
242
+
243
+ head_dim = int(getattr(hf_config, "head_dim", 512) or 512)
244
+ rope_dim = int(getattr(hf_config, "qk_rope_head_dim", 64) or 64)
245
+ return self._trace_dispatch(self.dispatch)(
246
+ heads_per_group=local_heads // local_groups,
247
+ head_dim=head_dim,
248
+ nope_dim=head_dim - rope_dim,
249
+ rope_dim=rope_dim,
250
+ quant_group_size=128,
251
+ tma_aligned_scales=capability.major >= 10,
252
+ launch_pdl=(False, True),
253
)
143
- tl.store(scale_addr, packed_val)
144
- else:
145
- scale_addrs = (
146
- scale_ptr + g * scale_stride_group + pid_token + qb_indices * scale_stride_k
254
+
255
+ def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
256
+ head_dim = compile_key.chunks_per_head * compile_key.quant_group_size
257
+ fp8_dim = compile_key.heads_per_group * head_dim
258
+ scale_dtype = torch.int32 if compile_key.tma_aligned_scales else torch.float32
259
+ return dict(
260
+ o=TritonWarmupTensor(
261
+ torch.bfloat16,
262
+ shape=(1, compile_key.heads_per_group, head_dim),
263
+ ),
264
+ positions=TritonWarmupTensor(torch.int64),
265
+ cos_sin_cache=TritonWarmupTensor(
266
+ torch.float32,
267
+ shape=(1, compile_key.half_rope * 2),
268
+ ),
269
+ fp8_buf=TritonWarmupTensor(
270
+ torch.float8_e4m3fn,
271
+ shape=(compile_key.heads_per_group, 1, fp8_dim),
272
+ strides=(fp8_dim, fp8_dim, 1),
273
+ ),
274
+ scale_buf=TritonWarmupTensor(
275
+ scale_dtype,
276
+ shape=(compile_key.heads_per_group, 1, 16),
277
+ strides=(16, 16, 16),
278
+ ),
279
+ num_tokens=1,
280
+ heads_per_group=compile_key.heads_per_group,
281
+ quant_group_size=compile_key.quant_group_size,
282
+ chunks_per_head=compile_key.chunks_per_head,
283
+ rope_start=compile_key.rope_start,
284
+ half_rope=compile_key.half_rope,
285
+ tma_aligned_scales=compile_key.tma_aligned_scales,
286
+ fp8_max=compile_key.fp8_max,
287
+ launch_pdl=compile_key.launch_pdl,
288
+ grid=(1, compile_key.heads_per_group),
289
+ )
290
+
291
+ @kernel_launcher
292
+ def __call__(
293
+ self,
294
+ o: torch.Tensor,
295
+ positions: torch.Tensor,
296
+ cos_sin_cache: torch.Tensor,
297
+ fp8_buf: torch.Tensor,
298
+ scale_buf: torch.Tensor,
299
+ num_tokens: int,
300
+ *,
301
+ heads_per_group: int,
302
+ quant_group_size: int,
303
+ chunks_per_head: int,
304
+ rope_start: int,
305
+ half_rope: int,
306
+ tma_aligned_scales: bool,
307
+ fp8_max: float,
308
+ launch_pdl: bool,
309
+ grid: tuple[int, int],
310
+ ) -> LaunchSpec:
311
+ return grid, dict(
312
+ fp8_ptr=fp8_buf,
313
+ scale_ptr=scale_buf,
314
+ o_stride_token=o.stride(0),
315
+ o_stride_head=o.stride(1),
316
+ cache_stride_pos=cos_sin_cache.stride(0),
317
+ fp8_stride_group=fp8_buf.stride(0),
318
+ fp8_stride_token=fp8_buf.stride(1),
319
+ scale_stride_group=scale_buf.stride(0),
320
+ scale_stride_k=scale_buf.stride(2),
321
+ eps=1e-10,
322
+ QUANT_GROUP_SIZE=quant_group_size,
323
+ CHUNKS_PER_HEAD=chunks_per_head,
324
+ ROPE_START=rope_start,
325
+ HALF_ROPE=half_rope,
326
+ TMA_ALIGNED_SCALES=tma_aligned_scales,
327
+ launch_pdl=launch_pdl,
328
+ num_stages=1,
329
+ num_warps=1,
330
)
148
- tl.store(scale_addrs, scales)
331
332
333
def fused_inv_rope_fp8_quant(
@@ -256,8 +438,8 @@ def _fused_inv_rope_fp8_quant_kernel_impl(
438
(scale_inner * tma_aligned_T, 1, tma_aligned_T),
439
)
440
grid = (tma_aligned_T, n_groups * heads_per_group)
259
- use_gdc = current_platform.is_arch_support_pdl()
260
- _fused_inv_rope_fp8_quant_per_head[grid](
441
+ launch_pdl = current_platform.is_arch_support_pdl()
442
+ _FUSED_INV_ROPE_FP8_QUANT_KERNEL(
443
o,
444
positions,
445
cos_sin_cache,
@@ -265,24 +447,14 @@ def _fused_inv_rope_fp8_quant_kernel_impl(
447
scale_buf,
448
num_tokens,
449
heads_per_group=heads_per_group,
268
- o_stride_token=o.stride(0),
269
- o_stride_head=o.stride(1),
270
- cache_stride_pos=cos_sin_cache.stride(0),
271
- fp8_stride_group=fp8_buf.stride(0),
272
- fp8_stride_token=fp8_buf.stride(1),
273
- scale_stride_group=scale_buf.stride(0),
274
- scale_stride_k=scale_buf.stride(2),
450
+ quant_group_size=quant_group_size,
451
+ chunks_per_head=chunks_per_head,
452
+ rope_start=rope_start,
453
+ half_rope=half_rope,
454
+ tma_aligned_scales=tma_aligned_scales,
455
fp8_max=fp8_max,
276
- eps=1e-10,
277
- QUANT_GROUP_SIZE=quant_group_size,
278
- CHUNKS_PER_HEAD=chunks_per_head,
279
- ROPE_START=rope_start,
280
- HALF_ROPE=half_rope,
281
- TMA_ALIGNED_SCALES=tma_aligned_scales,
282
- USE_GDC=use_gdc,
283
- launch_pdl=use_gdc,
284
- num_stages=1,
285
- num_warps=1,
456
+ launch_pdl=launch_pdl,
457
+ grid=grid,
458
)
459
return fp8_buf, scale_buf
460
@@ -326,3 +498,6 @@ direct_register_custom_op(
498
op_func=_fused_inv_rope_fp8_quant_kernel_impl,
499
fake_impl=_fused_inv_rope_fp8_quant_kernel_fake,
500
)
501
+
502
+
503
+_FUSED_INV_ROPE_FP8_QUANT_KERNEL = FusedInvRopeFP8QuantKernel()
vllm/models/deepseek_v4/common/ops/fused_mtp_input_rmsnorm.py
+286
-149
@@ -17,9 +17,19 @@ A single grid (T, hc_mult+1) drives both norms: task 0 is enorm on
17
inputs_embeds[token, :], task k+1 is hnorm on previous_hidden_states[token, k, :].
18
"""
19
20
+from dataclasses import dataclass
21
+from typing import Any
22
+
23
import torch
24
25
+from vllm.model_executor.warmup.jit_warmup_triton_helper import (
26
+ LaunchSpec,
27
+ TritonWarmupTensor,
28
+ VllmTritonJitKernel,
29
+ kernel_launcher,
30
+)
31
from vllm.triton_utils import tl, triton
32
+from vllm.utils.math_utils import next_power_of_2
33
34
35
@triton.jit
@@ -40,164 +50,291 @@ def _rmsnorm_row(
50
tl.store(out_row_ptr + block, y.to(out_row_ptr.dtype.element_ty), mask=mask)
51
52
43
-@triton.jit
44
-def _fused_mtp_input_rmsnorm_kernel(
45
- inputs_embeds_ptr,
46
- positions_ptr,
47
- prev_hidden_ptr,
48
- enorm_weight_ptr,
49
- hnorm_weight_ptr,
50
- enorm_out_ptr,
51
- hnorm_out_ptr,
52
- eps,
53
- HIDDEN: tl.constexpr,
54
- HC_MULT: tl.constexpr,
55
- BLOCK_SIZE: tl.constexpr,
53
+class FusedMTPInputRMSNormKernel(
54
+ VllmTritonJitKernel["FusedMTPInputRMSNormKernel.CompileKey"]
55
):
57
- # int64 token index so per-token offsets don't overflow int32 at
58
- # large num_tokens (matches the convention in fused_q_kv_rmsnorm).
59
- token_idx = tl.program_id(0).to(tl.int64)
60
- pid_task = tl.program_id(1)
61
-
62
- block = tl.arange(0, BLOCK_SIZE)
63
- mask = block < HIDDEN
64
-
65
- if pid_task == 0:
66
- # enorm path: load inputs_embeds[token, :] then zero-mask at pos==0.
67
- # Math is preserved: pos==0 → x=0 → variance=0 → RMSNorm output is 0
68
- # regardless of weight, matching torch.where(pos==0, 0, x) + RMSNorm.
69
- pos = tl.load(positions_ptr + token_idx)
70
- keep = pos != 0
71
- x = tl.load(
72
- inputs_embeds_ptr + token_idx * HIDDEN + block, mask=mask, other=0.0
56
+ @dataclass(frozen=True)
57
+ class CompileKey:
58
+ dtype: torch.dtype
59
+ hidden: int
60
+ hc_mult: int
61
+ block_size: int
62
+ eps: float
63
+
64
+ @staticmethod
65
+ @triton.jit
66
+ def kernel(
67
+ inputs_embeds_ptr,
68
+ positions_ptr,
69
+ prev_hidden_ptr,
70
+ enorm_weight_ptr,
71
+ hnorm_weight_ptr,
72
+ enorm_out_ptr,
73
+ hnorm_out_ptr,
74
+ eps,
75
+ HIDDEN: tl.constexpr,
76
+ HC_MULT: tl.constexpr,
77
+ BLOCK_SIZE: tl.constexpr,
78
+ ):
79
+ # int64 token index so per-token offsets don't overflow int32 at
80
+ # large num_tokens (matches the convention in fused_q_kv_rmsnorm).
81
+ token_idx = tl.program_id(0).to(tl.int64)
82
+ pid_task = tl.program_id(1)
83
+
84
+ block = tl.arange(0, BLOCK_SIZE)
85
+ mask = block < HIDDEN
86
+
87
+ if pid_task == 0:
88
+ # enorm path: load inputs_embeds[token, :] then zero-mask at pos==0.
89
+ # Math is preserved: pos==0 → x=0 → variance=0 → RMSNorm output is 0
90
+ # regardless of weight, matching torch.where(pos==0, 0, x) + RMSNorm.
91
+ pos = tl.load(positions_ptr + token_idx)
92
+ keep = pos != 0
93
+ x = tl.load(
94
+ inputs_embeds_ptr + token_idx * HIDDEN + block, mask=mask, other=0.0
95
+ )
96
+ x = tl.where(keep, x, 0.0)
97
+ _rmsnorm_row(
98
+ x,
99
+ enorm_weight_ptr,
100
+ enorm_out_ptr + token_idx * HIDDEN,
101
+ block,
102
+ mask,
103
+ eps,
104
+ HIDDEN,
105
+ )
106
+ else:
107
+ # hnorm path: load prev_hidden[token, slot, :].
108
+ slot = pid_task - 1
109
+ row_offset = (token_idx * HC_MULT + slot) * HIDDEN
110
+ x = tl.load(prev_hidden_ptr + row_offset + block, mask=mask, other=0.0)
111
+ _rmsnorm_row(
112
+ x,
113
+ hnorm_weight_ptr,
114
+ hnorm_out_ptr + row_offset,
115
+ block,
116
+ mask,
117
+ eps,
118
+ HIDDEN,
119
+ )
120
+
121
+ def dispatch( # type: ignore[override]
122
+ self,
123
+ *,
124
+ dtype: torch.dtype,
125
+ hidden: int,
126
+ hc_mult: int,
127
+ eps: float,
128
+ ) -> CompileKey:
129
+ return self.CompileKey(
130
+ dtype=dtype,
131
+ hidden=hidden,
132
+ hc_mult=hc_mult,
133
+ block_size=next_power_of_2(hidden),
134
+ eps=eps,
135
)
74
- x = tl.where(keep, x, 0.0)
75
- _rmsnorm_row(
76
- x,
77
- enorm_weight_ptr,
78
- enorm_out_ptr + token_idx * HIDDEN,
79
- block,
80
- mask,
81
- eps,
82
- HIDDEN,
136
+
137
+ def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]:
138
+ hf_config = vllm_config.model_config.hf_config
139
+ if int(getattr(hf_config, "num_nextn_predict_layers", 0) or 0) <= 0:
140
+ return []
141
+
142
+ return self._trace_dispatch(self.dispatch)(
143
+ dtype=vllm_config.model_config.dtype,
144
+ hidden=int(hf_config.hidden_size),
145
+ hc_mult=int(hf_config.hc_mult),
146
+ eps=float(hf_config.rms_norm_eps),
147
)
84
- else:
85
- # hnorm path: load prev_hidden[token, slot, :].
86
- slot = pid_task - 1
87
- row_offset = (token_idx * HC_MULT + slot) * HIDDEN
88
- x = tl.load(prev_hidden_ptr + row_offset + block, mask=mask, other=0.0)
148
+
149
+ def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
150
+ return dict(
151
+ inputs_embeds=TritonWarmupTensor(
152
+ compile_key.dtype,
153
+ shape=(1, compile_key.hidden),
154
+ ),
155
+ positions=TritonWarmupTensor(torch.int64),
156
+ previous_hidden_states=TritonWarmupTensor(
157
+ compile_key.dtype,
158
+ shape=(1, compile_key.hc_mult, compile_key.hidden),
159
+ ),
160
+ enorm_weight=TritonWarmupTensor(
161
+ compile_key.dtype,
162
+ shape=(compile_key.hidden,),
163
+ ),
164
+ hnorm_weight=TritonWarmupTensor(
165
+ compile_key.dtype,
166
+ shape=(compile_key.hidden,),
167
+ ),
168
+ _outputs=(
169
+ TritonWarmupTensor(
170
+ compile_key.dtype,
171
+ shape=(1, compile_key.hidden),
172
+ ),
173
+ TritonWarmupTensor(
174
+ compile_key.dtype,
175
+ shape=(1, compile_key.hc_mult, compile_key.hidden),
176
+ ),
177
+ ),
178
+ eps=compile_key.eps,
179
+ hc_mult=compile_key.hc_mult,
180
+ )
181
+
182
+ @kernel_launcher
183
+ def __call__(
184
+ self,
185
+ inputs_embeds: torch.Tensor,
186
+ positions: torch.Tensor,
187
+ previous_hidden_states: torch.Tensor,
188
+ enorm_weight: torch.Tensor,
189
+ hnorm_weight: torch.Tensor,
190
+ eps: float,
191
+ hc_mult: int,
192
+ *,
193
+ _outputs: tuple[Any, Any] | None = None,
194
+ ) -> LaunchSpec:
195
+ num_tokens, hidden = inputs_embeds.shape
196
+ if _outputs is None:
197
+ assert inputs_embeds.ndim == 2
198
+ assert previous_hidden_states.ndim == 3
199
+ assert previous_hidden_states.shape[1] == hc_mult
200
+ assert inputs_embeds.shape[0] == previous_hidden_states.shape[0], (
201
+ "token dim mismatch"
202
+ )
203
+ assert (
204
+ hidden
205
+ == previous_hidden_states.shape[2]
206
+ == enorm_weight.shape[0]
207
+ == hnorm_weight.shape[0]
208
+ )
209
+ assert (
210
+ inputs_embeds.is_contiguous() and previous_hidden_states.is_contiguous()
211
+ )
212
+ assert enorm_weight.is_contiguous() and hnorm_weight.is_contiguous()
213
+ _outputs = (
214
+ torch.empty_like(inputs_embeds),
215
+ torch.empty_like(previous_hidden_states),
216
+ )
217
+
218
+ enorm_out, hnorm_out = _outputs
219
+ grid = None if num_tokens == 0 else (num_tokens, hc_mult + 1)
220
+ return (
221
+ grid,
222
+ dict(
223
+ prev_hidden_ptr=previous_hidden_states,
224
+ enorm_out_ptr=enorm_out,
225
+ hnorm_out_ptr=hnorm_out,
226
+ HIDDEN=hidden,
227
+ HC_MULT=hc_mult,
228
+ BLOCK_SIZE=next_power_of_2(hidden),
229
+ ),
230
+ _outputs,
231
+ )
232
+
233
+
234
+class MTPSharedHeadRMSNormKernel(
235
+ VllmTritonJitKernel["MTPSharedHeadRMSNormKernel.CompileKey"]
236
+):
237
+ @dataclass(frozen=True)
238
+ class CompileKey:
239
+ dtype: torch.dtype
240
+ hidden: int
241
+ block_size: int
242
+ eps: float
243
+
244
+ @staticmethod
245
+ @triton.jit
246
+ def kernel(
247
+ x_ptr,
248
+ weight_ptr,
249
+ out_ptr,
250
+ eps,
251
+ HIDDEN: tl.constexpr,
252
+ BLOCK_SIZE: tl.constexpr,
253
+ ):
254
+ token_idx = tl.program_id(0).to(tl.int64)
255
+ block = tl.arange(0, BLOCK_SIZE)
256
+ mask = block < HIDDEN
257
+ x = tl.load(x_ptr + token_idx * HIDDEN + block, mask=mask, other=0.0)
258
_rmsnorm_row(
259
x,
91
- hnorm_weight_ptr,
92
- hnorm_out_ptr + row_offset,
260
+ weight_ptr,
261
+ out_ptr + token_idx * HIDDEN,
262
block,
263
mask,
264
eps,
265
HIDDEN,
266
)
267
268
+ def dispatch( # type: ignore[override]
269
+ self,
270
+ *,
271
+ dtype: torch.dtype,
272
+ hidden: int,
273
+ eps: float,
274
+ ) -> CompileKey:
275
+ return self.CompileKey(
276
+ dtype=dtype,
277
+ hidden=hidden,
278
+ block_size=next_power_of_2(hidden),
279
+ eps=eps,
280
+ )
281
100
-@triton.jit
101
-def _mtp_shared_head_rmsnorm_kernel(
102
- x_ptr,
103
- weight_ptr,
104
- out_ptr,
105
- eps,
106
- HIDDEN: tl.constexpr,
107
- BLOCK_SIZE: tl.constexpr,
108
-):
109
- token_idx = tl.program_id(0).to(tl.int64)
110
- block = tl.arange(0, BLOCK_SIZE)
111
- mask = block < HIDDEN
112
- x = tl.load(x_ptr + token_idx * HIDDEN + block, mask=mask, other=0.0)
113
- _rmsnorm_row(
114
- x,
115
- weight_ptr,
116
- out_ptr + token_idx * HIDDEN,
117
- block,
118
- mask,
119
- eps,
120
- HIDDEN,
121
- )
122
-
123
-
124
-def mtp_shared_head_rmsnorm(
125
- hidden_states: torch.Tensor,
126
- weight: torch.Tensor,
127
- eps: float,
128
-) -> torch.Tensor:
129
- """RMSNorm for MTP's SharedHead.norm, on (T, H) bf16 input.
130
-
131
- Uses the same ``_rmsnorm_row`` body as ``fused_mtp_input_rmsnorm`` so the
132
- MTP draft path runs one consistent RMSNorm implementation end to end.
133
- """
134
- assert hidden_states.ndim == 2
135
- assert hidden_states.is_contiguous()
136
- assert weight.is_contiguous()
137
- num_tokens, hidden = hidden_states.shape
138
- out = torch.empty_like(hidden_states)
139
- if num_tokens == 0:
140
- return out
141
- block_size = triton.next_power_of_2(hidden)
142
- _mtp_shared_head_rmsnorm_kernel[(num_tokens,)](
143
- hidden_states,
144
- weight,
145
- out,
146
- eps,
147
- HIDDEN=hidden,
148
- BLOCK_SIZE=block_size,
149
- )
150
- return out
151
-
152
-
153
-def fused_mtp_input_rmsnorm(
154
- inputs_embeds: torch.Tensor,
155
- positions: torch.Tensor,
156
- previous_hidden_states: torch.Tensor,
157
- enorm_weight: torch.Tensor,
158
- hnorm_weight: torch.Tensor,
159
- eps: float,
160
- hc_mult: int,
161
-) -> tuple[torch.Tensor, torch.Tensor]:
162
- """Returns (enorm_out, hnorm_out).
163
-
164
- enorm_out has the same shape as inputs_embeds (2D, [T, H]).
165
- hnorm_out has the same shape as previous_hidden_states (3D, [T, hc_mult, H]).
166
- previous_hidden_states must already be reshaped to 3D.
167
- """
168
- assert inputs_embeds.ndim == 2
169
- assert previous_hidden_states.ndim == 3
170
- assert previous_hidden_states.shape[1] == hc_mult
171
- assert inputs_embeds.shape[0] == previous_hidden_states.shape[0], (
172
- "token dim mismatch"
173
- )
174
- assert (
175
- inputs_embeds.shape[1]
176
- == previous_hidden_states.shape[2]
177
- == enorm_weight.shape[0]
178
- == hnorm_weight.shape[0]
179
- )
180
- assert inputs_embeds.is_contiguous() and previous_hidden_states.is_contiguous()
181
- assert enorm_weight.is_contiguous() and hnorm_weight.is_contiguous()
182
-
183
- num_tokens, hidden = inputs_embeds.shape
184
- enorm_out = torch.empty_like(inputs_embeds)
185
- hnorm_out = torch.empty_like(previous_hidden_states)
186
- if num_tokens == 0:
187
- return enorm_out, hnorm_out
188
-
189
- block_size = triton.next_power_of_2(hidden)
190
- _fused_mtp_input_rmsnorm_kernel[(num_tokens, hc_mult + 1)](
191
- inputs_embeds,
192
- positions,
193
- previous_hidden_states,
194
- enorm_weight,
195
- hnorm_weight,
196
- enorm_out,
197
- hnorm_out,
198
- eps,
199
- HIDDEN=hidden,
200
- HC_MULT=hc_mult,
201
- BLOCK_SIZE=block_size,
202
- )
203
- return enorm_out, hnorm_out
282
+ def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]:
283
+ hf_config = vllm_config.model_config.hf_config
284
+ if int(getattr(hf_config, "num_nextn_predict_layers", 0) or 0) <= 0:
285
+ return []
286
+
287
+ return self._trace_dispatch(self.dispatch)(
288
+ dtype=vllm_config.model_config.dtype,
289
+ hidden=int(hf_config.hidden_size),
290
+ eps=float(hf_config.rms_norm_eps),
291
+ )
292
+
293
+ def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
294
+ return dict(
295
+ hidden_states=TritonWarmupTensor(
296
+ compile_key.dtype,
297
+ shape=(1, compile_key.hidden),
298
+ ),
299
+ weight=TritonWarmupTensor(
300
+ compile_key.dtype,
301
+ shape=(compile_key.hidden,),
302
+ ),
303
+ _output=TritonWarmupTensor(
304
+ compile_key.dtype,
305
+ shape=(1, compile_key.hidden),
306
+ ),
307
+ eps=compile_key.eps,
308
+ )
309
+
310
+ @kernel_launcher
311
+ def __call__(
312
+ self,
313
+ hidden_states: torch.Tensor,
314
+ weight: torch.Tensor,
315
+ eps: float,
316
+ *,
317
+ _output: Any | None = None,
318
+ ) -> LaunchSpec:
319
+ num_tokens, hidden = hidden_states.shape
320
+ if _output is None:
321
+ assert hidden_states.ndim == 2
322
+ assert hidden_states.is_contiguous()
323
+ assert weight.is_contiguous()
324
+ _output = torch.empty_like(hidden_states)
325
+
326
+ grid = None if num_tokens == 0 else (num_tokens,)
327
+ return (
328
+ grid,
329
+ dict(
330
+ x_ptr=hidden_states,
331
+ out_ptr=_output,
332
+ HIDDEN=hidden,
333
+ BLOCK_SIZE=next_power_of_2(hidden),
334
+ ),
335
+ _output,
336
+ )
337
+
338
+
339
+_FUSED_MTP_INPUT_RMSNORM_KERNEL = FusedMTPInputRMSNormKernel()
340
+_MTP_SHARED_HEAD_RMSNORM_KERNEL = MTPSharedHeadRMSNormKernel()
vllm/models/deepseek_v4/common/ops/save_partial_states.py
+184
-92
@@ -1,101 +1,193 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+from dataclasses import dataclass
5
+from typing import Any
6
+
7
import torch
8
9
+from vllm.model_executor.warmup.jit_warmup_triton_helper import (
10
+ LaunchSpec,
11
+ TritonWarmupTensor,
12
+ VllmTritonJitKernel,
13
+ kernel_launcher,
14
+)
15
from vllm.triton_utils import tl, triton
16
+from vllm.utils.math_utils import next_power_of_2
17
18
9
-def save_partial_states(
10
- kv: torch.Tensor,
11
- score: torch.Tensor,
12
- ape: torch.Tensor,
13
- positions: torch.Tensor,
14
- state_cache: torch.Tensor,
15
- slot_mapping: torch.Tensor,
16
- block_size: int,
17
- state_width: int,
18
- compress_ratio: int,
19
- pdl_kwargs: dict | None = None,
20
-) -> None:
21
- """Write packed [kv, score+ape] partial states into the compressor cache.
22
-
23
- One program per token; pads (slot_id == -1) are skipped.
24
- """
25
- num_actual = slot_mapping.shape[0]
26
- head_size = kv.shape[-1]
27
- _save_partial_states_kernel[(num_actual,)](
28
- kv,
29
- kv.stride(0),
30
- score,
31
- score.stride(0),
32
- ape,
33
- ape.stride(0),
34
- positions,
35
- state_cache,
36
- state_cache.stride(0),
37
- state_cache.stride(1),
38
- slot_mapping,
39
- block_size,
40
- HEAD_SIZE=head_size,
41
- TRITON_BLOCK_SIZE=triton.next_power_of_2(head_size),
42
- STATE_WIDTH=state_width,
43
- COMPRESS_RATIO=compress_ratio,
44
- **(pdl_kwargs or {}),
45
- )
46
-
47
-
48
-@triton.jit
49
-def _save_partial_states_kernel(
50
- kv_ptr,
51
- kv_stride,
52
- score_ptr,
53
- score_stride,
54
- ape_ptr,
55
- ape_stride,
56
- positions_ptr,
57
- state_cache_ptr,
58
- state_cache_stride0,
59
- state_cache_stride1,
60
- slot_mapping_ptr,
61
- block_size,
62
- HEAD_SIZE: tl.constexpr,
63
- TRITON_BLOCK_SIZE: tl.constexpr,
64
- # state_cache last dim packs [kv_state, score_state], each STATE_WIDTH wide.
65
- STATE_WIDTH: tl.constexpr,
66
- COMPRESS_RATIO: tl.constexpr,
19
+class SavePartialStatesKernel(
20
+ VllmTritonJitKernel["SavePartialStatesKernel.CompileKey"]
21
):
68
- token_idx = tl.program_id(0)
69
- slot_id = tl.load(slot_mapping_ptr + token_idx)
70
-
71
- # Skip padded / invalid tokens (slot_id == -1 is the PAD sentinel used
72
- # by vLLM). During CUDA graph replay the batch may contain padding
73
- # tokens whose slot_mapping is -1; writing to kv_state[-1] would be an
74
- # illegal memory access.
75
- if slot_id < 0:
76
- return
77
-
78
- block_idx = slot_id // block_size
79
- pos_in_block = slot_id % block_size
80
- base_ptr = (
81
- state_cache_ptr
82
- + block_idx * state_cache_stride0
83
- + pos_in_block * state_cache_stride1
84
- )
85
-
86
- block = tl.arange(0, TRITON_BLOCK_SIZE)
87
- mask = block < HEAD_SIZE
88
-
89
- kv = tl.load(kv_ptr + token_idx * kv_stride + block, mask=mask)
90
- tl.store(base_ptr + block, kv, mask=mask)
91
-
92
- # Fused: score += ape[position % compress_ratio]
93
- position = tl.load(positions_ptr + token_idx)
94
- ape_row = position % COMPRESS_RATIO
95
- ape = tl.load(ape_ptr + ape_row * ape_stride + block, mask=mask)
96
- score = tl.load(score_ptr + token_idx * score_stride + block, mask=mask)
97
- tl.store(
98
- base_ptr + STATE_WIDTH + block,
99
- score + ape,
100
- mask=mask,
101
- )
22
+ @dataclass(frozen=True)
23
+ class CompileKey:
24
+ head_size: int
25
+ triton_block_size: int
26
+ state_width: int
27
+ compress_ratio: int
28
+ kv_stride: int
29
+ score_stride: int
30
+ ape_stride: int
31
+ state_cache_stride0: int
32
+ state_cache_stride1: int
33
+ block_size: int
34
+
35
+ @staticmethod
36
+ @triton.jit
37
+ def kernel(
38
+ kv_ptr,
39
+ kv_stride,
40
+ score_ptr,
41
+ score_stride,
42
+ ape_ptr,
43
+ ape_stride,
44
+ positions_ptr,
45
+ state_cache_ptr,
46
+ state_cache_stride0,
47
+ state_cache_stride1,
48
+ slot_mapping_ptr,
49
+ block_size,
50
+ HEAD_SIZE: tl.constexpr,
51
+ TRITON_BLOCK_SIZE: tl.constexpr,
52
+ # state_cache last dim packs [kv_state, score_state], each STATE_WIDTH wide.
53
+ STATE_WIDTH: tl.constexpr,
54
+ COMPRESS_RATIO: tl.constexpr,
55
+ ):
56
+ token_idx = tl.program_id(0)
57
+ slot_id = tl.load(slot_mapping_ptr + token_idx)
58
+
59
+ # Skip padded / invalid tokens (slot_id == -1 is the PAD sentinel used
60
+ # by vLLM). During CUDA graph replay the batch may contain padding
61
+ # tokens whose slot_mapping is -1; writing to kv_state[-1] would be an
62
+ # illegal memory access.
63
+ if slot_id < 0:
64
+ return
65
+
66
+ block_idx = slot_id // block_size
67
+ pos_in_block = slot_id % block_size
68
+ base_ptr = (
69
+ state_cache_ptr
70
+ + block_idx * state_cache_stride0
71
+ + pos_in_block * state_cache_stride1
72
+ )
73
+
74
+ block = tl.arange(0, TRITON_BLOCK_SIZE)
75
+ mask = block < HEAD_SIZE
76
+
77
+ kv = tl.load(kv_ptr + token_idx * kv_stride + block, mask=mask)
78
+ tl.store(base_ptr + block, kv, mask=mask)
79
+
80
+ # Fused: score += ape[position % compress_ratio]
81
+ position = tl.load(positions_ptr + token_idx)
82
+ ape_row = position % COMPRESS_RATIO
83
+ ape = tl.load(ape_ptr + ape_row * ape_stride + block, mask=mask)
84
+ score = tl.load(score_ptr + token_idx * score_stride + block, mask=mask)
85
+ tl.store(
86
+ base_ptr + STATE_WIDTH + block,
87
+ score + ape,
88
+ mask=mask,
89
+ )
90
+
91
+ def dispatch( # type: ignore[override]
92
+ self,
93
+ *,
94
+ head_size: int,
95
+ **compile_key_fields: int,
96
+ ) -> CompileKey:
97
+ return self.CompileKey(
98
+ **compile_key_fields,
99
+ head_size=head_size,
100
+ triton_block_size=next_power_of_2(head_size),
101
+ )
102
+
103
+ def get_warmup_keys(
104
+ self,
105
+ *,
106
+ head_dim: int,
107
+ compress_ratio: int,
108
+ ) -> list[CompileKey]:
109
+ if head_dim <= 0 or compress_ratio not in (4, 128):
110
+ return []
111
+
112
+ coefficient = 2 if compress_ratio == 4 else 1
113
+ return self._trace_dispatch(self.dispatch)(
114
+ head_size=coefficient * head_dim,
115
+ state_width=coefficient * head_dim,
116
+ compress_ratio=compress_ratio,
117
+ kv_stride=2 * coefficient * head_dim,
118
+ score_stride=2 * coefficient * head_dim,
119
+ ape_stride=coefficient * head_dim,
120
+ state_cache_stride0=8 * coefficient * head_dim,
121
+ state_cache_stride1=2 * coefficient * head_dim,
122
+ block_size=4 if compress_ratio == 4 else 8,
123
+ )
124
+
125
+ def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
126
+ return dict(
127
+ kv=TritonWarmupTensor(
128
+ torch.float32,
129
+ shape=(1, compile_key.head_size),
130
+ strides=(compile_key.kv_stride, 1),
131
+ ),
132
+ score=TritonWarmupTensor(
133
+ torch.float32,
134
+ shape=(1, 1),
135
+ strides=(compile_key.score_stride, 1),
136
+ ),
137
+ ape=TritonWarmupTensor(
138
+ torch.float32,
139
+ shape=(1, 1),
140
+ strides=(compile_key.ape_stride, 1),
141
+ ),
142
+ positions=TritonWarmupTensor(torch.int64),
143
+ state_cache=TritonWarmupTensor(
144
+ torch.float32,
145
+ shape=(1, 1, compile_key.state_width),
146
+ strides=(
147
+ compile_key.state_cache_stride0,
148
+ compile_key.state_cache_stride1,
149
+ 1,
150
+ ),
151
+ ),
152
+ slot_mapping=TritonWarmupTensor(torch.int64),
153
+ block_size=compile_key.block_size,
154
+ state_width=compile_key.state_width,
155
+ compress_ratio=compile_key.compress_ratio,
156
+ )
157
+
158
+ @kernel_launcher
159
+ def __call__(
160
+ self,
161
+ kv: torch.Tensor,
162
+ score: torch.Tensor,
163
+ ape: torch.Tensor,
164
+ positions: torch.Tensor,
165
+ state_cache: torch.Tensor,
166
+ slot_mapping: torch.Tensor,
167
+ block_size: int,
168
+ state_width: int,
169
+ compress_ratio: int,
170
+ pdl_kwargs: dict | None = None,
171
+ ) -> LaunchSpec:
172
+ """Write packed [kv, score+ape] partial states into the compressor cache.
173
+
174
+ One program per token; pads (slot_id == -1) are skipped.
175
+ """
176
+ num_actual = slot_mapping.shape[0]
177
+ head_size = kv.shape[-1]
178
+ return (num_actual,), dict(
179
+ kv_stride=kv.stride(0),
180
+ score_stride=score.stride(0),
181
+ ape_stride=ape.stride(0),
182
+ state_cache_stride0=state_cache.stride(0),
183
+ state_cache_stride1=state_cache.stride(1),
184
+ block_size=block_size,
185
+ HEAD_SIZE=head_size,
186
+ TRITON_BLOCK_SIZE=next_power_of_2(head_size),
187
+ STATE_WIDTH=state_width,
188
+ COMPRESS_RATIO=compress_ratio,
189
+ **(pdl_kwargs or {}),
190
+ )
191
+
192
+
193
+_SAVE_PARTIAL_STATES_KERNEL = SavePartialStatesKernel()
vllm/models/deepseek_v4/compressor.py
+14
-2
@@ -18,7 +18,7 @@ from vllm.models.deepseek_v4.common.ops.fused_compress_quant_cache import (
18
)
19
from vllm.models.deepseek_v4.common.ops.fused_indexer_q import MXFP4_BLOCK_SIZE
20
from vllm.models.deepseek_v4.common.ops.save_partial_states import (
21
- save_partial_states,
21
+ _SAVE_PARTIAL_STATES_KERNEL,
22
)
23
from vllm.platforms import current_platform
24
from vllm.v1.attention.backend import (
@@ -306,6 +306,18 @@ class DeepseekCompressor(nn.Module):
306
f"Unsupported head_dim for fused quant+cache: {self.head_dim}"
307
)
308
309
+ if vllm_config.kernel_config.enable_jit_warmup:
310
+ _SAVE_PARTIAL_STATES_KERNEL.register_warmup(
311
+ head_dim=self.head_dim,
312
+ compress_ratio=self.compress_ratio,
313
+ )
314
+ if self.head_dim != 512:
315
+ from vllm.models.deepseek_v4.common.ops.fused_compress_quant_cache import ( # noqa: E501
316
+ _FUSED_KV_COMPRESS_NORM_ROPE_INSERT_INDEXER_TRITON_KERNEL,
317
+ )
318
+
319
+ _FUSED_KV_COMPRESS_NORM_ROPE_INSERT_INDEXER_TRITON_KERNEL.register_warmup()
320
+
321
def forward(
322
self,
323
# [num_tokens, 2 * self.coff * self.head_dim]
@@ -351,7 +363,7 @@ class DeepseekCompressor(nn.Module):
363
# GEMM; state_cache from this kernel) but neither emits/waits on PDL
364
# grid dependency primitives, so launch_pdl=True caused a
365
# read-after-write race and non-deterministic output.
354
- save_partial_states(
366
+ _SAVE_PARTIAL_STATES_KERNEL(
367
kv=kv,
368
score=score,
369
ape=self.ape,
vllm/models/deepseek_v4/nvidia/mtp.py
+9
-5
@@ -51,9 +51,9 @@ from vllm.models.common.ops.sequence_parallel import (
51
sp_padding_mask,
52
sp_shard,
53
)
54
-from vllm.models.deepseek_v4.common.ops import (
55
- fused_mtp_input_rmsnorm,
56
- mtp_shared_head_rmsnorm,
54
+from vllm.models.deepseek_v4.common.ops.fused_mtp_input_rmsnorm import (
55
+ _FUSED_MTP_INPUT_RMSNORM_KERNEL,
56
+ _MTP_SHARED_HEAD_RMSNORM_KERNEL,
57
)
58
from vllm.sequence import IntermediateTensors
59
@@ -141,6 +141,10 @@ class DeepSeekV4MultiTokenPredictorLayer(nn.Module):
141
aux_stream_list=aux_stream_list,
142
)
143
144
+ if vllm_config.kernel_config.enable_jit_warmup:
145
+ _FUSED_MTP_INPUT_RMSNORM_KERNEL.register_warmup()
146
+ _MTP_SHARED_HEAD_RMSNORM_KERNEL.register_warmup()
147
+
148
def forward(
149
self,
150
input_ids: torch.Tensor,
@@ -157,7 +161,7 @@ class DeepSeekV4MultiTokenPredictorLayer(nn.Module):
161
-1, self.hc_mult, self.config.hidden_size
162
)
163
# Fused: mask inputs at position 0 (not needed by MTP), enorm, hnorm.
160
- inputs_embeds, previous_hidden_states = fused_mtp_input_rmsnorm(
164
+ inputs_embeds, previous_hidden_states = _FUSED_MTP_INPUT_RMSNORM_KERNEL(
165
inputs_embeds,
166
positions,
167
previous_hidden_states,
@@ -270,7 +274,7 @@ class DeepSeekV4MultiTokenPredictor(nn.Module):
274
mtp_layer.rms_norm_eps,
275
mtp_layer.hc_eps,
276
)
273
- hidden_states = mtp_shared_head_rmsnorm(
277
+ hidden_states = _MTP_SHARED_HEAD_RMSNORM_KERNEL(
278
hidden_states,
279
mtp_layer.shared_head.norm.weight.data,
280
mtp_layer.shared_head.norm.variance_epsilon,
vllm/models/deepseek_v4/sparse_mla.py
+163
-67
@@ -9,6 +9,12 @@ import torch
9
10
from vllm.config import VllmConfig
11
from vllm.config.cache import CacheDType
12
+from vllm.model_executor.warmup.jit_warmup_triton_helper import (
13
+ LaunchSpec,
14
+ TritonWarmupTensor,
15
+ VllmTritonJitKernel,
16
+ kernel_launcher,
17
+)
18
from vllm.platforms import current_platform
19
from vllm.platforms.interface import DeviceCapability
20
from vllm.triton_utils import tl, triton
@@ -338,90 +344,180 @@ def build_c128a_topk_metadata(
344
if num_tokens == 0:
345
return global_decode, decode_lens, prefill_local
346
341
- _build_c128a_topk_metadata_kernel[(num_tokens,)](
347
+ _BUILD_C128A_TOPK_METADATA_KERNEL(
348
global_decode_buffer,
343
- global_decode_buffer.stride(0),
349
decode_lens_buffer,
350
prefill_buffer,
346
- prefill_buffer.stride(0),
351
positions,
352
compress_ratio,
353
max_compressed_tokens,
354
num_decode_tokens,
355
token_to_req_indices,
356
block_table,
353
- block_table.stride(0),
357
block_size,
358
slot_mapping,
356
- BLOCK_SIZE=1024,
359
)
360
return global_decode, decode_lens, prefill_local
361
362
361
-@triton.jit
362
-def _build_c128a_topk_metadata_kernel(
363
- # Decode outputs
364
- global_decode_ptr,
365
- global_decode_stride,
366
- decode_lens_ptr,
367
- # Prefill output
368
- prefill_local_ptr,
369
- prefill_local_stride,
370
- # Inputs
371
- positions_ptr,
372
- compress_ratio,
373
- max_compressed_tokens,
374
- num_decode_tokens,
375
- token_to_req_indices_ptr,
376
- block_table_ptr,
377
- block_table_stride,
378
- block_size,
379
- slot_mapping_ptr,
380
- BLOCK_SIZE: tl.constexpr,
363
+class BuildC128ATopkMetadataKernel(
364
+ VllmTritonJitKernel["BuildC128ATopkMetadataKernel.CompileKey"]
365
):
382
- token_idx = tl.program_id(0)
383
- position = tl.load(positions_ptr + token_idx)
384
- num_compressed = (position + 1) // compress_ratio
385
- num_compressed = tl.minimum(num_compressed, max_compressed_tokens)
386
- is_decode = token_idx < num_decode_tokens
387
-
388
- if is_decode:
389
- # --- Decode: block-table lookup → global slot ids + count ---
390
- is_valid_token = tl.load(slot_mapping_ptr + token_idx) >= 0
391
- req_idx = tl.load(token_to_req_indices_ptr + token_idx)
392
- count = tl.zeros((), dtype=tl.int32)
393
- for i in range(0, max_compressed_tokens, BLOCK_SIZE):
394
- offset = i + tl.arange(0, BLOCK_SIZE)
395
- mask = offset < max_compressed_tokens
396
- is_valid = offset < num_compressed
397
-
398
- block_indices = offset // block_size
399
- block_numbers = tl.load(
400
- block_table_ptr + req_idx * block_table_stride + block_indices,
401
- mask=mask & is_valid,
402
- )
403
- block_offsets = offset % block_size
404
- slot_ids = block_numbers * block_size + block_offsets
405
- slot_ids = tl.where(is_valid, slot_ids, -1)
366
+ @dataclass(frozen=True)
367
+ class CompileKey:
368
+ compress_ratio: int
369
+ max_compressed_tokens: int
370
+ block_size: int
371
+ triton_block_size: int
372
+
373
+ @staticmethod
374
+ @triton.jit(
375
+ do_not_specialize=[
376
+ "global_decode_stride",
377
+ "prefill_local_stride",
378
+ "num_decode_tokens",
379
+ "block_table_stride",
380
+ ]
381
+ )
382
+ def kernel(
383
+ # Decode outputs
384
+ global_decode_ptr,
385
+ global_decode_stride,
386
+ decode_lens_ptr,
387
+ # Prefill output
388
+ prefill_local_ptr,
389
+ prefill_local_stride,
390
+ # Inputs
391
+ positions_ptr,
392
+ compress_ratio,
393
+ max_compressed_tokens,
394
+ num_decode_tokens,
395
+ token_to_req_indices_ptr,
396
+ block_table_ptr,
397
+ block_table_stride,
398
+ block_size,
399
+ slot_mapping_ptr,
400
+ BLOCK_SIZE: tl.constexpr,
401
+ ):
402
+ token_idx = tl.program_id(0)
403
+ position = tl.load(positions_ptr + token_idx)
404
+ num_compressed = (position + 1) // compress_ratio
405
+ num_compressed = tl.minimum(num_compressed, max_compressed_tokens)
406
+ is_decode = token_idx < num_decode_tokens
407
+
408
+ if is_decode:
409
+ # --- Decode: block-table lookup → global slot ids + count ---
410
+ is_valid_token = tl.load(slot_mapping_ptr + token_idx) >= 0
411
+ req_idx = tl.load(token_to_req_indices_ptr + token_idx)
412
+ count = tl.zeros((), dtype=tl.int32)
413
+ for i in range(0, max_compressed_tokens, BLOCK_SIZE):
414
+ offset = i + tl.arange(0, BLOCK_SIZE)
415
+ mask = offset < max_compressed_tokens
416
+ is_valid = offset < num_compressed
417
+
418
+ block_indices = offset // block_size
419
+ block_numbers = tl.load(
420
+ block_table_ptr + req_idx * block_table_stride + block_indices,
421
+ mask=mask & is_valid,
422
+ )
423
+ block_offsets = offset % block_size
424
+ slot_ids = block_numbers * block_size + block_offsets
425
+ slot_ids = tl.where(is_valid, slot_ids, -1)
426
+ tl.store(
427
+ global_decode_ptr + token_idx * global_decode_stride + offset,
428
+ slot_ids,
429
+ mask=mask,
430
+ )
431
+ count += tl.sum(is_valid.to(tl.int32), axis=0)
432
+
433
tl.store(
407
- global_decode_ptr + token_idx * global_decode_stride + offset,
408
- slot_ids,
409
- mask=mask,
434
+ decode_lens_ptr + token_idx,
435
+ tl.where(is_valid_token, count, 0),
436
)
411
- count += tl.sum(is_valid.to(tl.int32), axis=0)
437
+ else:
438
+ # --- Prefill: write local indices ---
439
+ pfx_idx = token_idx - num_decode_tokens
440
+ for i in range(0, max_compressed_tokens, BLOCK_SIZE):
441
+ offset = i + tl.arange(0, BLOCK_SIZE)
442
+ mask = offset < max_compressed_tokens
443
+ tl.store(
444
+ prefill_local_ptr + pfx_idx * prefill_local_stride + offset,
445
+ tl.where(offset < num_compressed, offset, -1),
446
+ mask=mask,
447
+ )
448
+
449
+ def dispatch( # type: ignore[override]
450
+ self,
451
+ **compile_key_fields: int,
452
+ ) -> CompileKey:
453
+ return self.CompileKey(**compile_key_fields)
454
+
455
+ def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]:
456
+ hf_config = vllm_config.model_config.hf_config
457
+ if getattr(hf_config, "model_type", None) != "deepseek_v4":
458
+ return []
459
+
460
+ compress_ratio = 128
461
+ max_model_len = vllm_config.model_config.max_model_len
462
+ if max_model_len <= 0:
463
+ return []
464
+ max_compressed_tokens = cdiv(max_model_len, compress_ratio)
465
+ max_compressed_tokens = (
466
+ cdiv(max_compressed_tokens, _C128A_TOPK_ALIGNMENT) * _C128A_TOPK_ALIGNMENT
467
+ )
468
+ return self._trace_dispatch(self.dispatch)(
469
+ compress_ratio=compress_ratio,
470
+ max_compressed_tokens=max_compressed_tokens,
471
+ # DeepSeek V4 sparse MLA uses 256-token KV pages; C128A metadata
472
+ # works in compressed-token units.
473
+ block_size=256 // compress_ratio,
474
+ triton_block_size=1024,
475
+ )
476
413
- tl.store(
414
- decode_lens_ptr + token_idx,
415
- tl.where(is_valid_token, count, 0),
477
+ def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
478
+ int32_ptr = TritonWarmupTensor(torch.int32)
479
+ return dict(
480
+ global_decode_buffer=TritonWarmupTensor(
481
+ torch.int32, shape=(1, 1), strides=(1, 1)
482
+ ),
483
+ decode_lens_buffer=int32_ptr,
484
+ prefill_buffer=TritonWarmupTensor(
485
+ torch.int32, shape=(1, 1), strides=(1, 1)
486
+ ),
487
+ positions=TritonWarmupTensor(torch.int64),
488
+ compress_ratio=compile_key.compress_ratio,
489
+ max_compressed_tokens=compile_key.max_compressed_tokens,
490
+ num_decode_tokens=0,
491
+ token_to_req_indices=int32_ptr,
492
+ block_table=TritonWarmupTensor(torch.int32, shape=(1, 1), strides=(1, 1)),
493
+ block_size=compile_key.block_size,
494
+ slot_mapping=TritonWarmupTensor(torch.int64),
495
)
417
- else:
418
- # --- Prefill: write local indices ---
419
- pfx_idx = token_idx - num_decode_tokens
420
- for i in range(0, max_compressed_tokens, BLOCK_SIZE):
421
- offset = i + tl.arange(0, BLOCK_SIZE)
422
- mask = offset < max_compressed_tokens
423
- tl.store(
424
- prefill_local_ptr + pfx_idx * prefill_local_stride + offset,
425
- tl.where(offset < num_compressed, offset, -1),
426
- mask=mask,
427
- )
496
+
497
+ @kernel_launcher
498
+ def __call__(
499
+ self,
500
+ global_decode_buffer: torch.Tensor,
501
+ decode_lens_buffer: torch.Tensor,
502
+ prefill_buffer: torch.Tensor,
503
+ positions: torch.Tensor,
504
+ compress_ratio: int,
505
+ max_compressed_tokens: int,
506
+ num_decode_tokens: int,
507
+ token_to_req_indices: torch.Tensor,
508
+ block_table: torch.Tensor,
509
+ block_size: int,
510
+ slot_mapping: torch.Tensor,
511
+ ) -> LaunchSpec:
512
+ return (positions.shape[0],), dict(
513
+ global_decode_ptr=global_decode_buffer,
514
+ global_decode_stride=global_decode_buffer.stride(0),
515
+ decode_lens_ptr=decode_lens_buffer,
516
+ prefill_local_ptr=prefill_buffer,
517
+ prefill_local_stride=prefill_buffer.stride(0),
518
+ block_table_stride=block_table.stride(0),
519
+ BLOCK_SIZE=1024,
520
+ )
521
+
522
+
523
+_BUILD_C128A_TOPK_METADATA_KERNEL = BuildC128ATopkMetadataKernel()
vllm/models/deepseek_v4/xpu/mtp.py
+9
-5
@@ -39,9 +39,9 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader
39
from vllm.model_executor.models.deepseek_mtp import SharedHead
40
from vllm.model_executor.models.deepseek_v2 import get_spec_layer_idx_from_weight_name
41
from vllm.model_executor.models.utils import maybe_prefix
42
-from vllm.models.deepseek_v4.common.ops import (
43
- fused_mtp_input_rmsnorm,
44
- mtp_shared_head_rmsnorm,
42
+from vllm.models.deepseek_v4.common.ops.fused_mtp_input_rmsnorm import (
43
+ _FUSED_MTP_INPUT_RMSNORM_KERNEL,
44
+ _MTP_SHARED_HEAD_RMSNORM_KERNEL,
45
)
46
from vllm.platforms import current_platform
47
from vllm.sequence import IntermediateTensors
@@ -130,6 +130,10 @@ class DeepSeekV4MultiTokenPredictorLayer(nn.Module):
130
131
self.hc_head_op = HCHeadOp()
132
133
+ if vllm_config.kernel_config.enable_jit_warmup:
134
+ _FUSED_MTP_INPUT_RMSNORM_KERNEL.register_warmup()
135
+ _MTP_SHARED_HEAD_RMSNORM_KERNEL.register_warmup()
136
+
137
def forward(
138
self,
139
input_ids: torch.Tensor,
@@ -146,7 +150,7 @@ class DeepSeekV4MultiTokenPredictorLayer(nn.Module):
150
-1, self.hc_mult, self.config.hidden_size
151
)
152
# Fused: mask inputs at position 0 (not needed by MTP), enorm, hnorm.
149
- inputs_embeds, previous_hidden_states = fused_mtp_input_rmsnorm(
153
+ inputs_embeds, previous_hidden_states = _FUSED_MTP_INPUT_RMSNORM_KERNEL(
154
inputs_embeds,
155
positions,
156
previous_hidden_states,
@@ -253,7 +257,7 @@ class DeepSeekV4MultiTokenPredictor(nn.Module):
257
mtp_layer.rms_norm_eps,
258
mtp_layer.hc_eps,
259
)
256
- hidden_states = mtp_shared_head_rmsnorm(
260
+ hidden_states = _MTP_SHARED_HEAD_RMSNORM_KERNEL(
261
hidden_states,
262
mtp_layer.shared_head.norm.weight.data,
263
mtp_layer.shared_head.norm.variance_epsilon,