[5/N][warmup][DSv4] Migrate NVIDIA CuTeDSL attention kernels (#53566)
Signed-off-by: LopezCastroRoberto <rocastro@redhat.com> Co-authored-by: Codex <codex@openai.com>
Roberto L. Castro committed
Sep 12, 2026 at 06:05 UTC
120ec4ebd2803c3cc63729e18fcfd42b5ba47447
28 files changed
+3108
-2535
tests/model_executor/test_jit_warmup_cutedsl_launcher.py
+38
-3
@@ -5,10 +5,10 @@ from dataclasses import dataclass
5
from typing import Any
6
7
from vllm.model_executor.warmup import jit_warmup_cutedsl_helper
8
+from vllm.model_executor.warmup.jit_warmup import kernel_launcher
9
from vllm.model_executor.warmup.jit_warmup_cutedsl_helper import (
10
CuTeDSLLaunchSpec,
11
VllmCuTeDSLJitKernel,
11
- cutedsl_kernel_launcher,
12
)
13
14
@@ -30,7 +30,7 @@ class _TestCuTeDSLKernel(VllmCuTeDSLJitKernel["_TestCuTeDSLKernel.CompileKey"]):
30
def warmup_inputs(self, compile_key: CompileKey) -> tuple[Any, ...]:
31
return (f"fake-{compile_key.variant}",)
32
33
- @cutedsl_kernel_launcher
33
+ @kernel_launcher
34
def __call__(
35
self,
36
payload: str,
@@ -38,7 +38,7 @@ class _TestCuTeDSLKernel(VllmCuTeDSLJitKernel["_TestCuTeDSLKernel.CompileKey"]):
38
variant: int,
39
) -> CuTeDSLLaunchSpec[CompileKey]:
40
compile_key = self.dispatch(variant=variant)
41
- return compile_key, (payload,), {"variant": variant}
41
+ return compile_key, (payload,)
42
43
44
def test_cutedsl_launcher_reuses_compiled_executor(monkeypatch) -> None:
@@ -74,3 +74,38 @@ def test_cutedsl_launcher_reuses_compiled_executor(monkeypatch) -> None:
74
(("entry", 2), ("fake-2",)),
75
]
76
assert launch_calls == [(1, ("runtime",)), (2, ("other",))]
77
+
78
+
79
+def test_cutedsl_launcher_runs_post_launch_finalizer(monkeypatch) -> None:
80
+ events: list[str] = []
81
+
82
+ def compile_cutedsl(entry: tuple[str, int], *args: Any) -> Any:
83
+ def executor(*runtime_args: Any) -> None:
84
+ events.append("launch")
85
+
86
+ return executor
87
+
88
+ class FinalizedKernel(_TestCuTeDSLKernel):
89
+ @kernel_launcher
90
+ def __call__(
91
+ self,
92
+ payload: str,
93
+ *,
94
+ variant: int,
95
+ ) -> CuTeDSLLaunchSpec[_TestCuTeDSLKernel.CompileKey]:
96
+ compile_key = self.dispatch(variant=variant)
97
+
98
+ def finalize() -> str:
99
+ events.append("finalize")
100
+ return "output"
101
+
102
+ return compile_key, (payload,), "unused", finalize
103
+
104
+ monkeypatch.setattr(
105
+ jit_warmup_cutedsl_helper,
106
+ "compile_cutedsl",
107
+ compile_cutedsl,
108
+ )
109
+
110
+ assert FinalizedKernel()("runtime", variant=1) == "output"
111
+ assert events == ["launch", "finalize"]
tests/model_executor/test_jit_warmup_triton_launcher.py
+4
-4
@@ -4,10 +4,10 @@
4
from dataclasses import dataclass
5
from typing import Any
6
7
+from vllm.model_executor.warmup.jit_warmup import kernel_launcher
8
from vllm.model_executor.warmup.jit_warmup_triton_helper import (
9
LaunchSpec,
10
VllmTritonJitKernel,
10
- kernel_launcher,
11
)
12
13
@@ -107,13 +107,13 @@ def test_compute_slot_mapping_uses_named_launcher_inputs(monkeypatch) -> None:
107
)
108
launches: list[tuple[Any, ...]] = []
109
110
- def launch(grid: Any, inputs: Any, **kwargs: Any) -> None:
111
- launches.append((grid, inputs, kwargs))
110
+ def launch(launch_spec: Any, inputs: Any) -> None:
111
+ launches.append((launch_spec, inputs))
112
113
monkeypatch.setattr(owner, "launch", launch)
114
owner.compile(compile_key)
115
116
- grid, inputs, kwargs = launches[0]
116
+ (grid, kwargs), inputs = launches[0]
117
assert grid == (2,)
118
assert inputs["num_tokens"] == 2
119
assert inputs["block_table_stride"] == 128
vllm/cute_utils/__init__.py
+8
@@ -22,6 +22,14 @@ _TORCH_TO_CUTE_DTYPE = {
22
torch.float32: Float32,
23
}
24
25
+
26
+def torch_to_cute_dtype(dtype: torch.dtype):
27
+ try:
28
+ return _TORCH_TO_CUTE_DTYPE[dtype]
29
+ except KeyError as exc:
30
+ raise TypeError(f"Unsupported CuTe dtype {dtype}") from exc
31
+
32
+
33
_CUTE_TO_PTX_DTYPE = {
34
BFloat16: "bf16",
35
Float16: "f16",
vllm/model_executor/kernels/attention/dsa/dcp_indexer_cutedsl.py
+6
-12
@@ -12,16 +12,15 @@ from cutlass import Float32, Int32, Uint32, Uint64
12
from quack.compile_utils import make_fake_tensor
13
14
from vllm.cute_utils import recast_val
15
+from vllm.model_executor.warmup.jit_warmup import kernel_launcher
16
from vllm.model_executor.warmup.jit_warmup_cutedsl_helper import (
17
CuTeDSLLaunchSpec,
18
VllmCuTeDSLJitKernel,
18
- cutedsl_kernel_launcher,
19
)
20
from vllm.model_executor.warmup.jit_warmup_triton_helper import (
21
LaunchSpec,
22
TritonWarmupTensor,
23
VllmTritonJitKernel,
24
- kernel_launcher,
24
)
25
from vllm.triton_utils import tl, triton
26
@@ -569,7 +568,7 @@ class StableTopKFromGatheredCandidatesKernel(
568
)
569
return gathered, out
570
572
- @cutedsl_kernel_launcher
571
+ @kernel_launcher
572
def __call__(
573
self,
574
gathered: torch.Tensor,
@@ -577,16 +576,11 @@ class StableTopKFromGatheredCandidatesKernel(
576
*,
577
topk: int,
578
) -> CuTeDSLLaunchSpec[CompileKey]:
580
- compile_key = self.dispatch(topk=topk, num_candidates=gathered.shape[1])
581
- return (
582
- compile_key,
583
- (gathered, out),
584
- {
585
- "gathered_shape": tuple(gathered.shape),
586
- "out_shape": tuple(out.shape),
587
- "topk": topk,
588
- },
579
+ compile_key = self.dispatch(
580
+ topk=topk,
581
+ num_candidates=gathered.shape[1],
582
)
583
+ return compile_key, (gathered, out)
584
585
586
_PACK_DCP_TOPK_CANDIDATES_KERNEL = PackDCPTopkCandidatesKernel()
vllm/model_executor/warmup/jit_warmup.py
+23
@@ -15,6 +15,7 @@ from collections.abc import Callable, Iterable, Iterator, Mapping
15
from contextlib import contextmanager
16
from contextvars import ContextVar
17
from dataclasses import dataclass, fields
18
+from functools import wraps
19
from typing import Any, Generic, TypeVar, cast
20
21
__all__ = [
@@ -23,6 +24,7 @@ __all__ = [
24
"WarmupIntRange",
25
"get_ast_full_name",
26
"get_function_source_node",
27
+ "kernel_launcher",
28
"zip_inputs",
29
]
30
@@ -30,6 +32,26 @@ __all__ = [
32
CompileKeyT = TypeVar("CompileKeyT")
33
34
35
+def kernel_launcher(call_fn: Callable[..., Any]) -> Callable[..., Any]:
36
+ """Delegate a declarative launch specification to its backend owner."""
37
+ signature = inspect.signature(call_fn)
38
+
39
+ @wraps(call_fn)
40
+ def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
41
+ launch_spec = call_fn(self, *args, **kwargs)
42
+ if not self.bind_launch_inputs:
43
+ return self.launch(launch_spec, {})
44
+
45
+ bound = signature.bind(self, *args, **kwargs)
46
+ bound.apply_defaults()
47
+ inputs = {
48
+ name: value for name, value in bound.arguments.items() if name != "self"
49
+ }
50
+ return self.launch(launch_spec, inputs)
51
+
52
+ return wrapper
53
+
54
+
55
@dataclass(frozen=True)
56
class WarmupIntRange:
57
"""Expand integers with range semantics or a custom monotonic progression."""
@@ -549,6 +571,7 @@ class VllmJitKernel(Generic[CompileKeyT], ABC):
571
"""Kernel wrapper that owns dispatch, warmup keys, and compilation."""
572
573
CompileKey: type[CompileKeyT]
574
+ bind_launch_inputs = True
575
576
def __init__(self) -> None:
577
self._dispatch_trace = _trace_compile_key_dispatch(self.dispatch)
vllm/model_executor/warmup/jit_warmup_cutedsl_helper.py
+25
-26
@@ -3,19 +3,25 @@
3
4
from abc import abstractmethod
5
from collections.abc import Callable, Mapping
6
-from functools import wraps
6
from typing import Any, ClassVar, Generic, TypeAlias, TypeVar
7
9
-from vllm.model_executor.warmup.jit_warmup import VllmJitKernel
8
+from vllm.model_executor.warmup.jit_warmup import (
9
+ VllmJitKernel,
10
+)
11
12
DEFAULT_CUTEDSL_COMPILE_OPTIONS = "--enable-tvm-ffi"
13
14
CompileKeyT = TypeVar("CompileKeyT")
14
-CuTeDSLLaunchSpec: TypeAlias = tuple[
15
- CompileKeyT,
16
- tuple[Any, ...],
17
- Mapping[str, Any] | None,
18
-]
15
+CuTeDSLLaunchSpec: TypeAlias = (
16
+ tuple[CompileKeyT, tuple[Any, ...]]
17
+ | tuple[CompileKeyT, tuple[Any, ...], Any]
18
+ | tuple[
19
+ CompileKeyT,
20
+ tuple[Any, ...],
21
+ Any,
22
+ Callable[[], Any],
23
+ ]
24
+)
25
26
27
def cutedsl_fake_stream(*, use_tvm_ffi_env_stream: bool = True) -> Any:
@@ -44,6 +50,7 @@ class VllmCuTeDSLJitKernel(VllmJitKernel[CompileKeyT], Generic[CompileKeyT]):
50
"""CuTeDSL owner whose compiled executor is shared by warmup and runtime."""
51
52
kernel: ClassVar[Any]
53
+ bind_launch_inputs = False
54
55
@abstractmethod
56
def warmup_inputs(self, compile_key: CompileKeyT) -> tuple[Any, ...]:
@@ -58,23 +65,15 @@ class VllmCuTeDSLJitKernel(VllmJitKernel[CompileKeyT], Generic[CompileKeyT]):
65
*self.warmup_inputs(compile_key),
66
)
67
61
-
62
-def cutedsl_kernel_launcher(
63
- call_fn: Callable[..., CuTeDSLLaunchSpec[CompileKeyT]],
64
-) -> Callable[..., Any]:
65
- """Invoke a cached CuTeDSL executor from a declarative ``__call__``."""
66
-
67
- @wraps(call_fn)
68
- def wrapper(
69
- self: VllmCuTeDSLJitKernel[CompileKeyT],
70
- *args: Any,
71
- **kwargs: Any,
68
+ def launch(
69
+ self,
70
+ launch_spec: CuTeDSLLaunchSpec[CompileKeyT],
71
+ _inputs: Mapping[str, Any],
72
) -> Any:
73
- compile_key, launch_args, runtime_context = call_fn(self, *args, **kwargs)
74
- executor = self._get_or_compile(
75
- compile_key,
76
- runtime_context=runtime_context,
77
- )
78
- return executor(*launch_args)
79
-
80
- return wrapper
73
+ # (compile_key, args), optionally followed by output and epilogue.
74
+ compile_key, launch_args = launch_spec[:2]
75
+ executor = self._get_or_compile(compile_key)
76
+ result = executor(*launch_args)
77
+ if len(launch_spec) == 4:
78
+ return launch_spec[3]()
79
+ return launch_spec[2] if len(launch_spec) == 3 else result
vllm/model_executor/warmup/jit_warmup_triton_helper.py
+14
-44
@@ -5,7 +5,7 @@ import inspect
5
from abc import abstractmethod
6
from collections.abc import Callable, Mapping
7
from dataclasses import dataclass
8
-from functools import cached_property, wraps
8
+from functools import cached_property
9
from typing import Any, ClassVar, Generic, TypeVar, cast
10
11
from vllm.model_executor.warmup.jit_warmup import (
@@ -128,14 +128,16 @@ class VllmTritonJitKernel(VllmJitKernel[CompileKeyT], Generic[CompileKeyT]):
128
129
def launch(
130
self,
131
- grid: tuple[int, ...] | None,
131
+ launch_spec: LaunchSpec,
132
inputs: Mapping[str, Any],
133
- /,
134
- *,
135
- kernel: Any = None,
136
- **kwargs: Any,
133
) -> Any:
138
- kernel = kernel if kernel is not None else self.kernel
134
+ if len(launch_spec) == 3:
135
+ grid, launch_kwargs, outputs = launch_spec
136
+ else:
137
+ grid, launch_kwargs = launch_spec
138
+ outputs = None
139
+ kwargs = dict(launch_kwargs)
140
+ kernel = kwargs.pop("kernel", None) or self.kernel
141
runtime_launcher = kwargs.pop("_runtime_launcher", None)
142
runtime_launcher_arg_count = kwargs.pop("_runtime_launcher_arg_count", 0)
143
for name, value in inputs.items():
@@ -151,52 +153,20 @@ class VllmTritonJitKernel(VllmJitKernel[CompileKeyT], Generic[CompileKeyT]):
153
kwargs["launch_pdl"] = self._warming_compile_key.launch_pdl
154
warmup = getattr(kernel, "warmup", None)
155
assert warmup is not None
154
- return warmup(grid=(1,), **kwargs)
156
+ warmup(grid=(1,), **kwargs)
157
+ return outputs
158
if grid is None:
156
- return None
159
+ return outputs
160
if runtime_launcher is not None:
161
regular_args = [
162
kwargs.pop(name)
163
for name in self._kernel_arg_names[:runtime_launcher_arg_count]
164
]
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],
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)
179
- def wrapper(
180
- self: VllmTritonJitKernel[Any],
181
- *args: Any,
182
- **kwargs: Any,
183
- ) -> Any:
184
- spec = call_fn(self, *args, **kwargs)
185
- if len(spec) == 3:
186
- grid, launch_kwargs, outputs = spec
165
+ runtime_launcher(kernel, grid, *regular_args, **kwargs)
166
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)
167
+ kernel[grid](**kwargs)
168
return outputs
169
198
- return wrapper
199
-
170
171
@dataclass(frozen=True)
172
class TritonPointerInputVariant:
vllm/models/common/ops/fused_qk_rmsnorm.py
+1
-1
@@ -5,11 +5,11 @@ from typing import Any
5
6
import torch
7
8
+from vllm.model_executor.warmup.jit_warmup import kernel_launcher
9
from vllm.model_executor.warmup.jit_warmup_triton_helper import (
10
LaunchSpec,
11
TritonWarmupTensor,
12
VllmTritonJitKernel,
12
- kernel_launcher,
13
)
14
from vllm.platforms import current_platform
15
from vllm.triton_utils import tl, triton
vllm/models/deepseek_v4/attention.py
+23
-3
@@ -50,6 +50,7 @@ from vllm.model_executor.layers.quantization import QuantizationConfig
50
from vllm.model_executor.models.utils import extract_layer_index
51
from vllm.models.deepseek_v4.common.rope import build_deepseek_v4_rope
52
from vllm.models.deepseek_v4.compressor import DeepseekCompressor
53
+from vllm.platforms import current_platform
54
from vllm.triton_utils import tl, triton
55
from vllm.utils.multi_stream_utils import (
56
execute_in_parallel,
@@ -420,7 +421,6 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
421
422
_COMBINE_TOPK_SWA_INDICES_KERNEL.register_warmup()
423
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()
@@ -447,7 +447,13 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
447
)
448
449
_COMPUTE_GLOBAL_TOPK_INDICES_AND_LENS_KERNEL.register_warmup()
450
- if not has_cutedsl():
450
+ if has_cutedsl():
451
+ from vllm.models.deepseek_v4.nvidia.ops.dequant_gather_k_cutedsl import ( # noqa: E501
452
+ _DEQUANT_GATHER_K_CACHE_CUTEDSL_KERNEL,
453
+ )
454
+
455
+ _DEQUANT_GATHER_K_CACHE_CUTEDSL_KERNEL.register_warmup()
456
+ else:
457
_DEQUANTIZE_AND_GATHER_K_CACHE_KERNEL.register_warmup()
458
elif backend_name == "FLASHINFER_MLA_SPARSE_DSV4":
459
from vllm.models.deepseek_v4.common.ops.cache_utils import (
@@ -934,6 +940,21 @@ class DeepseekV4Indexer(nn.Module):
940
"Using %s indexer cache for Lightning Indexer.",
941
"MXFP4" if self.use_fp4_kv else "FP8",
942
)
943
+ if vllm_config.kernel_config.enable_jit_warmup:
944
+ from vllm.utils.import_utils import has_cutedsl
945
+
946
+ if current_platform.is_cuda() and has_cutedsl():
947
+ from vllm.models.deepseek_v4.nvidia.ops.fused_indexer_q_cutedsl import ( # noqa: E501
948
+ _INDEXER_Q_FP8_KERNEL,
949
+ _INDEXER_Q_MXFP4_KERNEL,
950
+ )
951
+
952
+ indexer_q_kernel = (
953
+ _INDEXER_Q_MXFP4_KERNEL
954
+ if self.use_fp4_kv
955
+ else _INDEXER_Q_FP8_KERNEL
956
+ )
957
+ indexer_q_kernel.register_warmup()
958
959
# no tensor parallel, just replicated
960
self.wq_b = ReplicatedLinear(
@@ -1016,7 +1037,6 @@ class DeepseekV4Indexer(nn.Module):
1037
]
1038
1039
if vllm_config.kernel_config.enable_jit_warmup:
1019
- from vllm.platforms import current_platform
1040
from vllm.utils.import_utils import has_cutedsl
1041
1042
if not has_cutedsl() and not current_platform.is_xpu():
vllm/models/deepseek_v4/common/ops/cache_utils.py
+10
-4
@@ -24,6 +24,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
24
)
25
from vllm.model_executor.warmup.jit_warmup import (
26
WarmupIntRange,
27
+ kernel_launcher,
28
zip_inputs,
29
)
30
from vllm.model_executor.warmup.jit_warmup_triton_helper import (
@@ -31,7 +32,6 @@ from vllm.model_executor.warmup.jit_warmup_triton_helper import (
32
TritonPointerInputVariant,
33
TritonWarmupTensor,
34
VllmTritonJitKernel,
34
- kernel_launcher,
35
triton_scalar_specialization_rep,
36
)
37
from vllm.platforms import current_platform
@@ -521,11 +521,17 @@ def dequantize_and_gather_k_cache(
521
if has_cutedsl():
522
# lazily import, otherwise some tests fail due to CUDA driver init failure.
523
from vllm.models.deepseek_v4.nvidia.ops.dequant_gather_k_cutedsl import (
524
- dequantize_and_gather_k_cache_cutedsl,
524
+ _DEQUANT_GATHER_K_CACHE_CUTEDSL_KERNEL,
525
)
526
527
- dequantize_and_gather_k_cache_cutedsl(
528
- out, k_cache, seq_lens, gather_lens, block_table, block_size, offset
527
+ _DEQUANT_GATHER_K_CACHE_CUTEDSL_KERNEL(
528
+ out=out,
529
+ k_cache=k_cache,
530
+ seq_lens=seq_lens,
531
+ gather_lens=gather_lens,
532
+ block_table=block_table,
533
+ block_size=block_size,
534
+ offset=offset,
535
)
536
return
537
vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py
+1
-1
@@ -26,11 +26,11 @@ from typing import Any
26
27
import torch
28
29
+from vllm.model_executor.warmup.jit_warmup import kernel_launcher
30
from vllm.model_executor.warmup.jit_warmup_triton_helper import (
31
LaunchSpec,
32
TritonWarmupTensor,
33
VllmTritonJitKernel,
33
- kernel_launcher,
34
)
35
from vllm.platforms import current_platform
36
from vllm.triton_utils import tl, triton
vllm/models/deepseek_v4/common/ops/fused_indexer_q.py
+22
-22
@@ -6,11 +6,11 @@ from typing import Any
6
7
import torch
8
9
+from vllm.model_executor.warmup.jit_warmup import kernel_launcher
10
from vllm.model_executor.warmup.jit_warmup_triton_helper import (
11
LaunchSpec,
12
TritonWarmupTensor,
13
VllmTritonJitKernel,
13
- kernel_launcher,
14
)
15
from vllm.platforms import current_platform
16
from vllm.triton_utils import tl, triton
@@ -598,19 +598,19 @@ def fused_indexer_q_rope_quant(
598
if has_cutedsl():
599
# lazily import, otherwise some tests fail due to CUDA driver init failure.
600
from vllm.models.deepseek_v4.nvidia.ops.fused_indexer_q_cutedsl import (
601
- fused_indexer_q_rope_quant_mxfp4_cutedsl,
601
+ _INDEXER_Q_MXFP4_KERNEL,
602
)
603
604
- fused_indexer_q_rope_quant_mxfp4_cutedsl(
605
- positions,
606
- index_q,
607
- index_q_cos_sin_cache,
608
- index_weights,
609
- index_weights_softmax_scale,
610
- index_weights_head_scale,
611
- index_q_packed,
612
- index_q_scale,
613
- index_weights_out,
604
+ _INDEXER_Q_MXFP4_KERNEL(
605
+ positions=positions,
606
+ q=index_q,
607
+ cos_sin_cache=index_q_cos_sin_cache,
608
+ weights=index_weights,
609
+ weights_softmax_scale=index_weights_softmax_scale,
610
+ weights_head_scale=index_weights_head_scale,
611
+ q_packed=index_q_packed,
612
+ q_scale=index_q_scale,
613
+ weights_out=index_weights_out,
614
)
615
elif current_platform.is_xpu():
616
torch.ops.vllm.xpu_deepseek_fused_indexer_q_rope_mxfp4(
@@ -654,18 +654,18 @@ def fused_indexer_q_rope_quant(
654
if has_cutedsl():
655
# lazily import, otherwise some tests fail due to CUDA driver init failure.
656
from vllm.models.deepseek_v4.nvidia.ops.fused_indexer_q_cutedsl import (
657
- fused_indexer_q_rope_quant_fp8_cutedsl,
657
+ _INDEXER_Q_FP8_KERNEL,
658
)
659
660
- fused_indexer_q_rope_quant_fp8_cutedsl(
661
- positions,
662
- index_q,
663
- index_q_cos_sin_cache,
664
- index_weights,
665
- index_weights_softmax_scale,
666
- index_weights_head_scale,
667
- index_q_fp8,
668
- index_weights_out,
660
+ _INDEXER_Q_FP8_KERNEL(
661
+ positions=positions,
662
+ q=index_q,
663
+ cos_sin_cache=index_q_cos_sin_cache,
664
+ weights=index_weights,
665
+ weights_softmax_scale=index_weights_softmax_scale,
666
+ weights_head_scale=index_weights_head_scale,
667
+ q_fp8=index_q_fp8.view(torch.uint8),
668
+ weights_out=index_weights_out,
669
)
670
elif current_platform.is_xpu():
671
torch.ops.vllm.xpu_deepseek_fused_indexer_q_rope_fp8(
vllm/models/deepseek_v4/common/ops/fused_inv_rope_fp8_quant.py
+1
-1
@@ -12,11 +12,11 @@ from typing import Any
12
13
import torch
14
15
+from vllm.model_executor.warmup.jit_warmup import kernel_launcher
16
from vllm.model_executor.warmup.jit_warmup_triton_helper import (
17
LaunchSpec,
18
TritonWarmupTensor,
19
VllmTritonJitKernel,
19
- kernel_launcher,
20
)
21
from vllm.platforms import current_platform
22
from vllm.triton_utils import tl, triton
vllm/models/deepseek_v4/common/ops/fused_mtp_input_rmsnorm.py
+1
-1
@@ -22,11 +22,11 @@ from typing import Any
22
23
import torch
24
25
+from vllm.model_executor.warmup.jit_warmup import kernel_launcher
26
from vllm.model_executor.warmup.jit_warmup_triton_helper import (
27
LaunchSpec,
28
TritonWarmupTensor,
29
VllmTritonJitKernel,
29
- kernel_launcher,
30
)
31
from vllm.triton_utils import tl, triton
32
from vllm.utils.math_utils import next_power_of_2
vllm/models/deepseek_v4/common/ops/save_partial_states.py
+1
-1
@@ -6,11 +6,11 @@ from typing import Any
6
7
import torch
8
9
+from vllm.model_executor.warmup.jit_warmup import kernel_launcher
10
from vllm.model_executor.warmup.jit_warmup_triton_helper import (
11
LaunchSpec,
12
TritonWarmupTensor,
13
VllmTritonJitKernel,
13
- kernel_launcher,
14
)
15
from vllm.triton_utils import tl, triton
16
from vllm.utils.math_utils import next_power_of_2
vllm/models/deepseek_v4/compressor.py
+31
-5
@@ -203,8 +203,8 @@ class DeepseekCompressor(nn.Module):
203
prologue (kv/score split, save_partial_states launch). The
204
compress → norm → RoPE → store step is dispatched to a triton kernel
205
(``compress_norm_rope_store_triton``) by default, except for the NVIDIA
206
- head_dim=128 indexer path which uses the cutedsl kernel
207
- (``compress_norm_rope_store_cutedsl``) for better performance.
206
+ head_dim=512 path which uses the CuTeDSL compressor kernels for better
207
+ performance.
208
"""
209
210
def __init__(
@@ -317,7 +317,33 @@ class DeepseekCompressor(nn.Module):
317
head_dim=self.head_dim,
318
compress_ratio=self.compress_ratio,
319
)
320
- if self.head_dim != 512:
320
+ if current_platform.is_cuda() and self.head_dim == 512:
321
+ from vllm.models.deepseek_v4.nvidia.ops.sparse_attn_compress_cutedsl import ( # noqa: E501
322
+ _SPARSE_ATTN_COMPRESS_C128_BLOCK8_KERNEL,
323
+ _SPARSE_ATTN_COMPRESS_NORM_ROPE_STORE_C4_KERNEL,
324
+ _SPARSE_ATTN_COMPRESS_NORM_ROPE_STORE_FULL_C4_KERNEL,
325
+ _SPARSE_ATTN_NORM_ROPE_STORE_FULL_KERNEL,
326
+ _SPARSE_ATTN_NORM_ROPE_STORE_KERNEL,
327
+ )
328
+
329
+ store_full_kv = vllm_config.cache_config.cache_dtype != "fp8_ds_mla"
330
+ if self.compress_ratio == 4:
331
+ (
332
+ _SPARSE_ATTN_COMPRESS_NORM_ROPE_STORE_FULL_C4_KERNEL
333
+ if store_full_kv
334
+ else _SPARSE_ATTN_COMPRESS_NORM_ROPE_STORE_C4_KERNEL
335
+ ).register_warmup()
336
+ else:
337
+ _SPARSE_ATTN_COMPRESS_C128_BLOCK8_KERNEL.register_warmup()
338
+ if store_full_kv:
339
+ _SPARSE_ATTN_NORM_ROPE_STORE_FULL_KERNEL.register_warmup()
340
+ else:
341
+ _SPARSE_ATTN_NORM_ROPE_STORE_KERNEL.register_warmup(
342
+ vllm_config,
343
+ k_cache_prefix=self.k_cache_prefix,
344
+ compress_ratio=self.compress_ratio,
345
+ )
346
+ else:
347
from vllm.models.deepseek_v4.common.ops.fused_compress_quant_cache import ( # noqa: E501
348
_FUSED_KV_COMPRESS_NORM_ROPE_INSERT_INDEXER_TRITON_KERNEL,
349
)
@@ -419,13 +445,13 @@ class DeepseekCompressor(nn.Module):
445
compress_norm_rope_store_fn: Any
446
if current_platform.is_cuda() and self.head_dim == 512:
447
from .nvidia.ops.sparse_attn_compress_cutedsl import (
422
- compress_norm_rope_store_cutedsl,
448
+ _SPARSE_ATTN_COMPRESSOR_CUTEDSL_KERNEL,
449
)
450
451
# head=512 on CUDA always uses cutedsl, for both the fp8_ds_mla
452
# layout and the plain full-cache layout. The full-cache flags
453
# are consumed only here.
428
- compress_norm_rope_store_fn = compress_norm_rope_store_cutedsl
454
+ compress_norm_rope_store_fn = _SPARSE_ATTN_COMPRESSOR_CUTEDSL_KERNEL
455
extra_kwargs: dict[str, Any] = dict(
456
store_full_kv=store_full_kv,
457
store_full_fp8=store_full_fp8,
vllm/models/deepseek_v4/nvidia/ops/dequant_gather_k_cutedsl.py
+346
-279
@@ -1,7 +1,8 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-from functools import cache
4
+from dataclasses import dataclass
5
+from typing import Any
6
7
import cutlass
8
import cutlass.cute as cute
@@ -12,302 +13,340 @@ from cutlass.cute.nvgpu import cpasync
13
from quack.compile_utils import make_fake_tensor
14
15
from vllm.cute_utils import _bf16x2_mul, cvt
16
+from vllm.model_executor.warmup.jit_warmup import kernel_launcher, zip_inputs
17
+from vllm.model_executor.warmup.jit_warmup_cutedsl_helper import (
18
+ CuTeDSLLaunchSpec,
19
+ VllmCuTeDSLJitKernel,
20
+)
21
22
17
-def dequantize_and_gather_k_cache_cutedsl(
18
- out: torch.Tensor,
19
- k_cache: torch.Tensor,
20
- seq_lens: torch.Tensor,
21
- gather_lens: torch.Tensor | None,
22
- block_table: torch.Tensor,
23
- block_size: int,
24
- offset: int,
25
-) -> None:
26
- DequantGatherKCacheKernel.compile(
27
- block_size=block_size,
28
- has_gather_lens=gather_lens is not None,
29
- )(out, k_cache, seq_lens, gather_lens, block_table, offset)
30
-
31
-
32
-class DequantGatherKCacheKernel:
23
+class DequantGatherKCacheKernel(
24
+ VllmCuTeDSLJitKernel["DequantGatherKCacheKernel.CompileKey"]
25
+):
26
# Hard-coded for DSv4.
27
head_dim = 512
28
+ fp8_dim = 448
29
group_size = 64 # 1 scale per 64 elems
30
+ num_warps = 4
31
+ tb_size = num_warps * 32
32
+ num_stages = 4
33
37
- def __init__(self, fp8_dim: int = 448, block_size: int = 64):
38
- self.fp8_dim = fp8_dim
39
- self.bf16_dim = self.head_dim - fp8_dim
40
- self.data_dim = fp8_dim + self.bf16_dim * 2
41
- self.block_size = block_size
42
-
43
- self.num_warps = 4
44
- self.tb_size = self.num_warps * 32
45
- self.num_stages = 4
46
-
47
- @cute.jit
48
- def __call__(
49
- self,
50
- out: cute.Tensor,
51
- k_cache: cute.Tensor,
52
- seq_lens: cute.Tensor,
53
- gather_lens: cute.Tensor | None,
54
- block_table: cute.Tensor,
55
- offset: Int32,
56
- stream: CUstream,
57
- ):
58
- # Split k_cache into k_data and k_scale. Each [block_size, head_bytes]
59
- # block is actually a concat of
60
- # [block_size, fp8_dim + bf16_dim * 2] and [block_size, 8].
61
- k_data = cute.make_tensor(
62
- k_cache.iterator,
63
- layout=cute.make_layout(
64
- (k_cache.shape[0], self.block_size, self.data_dim),
65
- stride=(k_cache.stride[0], self.data_dim, 1),
66
- ),
67
- )
68
- k_scale = cute.make_tensor(
69
- k_cache.iterator + (self.block_size * self.data_dim),
70
- layout=cute.make_layout(
71
- (k_cache.shape[0], self.block_size, 8),
72
- stride=(k_cache.stride[0], 8, 1),
73
- ),
74
- )
75
-
76
- num_reqs = out.shape[0]
77
- max_rows = out.shape[1] - offset
78
- num_workers = cutlass.max(
79
- 1, cutlass.min(cute.ceil_div(max_rows, 4), cute.ceil_div(8192, num_reqs))
80
- )
81
- grid = (num_reqs, num_workers, 1)
82
- self.kernel(
83
- out,
84
- k_data,
85
- k_scale,
86
- seq_lens,
87
- gather_lens,
88
- block_table,
89
- offset,
90
- ).launch(grid=grid, block=(self.tb_size, 1, 1), stream=stream)
34
+ @dataclass(frozen=True)
35
+ class CompileKey:
36
+ block_size: int
37
+ has_gather_lens: bool
38
92
- @cute.jit
93
- def load_g2s(
94
- self,
95
- k_data_slice: cute.Tensor,
96
- k_scale: cute.Tensor,
97
- block_table: cute.Tensor,
98
- s_kdata_slice: cute.Tensor,
99
- s_kscale: cute.Tensor,
100
- req_id,
101
- pos,
102
- lane_id,
103
- stage_id,
104
- ):
105
- # k_data_slice: [num_blocks, block_size, (16, data_dim/16)]
106
- # s_kdata_slice: [(4, data_dim/16), num_stages]
107
-
108
- op = cpasync.CopyG2SOp(cute.nvgpu.LoadCacheMode.GLOBAL)
109
- cp16_atom = cute.make_copy_atom(op, Uint32, num_bits_per_copy=128)
110
- cp8_atom = cute.make_copy_atom(cpasync.CopyG2SOp(), Uint8, num_bits_per_copy=64)
111
- page_id = block_table[req_id, pos // self.block_size]
112
- block_offset = pos % self.block_size
113
-
114
- # Load the first 512 bytes (32x16B).
115
- idx = lane_id
116
- src = k_data_slice[page_id, block_offset, (None, idx)]
117
- cute.copy(
118
- cp16_atom,
119
- cute.recast_tensor(src, Uint32),
120
- s_kdata_slice[(None, idx), stage_id],
121
- )
39
+ @staticmethod
40
+ def kernel(compile_key: CompileKey) -> Any:
41
+ head_dim = DequantGatherKCacheKernel.head_dim
42
+ fp8_dim = DequantGatherKCacheKernel.fp8_dim
43
+ group_size = DequantGatherKCacheKernel.group_size
44
+ num_warps = DequantGatherKCacheKernel.num_warps
45
+ tb_size = DequantGatherKCacheKernel.tb_size
46
+ num_stages = DequantGatherKCacheKernel.num_stages
47
+ bf16_dim = head_dim - fp8_dim
48
+ data_dim = fp8_dim + bf16_dim * 2
49
+ block_size = compile_key.block_size
50
+
51
+ @cute.jit
52
+ def load_g2s(
53
+ k_data_slice: cute.Tensor,
54
+ k_scale: cute.Tensor,
55
+ block_table: cute.Tensor,
56
+ s_kdata_slice: cute.Tensor,
57
+ s_kscale: cute.Tensor,
58
+ req_id,
59
+ pos,
60
+ lane_id,
61
+ stage_id,
62
+ ):
63
+ op = cpasync.CopyG2SOp(cute.nvgpu.LoadCacheMode.GLOBAL)
64
+ cp16_atom = cute.make_copy_atom(op, Uint32, num_bits_per_copy=128)
65
+ cp8_atom = cute.make_copy_atom(
66
+ cpasync.CopyG2SOp(), Uint8, num_bits_per_copy=64
67
+ )
68
+ page_id = block_table[req_id, pos // block_size]
69
+ block_offset = pos % block_size
70
123
- # Load the tail 64 bytes.
124
- idx += 32
125
- if idx < cutlass.const_expr(self.data_dim // 16):
71
+ # Load the first 512 bytes (32x16B).
72
+ idx = lane_id
73
src = k_data_slice[page_id, block_offset, (None, idx)]
74
cute.copy(
75
cp16_atom,
76
cute.recast_tensor(src, Uint32),
77
s_kdata_slice[(None, idx), stage_id],
78
)
132
- elif idx == cutlass.const_expr(self.data_dim // 16):
133
- cute.copy(
134
- cp8_atom,
135
- k_scale[page_id, block_offset, None],
136
- s_kscale[None, stage_id],
137
- )
79
139
- @cute.kernel
140
- def kernel(
141
- self,
142
- out: cute.Tensor,
143
- k_data: cute.Tensor,
144
- k_scale: cute.Tensor,
145
- seq_lens: cute.Tensor,
146
- gather_lens: cute.Tensor | None,
147
- block_table: cute.Tensor,
148
- offset: Int32,
149
- ):
150
- req_id, worker_id, _ = cute.arch.block_idx()
151
- tid, _, _ = cute.arch.thread_idx()
152
- warp_id = cute.arch.make_warp_uniform(tid // 32)
153
- lane_id = tid % 32
154
-
155
- _, num_workers, _ = cute.arch.grid_dim()
156
-
157
- # Prepare smem.
158
- smem = cutlass.utils.SmemAllocator()
159
- s_kdata = smem.allocate_tensor(
160
- Uint32,
161
- cute.make_layout((self.data_dim // 4, self.num_warps, self.num_stages)),
162
- byte_alignment=16,
163
- )[None, warp_id, None]
164
- s_kscale = smem.allocate_tensor(
165
- Uint8,
166
- cute.make_layout((8, self.num_warps, self.num_stages)),
167
- byte_alignment=8,
168
- )[None, warp_id, None]
169
-
170
- # Prepare for 16B cp.async, also for BF16 smem loads later.
171
- k_data_slice = cute.logical_divide(k_data, (None, None, 16))
172
- s_kdata_16B_slice = cute.logical_divide(s_kdata, (4, None))
173
-
174
- # Load FP8 elems in 8B units, so once dequantized, they are 16B units.
175
- s_kdata_8B_slice = cute.logical_divide(s_kdata, (2, None))
176
-
177
- # 16B st.global.
178
- out_slice = cute.logical_divide(out, (None, None, 8))
179
-
180
- cp_op = cute.nvgpu.CopyUniversalOp()
181
- cp8_atom = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=64)
182
- cp16_atom = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=128)
183
-
184
- seq_len = seq_lens[req_id]
185
- gather_len = seq_len
186
- if cutlass.const_expr(gather_lens is not None):
187
- gather_len = gather_lens[req_id] # type: ignore[index]
188
- start_pos = seq_len - gather_len
189
-
190
- # Start prefetch.
191
- for i in cutlass.range_constexpr(self.num_stages - 1):
192
- next_pos = (
193
- start_pos
194
- + worker_id * self.num_warps
195
- + warp_id
196
- + i * num_workers * self.num_warps
197
- )
198
- if next_pos < seq_len:
199
- self.load_g2s(
200
- k_data_slice,
201
- k_scale,
202
- block_table,
203
- s_kdata_16B_slice,
204
- s_kscale,
205
- req_id,
206
- next_pos,
207
- lane_id,
208
- i,
80
+ # Load the tail 64 bytes.
81
+ idx += 32
82
+ if idx < cutlass.const_expr(data_dim // 16):
83
+ src = k_data_slice[page_id, block_offset, (None, idx)]
84
+ cute.copy(
85
+ cp16_atom,
86
+ cute.recast_tensor(src, Uint32),
87
+ s_kdata_slice[(None, idx), stage_id],
88
)
210
- cute.arch.cp_async_commit_group()
211
- prefetch_stage = self.num_stages - 1
212
- compute_stage = 0
213
-
214
- # Main loop.
215
- for i in range(
216
- worker_id * self.num_warps + warp_id,
217
- gather_len,
218
- num_workers * self.num_warps,
219
- ):
220
- pos = start_pos + i
221
-
222
- # Prefetch next stage.
223
- next_pos = pos + num_workers * self.num_warps * (self.num_stages - 1)
224
- if next_pos < seq_len:
225
- self.load_g2s(
226
- k_data_slice,
227
- k_scale,
228
- block_table,
229
- s_kdata_16B_slice,
230
- s_kscale,
231
- req_id,
232
- next_pos,
233
- lane_id,
234
- prefetch_stage,
89
+ elif idx == cutlass.const_expr(data_dim // 16):
90
+ cute.copy(
91
+ cp8_atom,
92
+ k_scale[page_id, block_offset, None],
93
+ s_kscale[None, stage_id],
94
)
236
- prefetch_stage = (prefetch_stage + 1) % self.num_stages
237
- cute.arch.cp_async_commit_group()
238
-
239
- # Wait for gmem->smem to finish.
240
- cute.arch.cp_async_wait_group(self.num_stages - 1)
241
- cute.arch.sync_warp()
242
-
243
- # There are 512 elems per token. As a warp, data0 holds the first
244
- # 256 elems and data1 holds the second 256 elems, i.e. each thread
245
- # holds 8 FP8 elems. This keeps the dequantized 8 BF16 elems as
246
- # contiguous 16B global stores. On Blackwell, this might not be
247
- # necessary as we have 32B global stores, but doing it this way
248
- # does not seem to be slower.
249
- data0 = cute.make_rmem_tensor((2,), Uint32)
250
- data1 = cute.make_rmem_tensor((2,), Uint32)
251
- cute.copy(cp8_atom, s_kdata_8B_slice[(None, lane_id), compute_stage], data0)
252
- cute.copy(
253
- cp8_atom,
254
- s_kdata_8B_slice[(None, lane_id + 32), compute_stage],
255
- data1,
256
- )
95
258
- # Convert to bf16x2 via bit manipulation. FP8 scales are per 64
259
- # elements. An 8-element chunk advances the scale index by
260
- # chunk_id * 8 // group_size.
261
- scale0_u32 = Uint32(s_kscale[lane_id * 8 // self.group_size, compute_stage])
262
- scale0_bf16x2 = (scale0_u32 << Uint32(23)) | (scale0_u32 << Uint32(7))
263
- scale1_u32 = Uint32(
264
- s_kscale[(lane_id + 32) * 8 // self.group_size, compute_stage]
265
- )
266
- scale1_bf16x2 = (scale1_u32 << Uint32(23)) | (scale1_u32 << Uint32(7))
267
-
268
- # cvt.rn.scaled::n2::ue8m0.bf16x2.e4m3x2 requires PTX 9.2
269
- # (CUDA 13.2).
270
- dequant0 = cute.make_rmem_tensor(4, Uint32)
271
- dequant1 = cute.make_rmem_tensor(4, Uint32)
272
- for j in cutlass.range_constexpr(2):
273
- tmp0 = cvt.fp8x4_to_bf16x4(data0[j])
274
- tmp1 = cvt.fp8x4_to_bf16x4(data1[j])
275
-
276
- # BF16 multiply is safe because the scales are exact powers of 2.
277
- dequant0[j * 2] = _bf16x2_mul(tmp0[0], scale0_bf16x2)
278
- dequant1[j * 2] = _bf16x2_mul(tmp1[0], scale1_bf16x2)
279
- dequant0[j * 2 + 1] = _bf16x2_mul(tmp0[1], scale0_bf16x2)
280
- dequant1[j * 2 + 1] = _bf16x2_mul(tmp1[1], scale1_bf16x2)
281
-
282
- # Last 64 elems are BF16 tail, corresponds to dequant1 of last
283
- # 8 threads. We have 448 FP8 + 64 BF16 -> 28x 16B for FP8 +
284
- # 8x 16B for BF16.
285
- if lane_id + 32 >= self.fp8_dim // 8:
286
- idx = self.fp8_dim // 16 + (lane_id + 32) - self.fp8_dim // 8
96
+ @cute.kernel
97
+ def device_kernel(
98
+ out: cute.Tensor,
99
+ k_data: cute.Tensor,
100
+ k_scale: cute.Tensor,
101
+ seq_lens: cute.Tensor,
102
+ gather_lens: cute.Tensor | None,
103
+ block_table: cute.Tensor,
104
+ offset: Int32,
105
+ ):
106
+ req_id, worker_id, _ = cute.arch.block_idx()
107
+ tid, _, _ = cute.arch.thread_idx()
108
+ warp_id = cute.arch.make_warp_uniform(tid // 32)
109
+ lane_id = tid % 32
110
+
111
+ _, num_workers, _ = cute.arch.grid_dim()
112
+
113
+ # Prepare smem.
114
+ smem = cutlass.utils.SmemAllocator()
115
+ s_kdata = smem.allocate_tensor(
116
+ Uint32,
117
+ cute.make_layout((data_dim // 4, num_warps, num_stages)),
118
+ byte_alignment=16,
119
+ )[None, warp_id, None]
120
+ s_kscale = smem.allocate_tensor(
121
+ Uint8,
122
+ cute.make_layout((8, num_warps, num_stages)),
123
+ byte_alignment=8,
124
+ )[None, warp_id, None]
125
+
126
+ # Prepare for 16B cp.async, also for BF16 smem loads later.
127
+ k_data_slice = cute.logical_divide(k_data, (None, None, 16))
128
+ s_kdata_16B_slice = cute.logical_divide(s_kdata, (4, None))
129
+
130
+ # Load FP8 elems in 8B units, so once dequantized, they are 16B units.
131
+ s_kdata_8B_slice = cute.logical_divide(s_kdata, (2, None))
132
+
133
+ # 16B st.global.
134
+ out_slice = cute.logical_divide(out, (None, None, 8))
135
+
136
+ cp_op = cute.nvgpu.CopyUniversalOp()
137
+ cp8_atom = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=64)
138
+ cp16_atom = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=128)
139
+
140
+ seq_len = seq_lens[req_id]
141
+ gather_len = seq_len
142
+ if cutlass.const_expr(gather_lens is not None):
143
+ gather_len = gather_lens[req_id] # type: ignore[index]
144
+ start_pos = seq_len - gather_len
145
+
146
+ # Start prefetch.
147
+ for i in cutlass.range_constexpr(num_stages - 1):
148
+ next_pos = (
149
+ start_pos
150
+ + worker_id * num_warps
151
+ + warp_id
152
+ + i * num_workers * num_warps
153
+ )
154
+ if next_pos < seq_len:
155
+ load_g2s(
156
+ k_data_slice,
157
+ k_scale,
158
+ block_table,
159
+ s_kdata_16B_slice,
160
+ s_kscale,
161
+ req_id,
162
+ next_pos,
163
+ lane_id,
164
+ i,
165
+ )
166
+ cute.arch.cp_async_commit_group()
167
+ prefetch_stage = num_stages - 1
168
+ compute_stage = 0
169
+
170
+ # Main loop.
171
+ for i in range(
172
+ worker_id * num_warps + warp_id,
173
+ gather_len,
174
+ num_workers * num_warps,
175
+ ):
176
+ pos = start_pos + i
177
+
178
+ # Prefetch next stage.
179
+ next_pos = pos + num_workers * num_warps * (num_stages - 1)
180
+ if next_pos < seq_len:
181
+ load_g2s(
182
+ k_data_slice,
183
+ k_scale,
184
+ block_table,
185
+ s_kdata_16B_slice,
186
+ s_kscale,
187
+ req_id,
188
+ next_pos,
189
+ lane_id,
190
+ prefetch_stage,
191
+ )
192
+ prefetch_stage = (prefetch_stage + 1) % num_stages
193
+ cute.arch.cp_async_commit_group()
194
+
195
+ # Wait for gmem->smem to finish.
196
+ cute.arch.cp_async_wait_group(num_stages - 1)
197
+ cute.arch.sync_warp()
198
+
199
+ # There are 512 elems per token. As a warp, data0 holds the first
200
+ # 256 elems and data1 holds the second 256 elems, i.e. each thread
201
+ # holds 8 FP8 elems. This keeps the dequantized 8 BF16 elems as
202
+ # contiguous 16B global stores. On Blackwell, this might not be
203
+ # necessary as we have 32B global stores, but doing it this way
204
+ # does not seem to be slower.
205
+ data0 = cute.make_rmem_tensor((2,), Uint32)
206
+ data1 = cute.make_rmem_tensor((2,), Uint32)
207
cute.copy(
288
- cp16_atom,
289
- s_kdata_16B_slice[(None, idx), compute_stage],
290
- dequant1,
208
+ cp8_atom, s_kdata_8B_slice[(None, lane_id), compute_stage], data0
209
+ )
210
+ cute.copy(
211
+ cp8_atom,
212
+ s_kdata_8B_slice[(None, lane_id + 32), compute_stage],
213
+ data1,
214
)
215
293
- # Store two 16B BF16 chunks per lane: first half, then second half.
294
- dst = out_slice[req_id, offset + i, (None, lane_id)]
295
- cute.copy(cp16_atom, dequant0, cute.recast_tensor(dst, Uint32))
216
+ # Convert to bf16x2 via bit manipulation. FP8 scales are per 64
217
+ # elements. An 8-element chunk advances the scale index by
218
+ # chunk_id * 8 // group_size.
219
+ scale0_u32 = Uint32(s_kscale[lane_id * 8 // group_size, compute_stage])
220
+ scale0_bf16x2 = (scale0_u32 << Uint32(23)) | (scale0_u32 << Uint32(7))
221
+ scale1_u32 = Uint32(
222
+ s_kscale[(lane_id + 32) * 8 // group_size, compute_stage]
223
+ )
224
+ scale1_bf16x2 = (scale1_u32 << Uint32(23)) | (scale1_u32 << Uint32(7))
225
+
226
+ # cvt.rn.scaled::n2::ue8m0.bf16x2.e4m3x2 requires PTX 9.2
227
+ # (CUDA 13.2).
228
+ dequant0 = cute.make_rmem_tensor(4, Uint32)
229
+ dequant1 = cute.make_rmem_tensor(4, Uint32)
230
+ for j in cutlass.range_constexpr(2):
231
+ tmp0 = cvt.fp8x4_to_bf16x4(data0[j])
232
+ tmp1 = cvt.fp8x4_to_bf16x4(data1[j])
233
+
234
+ # BF16 multiply is safe because the scales are exact powers of 2.
235
+ dequant0[j * 2] = _bf16x2_mul(tmp0[0], scale0_bf16x2)
236
+ dequant1[j * 2] = _bf16x2_mul(tmp1[0], scale1_bf16x2)
237
+ dequant0[j * 2 + 1] = _bf16x2_mul(tmp0[1], scale0_bf16x2)
238
+ dequant1[j * 2 + 1] = _bf16x2_mul(tmp1[1], scale1_bf16x2)
239
+
240
+ # Last 64 elems are BF16 tail, corresponds to dequant1 of last
241
+ # 8 threads. We have 448 FP8 + 64 BF16 -> 28x 16B for FP8 +
242
+ # 8x 16B for BF16.
243
+ if lane_id + 32 >= fp8_dim // 8:
244
+ idx = fp8_dim // 16 + (lane_id + 32) - fp8_dim // 8
245
+ cute.copy(
246
+ cp16_atom,
247
+ s_kdata_16B_slice[(None, idx), compute_stage],
248
+ dequant1,
249
+ )
250
+
251
+ # Store two 16B BF16 chunks per lane: first half, then second half.
252
+ dst = out_slice[req_id, offset + i, (None, lane_id)]
253
+ cute.copy(cp16_atom, dequant0, cute.recast_tensor(dst, Uint32))
254
+
255
+ dst = out_slice[req_id, offset + i, (None, lane_id + 32)]
256
+ cute.copy(cp16_atom, dequant1, cute.recast_tensor(dst, Uint32))
257
+
258
+ compute_stage = (compute_stage + 1) % num_stages
259
+
260
+ @cute.jit
261
+ def host_entrypoint(
262
+ out: cute.Tensor,
263
+ k_cache: cute.Tensor,
264
+ seq_lens: cute.Tensor,
265
+ gather_lens: cute.Tensor | None,
266
+ block_table: cute.Tensor,
267
+ offset: Int32,
268
+ stream: CUstream,
269
+ ):
270
+ # Split k_cache into k_data and k_scale. Each [block_size, head_bytes]
271
+ # block is actually a concat of
272
+ # [block_size, fp8_dim + bf16_dim * 2] and [block_size, 8].
273
+ k_data = cute.make_tensor(
274
+ k_cache.iterator,
275
+ layout=cute.make_layout(
276
+ (k_cache.shape[0], block_size, data_dim),
277
+ stride=(k_cache.stride[0], data_dim, 1),
278
+ ),
279
+ )
280
+ # k_data_slice: [num_blocks, block_size, (16, data_dim/16)]
281
+ # s_kdata_slice: [(4, data_dim/16), num_stages]
282
+ k_scale = cute.make_tensor(
283
+ k_cache.iterator + (block_size * data_dim),
284
+ layout=cute.make_layout(
285
+ (k_cache.shape[0], block_size, 8),
286
+ stride=(k_cache.stride[0], 8, 1),
287
+ ),
288
+ )
289
297
- dst = out_slice[req_id, offset + i, (None, lane_id + 32)]
298
- cute.copy(cp16_atom, dequant1, cute.recast_tensor(dst, Uint32))
290
+ num_reqs = out.shape[0]
291
+ max_rows = out.shape[1] - offset
292
+ num_workers = cutlass.max(
293
+ 1,
294
+ cutlass.min(cute.ceil_div(max_rows, 4), cute.ceil_div(8192, num_reqs)),
295
+ )
296
+ grid = (num_reqs, num_workers, 1)
297
+ device_kernel(
298
+ out,
299
+ k_data,
300
+ k_scale,
301
+ seq_lens,
302
+ gather_lens,
303
+ block_table,
304
+ offset,
305
+ ).launch(grid=grid, block=(tb_size, 1, 1), stream=stream)
306
+
307
+ return host_entrypoint
308
+
309
+ def dispatch( # type: ignore[override]
310
+ self,
311
+ *,
312
+ block_size: int,
313
+ has_gather_lens: bool,
314
+ ) -> CompileKey:
315
+ return self.CompileKey(
316
+ block_size=block_size,
317
+ has_gather_lens=has_gather_lens,
318
+ )
319
300
- compute_stage = (compute_stage + 1) % self.num_stages
320
+ def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]:
321
+ block_size = vllm_config.cache_config.block_size
322
+ compress_ratios = vllm_config.model_config.hf_config.compress_ratios
323
+ return self._trace_dispatch(self.dispatch)(
324
+ zip_inputs(
325
+ dict(
326
+ block_size=block_size,
327
+ has_gather_lens=True,
328
+ enabled=True,
329
+ ),
330
+ dict(
331
+ block_size=max(1, block_size // 4),
332
+ has_gather_lens=False,
333
+ enabled=4 in compress_ratios,
334
+ ),
335
+ dict(
336
+ block_size=max(1, block_size // 128),
337
+ has_gather_lens=False,
338
+ enabled=128 in compress_ratios,
339
+ ),
340
+ ),
341
+ valid_block_size=block_size > 0,
342
+ _when=lambda *, enabled, valid_block_size: enabled and valid_block_size,
343
+ )
344
302
- @cache
303
- @staticmethod
304
- def compile(
305
- fp8_dim: int = 448,
306
- block_size: int = 64,
307
- has_gather_lens: bool = True,
308
- ):
345
+ def warmup_inputs(self, compile_key: CompileKey) -> tuple[Any, ...]:
346
num_reqs = cute.sym_int()
310
- head_dim = DequantGatherKCacheKernel.head_dim
347
+ head_dim = self.head_dim
348
+ fp8_dim = self.fp8_dim
349
+ block_size = compile_key.block_size
350
head_bytes = fp8_dim + (head_dim - fp8_dim) * 2 + 8
351
352
out = make_fake_tensor(BFloat16, (num_reqs, cute.sym_int(), head_dim), 16)
@@ -318,19 +357,47 @@ class DequantGatherKCacheKernel:
357
assumed_align=32,
358
)
359
seq_lens = make_fake_tensor(Int32, (num_reqs,))
321
- gather_lens = make_fake_tensor(Int32, (num_reqs,)) if has_gather_lens else None
360
+ gather_lens = (
361
+ make_fake_tensor(Int32, (num_reqs,))
362
+ if compile_key.has_gather_lens
363
+ else None
364
+ )
365
block_table = make_fake_tensor(Int32, (num_reqs, cute.sym_int()))
366
324
- kernel = DequantGatherKCacheKernel(fp8_dim, block_size)
325
- stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
326
- return cute.compile(
327
- kernel,
367
+ return (
368
out,
369
k_cache,
370
seq_lens,
371
gather_lens,
372
block_table,
373
Int32(0),
334
- stream,
335
- options="--enable-tvm-ffi",
374
)
375
+
376
+ @kernel_launcher
377
+ def __call__(
378
+ self,
379
+ *,
380
+ out: torch.Tensor,
381
+ k_cache: torch.Tensor,
382
+ seq_lens: torch.Tensor,
383
+ gather_lens: torch.Tensor | None,
384
+ block_table: torch.Tensor,
385
+ block_size: int,
386
+ offset: int,
387
+ ) -> CuTeDSLLaunchSpec["DequantGatherKCacheKernel.CompileKey"]:
388
+ compile_key = self.dispatch(
389
+ block_size=block_size,
390
+ has_gather_lens=gather_lens is not None,
391
+ )
392
+ launch_args = (
393
+ out,
394
+ k_cache,
395
+ seq_lens,
396
+ gather_lens,
397
+ block_table,
398
+ offset,
399
+ )
400
+ return compile_key, launch_args
401
+
402
+
403
+_DEQUANT_GATHER_K_CACHE_CUTEDSL_KERNEL = DequantGatherKCacheKernel()
vllm/models/deepseek_v4/nvidia/ops/fused_indexer_q_cutedsl.py
+591
-510
@@ -1,6 +1,9 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-from functools import cache
3
+# ruff: noqa: E501
4
+
5
+from dataclasses import dataclass
6
+from typing import Any
7
8
import cutlass
9
import cutlass.cute as cute
@@ -14,594 +17,638 @@ from vllm.cute_utils import (
17
_bf16x2_max,
18
cvt,
19
recast_val,
20
+ torch_to_cute_dtype,
21
+)
22
+from vllm.model_executor.warmup.jit_warmup import kernel_launcher
23
+from vllm.model_executor.warmup.jit_warmup_cutedsl_helper import (
24
+ CuTeDSLLaunchSpec,
25
+ VllmCuTeDSLJitKernel,
26
)
27
from vllm.vllm_flash_attn.cute import utils as cute_utils
28
29
# MXFP4: 32 elements per block, packed 2 nibbles per byte, ue8m0 block scale.
30
MXFP4_BLOCK_SIZE = 32
31
23
-_TORCH_TO_CUTE = {
24
- torch.bfloat16: BFloat16,
25
- torch.float32: Float32,
26
-}
27
-
28
-
29
-def fused_indexer_q_rope_quant_mxfp4_cutedsl(
30
- positions: torch.Tensor,
31
- index_q: torch.Tensor,
32
- index_q_cos_sin_cache: torch.Tensor,
33
- index_weights: torch.Tensor,
34
- index_weights_softmax_scale: float,
35
- index_weights_head_scale: float,
36
- index_q_packed: torch.Tensor,
37
- index_q_scale: torch.Tensor,
38
- index_weights_out: torch.Tensor,
39
-) -> None:
40
- num_tokens, num_heads, head_dim = index_q.shape
41
- rope_dim = index_q_cos_sin_cache.shape[-1]
42
- rope_type = _TORCH_TO_CUTE[index_q_cos_sin_cache.dtype]
43
-
44
- # compile all variants at first invocation
45
- for coarsen in (1, 4):
46
- IndexerQMxFp4Kernel.compile(head_dim, rope_dim, num_heads, rope_type, coarsen)
47
-
48
- # heuristic
49
- coarsen = 1 if num_tokens < 512 else 4
50
- compiled = IndexerQMxFp4Kernel.compile(
51
- head_dim, rope_dim, num_heads, rope_type, coarsen
52
- )
53
- scale = float(index_weights_softmax_scale * index_weights_head_scale)
54
- compiled(
55
- positions,
56
- index_q,
57
- index_q_cos_sin_cache,
58
- index_weights,
59
- index_q_packed,
60
- index_q_scale,
61
- index_weights_out,
62
- scale,
63
- )
32
65
-
66
-def fused_indexer_q_rope_quant_fp8_cutedsl(
67
- positions: torch.Tensor,
68
- index_q: torch.Tensor,
69
- index_q_cos_sin_cache: torch.Tensor,
70
- index_weights: torch.Tensor,
71
- index_weights_softmax_scale: float,
72
- index_weights_head_scale: float,
73
- index_q_fp8: torch.Tensor,
74
- index_weights_out: torch.Tensor,
75
-) -> None:
76
- num_tokens, num_heads, head_dim = index_q.shape
77
- rope_dim = index_q_cos_sin_cache.shape[-1]
78
- rope_type = _TORCH_TO_CUTE[index_q_cos_sin_cache.dtype]
79
-
80
- for coarsen in (1, 4):
81
- IndexerQFp8Kernel.compile(head_dim, rope_dim, num_heads, rope_type, coarsen)
82
-
83
- coarsen = 1 if num_tokens < 512 else 4
84
- compiled = IndexerQFp8Kernel.compile(
85
- head_dim, rope_dim, num_heads, rope_type, coarsen
86
- )
87
- scale = float(index_weights_softmax_scale * index_weights_head_scale)
88
- # The cute kernel treats the FP8 buffer as raw bytes (Uint8).
89
- compiled(
90
- positions,
91
- index_q,
92
- index_q_cos_sin_cache,
93
- index_weights,
94
- index_q_fp8.view(torch.uint8),
95
- index_weights_out,
96
- scale,
33
+@cute.jit
34
+def _load_q_and_rope(
35
+ positions: cute.Tensor,
36
+ q: cute.Tensor,
37
+ cos_sin_cache: cute.Tensor,
38
+ num_heads: cutlass.Constexpr,
39
+ coarsen: cutlass.Constexpr,
40
+ tb_size: cutlass.Constexpr,
41
+ subwarp_size: cutlass.Constexpr,
42
+ threads_per_token: cutlass.Constexpr,
43
+ nope_dim: cutlass.Constexpr,
44
+ rope_dim: cutlass.Constexpr,
45
+ cos_sin_dtype: cutlass.Constexpr,
46
+):
47
+ """Compute thread indices, load Q (BF16), and apply interleaved RoPE.
48
+
49
+ Returns a tuple
50
+ (q_bf16x2, tid, global_tid, sublane, token_id, head_tile_id,
51
+ head_start, in_bounds, num_token_heads)
52
+ where ``q_bf16x2`` is a (coarsen, 8) rmem tile of Uint32 packed
53
+ bf16x2 pairs covering the 16 BF16 lanes owned by this thread for
54
+ each of ``coarsen`` heads. RoPE is applied in place to the
55
+ trailing ``rope_dim`` lanes; the leading nope lanes pass through.
56
+ """
57
+ block_id, _, _ = cute.arch.block_idx()
58
+ tid, _, _ = cute.arch.thread_idx()
59
+
60
+ num_tokens = q.shape[0]
61
+ num_token_heads = num_tokens * num_heads
62
+ global_tid = block_id * tb_size + tid
63
+
64
+ global_subwarp_id = global_tid // subwarp_size
65
+ sublane = tid % subwarp_size
66
+
67
+ token_id = global_subwarp_id // (num_heads // coarsen)
68
+ head_tile_id = global_subwarp_id % (num_heads // coarsen)
69
+ head_start = head_tile_id * coarsen
70
+
71
+ # NOTE: token_id may exceed bounds, hence we need to add load/store guards
72
+ # we can't do early exit because CuteDSL doesn't support it. and we also need
73
+ # all threads in a warp to be active since we utilize warp shuffle later.
74
+ # must_in_bounds is constexpr, True when a token spans a whole number of
75
+ # threadblocks. Then num_tokens * threads_per_token is always a multiple of
76
+ # tb_size, so the trailing block never straddles num_tokens and the compiler
77
+ # removes the bounds check. Otherwise the guard has to stay: e.g. 32 heads
78
+ # with coarsen=4 needs 64 threads per token, so an odd num_tokens leaves half
79
+ # of the last block addressing token num_tokens.
80
+ must_in_bounds = cutlass.const_expr(threads_per_token % tb_size == 0)
81
+ in_bounds = must_in_bounds or (token_id < num_tokens)
82
+
83
+ cp_op = cute.nvgpu.CopyUniversalOp()
84
+
85
+ _layout = cute.make_layout((coarsen, 8), stride=(8, 1))
86
+ q_bf16x2 = cute.make_rmem_tensor(_layout, Uint32)
87
+
88
+ if in_bounds:
89
+ # we can't do cute.copy() on the whole 2D tile directly because
90
+ # cute.copy() wants the 1st mode to be covered by the copy atom,
91
+ # and other modes as for loop. there is no fast way to
92
+ # "transpose" the tensor view.
93
+ q_tile = cute.local_tile(
94
+ q[token_id, None, None],
95
+ tiler=(coarsen, 16),
96
+ coord=(head_tile_id, sublane),
97
+ )
98
+ cp_u32x8 = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=256)
99
+ for i in cutlass.range_constexpr(coarsen):
100
+ src = cute.recast_tensor(q_tile[i, None], Uint32)
101
+ cute.copy(cp_u32x8, src, q_bf16x2[i, None])
102
+
103
+ # RoPE applies only to the trailing rope_dim values. We keep the rounded
104
+ # BF16 result in q_bits so the later amax and quantization see BF16.
105
+ # cos_sin_cache layout: [max_pos, rope_dim]
106
+ if in_bounds and sublane * 16 >= nope_dim:
107
+ cos_vals = cute.make_rmem_tensor((8,), Float32)
108
+ sin_vals = cute.make_rmem_tensor((8,), Float32)
109
+
110
+ pos = positions[token_id]
111
+
112
+ # select 8 elems from cos and sin
113
+ cos_id = sublane - nope_dim // 16
114
+ sin_id = cos_id + rope_dim // 16
115
+ cos_src = cute.local_tile(cos_sin_cache[pos, None], tiler=(8,), coord=(cos_id,))
116
+ sin_src = cute.local_tile(cos_sin_cache[pos, None], tiler=(8,), coord=(sin_id,))
117
+
118
+ cp_f32x8 = cute.make_copy_atom(cp_op, Float32, num_bits_per_copy=256)
119
+ cp_u32x4 = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=128)
120
+
121
+ if const_expr(cos_sin_dtype is Float32):
122
+ cute.copy(cp_f32x8, cos_src, cos_vals)
123
+ cute.copy(cp_f32x8, sin_src, sin_vals)
124
+ else:
125
+ cos_bf16x2 = cute.make_rmem_tensor((4,), Uint32)
126
+ sin_bf16x2 = cute.make_rmem_tensor((4,), Uint32)
127
+ cute.copy(cp_u32x4, cute.recast_tensor(cos_src, Uint32), cos_bf16x2)
128
+ cute.copy(cp_u32x4, cute.recast_tensor(sin_src, Uint32), sin_bf16x2)
129
+
130
+ for i in cutlass.range_constexpr(4):
131
+ cos0, cos1 = cvt.bf16x2_to_fp32x2(cos_bf16x2[i])
132
+ sin0, sin1 = cvt.bf16x2_to_fp32x2(sin_bf16x2[i])
133
+ cos_vals[i * 2] = cos0
134
+ cos_vals[i * 2 + 1] = cos1
135
+ sin_vals[i * 2] = sin0
136
+ sin_vals[i * 2 + 1] = sin1
137
+
138
+ for i in cutlass.range_constexpr(coarsen):
139
+ for j in cutlass.range_constexpr(8):
140
+ q0, q1 = cvt.bf16x2_to_fp32x2(q_bf16x2[i, j])
141
+ rot0 = q0 * cos_vals[j] - q1 * sin_vals[j]
142
+ rot1 = q0 * sin_vals[j] + q1 * cos_vals[j]
143
+ # convert back to BF16 to match numerics
144
+ q_bf16x2[i, j] = cvt.fp32x2_to_bf16x2(rot0, rot1)
145
+
146
+ return (
147
+ q_bf16x2,
148
+ tid,
149
+ global_tid,
150
+ sublane,
151
+ token_id,
152
+ head_tile_id,
153
+ head_start,
154
+ in_bounds,
155
+ num_token_heads,
156
)
157
158
100
-class IndexerQRopeQuantKernel:
101
- """Shared infrastructure for indexer-Q RoPE+quant fused kernels.
159
+class IndexerQMxFp4Kernel(VllmCuTeDSLJitKernel["IndexerQMxFp4Kernel.CompileKey"]):
160
+ tb_size = 128
161
103
- Subclasses implement ``kernel`` for a particular Q quantization scheme
104
- (MXFP4, FP8 e4m3, …). The base class owns the launch geometry and the
105
- common preamble: thread/token addressing, the BF16 Q load, and the
106
- interleaved-RoPE pass over the trailing ``rope_dim`` lanes.
107
- """
108
-
109
- def __init__(
110
- self,
111
- head_dim: int = 128,
112
- rope_dim: int = 64,
113
- num_heads: int = 64,
114
- cos_sin_dtype: type[cutlass.Numeric] = Float32,
115
- coarsen: int = 4,
116
- ):
117
- self.head_dim = head_dim
118
- self.rope_dim = rope_dim
119
- self.nope_dim = head_dim - rope_dim
120
- self.num_heads = num_heads
121
- self.cos_sin_dtype = cos_sin_dtype
162
+ @dataclass(frozen=True)
163
+ class CompileKey:
164
+ head_dim: int
165
+ rope_dim: int
166
+ num_heads: int
167
+ cos_sin_dtype: type[cutlass.Numeric]
168
+ coarsen: int
169
170
+ @staticmethod
171
+ def kernel(compile_key: CompileKey) -> Any:
172
+ head_dim = compile_key.head_dim
173
+ rope_dim = compile_key.rope_dim
174
+ nope_dim = head_dim - rope_dim
175
+ num_heads = compile_key.num_heads
176
+ cos_sin_dtype = compile_key.cos_sin_dtype
177
# process multiple heads at the same time to armotize RoPE load costs
178
+ coarsen = compile_key.coarsen
179
+ tb_size = IndexerQMxFp4Kernel.tb_size
180
assert num_heads % coarsen == 0
125
- self.coarsen = coarsen
181
182
# later we will use 32B load = 16 BF16 elems
183
# thus, head_dim=128 requires 8 threads to handle.
184
# let's call subwarp = 8 threads.
130
- self.subwarp_size = head_dim // 16
131
- self.tb_size = 128
132
- self.threads_per_token = (self.num_heads // self.coarsen) * self.subwarp_size
133
-
134
- @cute.jit
135
- def _load_q_and_rope(
136
- self,
137
- positions: cute.Tensor,
138
- q: cute.Tensor,
139
- cos_sin_cache: cute.Tensor,
140
- ):
141
- """Compute thread indices, load Q (BF16), and apply interleaved RoPE.
142
-
143
- Returns a tuple
144
- (q_bf16x2, tid, global_tid, sublane, token_id, head_tile_id,
145
- head_start, in_bounds, num_token_heads)
146
- where ``q_bf16x2`` is a (coarsen, 8) rmem tile of Uint32 packed
147
- bf16x2 pairs covering the 16 BF16 lanes owned by this thread for
148
- each of ``coarsen`` heads. RoPE is applied in place to the
149
- trailing ``rope_dim`` lanes; the leading nope lanes pass through.
150
- """
151
- block_id, _, _ = cute.arch.block_idx()
152
- tid, _, _ = cute.arch.thread_idx()
153
-
154
- num_tokens = q.shape[0]
155
- num_token_heads = num_tokens * self.num_heads
156
- global_tid = block_id * self.tb_size + tid
157
-
158
- global_subwarp_id = global_tid // self.subwarp_size
159
- sublane = tid % self.subwarp_size
160
-
161
- token_id = global_subwarp_id // (self.num_heads // self.coarsen)
162
- head_tile_id = global_subwarp_id % (self.num_heads // self.coarsen)
163
- head_start = head_tile_id * self.coarsen
164
-
165
- # NOTE: token_id may exceed bounds, hence we need to add load/store guards
166
- # we can't do early exit because CuteDSL doesn't support it. and we also need
167
- # all threads in a warp to be active since we utilize warp shuffle later.
168
- # must_in_bounds is constexpr, True when a token spans a whole number of
169
- # threadblocks. Then num_tokens * threads_per_token is always a multiple of
170
- # tb_size, so the trailing block never straddles num_tokens and the compiler
171
- # removes the bounds check. Otherwise the guard has to stay: e.g. 32 heads
172
- # with coarsen=4 needs 64 threads per token, so an odd num_tokens leaves half
173
- # of the last block addressing token num_tokens.
174
- must_in_bounds = cutlass.const_expr(self.threads_per_token % self.tb_size == 0)
175
- in_bounds = must_in_bounds or (token_id < num_tokens)
176
-
177
- cp_op = cute.nvgpu.CopyUniversalOp()
178
-
179
- _layout = cute.make_layout((self.coarsen, 8), stride=(8, 1))
180
- q_bf16x2 = cute.make_rmem_tensor(_layout, Uint32)
181
-
182
- if in_bounds:
183
- # we can't do cute.copy() on the whole 2D tile directly because
184
- # cute.copy() wants the 1st mode to be covered by the copy atom,
185
- # and other modes as for loop. there is no fast way to
186
- # "transpose" the tensor view.
187
- q_tile = cute.local_tile(
188
- q[token_id, None, None],
189
- tiler=(self.coarsen, 16),
190
- coord=(head_tile_id, sublane),
191
- )
192
- cp_u32x8 = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=256)
193
- for i in cutlass.range_constexpr(self.coarsen):
194
- src = cute.recast_tensor(q_tile[i, None], Uint32)
195
- cute.copy(cp_u32x8, src, q_bf16x2[i, None])
196
-
197
- # RoPE applies only to the trailing rope_dim values. We keep the rounded
198
- # BF16 result in q_bits so the later amax and quantization see BF16.
199
- # cos_sin_cache layout: [max_pos, rope_dim]
200
- if in_bounds and sublane * 16 >= self.nope_dim:
201
- cos_vals = cute.make_rmem_tensor((8,), Float32)
202
- sin_vals = cute.make_rmem_tensor((8,), Float32)
203
-
204
- pos = positions[token_id]
205
-
206
- # select 8 elems from cos and sin
207
- cos_id = sublane - self.nope_dim // 16
208
- sin_id = cos_id + self.rope_dim // 16
209
- cos_src = cute.local_tile(
210
- cos_sin_cache[pos, None], tiler=(8,), coord=(cos_id,)
211
- )
212
- sin_src = cute.local_tile(
213
- cos_sin_cache[pos, None], tiler=(8,), coord=(sin_id,)
185
+ subwarp_size = head_dim // 16
186
+ threads_per_token = (num_heads // coarsen) * subwarp_size
187
+
188
+ @cute.kernel
189
+ def device_kernel(
190
+ positions: cute.Tensor,
191
+ q: cute.Tensor,
192
+ cos_sin_cache: cute.Tensor,
193
+ weights: cute.Tensor,
194
+ q_quant: cute.Tensor,
195
+ q_scale: cute.Tensor,
196
+ weights_out: cute.Tensor,
197
+ scale: Float32,
198
+ ):
199
+ (
200
+ q_bf16x2,
201
+ tid,
202
+ global_tid,
203
+ sublane,
204
+ token_id,
205
+ head_tile_id,
206
+ head_start,
207
+ in_bounds,
208
+ num_token_heads,
209
+ ) = _load_q_and_rope(
210
+ positions,
211
+ q,
212
+ cos_sin_cache,
213
+ num_heads,
214
+ coarsen,
215
+ tb_size,
216
+ subwarp_size,
217
+ threads_per_token,
218
+ nope_dim,
219
+ rope_dim,
220
+ cos_sin_dtype,
221
)
222
216
- cp_f32x8 = cute.make_copy_atom(cp_op, Float32, num_bits_per_copy=256)
217
- cp_u32x4 = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=128)
218
-
219
- if const_expr(self.cos_sin_dtype is Float32):
220
- cute.copy(cp_f32x8, cos_src, cos_vals)
221
- cute.copy(cp_f32x8, sin_src, sin_vals)
222
- else:
223
- cos_bf16x2 = cute.make_rmem_tensor((4,), Uint32)
224
- sin_bf16x2 = cute.make_rmem_tensor((4,), Uint32)
225
- cute.copy(cp_u32x4, cute.recast_tensor(cos_src, Uint32), cos_bf16x2)
226
- cute.copy(cp_u32x4, cute.recast_tensor(sin_src, Uint32), sin_bf16x2)
227
-
228
- for i in cutlass.range_constexpr(4):
229
- cos0, cos1 = cvt.bf16x2_to_fp32x2(cos_bf16x2[i])
230
- sin0, sin1 = cvt.bf16x2_to_fp32x2(sin_bf16x2[i])
231
- cos_vals[i * 2] = cos0
232
- cos_vals[i * 2 + 1] = cos1
233
- sin_vals[i * 2] = sin0
234
- sin_vals[i * 2 + 1] = sin1
235
-
236
- for i in cutlass.range_constexpr(self.coarsen):
237
- for j in cutlass.range_constexpr(8):
238
- q0, q1 = cvt.bf16x2_to_fp32x2(q_bf16x2[i, j])
239
- rot0 = q0 * cos_vals[j] - q1 * sin_vals[j]
240
- rot1 = q0 * sin_vals[j] + q1 * cos_vals[j]
241
- # convert back to BF16 to match numerics
242
- q_bf16x2[i, j] = cvt.fp32x2_to_bf16x2(rot0, rot1)
243
-
244
- return (
245
- q_bf16x2,
246
- tid,
247
- global_tid,
248
- sublane,
249
- token_id,
250
- head_tile_id,
251
- head_start,
252
- in_bounds,
253
- num_token_heads,
254
- )
223
+ cp_op = cute.nvgpu.CopyUniversalOp()
224
225
+ # layout: [coarsen, 8]
226
+ q_fp4_tile = cute.local_tile(
227
+ q_quant[token_id, None, None],
228
+ tiler=(coarsen, 8),
229
+ coord=(head_tile_id, sublane),
230
+ )
231
257
-class IndexerQMxFp4Kernel(IndexerQRopeQuantKernel):
258
- """Eight-thread subwarps process one ``(token, head)`` row."""
259
-
260
- @cute.jit
261
- def __call__(
262
- self,
263
- positions: cute.Tensor,
264
- q: cute.Tensor,
265
- cos_sin_cache: cute.Tensor,
266
- weights: cute.Tensor,
267
- q_quant: cute.Tensor,
268
- q_scale: cute.Tensor,
269
- weights_out: cute.Tensor,
270
- scale: Float32,
271
- stream: CUstream,
272
- ):
273
- total_threads = q.shape[0] * self.threads_per_token
274
- grid = (cute.ceil_div(total_threads, self.tb_size), 1, 1)
275
- self.kernel(
276
- positions,
277
- q,
278
- cos_sin_cache,
279
- weights,
280
- q_quant,
281
- q_scale,
282
- weights_out,
283
- scale,
284
- ).launch(grid=grid, block=(self.tb_size, 1, 1), stream=stream)
232
+ for i in cutlass.range_constexpr(coarsen):
233
+ # compute amax in packed bf16x2 to save instructions
234
+ # Each thread holds 16 elems. Two adjacent threads form one 32-elem
235
+ # MXFP4 block, so a width-2 shuffle gives the block amax.
236
+ amax_bf16x2 = _bf16x2_abs(q_bf16x2[i, 0])
237
+ for j in cutlass.range_constexpr(1, 8):
238
+ amax_bf16x2 = _bf16x2_max(amax_bf16x2, _bf16x2_abs(q_bf16x2[i, j]))
239
+ amax_bf16x2 = cute_utils.warp_reduce(
240
+ amax_bf16x2,
241
+ _bf16x2_max,
242
+ width=MXFP4_BLOCK_SIZE // 16,
243
+ )
244
+ amax_pair = cvt.bf16x2_to_fp32x2(amax_bf16x2)
245
+ amax = cute_utils.fmax(amax_pair[0], amax_pair[1])
246
+
247
+ if in_bounds:
248
+ # compute block scale with bit manipulation
249
+ # UE8M0 stores ceil(log2(fp4_scale)) + 127. Adding the mantissa mask
250
+ # increments the exponent whenever fp4_scale is not exactly a power of 2
251
+ eps = cutlass.const_expr(float.fromhex("0x6p-126"))
252
+ fp4_scale = cute_utils.fmax(amax, eps) * Float32(1.0 / 6.0)
253
+ bits = recast_val(fp4_scale, Uint32)
254
+ ue8m0 = cute_utils.shr_u32(
255
+ bits + Uint32(0x7FFFFF), Uint32(23)
256
+ ) & Uint32(0xFF)
257
+
258
+ # Only one of the two threads in an MXFP4 block writes the shared scale.
259
+ if tid % 2 == 0:
260
+ mx_block = sublane // 2
261
+ q_scale[token_id, head_start + i, mx_block] = Uint8(ue8m0)
262
+
263
+ # If scale = 2^A and ue8m0 = A + 127, then inverse scale has exponent
264
+ # -A + 127 = 254 - ue8m0.
265
+ inv_scale_bits = (Uint32(254) - ue8m0) << Uint32(23)
266
+ inv_fp4_scale = recast_val(inv_scale_bits, Float32)
267
+
268
+ vals = cute.make_rmem_tensor(16, Float32)
269
+ for j in cutlass.range_constexpr(8):
270
+ q0, q1 = cvt.bf16x2_to_fp32x2(q_bf16x2[i, j])
271
+ vals[j * 2] = q0 * inv_fp4_scale
272
+ vals[j * 2 + 1] = q1 * inv_fp4_scale
273
+
274
+ # pack to FP4
275
+ packed = cute.make_rmem_tensor((2,), Uint32)
276
+ packed[0] = cvt.fp32x8_to_fp4x8(vals, 0)
277
+ packed[1] = cvt.fp32x8_to_fp4x8(vals, 8)
278
+
279
+ dst = q_fp4_tile[i, None]
280
+ cp_u32x2 = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=64)
281
+ cute.copy(cp_u32x2, packed, cute.recast_tensor(dst, Uint32))
282
+
283
+ # Weight scaling is independent of the Q subwarp work. The first
284
+ # num_tokens * num_heads logical threads cover one weight each.
285
+ if global_tid < num_token_heads:
286
+ weight_token_id = global_tid // num_heads
287
+ weight_head_id = global_tid % num_heads
288
+ weights_out[weight_token_id, weight_head_id] = (
289
+ weights[weight_token_id, weight_head_id].to(Float32) * scale
290
+ )
291
286
- @cute.kernel
287
- def kernel(
292
+ @cute.jit
293
+ def host_entrypoint(
294
+ positions: cute.Tensor,
295
+ q: cute.Tensor,
296
+ cos_sin_cache: cute.Tensor,
297
+ weights: cute.Tensor,
298
+ q_quant: cute.Tensor,
299
+ q_scale: cute.Tensor,
300
+ weights_out: cute.Tensor,
301
+ scale: Float32,
302
+ stream: CUstream,
303
+ ):
304
+ total_threads = q.shape[0] * threads_per_token
305
+ grid = (cute.ceil_div(total_threads, tb_size), 1, 1)
306
+ device_kernel(
307
+ positions,
308
+ q,
309
+ cos_sin_cache,
310
+ weights,
311
+ q_quant,
312
+ q_scale,
313
+ weights_out,
314
+ scale,
315
+ ).launch(grid=grid, block=(tb_size, 1, 1), stream=stream)
316
+
317
+ return host_entrypoint
318
+
319
+ def dispatch( # type: ignore[override]
320
self,
289
- positions: cute.Tensor,
290
- q: cute.Tensor,
291
- cos_sin_cache: cute.Tensor,
292
- weights: cute.Tensor,
293
- q_quant: cute.Tensor,
294
- q_scale: cute.Tensor,
295
- weights_out: cute.Tensor,
296
- scale: Float32,
297
- ):
298
- (
299
- q_bf16x2,
300
- tid,
301
- global_tid,
302
- sublane,
303
- token_id,
304
- head_tile_id,
305
- head_start,
306
- in_bounds,
307
- num_token_heads,
308
- ) = self._load_q_and_rope(positions, q, cos_sin_cache)
309
-
310
- cp_op = cute.nvgpu.CopyUniversalOp()
311
-
312
- # layout: [coarsen, 8]
313
- q_fp4_tile = cute.local_tile(
314
- q_quant[token_id, None, None],
315
- tiler=(self.coarsen, 8),
316
- coord=(head_tile_id, sublane),
321
+ *,
322
+ cos_sin_dtype: type[cutlass.Numeric],
323
+ **compile_key_fields: int,
324
+ ) -> CompileKey:
325
+ return self.CompileKey(
326
+ **compile_key_fields,
327
+ cos_sin_dtype=cos_sin_dtype,
328
)
329
319
- for i in cutlass.range_constexpr(self.coarsen):
320
- # compute amax in packed bf16x2 to save instructions
321
- # Each thread holds 16 elems. Two adjacent threads form one 32-elem
322
- # MXFP4 block, so a width-2 shuffle gives the block amax.
323
- amax_bf16x2 = _bf16x2_abs(q_bf16x2[i, 0])
324
- for j in cutlass.range_constexpr(1, 8):
325
- amax_bf16x2 = _bf16x2_max(amax_bf16x2, _bf16x2_abs(q_bf16x2[i, j]))
326
- amax_bf16x2 = cute_utils.warp_reduce(
327
- amax_bf16x2,
328
- _bf16x2_max,
329
- width=MXFP4_BLOCK_SIZE // 16,
330
- )
331
- amax_pair = cvt.bf16x2_to_fp32x2(amax_bf16x2)
332
- amax = cute_utils.fmax(amax_pair[0], amax_pair[1])
333
-
334
- if in_bounds:
335
- # compute block scale with bit manipulation
336
- # UE8M0 stores ceil(log2(fp4_scale)) + 127. Adding the mantissa mask
337
- # increments the exponent whenever fp4_scale is not exactly a power of 2
338
- eps = cutlass.const_expr(float.fromhex("0x6p-126"))
339
- fp4_scale = cute_utils.fmax(amax, eps) * Float32(1.0 / 6.0)
340
- bits = recast_val(fp4_scale, Uint32)
341
- ue8m0 = cute_utils.shr_u32(
342
- bits + Uint32(0x7FFFFF), Uint32(23)
343
- ) & Uint32(0xFF)
344
-
345
- # Only one of the two threads in an MXFP4 block writes the shared scale.
346
- if tid % 2 == 0:
347
- mx_block = sublane // 2
348
- q_scale[token_id, head_start + i, mx_block] = Uint8(ue8m0)
349
-
350
- # If scale = 2^A and ue8m0 = A + 127, then inverse scale has exponent
351
- # -A + 127 = 254 - ue8m0.
352
- inv_scale_bits = (Uint32(254) - ue8m0) << Uint32(23)
353
- inv_fp4_scale = recast_val(inv_scale_bits, Float32)
354
-
355
- vals = cute.make_rmem_tensor(16, Float32)
356
- for j in cutlass.range_constexpr(8):
357
- q0, q1 = cvt.bf16x2_to_fp32x2(q_bf16x2[i, j])
358
- vals[j * 2] = q0 * inv_fp4_scale
359
- vals[j * 2 + 1] = q1 * inv_fp4_scale
360
-
361
- # pack to FP4
362
- packed = cute.make_rmem_tensor((2,), Uint32)
363
- packed[0] = cvt.fp32x8_to_fp4x8(vals, 0)
364
- packed[1] = cvt.fp32x8_to_fp4x8(vals, 8)
365
-
366
- dst = q_fp4_tile[i, None]
367
- cp_u32x2 = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=64)
368
- cute.copy(cp_u32x2, packed, cute.recast_tensor(dst, Uint32))
369
-
370
- # Weight scaling is independent of the Q subwarp work. The first
371
- # num_tokens * num_heads logical threads cover one weight each.
372
- if global_tid < num_token_heads:
373
- weight_token_id = global_tid // self.num_heads
374
- weight_head_id = global_tid % self.num_heads
375
- weights_out[weight_token_id, weight_head_id] = (
376
- weights[weight_token_id, weight_head_id].to(Float32) * scale
377
- )
330
+ def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]:
331
+ from vllm.v1.attention.backends.mla.indexer import dsa_indexer_uses_fp4
332
+
333
+ hf_config = vllm_config.model_config.hf_config
334
+ if hf_config is None or not dsa_indexer_uses_fp4(vllm_config):
335
+ return []
336
+
337
+ num_heads = int(getattr(hf_config, "index_n_heads", 0) or 0)
338
+ head_dim = int(getattr(hf_config, "index_head_dim", 0) or 0)
339
+ rope_dim = int(getattr(hf_config, "qk_rope_head_dim", 0) or 0)
340
+ if num_heads <= 0 or head_dim <= 0 or rope_dim <= 0:
341
+ return []
342
+
343
+ return self._trace_dispatch(self.dispatch)(
344
+ head_dim=head_dim,
345
+ rope_dim=rope_dim,
346
+ num_heads=num_heads,
347
+ cos_sin_dtype=Float32,
348
+ coarsen=(1, 4),
349
+ )
350
379
- @cache
380
- @staticmethod
381
- def compile(
382
- head_dim: int = 128,
383
- rope_dim: int = 64,
384
- num_heads: int = 64,
385
- cos_sin_dtype: type[cutlass.Numeric] = Float32,
386
- coarsen: int = 4,
387
- ):
351
+ def warmup_inputs(self, compile_key: CompileKey) -> tuple[Any, ...]:
352
num_tokens = cute.sym_int()
353
max_pos = cute.sym_int()
390
-
354
q = make_fake_tensor(
392
- BFloat16, (num_tokens, num_heads, head_dim), divisibility=16
355
+ BFloat16,
356
+ (num_tokens, compile_key.num_heads, compile_key.head_dim),
357
+ divisibility=16,
358
)
359
positions = make_fake_tensor(Int64, (num_tokens,), divisibility=1)
360
cos_sin_cache = make_fake_tensor(
396
- cos_sin_dtype,
397
- (max_pos, rope_dim),
361
+ compile_key.cos_sin_dtype,
362
+ (max_pos, compile_key.rope_dim),
363
divisibility=8,
364
)
400
- weights = make_fake_tensor(BFloat16, (num_tokens, num_heads), divisibility=8)
401
- q_fp4 = make_fake_tensor(
365
+ weights = make_fake_tensor(
366
+ BFloat16, (num_tokens, compile_key.num_heads), divisibility=8
367
+ )
368
+ q_packed = make_fake_tensor(
369
Uint8,
403
- (num_tokens, num_heads, head_dim // 2),
370
+ (num_tokens, compile_key.num_heads, compile_key.head_dim // 2),
371
divisibility=16,
372
)
373
q_scale = make_fake_tensor(
374
Uint8,
408
- (num_tokens, num_heads, head_dim // MXFP4_BLOCK_SIZE),
375
+ (
376
+ num_tokens,
377
+ compile_key.num_heads,
378
+ compile_key.head_dim // MXFP4_BLOCK_SIZE,
379
+ ),
380
divisibility=4,
381
)
411
- weights_out = make_fake_tensor(Float32, (num_tokens, num_heads), divisibility=4)
412
-
413
- kernel = IndexerQMxFp4Kernel(
414
- head_dim, rope_dim, num_heads, cos_sin_dtype, coarsen
382
+ weights_out = make_fake_tensor(
383
+ Float32, (num_tokens, compile_key.num_heads), divisibility=4
384
)
416
- stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
417
- return cute.compile(
418
- kernel,
385
+ return (
386
positions,
387
q,
388
cos_sin_cache,
389
weights,
423
- q_fp4,
390
+ q_packed,
391
q_scale,
392
weights_out,
393
Float32(0.0),
427
- stream,
428
- options="--enable-tvm-ffi",
394
)
395
431
-
432
-class IndexerQFp8Kernel(IndexerQRopeQuantKernel):
433
- """Eight-thread subwarps process one ``(token, head)`` row and emit
434
- float8 e4m3fn with a single per-(token, head) scalar scale folded
435
- into the per-token weight (mirrors ``_fused_indexer_q_rope_quant_kernel``).
436
- """
437
-
438
- def __init__(
439
- self,
440
- head_dim: int = 128,
441
- rope_dim: int = 64,
442
- num_heads: int = 64,
443
- cos_sin_dtype: type[cutlass.Numeric] = Float32,
444
- coarsen: int = 4,
445
- ):
446
- super().__init__(head_dim, rope_dim, num_heads, cos_sin_dtype, coarsen)
447
- # Each subwarp owns `coarsen` heads; we use the first `coarsen`
448
- # threads of the subwarp to write the per-head weights using the
449
- # fp8 scale computed in the matching loop iteration.
450
- assert self.coarsen <= self.subwarp_size, (
451
- f"FP8 kernel requires coarsen ({self.coarsen}) <= "
452
- f"subwarp_size ({self.subwarp_size}) for the weight-fold step"
453
- )
454
-
455
- @cute.jit
396
+ @kernel_launcher
397
def __call__(
398
self,
458
- positions: cute.Tensor,
459
- q: cute.Tensor,
460
- cos_sin_cache: cute.Tensor,
461
- weights: cute.Tensor,
462
- q_fp8: cute.Tensor,
463
- weights_out: cute.Tensor,
464
- scale: Float32,
465
- stream: CUstream,
466
- ):
467
- total_threads = q.shape[0] * self.threads_per_token
468
- grid = (cute.ceil_div(total_threads, self.tb_size), 1, 1)
469
- self.kernel(
399
+ *,
400
+ positions: torch.Tensor,
401
+ q: torch.Tensor,
402
+ cos_sin_cache: torch.Tensor,
403
+ weights: torch.Tensor,
404
+ weights_softmax_scale: float,
405
+ weights_head_scale: float,
406
+ q_packed: torch.Tensor,
407
+ q_scale: torch.Tensor,
408
+ weights_out: torch.Tensor,
409
+ ) -> CuTeDSLLaunchSpec["IndexerQMxFp4Kernel.CompileKey"]:
410
+ num_tokens, num_heads, head_dim = q.shape
411
+ compile_key = self.dispatch(
412
+ head_dim=head_dim,
413
+ rope_dim=cos_sin_cache.shape[-1],
414
+ num_heads=num_heads,
415
+ cos_sin_dtype=torch_to_cute_dtype(cos_sin_cache.dtype),
416
+ coarsen=1 if num_tokens < 512 else 4,
417
+ )
418
+ launch_args = (
419
positions,
420
q,
421
cos_sin_cache,
422
weights,
474
- q_fp8,
423
+ q_packed,
424
+ q_scale,
425
weights_out,
476
- scale,
477
- ).launch(grid=grid, block=(self.tb_size, 1, 1), stream=stream)
426
+ float(weights_softmax_scale * weights_head_scale),
427
+ )
428
+ return compile_key, launch_args
429
479
- @cute.kernel
480
- def kernel(
481
- self,
482
- positions: cute.Tensor,
483
- q: cute.Tensor,
484
- cos_sin_cache: cute.Tensor,
485
- weights: cute.Tensor,
486
- q_fp8: cute.Tensor,
487
- weights_out: cute.Tensor,
488
- scale: Float32,
489
- ):
490
- (
491
- q_bf16x2,
492
- _tid,
493
- _global_tid,
494
- sublane,
495
- token_id,
496
- head_tile_id,
497
- head_start,
498
- in_bounds,
499
- _num_token_heads,
500
- ) = self._load_q_and_rope(positions, q, cos_sin_cache)
501
-
502
- cp_op = cute.nvgpu.CopyUniversalOp()
503
-
504
- # layout: [coarsen, 16] bytes (one e4m3fn per element).
505
- q_fp8_tile = cute.local_tile(
506
- q_fp8[token_id, None, None],
507
- tiler=(self.coarsen, 16),
508
- coord=(head_tile_id, sublane),
430
+
431
+class IndexerQFp8Kernel(VllmCuTeDSLJitKernel["IndexerQFp8Kernel.CompileKey"]):
432
+ tb_size = 128
433
+
434
+ @dataclass(frozen=True)
435
+ class CompileKey:
436
+ head_dim: int
437
+ rope_dim: int
438
+ num_heads: int
439
+ cos_sin_dtype: type[cutlass.Numeric]
440
+ coarsen: int
441
+
442
+ @staticmethod
443
+ def kernel(compile_key: CompileKey) -> Any:
444
+ head_dim = compile_key.head_dim
445
+ rope_dim = compile_key.rope_dim
446
+ nope_dim = head_dim - rope_dim
447
+ num_heads = compile_key.num_heads
448
+ cos_sin_dtype = compile_key.cos_sin_dtype
449
+ # process multiple heads at the same time to armotize RoPE load costs
450
+ coarsen = compile_key.coarsen
451
+ tb_size = IndexerQFp8Kernel.tb_size
452
+ assert num_heads % coarsen == 0
453
+
454
+ # later we will use 32B load = 16 BF16 elems
455
+ # thus, head_dim=128 requires 8 threads to handle.
456
+ # let's call subwarp = 8 threads.
457
+ subwarp_size = head_dim // 16
458
+ threads_per_token = (num_heads // coarsen) * subwarp_size
459
+ # Each subwarp owns `coarsen` heads; we use the first `coarsen`
460
+ # threads of the subwarp to write the per-head weights using the
461
+ # fp8 scale computed in the matching loop iteration.
462
+ assert coarsen <= subwarp_size, (
463
+ f"FP8 kernel requires coarsen ({coarsen}) <= "
464
+ f"subwarp_size ({subwarp_size}) for the weight-fold step"
465
)
466
511
- for i in cutlass.range_constexpr(self.coarsen):
512
- # Reduce amax across the full head_dim: each thread already holds
513
- # the max over its 16 lanes; a width=subwarp_size warp shuffle
514
- # spreads the head-wide max to every lane in the subwarp.
515
- amax_bf16x2 = _bf16x2_abs(q_bf16x2[i, 0])
516
- for j in cutlass.range_constexpr(1, 8):
517
- amax_bf16x2 = _bf16x2_max(amax_bf16x2, _bf16x2_abs(q_bf16x2[i, j]))
518
- amax_bf16x2 = cute_utils.warp_reduce(
519
- amax_bf16x2,
520
- _bf16x2_max,
521
- width=self.subwarp_size,
467
+ @cute.kernel
468
+ def device_kernel(
469
+ positions: cute.Tensor,
470
+ q: cute.Tensor,
471
+ cos_sin_cache: cute.Tensor,
472
+ weights: cute.Tensor,
473
+ q_fp8: cute.Tensor,
474
+ weights_out: cute.Tensor,
475
+ scale: Float32,
476
+ ):
477
+ (
478
+ q_bf16x2,
479
+ _tid,
480
+ _global_tid,
481
+ sublane,
482
+ token_id,
483
+ head_tile_id,
484
+ head_start,
485
+ in_bounds,
486
+ _num_token_heads,
487
+ ) = _load_q_and_rope(
488
+ positions,
489
+ q,
490
+ cos_sin_cache,
491
+ num_heads,
492
+ coarsen,
493
+ tb_size,
494
+ subwarp_size,
495
+ threads_per_token,
496
+ nope_dim,
497
+ rope_dim,
498
+ cos_sin_dtype,
499
+ )
500
+
501
+ cp_op = cute.nvgpu.CopyUniversalOp()
502
+
503
+ # layout: [coarsen, 16] bytes (one e4m3fn per element).
504
+ q_fp8_tile = cute.local_tile(
505
+ q_fp8[token_id, None, None],
506
+ tiler=(coarsen, 16),
507
+ coord=(head_tile_id, sublane),
508
)
523
- amax_pair = cvt.bf16x2_to_fp32x2(amax_bf16x2)
524
- amax = cute_utils.fmax(amax_pair[0], amax_pair[1])
525
-
526
- # scale = max(amax, eps) / fp8_max, then rounded UP to the next
527
- # power of two. Adding the mantissa mask before shifting out the
528
- # mantissa bumps the exponent whenever s isn't a pure pow2.
529
- fp32_scale = cute_utils.fmax(amax, Float32(1e-4)) * Float32(1.0 / 448.0)
530
- bits = recast_val(fp32_scale, Uint32)
531
- scale_exp = cute_utils.shr_u32(
532
- bits + Uint32(0x7FFFFF), Uint32(23)
533
- ) & Uint32(0xFF)
534
-
535
- # rounded scale = 2^(scale_exp - 127); bit pattern is scale_exp << 23
536
- fp8_scale_bits = scale_exp << Uint32(23)
537
- fp8_scale = recast_val(fp8_scale_bits, Float32)
538
- # inverse = 2^-(scale_exp - 127); bit pattern is (254 - scale_exp) << 23
539
- inv_scale_bits = (Uint32(254) - scale_exp) << Uint32(23)
540
- inv_fp8_scale = recast_val(inv_scale_bits, Float32)
541
-
542
- # Weight fold: weights_out = weights * q_scale * scale_combined.
543
- # All threads in the subwarp share the same fp8_scale after the
544
- # warp_reduce above, so we let thread `sublane == i` write the
545
- # weight for head `head_start + i`.
546
- if in_bounds and sublane == i:
547
- head_id = head_start + i
548
- weights_out[token_id, head_id] = (
549
- weights[token_id, head_id].to(Float32) * scale * fp8_scale
509
+
510
+ for i in cutlass.range_constexpr(coarsen):
511
+ # Reduce amax across the full head_dim: each thread already holds
512
+ # the max over its 16 lanes; a width=subwarp_size warp shuffle
513
+ # spreads the head-wide max to every lane in the subwarp.
514
+ amax_bf16x2 = _bf16x2_abs(q_bf16x2[i, 0])
515
+ for j in cutlass.range_constexpr(1, 8):
516
+ amax_bf16x2 = _bf16x2_max(amax_bf16x2, _bf16x2_abs(q_bf16x2[i, j]))
517
+ amax_bf16x2 = cute_utils.warp_reduce(
518
+ amax_bf16x2,
519
+ _bf16x2_max,
520
+ width=subwarp_size,
521
)
522
+ amax_pair = cvt.bf16x2_to_fp32x2(amax_bf16x2)
523
+ amax = cute_utils.fmax(amax_pair[0], amax_pair[1])
524
+
525
+ # scale = max(amax, eps) / fp8_max, then rounded UP to the next
526
+ # power of two. Adding the mantissa mask before shifting out the
527
+ # mantissa bumps the exponent whenever s isn't a pure pow2.
528
+ fp32_scale = cute_utils.fmax(amax, Float32(1e-4)) * Float32(1.0 / 448.0)
529
+ bits = recast_val(fp32_scale, Uint32)
530
+ scale_exp = cute_utils.shr_u32(
531
+ bits + Uint32(0x7FFFFF), Uint32(23)
532
+ ) & Uint32(0xFF)
533
552
- if in_bounds:
553
- # 16 BF16 → 16 e4m3 bytes per thread, packed into 4 b32s
554
- # (one cp.async-shaped 128-bit store per row).
555
- packed = cute.make_rmem_tensor((4,), Uint32)
556
- for j in cutlass.range_constexpr(4):
557
- q0, q1 = cvt.bf16x2_to_fp32x2(q_bf16x2[i, j * 2])
558
- q2, q3 = cvt.bf16x2_to_fp32x2(q_bf16x2[i, j * 2 + 1])
559
- packed[j] = cvt.fp32x4_to_fp8x4(
560
- q0 * inv_fp8_scale,
561
- q1 * inv_fp8_scale,
562
- q2 * inv_fp8_scale,
563
- q3 * inv_fp8_scale,
534
+ # rounded scale = 2^(scale_exp - 127); bit pattern is scale_exp << 23
535
+ fp8_scale_bits = scale_exp << Uint32(23)
536
+ fp8_scale = recast_val(fp8_scale_bits, Float32)
537
+ # inverse = 2^-(scale_exp - 127); bit pattern is (254 - scale_exp) << 23
538
+ inv_scale_bits = (Uint32(254) - scale_exp) << Uint32(23)
539
+ inv_fp8_scale = recast_val(inv_scale_bits, Float32)
540
+
541
+ # Weight fold: weights_out = weights * q_scale * scale_combined.
542
+ # All threads in the subwarp share the same fp8_scale after the
543
+ # warp_reduce above, so we let thread `sublane == i` write the
544
+ # weight for head `head_start + i`.
545
+ if in_bounds and sublane == i:
546
+ head_id = head_start + i
547
+ weights_out[token_id, head_id] = (
548
+ weights[token_id, head_id].to(Float32) * scale * fp8_scale
549
)
550
566
- dst = q_fp8_tile[i, None]
567
- cp_u32x4 = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=128)
568
- cute.copy(cp_u32x4, packed, cute.recast_tensor(dst, Uint32))
551
+ if in_bounds:
552
+ # 16 BF16 → 16 e4m3 bytes per thread, packed into 4 b32s
553
+ # (one cp.async-shaped 128-bit store per row).
554
+ packed = cute.make_rmem_tensor((4,), Uint32)
555
+ for j in cutlass.range_constexpr(4):
556
+ q0, q1 = cvt.bf16x2_to_fp32x2(q_bf16x2[i, j * 2])
557
+ q2, q3 = cvt.bf16x2_to_fp32x2(q_bf16x2[i, j * 2 + 1])
558
+ packed[j] = cvt.fp32x4_to_fp8x4(
559
+ q0 * inv_fp8_scale,
560
+ q1 * inv_fp8_scale,
561
+ q2 * inv_fp8_scale,
562
+ q3 * inv_fp8_scale,
563
+ )
564
+
565
+ dst = q_fp8_tile[i, None]
566
+ cp_u32x4 = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=128)
567
+ cute.copy(cp_u32x4, packed, cute.recast_tensor(dst, Uint32))
568
+
569
+ @cute.jit
570
+ def host_entrypoint(
571
+ positions: cute.Tensor,
572
+ q: cute.Tensor,
573
+ cos_sin_cache: cute.Tensor,
574
+ weights: cute.Tensor,
575
+ q_fp8: cute.Tensor,
576
+ weights_out: cute.Tensor,
577
+ scale: Float32,
578
+ stream: CUstream,
579
+ ):
580
+ total_threads = q.shape[0] * threads_per_token
581
+ grid = (cute.ceil_div(total_threads, tb_size), 1, 1)
582
+ device_kernel(
583
+ positions,
584
+ q,
585
+ cos_sin_cache,
586
+ weights,
587
+ q_fp8,
588
+ weights_out,
589
+ scale,
590
+ ).launch(grid=grid, block=(tb_size, 1, 1), stream=stream)
591
+
592
+ return host_entrypoint
593
+
594
+ def dispatch( # type: ignore[override]
595
+ self,
596
+ *,
597
+ cos_sin_dtype: type[cutlass.Numeric],
598
+ **compile_key_fields: int,
599
+ ) -> CompileKey:
600
+ return self.CompileKey(
601
+ **compile_key_fields,
602
+ cos_sin_dtype=cos_sin_dtype,
603
+ )
604
570
- @cache
571
- @staticmethod
572
- def compile(
573
- head_dim: int = 128,
574
- rope_dim: int = 64,
575
- num_heads: int = 64,
576
- cos_sin_dtype: type[cutlass.Numeric] = Float32,
577
- coarsen: int = 4,
578
- ):
605
+ def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]:
606
+ from vllm.v1.attention.backends.mla.indexer import dsa_indexer_uses_fp4
607
+
608
+ hf_config = vllm_config.model_config.hf_config
609
+ if hf_config is None or dsa_indexer_uses_fp4(vllm_config):
610
+ return []
611
+
612
+ num_heads = int(getattr(hf_config, "index_n_heads", 0) or 0)
613
+ head_dim = int(getattr(hf_config, "index_head_dim", 0) or 0)
614
+ rope_dim = int(getattr(hf_config, "qk_rope_head_dim", 0) or 0)
615
+ if num_heads <= 0 or head_dim <= 0 or rope_dim <= 0:
616
+ return []
617
+
618
+ return self._trace_dispatch(self.dispatch)(
619
+ head_dim=head_dim,
620
+ rope_dim=rope_dim,
621
+ num_heads=num_heads,
622
+ cos_sin_dtype=Float32,
623
+ coarsen=(1, 4),
624
+ )
625
+
626
+ def warmup_inputs(self, compile_key: CompileKey) -> tuple[Any, ...]:
627
num_tokens = cute.sym_int()
628
max_pos = cute.sym_int()
581
-
629
q = make_fake_tensor(
583
- BFloat16, (num_tokens, num_heads, head_dim), divisibility=16
630
+ BFloat16,
631
+ (num_tokens, compile_key.num_heads, compile_key.head_dim),
632
+ divisibility=16,
633
)
634
positions = make_fake_tensor(Int64, (num_tokens,), divisibility=1)
635
cos_sin_cache = make_fake_tensor(
587
- cos_sin_dtype,
588
- (max_pos, rope_dim),
636
+ compile_key.cos_sin_dtype,
637
+ (max_pos, compile_key.rope_dim),
638
divisibility=8,
639
)
591
- weights = make_fake_tensor(BFloat16, (num_tokens, num_heads), divisibility=8)
640
+ weights = make_fake_tensor(
641
+ BFloat16, (num_tokens, compile_key.num_heads), divisibility=8
642
+ )
643
q_fp8 = make_fake_tensor(
644
Uint8,
594
- (num_tokens, num_heads, head_dim),
645
+ (num_tokens, compile_key.num_heads, compile_key.head_dim),
646
divisibility=16,
647
)
597
- weights_out = make_fake_tensor(Float32, (num_tokens, num_heads), divisibility=4)
598
-
599
- kernel = IndexerQFp8Kernel(
600
- head_dim, rope_dim, num_heads, cos_sin_dtype, coarsen
648
+ weights_out = make_fake_tensor(
649
+ Float32, (num_tokens, compile_key.num_heads), divisibility=4
650
)
602
- stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
603
- return cute.compile(
604
- kernel,
651
+ return (
652
positions,
653
q,
654
cos_sin_cache,
@@ -609,6 +656,40 @@ class IndexerQFp8Kernel(IndexerQRopeQuantKernel):
656
q_fp8,
657
weights_out,
658
Float32(0.0),
612
- stream,
613
- options="--enable-tvm-ffi",
659
)
660
+
661
+ @kernel_launcher
662
+ def __call__(
663
+ self,
664
+ *,
665
+ positions: torch.Tensor,
666
+ q: torch.Tensor,
667
+ cos_sin_cache: torch.Tensor,
668
+ weights: torch.Tensor,
669
+ weights_softmax_scale: float,
670
+ weights_head_scale: float,
671
+ q_fp8: torch.Tensor,
672
+ weights_out: torch.Tensor,
673
+ ) -> CuTeDSLLaunchSpec["IndexerQFp8Kernel.CompileKey"]:
674
+ num_tokens, num_heads, head_dim = q.shape
675
+ compile_key = self.dispatch(
676
+ head_dim=head_dim,
677
+ rope_dim=cos_sin_cache.shape[-1],
678
+ num_heads=num_heads,
679
+ cos_sin_dtype=torch_to_cute_dtype(cos_sin_cache.dtype),
680
+ coarsen=1 if num_tokens < 512 else 4,
681
+ )
682
+ launch_args = (
683
+ positions,
684
+ q,
685
+ cos_sin_cache,
686
+ weights,
687
+ q_fp8,
688
+ weights_out,
689
+ float(weights_softmax_scale * weights_head_scale),
690
+ )
691
+ return compile_key, launch_args
692
+
693
+
694
+_INDEXER_Q_MXFP4_KERNEL = IndexerQMxFp4Kernel()
695
+_INDEXER_Q_FP8_KERNEL = IndexerQFp8Kernel()
vllm/models/deepseek_v4/nvidia/ops/sparse_attn_compress_cutedsl.py
+1945
-1606
@@ -7,22 +7,25 @@ The public wrappers provide the C4 fused and C128 split kernels.
7
8
from __future__ import annotations
9
10
-from functools import cache
10
+from dataclasses import dataclass
11
from typing import Any
12
13
import cutlass
14
import cutlass.cute as cute
15
import torch
16
from cuda.bindings.driver import CUstream
17
-from cutlass import BFloat16, Float32, Int32, Int64, Uint8, Uint16, Uint32, const_expr
17
+from cutlass import Float32, Int32, Int64, Uint8, Uint16, Uint32, const_expr
18
from cutlass._mlir.dialects import llvm
19
from cutlass.cutlass_dsl import T, dsl_user_op
20
from quack.compile_utils import make_fake_tensor
21
22
-_TORCH_TO_CUTE = {
23
- torch.bfloat16: BFloat16,
24
- torch.float32: Float32,
25
-}
22
+from vllm.cute_utils import torch_to_cute_dtype
23
+from vllm.model_executor.warmup.jit_warmup import kernel_launcher
24
+from vllm.model_executor.warmup.jit_warmup_cutedsl_helper import (
25
+ CuTeDSLLaunchSpec,
26
+ VllmCuTeDSLJitKernel,
27
+)
28
+from vllm.utils.math_utils import round_up
29
30
31
@dsl_user_op
@@ -72,361 +75,406 @@ def _fp32x2_to_fp8e4m3x2(a: Float32, b: Float32, *, loc=None, ip=None) -> Uint16
75
return Uint16(out)
76
77
75
-class SparseAttnCompressNormRopeStoreC4Kernel:
78
+class SparseAttnCompressNormRopeStoreC4Kernel(
79
+ VllmCuTeDSLJitKernel["SparseAttnCompressNormRopeStoreC4Kernel.CompileKey"]
80
+):
81
min_scale = 1.0e-4
82
rcp_ln2 = 1.4426950408889634
83
79
- def __init__(
80
- self,
81
- head_size: int,
82
- state_width: int,
83
- rope_head_dim: int,
84
- fp8_max: float,
85
- quant_block: int,
86
- token_stride: int,
87
- scale_dim: int,
88
- compress_ratio: int,
89
- overlap: bool,
90
- ):
91
- self.head_dim = head_size
92
- self.state_width = state_width
93
- self.rope_dim = rope_head_dim
94
- self.nope_dim = head_size - rope_head_dim
95
- self.fp8_max = fp8_max
96
- self.quant_block = quant_block
97
- self.token_stride = token_stride
98
- self.scale_dim = scale_dim
99
- self.elems_per_lane = 8
100
- self.copy_elems = 4
101
- self.copy_chunks = self.elems_per_lane // self.copy_elems
102
- self.lanes_per_group = quant_block // self.elems_per_lane
103
- self.groups_per_warp = 32 // self.lanes_per_group
104
- self.scale_reduce_steps = self.lanes_per_group.bit_length() - 1
105
- self.scale_reduce_offset = self.lanes_per_group // 2
106
- self.num_warps = (head_size // quant_block) // self.groups_per_warp
107
- self.nope_blocks = self.nope_dim // quant_block
108
- self.tb_size = self.num_warps * 32
109
- self.compress_ratio = compress_ratio
110
- self.overlap = overlap
111
- self.window = (1 + int(overlap)) * compress_ratio
112
-
113
- @cute.jit
114
- def __call__(
115
- self,
116
- state_cache: cute.Tensor,
117
- token_to_req_indices: cute.Tensor,
118
- positions: cute.Tensor,
119
- slot_mapping: cute.Tensor,
120
- block_table: cute.Tensor,
121
- block_size: Int64,
122
- rms_norm_weight: cute.Tensor,
123
- rms_norm_eps: Float32,
124
- cos_sin_cache: cute.Tensor,
125
- k_cache: cute.Tensor,
126
- kv_slot_mapping: cute.Tensor,
127
- kv_cache_block_size: Int64,
128
- stream: CUstream,
129
- ):
130
- grid = (slot_mapping.shape[0], 1, 1)
131
- self.kernel(
132
- state_cache,
133
- token_to_req_indices,
134
- positions,
135
- slot_mapping,
136
- block_table,
137
- block_size,
138
- rms_norm_weight,
139
- rms_norm_eps,
140
- cos_sin_cache,
141
- k_cache,
142
- kv_slot_mapping,
143
- kv_cache_block_size,
144
- ).launch(grid=grid, block=(self.tb_size, 1, 1), stream=stream)
145
-
146
- @cute.kernel
147
- def kernel(
148
- self,
149
- state_cache: cute.Tensor,
150
- token_to_req_indices: cute.Tensor,
151
- positions: cute.Tensor,
152
- slot_mapping: cute.Tensor,
153
- block_table: cute.Tensor,
154
- block_size: Int64,
155
- rms_norm_weight: cute.Tensor,
156
- rms_norm_eps: Float32,
157
- cos_sin_cache: cute.Tensor,
158
- k_cache: cute.Tensor,
159
- kv_slot_mapping: cute.Tensor,
160
- kv_cache_block_size: Int64,
161
- ):
162
- token_idx, _, _ = cute.arch.block_idx()
163
- tid, _, _ = cute.arch.thread_idx()
164
- warp_id = cute.arch.make_warp_uniform(tid // 32)
165
- lane_id = tid % 32
166
- group_lane = lane_id % self.lanes_per_group
167
- group_idx = warp_id * self.groups_per_warp + lane_id // self.lanes_per_group
168
- elem_base = group_idx * self.quant_block + group_lane * self.elems_per_lane
169
-
170
- slot_id = slot_mapping[token_idx]
171
- has_position = token_idx < positions.shape[0]
172
- position = Int64(0)
173
- if has_position:
174
- position = positions[token_idx]
175
- boundary = has_position and (
176
- (position + Int64(1)) % Int64(self.compress_ratio) == Int64(0)
177
- )
178
- has_req_idx = token_idx < token_to_req_indices.shape[0]
179
- has_kv_slot_idx = token_idx < kv_slot_mapping.shape[0]
180
- kv_slot_idx = Int64(-1)
181
- if has_kv_slot_idx:
182
- kv_slot_idx = kv_slot_mapping[token_idx]
183
- active = (
184
- slot_id >= Int64(0) and has_req_idx and boundary and kv_slot_idx >= Int64(0)
185
- )
84
+ @dataclass(frozen=True)
85
+ class CompileKey:
86
+ head_size: int
87
+ state_width: int
88
+ rope_head_dim: int
89
+ fp8_max: float
90
+ quant_block: int
91
+ token_stride: int
92
+ scale_dim: int
93
+ compress_ratio: int
94
+ overlap: bool
95
+ norm_weight_dtype: type[cutlass.Numeric]
96
187
- if active:
188
- req_idx = token_to_req_indices[token_idx]
189
- start = position - Int64(self.window - 1)
97
+ @staticmethod
98
+ def kernel(compile_key: CompileKey) -> Any:
99
+ min_scale = SparseAttnCompressNormRopeStoreC4Kernel.min_scale
100
+ rcp_ln2 = SparseAttnCompressNormRopeStoreC4Kernel.rcp_ln2
101
+ head_dim = compile_key.head_size
102
+ state_width = compile_key.state_width
103
+ rope_dim = compile_key.rope_head_dim
104
+ nope_dim = head_dim - rope_dim
105
+ fp8_max = compile_key.fp8_max
106
+ quant_block = compile_key.quant_block
107
+ token_stride = compile_key.token_stride
108
+ scale_dim = compile_key.scale_dim
109
+ elems_per_lane = 8
110
+ copy_elems = 4
111
+ copy_chunks = elems_per_lane // copy_elems
112
+ lanes_per_group = quant_block // elems_per_lane
113
+ groups_per_warp = 32 // lanes_per_group
114
+ scale_reduce_steps = lanes_per_group.bit_length() - 1
115
+ scale_reduce_offset = lanes_per_group // 2
116
+ num_warps = (head_dim // quant_block) // groups_per_warp
117
+ nope_blocks = nope_dim // quant_block
118
+ tb_size = num_warps * 32
119
+ compress_ratio = compile_key.compress_ratio
120
+ overlap = compile_key.overlap
121
+ window = (1 + int(overlap)) * compress_ratio
122
+
123
+ @cute.kernel
124
+ def device_kernel(
125
+ state_cache: cute.Tensor,
126
+ token_to_req_indices: cute.Tensor,
127
+ positions: cute.Tensor,
128
+ slot_mapping: cute.Tensor,
129
+ block_table: cute.Tensor,
130
+ block_size: Int64,
131
+ rms_norm_weight: cute.Tensor,
132
+ rms_norm_eps: Float32,
133
+ cos_sin_cache: cute.Tensor,
134
+ k_cache: cute.Tensor,
135
+ kv_slot_mapping: cute.Tensor,
136
+ kv_cache_block_size: Int64,
137
+ ):
138
+ token_idx, _, _ = cute.arch.block_idx()
139
+ tid, _, _ = cute.arch.thread_idx()
140
+ warp_id = cute.arch.make_warp_uniform(tid // 32)
141
+ lane_id = tid % 32
142
+ group_lane = lane_id % lanes_per_group
143
+ group_idx = warp_id * groups_per_warp + lane_id // lanes_per_group
144
+ elem_base = group_idx * quant_block + group_lane * elems_per_lane
145
191
- smem = cutlass.utils.SmemAllocator()
192
- s_block_numbers = smem.allocate_tensor(
193
- Int32, cute.make_layout((self.window,)), byte_alignment=4
194
- )
195
- partial_sums = smem.allocate_tensor(
196
- Float32, cute.make_layout((self.num_warps,)), byte_alignment=4
146
+ slot_id = slot_mapping[token_idx]
147
+ has_position = token_idx < positions.shape[0]
148
+ position = Int64(0)
149
+ if has_position:
150
+ position = positions[token_idx]
151
+ boundary = has_position and (
152
+ (position + Int64(1)) % Int64(compress_ratio) == Int64(0)
153
)
198
- rrms_shared = smem.allocate_tensor(
199
- Float32, cute.make_layout((1,)), byte_alignment=4
154
+ has_req_idx = token_idx < token_to_req_indices.shape[0]
155
+ has_kv_slot_idx = token_idx < kv_slot_mapping.shape[0]
156
+ kv_slot_idx = Int64(-1)
157
+ if has_kv_slot_idx:
158
+ kv_slot_idx = kv_slot_mapping[token_idx]
159
+ active = (
160
+ slot_id >= Int64(0)
161
+ and has_req_idx
162
+ and boundary
163
+ and kv_slot_idx >= Int64(0)
164
)
165
202
- for row in cutlass.range_constexpr(self.window):
203
- pos = start + Int64(row)
204
- if tid == row:
205
- block_number_i32 = Int32(0)
206
- if pos >= Int64(0):
207
- block_index = pos // block_size
208
- block_number_i32 = block_table[req_idx, block_index]
209
- s_block_numbers[row] = block_number_i32
210
- cute.arch.sync_threads()
166
+ if active:
167
+ req_idx = token_to_req_indices[token_idx]
168
+ start = position - Int64(window - 1)
169
212
- local_max = cute.make_rmem_tensor((self.elems_per_lane,), Float32)
213
- local_sum = cute.make_rmem_tensor((self.elems_per_lane,), Float32)
214
- local_product = cute.make_rmem_tensor((self.elems_per_lane,), Float32)
170
+ smem = cutlass.utils.SmemAllocator()
171
+ s_block_numbers = smem.allocate_tensor(
172
+ Int32, cute.make_layout((window,)), byte_alignment=4
173
+ )
174
+ partial_sums = smem.allocate_tensor(
175
+ Float32, cute.make_layout((num_warps,)), byte_alignment=4
176
+ )
177
+ rrms_shared = smem.allocate_tensor(
178
+ Float32, cute.make_layout((1,)), byte_alignment=4
179
+ )
180
216
- for e in cutlass.range_constexpr(self.elems_per_lane):
217
- local_max[e] = -Float32.inf
218
- local_sum[e] = Float32(0.0)
219
- local_product[e] = Float32(0.0)
181
+ for row in cutlass.range_constexpr(window):
182
+ pos = start + Int64(row)
183
+ if tid == row:
184
+ block_number_i32 = Int32(0)
185
+ if pos >= Int64(0):
186
+ block_index = pos // block_size
187
+ block_number_i32 = block_table[req_idx, block_index]
188
+ s_block_numbers[row] = block_number_i32
189
+ cute.arch.sync_threads()
190
+
191
+ local_max = cute.make_rmem_tensor((elems_per_lane,), Float32)
192
+ local_sum = cute.make_rmem_tensor((elems_per_lane,), Float32)
193
+ local_product = cute.make_rmem_tensor((elems_per_lane,), Float32)
194
+
195
+ for e in cutlass.range_constexpr(elems_per_lane):
196
+ local_max[e] = -Float32.inf
197
+ local_sum[e] = Float32(0.0)
198
+ local_product[e] = Float32(0.0)
199
+
200
+ cp_f32x4 = cute.make_copy_atom(
201
+ cute.nvgpu.CopyUniversalOp(), Float32, num_bits_per_copy=128
202
+ )
203
+ copy_layout = cute.make_layout(
204
+ (copy_chunks, copy_elems),
205
+ stride=(copy_elems, 1),
206
+ )
207
+ kv_vals = cute.make_rmem_tensor(copy_layout, Float32)
208
+ score_vals = cute.make_rmem_tensor(copy_layout, Float32)
209
221
- cp_f32x4 = cute.make_copy_atom(
222
- cute.nvgpu.CopyUniversalOp(), Float32, num_bits_per_copy=128
223
- )
224
- copy_layout = cute.make_layout(
225
- (self.copy_chunks, self.copy_elems),
226
- stride=(self.copy_elems, 1),
227
- )
228
- kv_vals = cute.make_rmem_tensor(copy_layout, Float32)
229
- score_vals = cute.make_rmem_tensor(copy_layout, Float32)
230
-
231
- for row in cutlass.range_constexpr(self.window):
232
- pos = start + Int64(row)
233
- if pos >= Int64(0):
234
- block_index = pos // block_size
235
- block_offset = pos - block_index * block_size
236
- block_number = s_block_numbers[row].to(Int64)
237
- head_offset = Int64((row // self.compress_ratio) * self.head_dim)
238
- row_tensor = state_cache[block_number, block_offset, None]
239
- for chunk in cutlass.range_constexpr(self.copy_chunks):
240
- copy_elem = const_expr(chunk * self.copy_elems)
241
- col_tile = (
242
- head_offset + (elem_base + Int32(copy_elem)).to(Int64)
243
- ) // Int64(self.copy_elems)
244
- kv_src = cute.local_tile(
245
- row_tensor,
246
- tiler=(self.copy_elems,),
247
- coord=(col_tile,),
248
- )
249
- score_src = cute.local_tile(
250
- row_tensor,
251
- tiler=(self.copy_elems,),
252
- coord=(
253
- col_tile + Int64(self.state_width // self.copy_elems),
254
- ),
255
- )
256
- cute.copy(cp_f32x4, kv_src, kv_vals[chunk, None])
257
- cute.copy(cp_f32x4, score_src, score_vals[chunk, None])
258
-
259
- for e in cutlass.range_constexpr(self.elems_per_lane):
260
- chunk = const_expr(e // self.copy_elems)
261
- copy_elem = const_expr(e % self.copy_elems)
262
- score = score_vals[chunk, copy_elem]
263
- kv = kv_vals[chunk, copy_elem]
264
- new_max = cute.arch.fmax(local_max[e], score)
265
- old_scale = cute.math.exp2(
266
- (local_max[e] - new_max) * Float32(self.rcp_ln2),
267
- fastmath=True,
268
- )
269
- new_scale = cute.math.exp2(
270
- (score - new_max) * Float32(self.rcp_ln2),
271
- fastmath=True,
272
- )
273
- local_sum[e] = local_sum[e] * old_scale + new_scale
274
- local_product[e] = local_product[e] * old_scale + kv * new_scale
275
- local_max[e] = new_max
210
+ for row in cutlass.range_constexpr(window):
211
+ pos = start + Int64(row)
212
+ if pos >= Int64(0):
213
+ block_index = pos // block_size
214
+ block_offset = pos - block_index * block_size
215
+ block_number = s_block_numbers[row].to(Int64)
216
+ head_offset = Int64((row // compress_ratio) * head_dim)
217
+ row_tensor = state_cache[block_number, block_offset, None]
218
+ for chunk in cutlass.range_constexpr(copy_chunks):
219
+ copy_elem = const_expr(chunk * copy_elems)
220
+ col_tile = (
221
+ head_offset + (elem_base + Int32(copy_elem)).to(Int64)
222
+ ) // Int64(copy_elems)
223
+ kv_src = cute.local_tile(
224
+ row_tensor,
225
+ tiler=(copy_elems,),
226
+ coord=(col_tile,),
227
+ )
228
+ score_src = cute.local_tile(
229
+ row_tensor,
230
+ tiler=(copy_elems,),
231
+ coord=(col_tile + Int64(state_width // copy_elems),),
232
+ )
233
+ cute.copy(cp_f32x4, kv_src, kv_vals[chunk, None])
234
+ cute.copy(cp_f32x4, score_src, score_vals[chunk, None])
235
+
236
+ for e in cutlass.range_constexpr(elems_per_lane):
237
+ chunk = const_expr(e // copy_elems)
238
+ copy_elem = const_expr(e % copy_elems)
239
+ score = score_vals[chunk, copy_elem]
240
+ kv = kv_vals[chunk, copy_elem]
241
+ new_max = cute.arch.fmax(local_max[e], score)
242
+ old_scale = cute.math.exp2(
243
+ (local_max[e] - new_max) * Float32(rcp_ln2),
244
+ fastmath=True,
245
+ )
246
+ new_scale = cute.math.exp2(
247
+ (score - new_max) * Float32(rcp_ln2),
248
+ fastmath=True,
249
+ )
250
+ local_sum[e] = local_sum[e] * old_scale + new_scale
251
+ local_product[e] = (
252
+ local_product[e] * old_scale + kv * new_scale
253
+ )
254
+ local_max[e] = new_max
255
+
256
+ x = cute.make_rmem_tensor((elems_per_lane,), Float32)
257
+ local_sumsq = Float32(0.0)
258
+ for e in cutlass.range_constexpr(elems_per_lane):
259
+ x[e] = local_product[e] / local_sum[e]
260
+ local_sumsq += x[e] * x[e]
261
+
262
+ warp_sum = local_sumsq
263
+ for step in cutlass.range_constexpr(5):
264
+ offset = const_expr(16 >> step)
265
+ warp_sum += cute.arch.shuffle_sync_bfly(warp_sum, offset)
266
277
- x = cute.make_rmem_tensor((self.elems_per_lane,), Float32)
278
- local_sumsq = Float32(0.0)
279
- for e in cutlass.range_constexpr(self.elems_per_lane):
280
- x[e] = local_product[e] / local_sum[e]
281
- local_sumsq += x[e] * x[e]
267
+ if lane_id == 0:
268
+ partial_sums[warp_id] = warp_sum
269
+ cute.arch.sync_threads()
270
+ if tid == 0:
271
+ total = Float32(0.0)
272
+ for i in cutlass.range_constexpr(num_warps):
273
+ total += partial_sums[i]
274
+ rrms_shared[0] = cute.math.rsqrt(
275
+ total / Float32(head_dim) + rms_norm_eps, fastmath=True
276
+ )
277
+ cute.arch.sync_threads()
278
283
- warp_sum = local_sumsq
284
- for step in cutlass.range_constexpr(5):
285
- offset = const_expr(16 >> step)
286
- warp_sum += cute.arch.shuffle_sync_bfly(warp_sum, offset)
279
+ rrms = rrms_shared[0]
280
+ for e in cutlass.range_constexpr(elems_per_lane):
281
+ elem = elem_base + e
282
+ x[e] = x[e] * rrms * rms_norm_weight[elem].to(Float32)
283
288
- if lane_id == 0:
289
- partial_sums[warp_id] = warp_sum
290
- cute.arch.sync_threads()
291
- if tid == 0:
292
- total = Float32(0.0)
293
- for i in cutlass.range_constexpr(self.num_warps):
294
- total += partial_sums[i]
295
- rrms_shared[0] = cute.math.rsqrt(
296
- total / Float32(self.head_dim) + rms_norm_eps, fastmath=True
284
+ k_cache_u16 = cute.recast_tensor(k_cache, Uint16)
285
+ k_cache_u32 = cute.recast_tensor(k_cache, Uint32)
286
+ page = kv_slot_idx // kv_cache_block_size
287
+ kv_offset = kv_slot_idx - page * kv_cache_block_size
288
+ value_base = page * k_cache.stride[0] + kv_offset * Int64(token_stride)
289
+ scale_base = (
290
+ page * k_cache.stride[0]
291
+ + kv_cache_block_size * Int64(token_stride)
292
+ + kv_offset * Int64(scale_dim)
293
)
298
- cute.arch.sync_threads()
299
-
300
- rrms = rrms_shared[0]
301
- for e in cutlass.range_constexpr(self.elems_per_lane):
302
- elem = elem_base + e
303
- x[e] = x[e] * rrms * rms_norm_weight[elem].to(Float32)
304
-
305
- k_cache_u16 = cute.recast_tensor(k_cache, Uint16)
306
- k_cache_u32 = cute.recast_tensor(k_cache, Uint32)
307
- page = kv_slot_idx // kv_cache_block_size
308
- kv_offset = kv_slot_idx - page * kv_cache_block_size
309
- value_base = page * k_cache.stride[0] + kv_offset * Int64(self.token_stride)
310
- scale_base = (
311
- page * k_cache.stride[0]
312
- + kv_cache_block_size * Int64(self.token_stride)
313
- + kv_offset * Int64(self.scale_dim)
314
- )
294
316
- if group_idx == self.nope_blocks:
317
- compressed_pos = (position // Int64(self.compress_ratio)) * Int64(
318
- self.compress_ratio
319
- )
320
- for pair in cutlass.range_constexpr(self.elems_per_lane // 2):
321
- elem = const_expr(pair * 2)
322
- pair_idx = (elem_base - self.nope_dim) // 2 + Int32(pair)
323
- cos_v = cos_sin_cache[compressed_pos, pair_idx]
324
- sin_v = cos_sin_cache[
325
- compressed_pos, pair_idx + Int32(self.rope_dim // 2)
326
- ]
327
- real = x[elem] * cos_v - x[elem + 1] * sin_v
328
- imag = x[elem] * sin_v + x[elem + 1] * cos_v
329
- packed = _fp32x2_to_bf16x2(real, imag)
330
- out_base = (
331
- value_base
332
- + Int64(self.nope_dim)
333
- + ((elem_base - self.nope_dim + Int32(elem)) * 2).to(Int64)
295
+ if group_idx == nope_blocks:
296
+ compressed_pos = (position // Int64(compress_ratio)) * Int64(
297
+ compress_ratio
298
)
335
- k_cache_u32.iterator[out_base // Int64(4)] = packed
336
- else:
337
- q = cute.make_rmem_tensor((self.elems_per_lane,), Float32)
338
- local_absmax = Float32(0.0)
339
- for pair in cutlass.range_constexpr(self.elems_per_lane // 2):
340
- elem = const_expr(pair * 2)
341
- q_packed = _fp32x2_to_bf16x2(x[elem], x[elem + 1])
342
- q0, q1 = _bf16x2_to_fp32(q_packed)
343
- q[elem] = q0
344
- q[elem + 1] = q1
345
- local_absmax = cute.arch.fmax(
346
- local_absmax,
347
- cute.arch.fmax(cute.math.absf(q0), cute.math.absf(q1)),
299
+ for pair in cutlass.range_constexpr(elems_per_lane // 2):
300
+ elem = const_expr(pair * 2)
301
+ pair_idx = (elem_base - nope_dim) // 2 + Int32(pair)
302
+ cos_v = cos_sin_cache[compressed_pos, pair_idx]
303
+ sin_v = cos_sin_cache[
304
+ compressed_pos, pair_idx + Int32(rope_dim // 2)
305
+ ]
306
+ real = x[elem] * cos_v - x[elem + 1] * sin_v
307
+ imag = x[elem] * sin_v + x[elem + 1] * cos_v
308
+ packed = _fp32x2_to_bf16x2(real, imag)
309
+ out_base = (
310
+ value_base
311
+ + Int64(nope_dim)
312
+ + ((elem_base - nope_dim + Int32(elem)) * 2).to(Int64)
313
+ )
314
+ k_cache_u32.iterator[out_base // Int64(4)] = packed
315
+ else:
316
+ q = cute.make_rmem_tensor((elems_per_lane,), Float32)
317
+ local_absmax = Float32(0.0)
318
+ for pair in cutlass.range_constexpr(elems_per_lane // 2):
319
+ elem = const_expr(pair * 2)
320
+ q_packed = _fp32x2_to_bf16x2(x[elem], x[elem + 1])
321
+ q0, q1 = _bf16x2_to_fp32(q_packed)
322
+ q[elem] = q0
323
+ q[elem + 1] = q1
324
+ local_absmax = cute.arch.fmax(
325
+ local_absmax,
326
+ cute.arch.fmax(cute.math.absf(q0), cute.math.absf(q1)),
327
+ )
328
+ absmax = local_absmax
329
+ group_mask_and_clamp = const_expr(
330
+ (cute.arch.WARP_SIZE - lanes_per_group) << 8
331
+ | (cute.arch.WARP_SIZE - 1)
332
)
349
- absmax = local_absmax
350
- group_mask_and_clamp = const_expr(
351
- (cute.arch.WARP_SIZE - self.lanes_per_group) << 8
352
- | (cute.arch.WARP_SIZE - 1)
353
- )
354
- for step in cutlass.range_constexpr(self.scale_reduce_steps):
355
- offset = const_expr(self.scale_reduce_offset >> step)
356
- absmax = cute.arch.fmax(
357
- absmax,
358
- cute.arch.shuffle_sync_bfly(
333
+ for step in cutlass.range_constexpr(scale_reduce_steps):
334
+ offset = const_expr(scale_reduce_offset >> step)
335
+ absmax = cute.arch.fmax(
336
absmax,
360
- offset=offset,
361
- mask_and_clamp=group_mask_and_clamp,
362
- ),
363
- )
364
- scale_raw = cute.arch.fmax(
365
- Float32(self.min_scale),
366
- absmax / Float32(self.fp8_max),
367
- )
368
- bits = _recast_val(scale_raw, Uint32)
369
- ue8m0 = ((bits + Uint32(0x7FFFFF)) >> Uint32(23)) & Uint32(0xFF)
370
- inv_scale = _recast_val((Uint32(254) - ue8m0) << Uint32(23), Float32)
371
- for pair in cutlass.range_constexpr(self.elems_per_lane // 2):
372
- elem = const_expr(pair * 2)
373
- y0 = cutlass.min(
374
- cute.arch.fmax(q[elem] * inv_scale, Float32(-self.fp8_max)),
375
- Float32(self.fp8_max),
337
+ cute.arch.shuffle_sync_bfly(
338
+ absmax,
339
+ offset=offset,
340
+ mask_and_clamp=group_mask_and_clamp,
341
+ ),
342
+ )
343
+ scale_raw = cute.arch.fmax(
344
+ Float32(min_scale),
345
+ absmax / Float32(fp8_max),
346
)
377
- y1 = cutlass.min(
378
- cute.arch.fmax(q[elem + 1] * inv_scale, Float32(-self.fp8_max)),
379
- Float32(self.fp8_max),
347
+ bits = _recast_val(scale_raw, Uint32)
348
+ ue8m0 = ((bits + Uint32(0x7FFFFF)) >> Uint32(23)) & Uint32(0xFF)
349
+ inv_scale = _recast_val(
350
+ (Uint32(254) - ue8m0) << Uint32(23), Float32
351
)
381
- packed_fp8 = _fp32x2_to_fp8e4m3x2(y0, y1)
382
- out_base = value_base + (elem_base + Int32(elem)).to(Int64)
383
- k_cache_u16.iterator[out_base // Int64(2)] = packed_fp8
384
- if group_lane == 0:
385
- k_cache.iterator[scale_base + group_idx.to(Int64)] = ue8m0.to(Uint8)
386
- if group_idx == 0:
387
- k_cache.iterator[scale_base + Int64(self.nope_blocks)] = Uint8(
388
- 0
352
+ for pair in cutlass.range_constexpr(elems_per_lane // 2):
353
+ elem = const_expr(pair * 2)
354
+ y0 = cutlass.min(
355
+ cute.arch.fmax(q[elem] * inv_scale, Float32(-fp8_max)),
356
+ Float32(fp8_max),
357
+ )
358
+ y1 = cutlass.min(
359
+ cute.arch.fmax(q[elem + 1] * inv_scale, Float32(-fp8_max)),
360
+ Float32(fp8_max),
361
+ )
362
+ packed_fp8 = _fp32x2_to_fp8e4m3x2(y0, y1)
363
+ out_base = value_base + (elem_base + Int32(elem)).to(Int64)
364
+ k_cache_u16.iterator[out_base // Int64(2)] = packed_fp8
365
+ if group_lane == 0:
366
+ k_cache.iterator[scale_base + group_idx.to(Int64)] = ue8m0.to(
367
+ Uint8
368
)
369
+ if group_idx == 0:
370
+ k_cache.iterator[scale_base + Int64(nope_blocks)] = Uint8(0)
371
+
372
+ @cute.jit
373
+ def host_entrypoint(
374
+ state_cache: cute.Tensor,
375
+ token_to_req_indices: cute.Tensor,
376
+ positions: cute.Tensor,
377
+ slot_mapping: cute.Tensor,
378
+ block_table: cute.Tensor,
379
+ block_size: Int64,
380
+ rms_norm_weight: cute.Tensor,
381
+ rms_norm_eps: Float32,
382
+ cos_sin_cache: cute.Tensor,
383
+ k_cache: cute.Tensor,
384
+ kv_slot_mapping: cute.Tensor,
385
+ kv_cache_block_size: Int64,
386
+ stream: CUstream,
387
+ ):
388
+ grid = (slot_mapping.shape[0], 1, 1)
389
+ device_kernel(
390
+ state_cache,
391
+ token_to_req_indices,
392
+ positions,
393
+ slot_mapping,
394
+ block_table,
395
+ block_size,
396
+ rms_norm_weight,
397
+ rms_norm_eps,
398
+ cos_sin_cache,
399
+ k_cache,
400
+ kv_slot_mapping,
401
+ kv_cache_block_size,
402
+ ).launch(grid=grid, block=(tb_size, 1, 1), stream=stream)
403
+
404
+ return host_entrypoint
405
+
406
+ def dispatch( # type: ignore[override]
407
+ self,
408
+ *,
409
+ compress_ratio: int,
410
+ norm_weight_dtype: type[cutlass.Numeric],
411
+ head_size: int,
412
+ rope_head_dim: int,
413
+ ) -> CompileKey:
414
+ overlap = compress_ratio == 4
415
+ state_width = head_size * (2 if overlap else 1)
416
+ return self.CompileKey(
417
+ head_size=head_size,
418
+ state_width=state_width,
419
+ rope_head_dim=rope_head_dim,
420
+ fp8_max=448.0,
421
+ quant_block=64,
422
+ token_stride=576,
423
+ scale_dim=8,
424
+ compress_ratio=compress_ratio,
425
+ overlap=overlap,
426
+ norm_weight_dtype=norm_weight_dtype,
427
+ )
428
391
- @cache
392
- @staticmethod
393
- def compile(
394
- head_size: int = 512,
395
- state_width: int = 1024,
396
- rope_head_dim: int = 64,
397
- fp8_max: float = 448.0,
398
- quant_block: int = 64,
399
- token_stride: int = 576,
400
- scale_dim: int = 8,
401
- kv_block_stride: int = 74752,
402
- compress_ratio: int = 4,
403
- overlap: bool = True,
404
- norm_weight_dtype: type[cutlass.Numeric] = Float32,
405
- ):
429
+ def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]:
430
+ hf_config = vllm_config.model_config.hf_config
431
+ cache_config = vllm_config.cache_config
432
+ return self._trace_dispatch(self.dispatch)(
433
+ compress_ratio=4,
434
+ norm_weight_dtype=torch_to_cute_dtype(vllm_config.model_config.dtype),
435
+ head_size=hf_config.head_dim,
436
+ rope_head_dim=hf_config.qk_rope_head_dim,
437
+ enabled=(
438
+ cache_config.block_size > 0
439
+ and 4 in hf_config.compress_ratios
440
+ and cache_config.cache_dtype == "fp8_ds_mla"
441
+ ),
442
+ _when=lambda *, enabled: enabled,
443
+ )
444
+
445
+ def warmup_inputs(self, compile_key: CompileKey) -> tuple[Any, ...]:
446
+ head_size = compile_key.head_size
447
+ state_width = compile_key.state_width
448
+ rope_head_dim = compile_key.rope_head_dim
449
+ quant_block = compile_key.quant_block
450
+ token_stride = compile_key.token_stride
451
+ scale_dim = compile_key.scale_dim
452
+ compress_ratio = compile_key.compress_ratio
453
+ overlap = compile_key.overlap
454
+ norm_weight_dtype = compile_key.norm_weight_dtype
455
if compress_ratio != 4 or not overlap:
456
raise ValueError("CuTe DSL C4 fused sparse-attn requires C4 overlap.")
457
if head_size != 512:
458
raise ValueError(
410
- "CuTe DSL C4 fused sparse-attn currently requires head_size=512."
459
+ "CuTe DSL C4 fused sparse-attn currently requires head_size=512.",
460
)
461
if state_width != 2 * head_size:
462
raise ValueError(
414
- "CuTe DSL C4 fused sparse-attn requires state_width=2*head_size."
463
+ "CuTe DSL C4 fused sparse-attn requires state_width=2*head_size.",
464
)
465
if quant_block != 64:
466
raise ValueError(
418
- "CuTe DSL C4 fused sparse-attn currently requires quant_block=64."
467
+ "CuTe DSL C4 fused sparse-attn currently requires quant_block=64.",
468
)
469
if rope_head_dim != 64:
470
raise ValueError(
422
- "CuTe DSL C4 fused sparse-attn currently requires rope_head_dim=64."
471
+ "CuTe DSL C4 fused sparse-attn currently requires rope_head_dim=64.",
472
)
473
if token_stride < head_size + rope_head_dim:
474
raise ValueError("token_stride is too small for the packed FP8/BF16 row.")
475
expected_scale_dim = (head_size - rope_head_dim) // quant_block + 1
476
if scale_dim < expected_scale_dim:
477
raise ValueError("scale_dim is too small for the UE8M0 scale row.")
429
-
478
num_positions = cute.sym_int()
479
num_slots = cute.sym_int()
480
num_req_indices = cute.sym_int()
@@ -437,7 +485,6 @@ class SparseAttnCompressNormRopeStoreC4Kernel:
485
block_table_width = cute.sym_int()
486
max_pos = cute.sym_int()
487
state_cache_width = state_width * 2
440
-
488
state_cache = cute.runtime.make_fake_tensor(
489
Float32,
490
(num_state_blocks, state_cache_block_size, state_cache_width),
@@ -449,15 +496,21 @@ class SparseAttnCompressNormRopeStoreC4Kernel:
496
assumed_align=16,
497
)
498
token_to_req_indices = make_fake_tensor(
452
- Int32, (num_req_indices,), divisibility=4
499
+ Int32,
500
+ (num_req_indices,),
501
+ divisibility=4,
502
)
503
positions = make_fake_tensor(Int64, (num_positions,), divisibility=8)
504
slot_mapping = make_fake_tensor(Int64, (num_slots,), divisibility=8)
505
block_table = make_fake_tensor(
457
- Int32, (cute.sym_int(), block_table_width), divisibility=1
506
+ Int32,
507
+ (cute.sym_int(), block_table_width),
508
+ divisibility=1,
509
)
510
rms_norm_weight = make_fake_tensor(
460
- norm_weight_dtype, (head_size,), divisibility=4
511
+ norm_weight_dtype,
512
+ (head_size,),
513
+ divisibility=4,
514
)
515
cos_sin_cache = cute.runtime.make_fake_tensor(
516
Float32,
@@ -468,29 +521,12 @@ class SparseAttnCompressNormRopeStoreC4Kernel:
521
k_cache = cute.runtime.make_fake_tensor(
522
Uint8,
523
(num_kv_blocks, cute.sym_int(), cute.sym_int()),
471
- stride=(
472
- cute.sym_int64(divisibility=16),
473
- cute.sym_int64(divisibility=8),
474
- 1,
475
- ),
524
+ stride=(cute.sym_int64(divisibility=16), cute.sym_int64(divisibility=8), 1),
525
assumed_align=16,
526
)
527
kv_slot_mapping = make_fake_tensor(Int64, (num_kv_slots,), divisibility=8)
528
480
- kernel = SparseAttnCompressNormRopeStoreC4Kernel(
481
- head_size,
482
- state_width,
483
- rope_head_dim,
484
- fp8_max,
485
- quant_block,
486
- token_stride,
487
- scale_dim,
488
- compress_ratio,
489
- overlap,
490
- )
491
- stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
492
- return cute.compile(
493
- kernel,
529
+ return (
530
state_cache,
531
token_to_req_indices,
532
positions,
@@ -503,60 +539,34 @@ class SparseAttnCompressNormRopeStoreC4Kernel:
539
k_cache,
540
kv_slot_mapping,
541
Int64(0),
506
- stream,
507
- options="--enable-tvm-ffi",
542
)
543
510
-
511
-class SparseAttnCompressNormRopeStoreFullC4Kernel(
512
- SparseAttnCompressNormRopeStoreC4Kernel
513
-):
514
- def __init__(
544
+ @kernel_launcher
545
+ def __call__(
546
self,
516
- head_size: int,
517
- state_width: int,
518
- rope_head_dim: int,
519
- fp8_max: float,
520
- quant_block: int,
521
- token_stride: int,
522
- scale_dim: int,
547
+ *,
548
+ state_cache: torch.Tensor,
549
+ token_to_req_indices: torch.Tensor,
550
+ positions: torch.Tensor,
551
+ slot_mapping: torch.Tensor,
552
+ block_table: torch.Tensor,
553
+ block_size: int,
554
+ rms_norm_weight: torch.Tensor,
555
+ rms_norm_eps: float,
556
+ cos_sin_cache: torch.Tensor,
557
+ kv_cache: torch.Tensor,
558
+ kv_slot_mapping: torch.Tensor,
559
compress_ratio: int,
524
- overlap: bool,
525
- store_full_fp8: bool = False,
526
- ):
527
- super().__init__(
528
- head_size,
529
- state_width,
530
- rope_head_dim,
531
- fp8_max,
532
- quant_block,
533
- token_stride,
534
- scale_dim,
535
- compress_ratio,
536
- overlap,
560
+ head_dim: int,
561
+ rope_head_dim: int,
562
+ ) -> CuTeDSLLaunchSpec[SparseAttnCompressNormRopeStoreC4Kernel.CompileKey]:
563
+ compile_key = self.dispatch(
564
+ compress_ratio=compress_ratio,
565
+ norm_weight_dtype=torch_to_cute_dtype(rms_norm_weight.dtype),
566
+ head_size=head_dim,
567
+ rope_head_dim=rope_head_dim,
568
)
538
- self.store_full_fp8 = store_full_fp8
539
-
540
- @cute.jit
541
- def __call__(
542
- self,
543
- state_cache: cute.Tensor,
544
- token_to_req_indices: cute.Tensor,
545
- positions: cute.Tensor,
546
- slot_mapping: cute.Tensor,
547
- block_table: cute.Tensor,
548
- block_size: Int64,
549
- rms_norm_weight: cute.Tensor,
550
- rms_norm_eps: Float32,
551
- cos_sin_cache: cute.Tensor,
552
- k_cache: cute.Tensor,
553
- kv_slot_mapping: cute.Tensor,
554
- kv_cache_block_size: Int64,
555
- fp8_scale: cute.Tensor,
556
- stream: CUstream,
557
- ):
558
- grid = (slot_mapping.shape[0], 1, 1)
559
- self.kernel(
569
+ launch_args = (
570
state_cache,
571
token_to_req_indices,
572
positions,
@@ -566,283 +576,395 @@ class SparseAttnCompressNormRopeStoreFullC4Kernel(
576
rms_norm_weight,
577
rms_norm_eps,
578
cos_sin_cache,
569
- k_cache,
579
+ kv_cache,
580
kv_slot_mapping,
571
- kv_cache_block_size,
572
- fp8_scale,
573
- ).launch(grid=grid, block=(self.tb_size, 1, 1), stream=stream)
574
-
575
- @cute.kernel
576
- def kernel(
577
- self,
578
- state_cache: cute.Tensor,
579
- token_to_req_indices: cute.Tensor,
580
- positions: cute.Tensor,
581
- slot_mapping: cute.Tensor,
582
- block_table: cute.Tensor,
583
- block_size: Int64,
584
- rms_norm_weight: cute.Tensor,
585
- rms_norm_eps: Float32,
586
- cos_sin_cache: cute.Tensor,
587
- k_cache: cute.Tensor,
588
- kv_slot_mapping: cute.Tensor,
589
- kv_cache_block_size: Int64,
590
- fp8_scale: cute.Tensor,
591
- ):
592
- token_idx, _, _ = cute.arch.block_idx()
593
- tid, _, _ = cute.arch.thread_idx()
594
- warp_id = cute.arch.make_warp_uniform(tid // 32)
595
- lane_id = tid % 32
596
- group_lane = lane_id % self.lanes_per_group
597
- group_idx = warp_id * self.groups_per_warp + lane_id // self.lanes_per_group
598
- elem_base = group_idx * self.quant_block + group_lane * self.elems_per_lane
599
-
600
- slot_id = slot_mapping[token_idx]
601
- has_position = token_idx < positions.shape[0]
602
- position = Int64(0)
603
- if has_position:
604
- position = positions[token_idx]
605
- boundary = has_position and (
606
- (position + Int64(1)) % Int64(self.compress_ratio) == Int64(0)
607
- )
608
- has_req_idx = token_idx < token_to_req_indices.shape[0]
609
- has_kv_slot_idx = token_idx < kv_slot_mapping.shape[0]
610
- kv_slot_idx = Int64(-1)
611
- if has_kv_slot_idx:
612
- kv_slot_idx = kv_slot_mapping[token_idx]
613
- active = (
614
- slot_id >= Int64(0) and has_req_idx and boundary and kv_slot_idx >= Int64(0)
581
+ kv_cache.shape[1],
582
)
583
+ return compile_key, launch_args
584
617
- if active:
618
- req_idx = token_to_req_indices[token_idx]
619
- start = position - Int64(self.window - 1)
620
-
621
- smem = cutlass.utils.SmemAllocator()
622
- s_block_numbers = smem.allocate_tensor(
623
- Int32, cute.make_layout((self.window,)), byte_alignment=4
624
- )
625
- partial_sums = smem.allocate_tensor(
626
- Float32, cute.make_layout((self.num_warps,)), byte_alignment=4
627
- )
628
- rrms_shared = smem.allocate_tensor(
629
- Float32, cute.make_layout((1,)), byte_alignment=4
630
- )
585
632
- for row in cutlass.range_constexpr(self.window):
633
- pos = start + Int64(row)
634
- if tid == row:
635
- block_number_i32 = Int32(0)
636
- if pos >= Int64(0):
637
- block_index = pos // block_size
638
- block_number_i32 = block_table[req_idx, block_index]
639
- s_block_numbers[row] = block_number_i32
640
- cute.arch.sync_threads()
586
+class SparseAttnCompressNormRopeStoreFullC4Kernel(
587
+ VllmCuTeDSLJitKernel["SparseAttnCompressNormRopeStoreFullC4Kernel.CompileKey"]
588
+):
589
+ rcp_ln2 = 1.4426950408889634
590
642
- local_max = cute.make_rmem_tensor((self.elems_per_lane,), Float32)
643
- local_sum = cute.make_rmem_tensor((self.elems_per_lane,), Float32)
644
- local_product = cute.make_rmem_tensor((self.elems_per_lane,), Float32)
591
+ @dataclass(frozen=True)
592
+ class CompileKey:
593
+ head_size: int
594
+ state_width: int
595
+ rope_head_dim: int
596
+ fp8_max: float
597
+ quant_block: int
598
+ compress_ratio: int
599
+ overlap: bool
600
+ store_full_fp8: bool
601
+ norm_weight_dtype: type[cutlass.Numeric]
602
646
- for e in cutlass.range_constexpr(self.elems_per_lane):
647
- local_max[e] = -Float32.inf
648
- local_sum[e] = Float32(0.0)
649
- local_product[e] = Float32(0.0)
603
+ @staticmethod
604
+ def kernel(compile_key: CompileKey) -> Any:
605
+ rcp_ln2 = SparseAttnCompressNormRopeStoreFullC4Kernel.rcp_ln2
606
+ head_dim = compile_key.head_size
607
+ state_width = compile_key.state_width
608
+ rope_dim = compile_key.rope_head_dim
609
+ nope_dim = head_dim - rope_dim
610
+ fp8_max = compile_key.fp8_max
611
+ quant_block = compile_key.quant_block
612
+ elems_per_lane = 8
613
+ copy_elems = 4
614
+ copy_chunks = elems_per_lane // copy_elems
615
+ lanes_per_group = quant_block // elems_per_lane
616
+ groups_per_warp = 32 // lanes_per_group
617
+ num_warps = (head_dim // quant_block) // groups_per_warp
618
+ nope_blocks = nope_dim // quant_block
619
+ tb_size = num_warps * 32
620
+ compress_ratio = compile_key.compress_ratio
621
+ overlap = compile_key.overlap
622
+ window = (1 + int(overlap)) * compress_ratio
623
+ store_full_fp8 = compile_key.store_full_fp8
624
+
625
+ @cute.kernel
626
+ def device_kernel(
627
+ state_cache: cute.Tensor,
628
+ token_to_req_indices: cute.Tensor,
629
+ positions: cute.Tensor,
630
+ slot_mapping: cute.Tensor,
631
+ block_table: cute.Tensor,
632
+ block_size: Int64,
633
+ rms_norm_weight: cute.Tensor,
634
+ rms_norm_eps: Float32,
635
+ cos_sin_cache: cute.Tensor,
636
+ k_cache: cute.Tensor,
637
+ kv_slot_mapping: cute.Tensor,
638
+ kv_cache_block_size: Int64,
639
+ fp8_scale: cute.Tensor,
640
+ ):
641
+ token_idx, _, _ = cute.arch.block_idx()
642
+ tid, _, _ = cute.arch.thread_idx()
643
+ warp_id = cute.arch.make_warp_uniform(tid // 32)
644
+ lane_id = tid % 32
645
+ group_lane = lane_id % lanes_per_group
646
+ group_idx = warp_id * groups_per_warp + lane_id // lanes_per_group
647
+ elem_base = group_idx * quant_block + group_lane * elems_per_lane
648
651
- cp_f32x4 = cute.make_copy_atom(
652
- cute.nvgpu.CopyUniversalOp(), Float32, num_bits_per_copy=128
649
+ slot_id = slot_mapping[token_idx]
650
+ has_position = token_idx < positions.shape[0]
651
+ position = Int64(0)
652
+ if has_position:
653
+ position = positions[token_idx]
654
+ boundary = has_position and (
655
+ (position + Int64(1)) % Int64(compress_ratio) == Int64(0)
656
)
654
- copy_layout = cute.make_layout(
655
- (self.copy_chunks, self.copy_elems),
656
- stride=(self.copy_elems, 1),
657
+ has_req_idx = token_idx < token_to_req_indices.shape[0]
658
+ has_kv_slot_idx = token_idx < kv_slot_mapping.shape[0]
659
+ kv_slot_idx = Int64(-1)
660
+ if has_kv_slot_idx:
661
+ kv_slot_idx = kv_slot_mapping[token_idx]
662
+ active = (
663
+ slot_id >= Int64(0)
664
+ and has_req_idx
665
+ and boundary
666
+ and kv_slot_idx >= Int64(0)
667
)
658
- kv_vals = cute.make_rmem_tensor(copy_layout, Float32)
659
- score_vals = cute.make_rmem_tensor(copy_layout, Float32)
660
-
661
- for row in cutlass.range_constexpr(self.window):
662
- pos = start + Int64(row)
663
- if pos >= Int64(0):
664
- block_index = pos // block_size
665
- block_offset = pos - block_index * block_size
666
- block_number = s_block_numbers[row].to(Int64)
667
- head_offset = Int64((row // self.compress_ratio) * self.head_dim)
668
- row_tensor = state_cache[block_number, block_offset, None]
669
- for chunk in cutlass.range_constexpr(self.copy_chunks):
670
- copy_elem = const_expr(chunk * self.copy_elems)
671
- col_tile = (
672
- head_offset + (elem_base + Int32(copy_elem)).to(Int64)
673
- ) // Int64(self.copy_elems)
674
- kv_src = cute.local_tile(
675
- row_tensor,
676
- tiler=(self.copy_elems,),
677
- coord=(col_tile,),
678
- )
679
- score_src = cute.local_tile(
680
- row_tensor,
681
- tiler=(self.copy_elems,),
682
- coord=(
683
- col_tile + Int64(self.state_width // self.copy_elems),
684
- ),
685
- )
686
- cute.copy(cp_f32x4, kv_src, kv_vals[chunk, None])
687
- cute.copy(cp_f32x4, score_src, score_vals[chunk, None])
688
-
689
- for e in cutlass.range_constexpr(self.elems_per_lane):
690
- chunk = const_expr(e // self.copy_elems)
691
- copy_elem = const_expr(e % self.copy_elems)
692
- score = score_vals[chunk, copy_elem]
693
- kv = kv_vals[chunk, copy_elem]
694
- new_max = cute.arch.fmax(local_max[e], score)
695
- old_scale = cute.math.exp2(
696
- (local_max[e] - new_max) * Float32(self.rcp_ln2),
697
- fastmath=True,
698
- )
699
- new_scale = cute.math.exp2(
700
- (score - new_max) * Float32(self.rcp_ln2),
701
- fastmath=True,
702
- )
703
- local_sum[e] = local_sum[e] * old_scale + new_scale
704
- local_product[e] = local_product[e] * old_scale + kv * new_scale
705
- local_max[e] = new_max
668
707
- x = cute.make_rmem_tensor((self.elems_per_lane,), Float32)
708
- local_sumsq = Float32(0.0)
709
- for e in cutlass.range_constexpr(self.elems_per_lane):
710
- x[e] = local_product[e] / local_sum[e]
711
- local_sumsq += x[e] * x[e]
669
+ if active:
670
+ req_idx = token_to_req_indices[token_idx]
671
+ start = position - Int64(window - 1)
672
713
- warp_sum = local_sumsq
714
- for step in cutlass.range_constexpr(5):
715
- offset = const_expr(16 >> step)
716
- warp_sum += cute.arch.shuffle_sync_bfly(warp_sum, offset)
717
-
718
- if lane_id == 0:
719
- partial_sums[warp_id] = warp_sum
720
- cute.arch.sync_threads()
721
- if tid == 0:
722
- total = Float32(0.0)
723
- for i in cutlass.range_constexpr(self.num_warps):
724
- total += partial_sums[i]
725
- rrms_shared[0] = cute.math.rsqrt(
726
- total / Float32(self.head_dim) + rms_norm_eps, fastmath=True
673
+ smem = cutlass.utils.SmemAllocator()
674
+ s_block_numbers = smem.allocate_tensor(
675
+ Int32, cute.make_layout((window,)), byte_alignment=4
676
+ )
677
+ partial_sums = smem.allocate_tensor(
678
+ Float32, cute.make_layout((num_warps,)), byte_alignment=4
679
+ )
680
+ rrms_shared = smem.allocate_tensor(
681
+ Float32, cute.make_layout((1,)), byte_alignment=4
682
)
728
- cute.arch.sync_threads()
683
730
- rrms = rrms_shared[0]
731
- for e in cutlass.range_constexpr(self.elems_per_lane):
732
- elem = elem_base + e
733
- x[e] = x[e] * rrms * rms_norm_weight[elem].to(Float32)
684
+ for row in cutlass.range_constexpr(window):
685
+ pos = start + Int64(row)
686
+ if tid == row:
687
+ block_number_i32 = Int32(0)
688
+ if pos >= Int64(0):
689
+ block_index = pos // block_size
690
+ block_number_i32 = block_table[req_idx, block_index]
691
+ s_block_numbers[row] = block_number_i32
692
+ cute.arch.sync_threads()
693
+
694
+ local_max = cute.make_rmem_tensor((elems_per_lane,), Float32)
695
+ local_sum = cute.make_rmem_tensor((elems_per_lane,), Float32)
696
+ local_product = cute.make_rmem_tensor((elems_per_lane,), Float32)
697
+
698
+ for e in cutlass.range_constexpr(elems_per_lane):
699
+ local_max[e] = -Float32.inf
700
+ local_sum[e] = Float32(0.0)
701
+ local_product[e] = Float32(0.0)
702
+
703
+ cp_f32x4 = cute.make_copy_atom(
704
+ cute.nvgpu.CopyUniversalOp(), Float32, num_bits_per_copy=128
705
+ )
706
+ copy_layout = cute.make_layout(
707
+ (copy_chunks, copy_elems),
708
+ stride=(copy_elems, 1),
709
+ )
710
+ kv_vals = cute.make_rmem_tensor(copy_layout, Float32)
711
+ score_vals = cute.make_rmem_tensor(copy_layout, Float32)
712
735
- page = kv_slot_idx // kv_cache_block_size
736
- kv_offset = kv_slot_idx - page * kv_cache_block_size
737
- value_base = page * k_cache.stride[0] + kv_offset * k_cache.stride[1]
713
+ for row in cutlass.range_constexpr(window):
714
+ pos = start + Int64(row)
715
+ if pos >= Int64(0):
716
+ block_index = pos // block_size
717
+ block_offset = pos - block_index * block_size
718
+ block_number = s_block_numbers[row].to(Int64)
719
+ head_offset = Int64((row // compress_ratio) * head_dim)
720
+ row_tensor = state_cache[block_number, block_offset, None]
721
+ for chunk in cutlass.range_constexpr(copy_chunks):
722
+ copy_elem = const_expr(chunk * copy_elems)
723
+ col_tile = (
724
+ head_offset + (elem_base + Int32(copy_elem)).to(Int64)
725
+ ) // Int64(copy_elems)
726
+ kv_src = cute.local_tile(
727
+ row_tensor,
728
+ tiler=(copy_elems,),
729
+ coord=(col_tile,),
730
+ )
731
+ score_src = cute.local_tile(
732
+ row_tensor,
733
+ tiler=(copy_elems,),
734
+ coord=(col_tile + Int64(state_width // copy_elems),),
735
+ )
736
+ cute.copy(cp_f32x4, kv_src, kv_vals[chunk, None])
737
+ cute.copy(cp_f32x4, score_src, score_vals[chunk, None])
738
+
739
+ for e in cutlass.range_constexpr(elems_per_lane):
740
+ chunk = const_expr(e // copy_elems)
741
+ copy_elem = const_expr(e % copy_elems)
742
+ score = score_vals[chunk, copy_elem]
743
+ kv = kv_vals[chunk, copy_elem]
744
+ new_max = cute.arch.fmax(local_max[e], score)
745
+ old_scale = cute.math.exp2(
746
+ (local_max[e] - new_max) * Float32(rcp_ln2),
747
+ fastmath=True,
748
+ )
749
+ new_scale = cute.math.exp2(
750
+ (score - new_max) * Float32(rcp_ln2),
751
+ fastmath=True,
752
+ )
753
+ local_sum[e] = local_sum[e] * old_scale + new_scale
754
+ local_product[e] = (
755
+ local_product[e] * old_scale + kv * new_scale
756
+ )
757
+ local_max[e] = new_max
758
+
759
+ x = cute.make_rmem_tensor((elems_per_lane,), Float32)
760
+ local_sumsq = Float32(0.0)
761
+ for e in cutlass.range_constexpr(elems_per_lane):
762
+ x[e] = local_product[e] / local_sum[e]
763
+ local_sumsq += x[e] * x[e]
764
+
765
+ warp_sum = local_sumsq
766
+ for step in cutlass.range_constexpr(5):
767
+ offset = const_expr(16 >> step)
768
+ warp_sum += cute.arch.shuffle_sync_bfly(warp_sum, offset)
769
739
- if const_expr(self.store_full_fp8):
740
- k_cache_u16 = cute.recast_tensor(k_cache, Uint16)
741
- inv_fp8 = Float32(1.0) / fp8_scale[0]
742
- if group_idx == self.nope_blocks:
743
- compressed_pos = (position // Int64(self.compress_ratio)) * Int64(
744
- self.compress_ratio
745
- )
746
- for pair in cutlass.range_constexpr(self.elems_per_lane // 2):
747
- elem = const_expr(pair * 2)
748
- pair_idx = (elem_base - self.nope_dim) // 2 + Int32(pair)
749
- cos_v = cos_sin_cache[compressed_pos, pair_idx]
750
- sin_v = cos_sin_cache[
751
- compressed_pos, pair_idx + Int32(self.rope_dim // 2)
752
- ]
753
- real = x[elem] * cos_v - x[elem + 1] * sin_v
754
- imag = x[elem] * sin_v + x[elem + 1] * cos_v
755
- packed_bf16 = _fp32x2_to_bf16x2(real, imag)
756
- b0, b1 = _bf16x2_to_fp32(packed_bf16)
757
- y0 = cutlass.min(
758
- cutlass.max(b0 * inv_fp8, Float32(-self.fp8_max)),
759
- Float32(self.fp8_max),
760
- )
761
- y1 = cutlass.min(
762
- cutlass.max(b1 * inv_fp8, Float32(-self.fp8_max)),
763
- Float32(self.fp8_max),
764
- )
765
- packed_fp8 = _fp32x2_to_fp8e4m3x2(y0, y1)
766
- out_base = value_base + (elem_base + Int32(elem)).to(Int64)
767
- k_cache_u16.iterator[out_base // Int64(2)] = packed_fp8
768
- else:
769
- for pair in cutlass.range_constexpr(self.elems_per_lane // 2):
770
- elem = const_expr(pair * 2)
771
- packed_bf16 = _fp32x2_to_bf16x2(x[elem], x[elem + 1])
772
- b0, b1 = _bf16x2_to_fp32(packed_bf16)
773
- y0 = cutlass.min(
774
- cutlass.max(b0 * inv_fp8, Float32(-self.fp8_max)),
775
- Float32(self.fp8_max),
776
- )
777
- y1 = cutlass.min(
778
- cutlass.max(b1 * inv_fp8, Float32(-self.fp8_max)),
779
- Float32(self.fp8_max),
780
- )
781
- packed_fp8 = _fp32x2_to_fp8e4m3x2(y0, y1)
782
- out_base = value_base + (elem_base + Int32(elem)).to(Int64)
783
- k_cache_u16.iterator[out_base // Int64(2)] = packed_fp8
784
- else:
785
- k_cache_u32 = cute.recast_tensor(k_cache, Uint32)
786
- if group_idx == self.nope_blocks:
787
- compressed_pos = (position // Int64(self.compress_ratio)) * Int64(
788
- self.compress_ratio
770
+ if lane_id == 0:
771
+ partial_sums[warp_id] = warp_sum
772
+ cute.arch.sync_threads()
773
+ if tid == 0:
774
+ total = Float32(0.0)
775
+ for i in cutlass.range_constexpr(num_warps):
776
+ total += partial_sums[i]
777
+ rrms_shared[0] = cute.math.rsqrt(
778
+ total / Float32(head_dim) + rms_norm_eps, fastmath=True
779
)
790
- for pair in cutlass.range_constexpr(self.elems_per_lane // 2):
791
- elem = const_expr(pair * 2)
792
- pair_idx = (elem_base - self.nope_dim) // 2 + Int32(pair)
793
- cos_v = cos_sin_cache[compressed_pos, pair_idx]
794
- sin_v = cos_sin_cache[
795
- compressed_pos, pair_idx + Int32(self.rope_dim // 2)
796
- ]
797
- real = x[elem] * cos_v - x[elem + 1] * sin_v
798
- imag = x[elem] * sin_v + x[elem + 1] * cos_v
799
- packed_bf16 = _fp32x2_to_bf16x2(real, imag)
800
- out_base = value_base + ((elem_base + Int32(elem)) * 2).to(
801
- Int64
780
+ cute.arch.sync_threads()
781
+
782
+ rrms = rrms_shared[0]
783
+ for e in cutlass.range_constexpr(elems_per_lane):
784
+ elem = elem_base + e
785
+ x[e] = x[e] * rrms * rms_norm_weight[elem].to(Float32)
786
+
787
+ page = kv_slot_idx // kv_cache_block_size
788
+ kv_offset = kv_slot_idx - page * kv_cache_block_size
789
+ value_base = page * k_cache.stride[0] + kv_offset * k_cache.stride[1]
790
+
791
+ if const_expr(store_full_fp8):
792
+ k_cache_u16 = cute.recast_tensor(k_cache, Uint16)
793
+ inv_fp8 = Float32(1.0) / fp8_scale[0]
794
+ if group_idx == nope_blocks:
795
+ compressed_pos = (position // Int64(compress_ratio)) * Int64(
796
+ compress_ratio
797
)
803
- k_cache_u32.iterator[out_base // Int64(4)] = packed_bf16
798
+ for pair in cutlass.range_constexpr(elems_per_lane // 2):
799
+ elem = const_expr(pair * 2)
800
+ pair_idx = (elem_base - nope_dim) // 2 + Int32(pair)
801
+ cos_v = cos_sin_cache[compressed_pos, pair_idx]
802
+ sin_v = cos_sin_cache[
803
+ compressed_pos, pair_idx + Int32(rope_dim // 2)
804
+ ]
805
+ real = x[elem] * cos_v - x[elem + 1] * sin_v
806
+ imag = x[elem] * sin_v + x[elem + 1] * cos_v
807
+ packed_bf16 = _fp32x2_to_bf16x2(real, imag)
808
+ b0, b1 = _bf16x2_to_fp32(packed_bf16)
809
+ y0 = cutlass.min(
810
+ cutlass.max(b0 * inv_fp8, Float32(-fp8_max)),
811
+ Float32(fp8_max),
812
+ )
813
+ y1 = cutlass.min(
814
+ cutlass.max(b1 * inv_fp8, Float32(-fp8_max)),
815
+ Float32(fp8_max),
816
+ )
817
+ packed_fp8 = _fp32x2_to_fp8e4m3x2(y0, y1)
818
+ out_base = value_base + (elem_base + Int32(elem)).to(Int64)
819
+ k_cache_u16.iterator[out_base // Int64(2)] = packed_fp8
820
+ else:
821
+ for pair in cutlass.range_constexpr(elems_per_lane // 2):
822
+ elem = const_expr(pair * 2)
823
+ packed_bf16 = _fp32x2_to_bf16x2(x[elem], x[elem + 1])
824
+ b0, b1 = _bf16x2_to_fp32(packed_bf16)
825
+ y0 = cutlass.min(
826
+ cutlass.max(b0 * inv_fp8, Float32(-fp8_max)),
827
+ Float32(fp8_max),
828
+ )
829
+ y1 = cutlass.min(
830
+ cutlass.max(b1 * inv_fp8, Float32(-fp8_max)),
831
+ Float32(fp8_max),
832
+ )
833
+ packed_fp8 = _fp32x2_to_fp8e4m3x2(y0, y1)
834
+ out_base = value_base + (elem_base + Int32(elem)).to(Int64)
835
+ k_cache_u16.iterator[out_base // Int64(2)] = packed_fp8
836
else:
805
- for pair in cutlass.range_constexpr(self.elems_per_lane // 2):
806
- elem = const_expr(pair * 2)
807
- packed_bf16 = _fp32x2_to_bf16x2(x[elem], x[elem + 1])
808
- out_base = value_base + ((elem_base + Int32(elem)) * 2).to(
809
- Int64
837
+ k_cache_u32 = cute.recast_tensor(k_cache, Uint32)
838
+ if group_idx == nope_blocks:
839
+ compressed_pos = (position // Int64(compress_ratio)) * Int64(
840
+ compress_ratio
841
)
811
- k_cache_u32.iterator[out_base // Int64(4)] = packed_bf16
842
+ for pair in cutlass.range_constexpr(elems_per_lane // 2):
843
+ elem = const_expr(pair * 2)
844
+ pair_idx = (elem_base - nope_dim) // 2 + Int32(pair)
845
+ cos_v = cos_sin_cache[compressed_pos, pair_idx]
846
+ sin_v = cos_sin_cache[
847
+ compressed_pos, pair_idx + Int32(rope_dim // 2)
848
+ ]
849
+ real = x[elem] * cos_v - x[elem + 1] * sin_v
850
+ imag = x[elem] * sin_v + x[elem + 1] * cos_v
851
+ packed_bf16 = _fp32x2_to_bf16x2(real, imag)
852
+ out_base = value_base + ((elem_base + Int32(elem)) * 2).to(
853
+ Int64
854
+ )
855
+ k_cache_u32.iterator[out_base // Int64(4)] = packed_bf16
856
+ else:
857
+ for pair in cutlass.range_constexpr(elems_per_lane // 2):
858
+ elem = const_expr(pair * 2)
859
+ packed_bf16 = _fp32x2_to_bf16x2(x[elem], x[elem + 1])
860
+ out_base = value_base + ((elem_base + Int32(elem)) * 2).to(
861
+ Int64
862
+ )
863
+ k_cache_u32.iterator[out_base // Int64(4)] = packed_bf16
864
+
865
+ @cute.jit
866
+ def host_entrypoint(
867
+ state_cache: cute.Tensor,
868
+ token_to_req_indices: cute.Tensor,
869
+ positions: cute.Tensor,
870
+ slot_mapping: cute.Tensor,
871
+ block_table: cute.Tensor,
872
+ block_size: Int64,
873
+ rms_norm_weight: cute.Tensor,
874
+ rms_norm_eps: Float32,
875
+ cos_sin_cache: cute.Tensor,
876
+ k_cache: cute.Tensor,
877
+ kv_slot_mapping: cute.Tensor,
878
+ kv_cache_block_size: Int64,
879
+ fp8_scale: cute.Tensor,
880
+ stream: CUstream,
881
+ ):
882
+ grid = (slot_mapping.shape[0], 1, 1)
883
+ device_kernel(
884
+ state_cache,
885
+ token_to_req_indices,
886
+ positions,
887
+ slot_mapping,
888
+ block_table,
889
+ block_size,
890
+ rms_norm_weight,
891
+ rms_norm_eps,
892
+ cos_sin_cache,
893
+ k_cache,
894
+ kv_slot_mapping,
895
+ kv_cache_block_size,
896
+ fp8_scale,
897
+ ).launch(grid=grid, block=(tb_size, 1, 1), stream=stream)
898
+
899
+ return host_entrypoint
900
+
901
+ def dispatch( # type: ignore[override]
902
+ self,
903
+ *,
904
+ compress_ratio: int,
905
+ store_full_fp8: bool,
906
+ norm_weight_dtype: type[cutlass.Numeric],
907
+ head_size: int,
908
+ rope_head_dim: int,
909
+ ) -> CompileKey:
910
+ overlap = compress_ratio == 4
911
+ return self.CompileKey(
912
+ head_size=head_size,
913
+ state_width=head_size * (2 if overlap else 1),
914
+ rope_head_dim=rope_head_dim,
915
+ fp8_max=448.0,
916
+ quant_block=64,
917
+ compress_ratio=compress_ratio,
918
+ overlap=overlap,
919
+ store_full_fp8=store_full_fp8,
920
+ norm_weight_dtype=norm_weight_dtype,
921
+ )
922
813
- @cache
814
- @staticmethod
815
- def compile(
816
- head_size: int = 512,
817
- state_width: int = 1024,
818
- rope_head_dim: int = 64,
819
- fp8_max: float = 448.0,
820
- quant_block: int = 64,
821
- token_stride: int = 576,
822
- scale_dim: int = 8,
823
- kv_block_stride: int = 74752,
824
- compress_ratio: int = 4,
825
- overlap: bool = True,
826
- store_full_fp8: bool = False,
827
- norm_weight_dtype: type[cutlass.Numeric] = Float32,
828
- ):
923
+ def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]:
924
+ hf_config = vllm_config.model_config.hf_config
925
+ cache_config = vllm_config.cache_config
926
+ return self._trace_dispatch(self.dispatch)(
927
+ compress_ratio=4,
928
+ store_full_fp8=(
929
+ cache_config.cache_dtype.startswith("fp8")
930
+ and cache_config.cache_dtype != "fp8_ds_mla"
931
+ ),
932
+ norm_weight_dtype=torch_to_cute_dtype(vllm_config.model_config.dtype),
933
+ head_size=hf_config.head_dim,
934
+ rope_head_dim=hf_config.qk_rope_head_dim,
935
+ enabled=(
936
+ cache_config.block_size > 0
937
+ and 4 in hf_config.compress_ratios
938
+ and cache_config.cache_dtype != "fp8_ds_mla"
939
+ ),
940
+ _when=lambda *, enabled: enabled,
941
+ )
942
+
943
+ def warmup_inputs(self, compile_key: CompileKey) -> tuple[Any, ...]:
944
+ head_size = compile_key.head_size
945
+ state_width = compile_key.state_width
946
+ rope_head_dim = compile_key.rope_head_dim
947
+ quant_block = compile_key.quant_block
948
+ compress_ratio = compile_key.compress_ratio
949
+ overlap = compile_key.overlap
950
+ norm_weight_dtype = compile_key.norm_weight_dtype
951
if compress_ratio != 4 or not overlap:
952
raise ValueError("CuTe DSL C4 fused sparse-attn requires C4 overlap.")
953
if head_size != 512:
954
raise ValueError(
833
- "CuTe DSL C4 fused sparse-attn currently requires head_size=512."
955
+ "CuTe DSL C4 fused sparse-attn currently requires head_size=512.",
956
)
957
if state_width != 2 * head_size:
958
raise ValueError(
837
- "CuTe DSL C4 fused sparse-attn requires state_width=2*head_size."
959
+ "CuTe DSL C4 fused sparse-attn requires state_width=2*head_size.",
960
)
961
if quant_block != 64:
962
raise ValueError(
841
- "CuTe DSL C4 fused sparse-attn currently requires quant_block=64."
963
+ "CuTe DSL C4 fused sparse-attn currently requires quant_block=64.",
964
)
965
if rope_head_dim != 64:
966
raise ValueError(
845
- "CuTe DSL C4 fused sparse-attn currently requires rope_head_dim=64."
967
+ "CuTe DSL C4 fused sparse-attn currently requires rope_head_dim=64.",
968
)
969
num_positions = cute.sym_int()
970
num_slots = cute.sym_int()
@@ -854,7 +976,6 @@ class SparseAttnCompressNormRopeStoreFullC4Kernel(
976
block_table_width = cute.sym_int()
977
max_pos = cute.sym_int()
978
state_cache_width = state_width * 2
857
-
979
state_cache = cute.runtime.make_fake_tensor(
980
Float32,
981
(num_state_blocks, state_cache_block_size, state_cache_width),
@@ -866,15 +987,21 @@ class SparseAttnCompressNormRopeStoreFullC4Kernel(
987
assumed_align=16,
988
)
989
token_to_req_indices = make_fake_tensor(
869
- Int32, (num_req_indices,), divisibility=4
990
+ Int32,
991
+ (num_req_indices,),
992
+ divisibility=4,
993
)
994
positions = make_fake_tensor(Int64, (num_positions,), divisibility=8)
995
slot_mapping = make_fake_tensor(Int64, (num_slots,), divisibility=8)
996
block_table = make_fake_tensor(
874
- Int32, (cute.sym_int(), block_table_width), divisibility=1
997
+ Int32,
998
+ (cute.sym_int(), block_table_width),
999
+ divisibility=1,
1000
)
1001
rms_norm_weight = make_fake_tensor(
877
- norm_weight_dtype, (head_size,), divisibility=4
1002
+ norm_weight_dtype,
1003
+ (head_size,),
1004
+ divisibility=4,
1005
)
1006
cos_sin_cache = cute.runtime.make_fake_tensor(
1007
Float32,
@@ -885,31 +1012,12 @@ class SparseAttnCompressNormRopeStoreFullC4Kernel(
1012
k_cache = cute.runtime.make_fake_tensor(
1013
Uint8,
1014
(num_kv_blocks, cute.sym_int(), cute.sym_int()),
888
- stride=(
889
- cute.sym_int64(divisibility=16),
890
- cute.sym_int64(divisibility=8),
891
- 1,
892
- ),
1015
+ stride=(cute.sym_int64(divisibility=16), cute.sym_int64(divisibility=8), 1),
1016
assumed_align=16,
1017
)
1018
kv_slot_mapping = make_fake_tensor(Int64, (num_kv_slots,), divisibility=8)
1019
fp8_scale = make_fake_tensor(Float32, (1,), divisibility=1)
897
-
898
- kernel = SparseAttnCompressNormRopeStoreFullC4Kernel(
899
- head_size,
900
- state_width,
901
- rope_head_dim,
902
- fp8_max,
903
- quant_block,
904
- token_stride,
905
- scale_dim,
906
- compress_ratio,
907
- overlap,
908
- store_full_fp8,
909
- )
910
- stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
911
- return cute.compile(
912
- kernel,
1020
+ return (
1021
state_cache,
1022
token_to_req_indices,
1023
positions,
@@ -923,272 +1031,348 @@ class SparseAttnCompressNormRopeStoreFullC4Kernel(
1031
kv_slot_mapping,
1032
Int64(0),
1033
fp8_scale,
926
- stream,
927
- options="--enable-tvm-ffi",
1034
)
1035
930
-
931
-class SparseAttnCompressC128Block8Kernel:
932
- head_tile = 64
933
- rows_per_warp = 16
934
- elems_per_lane = 2
935
- lanes_per_row = head_tile // elems_per_lane
936
- num_warps = 8
937
- stats_lane_stride = lanes_per_row + 1
938
- final_reduce_steps = 3
939
- final_reduce_initial_offset = 4
940
- tb_size = num_warps * 32
941
- compress_ratio = 128
942
- state_block_size = 8
943
- rcp_ln2 = 1.4426950408889634
944
-
945
- def __init__(
1036
+ @kernel_launcher
1037
+ def __call__(
1038
self,
947
- head_size: int,
948
- state_width: int,
949
- ):
950
- self.head_dim = head_size
951
- self.num_splits = head_size // self.head_tile
952
- self.state_width = state_width
953
-
954
- @cute.jit
955
- def __call__(
956
- self,
957
- state_cache: cute.Tensor,
958
- token_to_req_indices: cute.Tensor,
959
- positions: cute.Tensor,
960
- slot_mapping: cute.Tensor,
961
- block_table: cute.Tensor,
962
- compressed_kv: cute.Tensor,
963
- stream: CUstream,
964
- ):
965
- grid = (slot_mapping.shape[0] * self.num_splits, 1, 1)
966
- self.kernel(
1039
+ *,
1040
+ state_cache: torch.Tensor,
1041
+ token_to_req_indices: torch.Tensor,
1042
+ positions: torch.Tensor,
1043
+ slot_mapping: torch.Tensor,
1044
+ block_table: torch.Tensor,
1045
+ block_size: int,
1046
+ rms_norm_weight: torch.Tensor,
1047
+ rms_norm_eps: float,
1048
+ cos_sin_cache: torch.Tensor,
1049
+ kv_cache: torch.Tensor,
1050
+ kv_slot_mapping: torch.Tensor,
1051
+ compress_ratio: int,
1052
+ store_full_fp8: bool,
1053
+ fp8_scale: torch.Tensor | None,
1054
+ head_dim: int,
1055
+ rope_head_dim: int,
1056
+ ) -> CuTeDSLLaunchSpec[SparseAttnCompressNormRopeStoreFullC4Kernel.CompileKey]:
1057
+ compile_key = self.dispatch(
1058
+ compress_ratio=compress_ratio,
1059
+ store_full_fp8=store_full_fp8,
1060
+ norm_weight_dtype=torch_to_cute_dtype(rms_norm_weight.dtype),
1061
+ head_size=head_dim,
1062
+ rope_head_dim=rope_head_dim,
1063
+ )
1064
+ if fp8_scale is None:
1065
+ fp8_scale = torch.ones(1, dtype=torch.float32, device=kv_cache.device)
1066
+ launch_args = (
1067
state_cache,
1068
token_to_req_indices,
1069
positions,
1070
slot_mapping,
1071
block_table,
972
- compressed_kv,
973
- ).launch(grid=grid, block=(self.tb_size, 1, 1), stream=stream)
1072
+ block_size,
1073
+ rms_norm_weight,
1074
+ rms_norm_eps,
1075
+ cos_sin_cache,
1076
+ kv_cache.view(torch.uint8),
1077
+ kv_slot_mapping,
1078
+ kv_cache.shape[1],
1079
+ fp8_scale,
1080
+ )
1081
+ return compile_key, launch_args
1082
975
- @cute.kernel
976
- def kernel(
977
- self,
978
- state_cache: cute.Tensor,
979
- token_to_req_indices: cute.Tensor,
980
- positions: cute.Tensor,
981
- slot_mapping: cute.Tensor,
982
- block_table: cute.Tensor,
983
- compressed_kv: cute.Tensor,
984
- ):
985
- block_id, _, _ = cute.arch.block_idx()
986
- tid, _, _ = cute.arch.thread_idx()
987
- warp_id = cute.arch.make_warp_uniform(tid // 32)
988
- lane_id = tid % 32
989
- col_group = lane_id % self.lanes_per_row
990
-
991
- token_idx = block_id // self.num_splits
992
- split_idx = block_id - token_idx * self.num_splits
993
- col_base = split_idx * self.head_tile + col_group * self.elems_per_lane
994
-
995
- position = Int64(0)
996
- req_idx = Int32(0)
997
- slot_id = Int64(-1)
998
- has_position = token_idx < positions.shape[0]
999
- has_req_idx = token_idx < token_to_req_indices.shape[0]
1000
- if lane_id == 0:
1001
- slot_id = slot_mapping[token_idx]
1002
- if lane_id == 0 and has_position:
1003
- position = positions[token_idx]
1004
- if lane_id == 0 and has_req_idx:
1005
- req_idx = token_to_req_indices[token_idx]
1006
- slot_id = cute.arch.shuffle_sync(slot_id, offset=0)
1007
- position = cute.arch.shuffle_sync(position, offset=0)
1008
- req_idx = cute.arch.shuffle_sync(req_idx, offset=0)
1009
- boundary = has_position and (
1010
- (position + Int64(1)) % Int64(self.compress_ratio) == Int64(0)
1083
+
1084
+class SparseAttnCompressC128Block8Kernel(
1085
+ VllmCuTeDSLJitKernel["SparseAttnCompressC128Block8Kernel.CompileKey"]
1086
+):
1087
+ head_tile = 64
1088
+ rows_per_warp = 16
1089
+ elems_per_lane = 2
1090
+ lanes_per_row = head_tile // elems_per_lane
1091
+ num_warps = 8
1092
+ stats_lane_stride = lanes_per_row + 1
1093
+ final_reduce_steps = 3
1094
+ final_reduce_initial_offset = 4
1095
+ tb_size = num_warps * 32
1096
+ compress_ratio = 128
1097
+ state_block_size = 8
1098
+ rcp_ln2 = 1.4426950408889634
1099
+
1100
+ @dataclass(frozen=True)
1101
+ class CompileKey:
1102
+ head_size: int
1103
+ state_width: int
1104
+
1105
+ @staticmethod
1106
+ def kernel(compile_key: CompileKey) -> Any:
1107
+ head_tile = SparseAttnCompressC128Block8Kernel.head_tile
1108
+ rows_per_warp = SparseAttnCompressC128Block8Kernel.rows_per_warp
1109
+ elems_per_lane = SparseAttnCompressC128Block8Kernel.elems_per_lane
1110
+ lanes_per_row = SparseAttnCompressC128Block8Kernel.lanes_per_row
1111
+ num_warps = SparseAttnCompressC128Block8Kernel.num_warps
1112
+ stats_lane_stride = SparseAttnCompressC128Block8Kernel.stats_lane_stride
1113
+ final_reduce_steps = SparseAttnCompressC128Block8Kernel.final_reduce_steps
1114
+ final_reduce_initial_offset = (
1115
+ SparseAttnCompressC128Block8Kernel.final_reduce_initial_offset
1116
)
1012
- start = position - Int64(self.compress_ratio - 1)
1013
- active = slot_id >= Int64(0) and has_req_idx and boundary
1014
-
1015
- if active:
1016
- smem = cutlass.utils.SmemAllocator()
1017
- s_max = smem.allocate_tensor(
1018
- Float32,
1019
- cute.make_layout(
1020
- (
1021
- self.num_warps,
1022
- self.lanes_per_row,
1023
- self.elems_per_lane,
1024
- ),
1025
- stride=(
1026
- self.stats_lane_stride * self.elems_per_lane,
1027
- self.elems_per_lane,
1028
- 1,
1029
- ),
1030
- ),
1031
- byte_alignment=4,
1117
+ tb_size = SparseAttnCompressC128Block8Kernel.tb_size
1118
+ compress_ratio = SparseAttnCompressC128Block8Kernel.compress_ratio
1119
+ state_block_size = SparseAttnCompressC128Block8Kernel.state_block_size
1120
+ rcp_ln2 = SparseAttnCompressC128Block8Kernel.rcp_ln2
1121
+ head_dim = compile_key.head_size
1122
+ num_splits = head_dim // head_tile
1123
+ state_width = compile_key.state_width
1124
+
1125
+ @cute.kernel
1126
+ def device_kernel(
1127
+ state_cache: cute.Tensor,
1128
+ token_to_req_indices: cute.Tensor,
1129
+ positions: cute.Tensor,
1130
+ slot_mapping: cute.Tensor,
1131
+ block_table: cute.Tensor,
1132
+ compressed_kv: cute.Tensor,
1133
+ ):
1134
+ block_id, _, _ = cute.arch.block_idx()
1135
+ tid, _, _ = cute.arch.thread_idx()
1136
+ warp_id = cute.arch.make_warp_uniform(tid // 32)
1137
+ lane_id = tid % 32
1138
+ col_group = lane_id % lanes_per_row
1139
+
1140
+ token_idx = block_id // num_splits
1141
+ split_idx = block_id - token_idx * num_splits
1142
+ col_base = split_idx * head_tile + col_group * elems_per_lane
1143
+
1144
+ position = Int64(0)
1145
+ req_idx = Int32(0)
1146
+ slot_id = Int64(-1)
1147
+ has_position = token_idx < positions.shape[0]
1148
+ has_req_idx = token_idx < token_to_req_indices.shape[0]
1149
+ if lane_id == 0:
1150
+ slot_id = slot_mapping[token_idx]
1151
+ if lane_id == 0 and has_position:
1152
+ position = positions[token_idx]
1153
+ if lane_id == 0 and has_req_idx:
1154
+ req_idx = token_to_req_indices[token_idx]
1155
+ slot_id = cute.arch.shuffle_sync(slot_id, offset=0)
1156
+ position = cute.arch.shuffle_sync(position, offset=0)
1157
+ req_idx = cute.arch.shuffle_sync(req_idx, offset=0)
1158
+ boundary = has_position and (
1159
+ (position + Int64(1)) % Int64(compress_ratio) == Int64(0)
1160
)
1033
- s_sum = smem.allocate_tensor(
1034
- Float32,
1035
- cute.make_layout(
1036
- (
1037
- self.num_warps,
1038
- self.lanes_per_row,
1039
- self.elems_per_lane,
1040
- ),
1041
- stride=(
1042
- self.stats_lane_stride * self.elems_per_lane,
1043
- self.elems_per_lane,
1044
- 1,
1161
+ start = position - Int64(compress_ratio - 1)
1162
+ active = slot_id >= Int64(0) and has_req_idx and boundary
1163
+
1164
+ if active:
1165
+ smem = cutlass.utils.SmemAllocator()
1166
+ s_max = smem.allocate_tensor(
1167
+ Float32,
1168
+ cute.make_layout(
1169
+ (
1170
+ num_warps,
1171
+ lanes_per_row,
1172
+ elems_per_lane,
1173
+ ),
1174
+ stride=(
1175
+ stats_lane_stride * elems_per_lane,
1176
+ elems_per_lane,
1177
+ 1,
1178
+ ),
1179
),
1046
- ),
1047
- byte_alignment=4,
1048
- )
1049
- s_product = smem.allocate_tensor(
1050
- Float32,
1051
- cute.make_layout(
1052
- (
1053
- self.num_warps,
1054
- self.lanes_per_row,
1055
- self.elems_per_lane,
1180
+ byte_alignment=4,
1181
+ )
1182
+ s_sum = smem.allocate_tensor(
1183
+ Float32,
1184
+ cute.make_layout(
1185
+ (
1186
+ num_warps,
1187
+ lanes_per_row,
1188
+ elems_per_lane,
1189
+ ),
1190
+ stride=(
1191
+ stats_lane_stride * elems_per_lane,
1192
+ elems_per_lane,
1193
+ 1,
1194
+ ),
1195
),
1057
- stride=(
1058
- self.stats_lane_stride * self.elems_per_lane,
1059
- self.elems_per_lane,
1060
- 1,
1196
+ byte_alignment=4,
1197
+ )
1198
+ s_product = smem.allocate_tensor(
1199
+ Float32,
1200
+ cute.make_layout(
1201
+ (
1202
+ num_warps,
1203
+ lanes_per_row,
1204
+ elems_per_lane,
1205
+ ),
1206
+ stride=(
1207
+ stats_lane_stride * elems_per_lane,
1208
+ elems_per_lane,
1209
+ 1,
1210
+ ),
1211
),
1062
- ),
1063
- byte_alignment=4,
1064
- )
1212
+ byte_alignment=4,
1213
+ )
1214
1066
- row_layout = cute.make_layout(
1067
- (self.rows_per_warp, self.elems_per_lane),
1068
- stride=(self.elems_per_lane, 1),
1069
- )
1070
- kv_vals = cute.make_rmem_tensor(row_layout, Float32)
1071
- score_vals = cute.make_rmem_tensor(row_layout, Float32)
1072
- local_max = cute.make_rmem_tensor((self.elems_per_lane,), Float32)
1073
- local_sum = cute.make_rmem_tensor((self.elems_per_lane,), Float32)
1074
- local_product = cute.make_rmem_tensor((self.elems_per_lane,), Float32)
1075
-
1076
- for e in cutlass.range_constexpr(self.elems_per_lane):
1077
- local_max[e] = -Float32.inf
1078
- local_sum[e] = Float32(0.0)
1079
- local_product[e] = Float32(0.0)
1080
-
1081
- first_block_index = start // Int64(self.state_block_size)
1082
- warp_block_index = first_block_index + (warp_id * 2).to(Int64)
1083
- block0_i32 = Int32(0)
1084
- block1_i32 = Int32(0)
1085
- if lane_id == 0:
1086
- block0_i32 = block_table[req_idx, warp_block_index]
1087
- block1_i32 = block_table[req_idx, warp_block_index + Int64(1)]
1088
- block0_i32 = cute.arch.shuffle_sync(block0_i32, offset=0)
1089
- block1_i32 = cute.arch.shuffle_sync(block1_i32, offset=0)
1215
+ row_layout = cute.make_layout(
1216
+ (rows_per_warp, elems_per_lane),
1217
+ stride=(elems_per_lane, 1),
1218
+ )
1219
+ kv_vals = cute.make_rmem_tensor(row_layout, Float32)
1220
+ score_vals = cute.make_rmem_tensor(row_layout, Float32)
1221
+ local_max = cute.make_rmem_tensor((elems_per_lane,), Float32)
1222
+ local_sum = cute.make_rmem_tensor((elems_per_lane,), Float32)
1223
+ local_product = cute.make_rmem_tensor((elems_per_lane,), Float32)
1224
+
1225
+ for e in cutlass.range_constexpr(elems_per_lane):
1226
+ local_max[e] = -Float32.inf
1227
+ local_sum[e] = Float32(0.0)
1228
+ local_product[e] = Float32(0.0)
1229
+
1230
+ first_block_index = start // Int64(state_block_size)
1231
+ warp_block_index = first_block_index + (warp_id * 2).to(Int64)
1232
+ block0_i32 = Int32(0)
1233
+ block1_i32 = Int32(0)
1234
+ if lane_id == 0:
1235
+ block0_i32 = block_table[req_idx, warp_block_index]
1236
+ block1_i32 = block_table[req_idx, warp_block_index + Int64(1)]
1237
+ block0_i32 = cute.arch.shuffle_sync(block0_i32, offset=0)
1238
+ block1_i32 = cute.arch.shuffle_sync(block1_i32, offset=0)
1239
1091
- cp_f32x2 = cute.make_copy_atom(
1092
- cute.nvgpu.CopyUniversalOp(), Float32, num_bits_per_copy=64
1093
- )
1094
- final_mask_and_clamp = const_expr(
1095
- (cute.arch.WARP_SIZE - self.num_warps) << 8 | (cute.arch.WARP_SIZE - 1)
1096
- )
1097
- col_tile = col_base.to(Int64) // Int64(self.elems_per_lane)
1098
- score_col_tile = col_tile + Int64(self.state_width // self.elems_per_lane)
1099
-
1100
- for i in cutlass.range_constexpr(self.rows_per_warp):
1101
- block_number_i32 = block0_i32
1102
- block_offset = Int64(i)
1103
- if const_expr(i >= self.state_block_size):
1104
- block_number_i32 = block1_i32
1105
- block_offset = Int64(i - self.state_block_size)
1106
- row_tensor = state_cache[block_number_i32.to(Int64), block_offset, None]
1107
- kv_src = cute.local_tile(
1108
- row_tensor,
1109
- tiler=(self.elems_per_lane,),
1110
- coord=(col_tile,),
1240
+ cp_f32x2 = cute.make_copy_atom(
1241
+ cute.nvgpu.CopyUniversalOp(), Float32, num_bits_per_copy=64
1242
)
1112
- score_src = cute.local_tile(
1113
- row_tensor,
1114
- tiler=(self.elems_per_lane,),
1115
- coord=(score_col_tile,),
1243
+ final_mask_and_clamp = const_expr(
1244
+ (cute.arch.WARP_SIZE - num_warps) << 8 | (cute.arch.WARP_SIZE - 1)
1245
)
1117
- cute.copy(cp_f32x2, kv_src, kv_vals[i, None])
1118
- cute.copy(cp_f32x2, score_src, score_vals[i, None])
1246
+ col_tile = col_base.to(Int64) // Int64(elems_per_lane)
1247
+ score_col_tile = col_tile + Int64(state_width // elems_per_lane)
1248
+
1249
+ for i in cutlass.range_constexpr(rows_per_warp):
1250
+ block_number_i32 = block0_i32
1251
+ block_offset = Int64(i)
1252
+ if const_expr(i >= state_block_size):
1253
+ block_number_i32 = block1_i32
1254
+ block_offset = Int64(i - state_block_size)
1255
+ row_tensor = state_cache[
1256
+ block_number_i32.to(Int64), block_offset, None
1257
+ ]
1258
+ kv_src = cute.local_tile(
1259
+ row_tensor,
1260
+ tiler=(elems_per_lane,),
1261
+ coord=(col_tile,),
1262
+ )
1263
+ score_src = cute.local_tile(
1264
+ row_tensor,
1265
+ tiler=(elems_per_lane,),
1266
+ coord=(score_col_tile,),
1267
+ )
1268
+ cute.copy(cp_f32x2, kv_src, kv_vals[i, None])
1269
+ cute.copy(cp_f32x2, score_src, score_vals[i, None])
1270
+
1271
+ for e in cutlass.range_constexpr(elems_per_lane):
1272
+ local_max[e] = cute.arch.fmax(local_max[e], score_vals[i, e])
1273
1120
- for e in cutlass.range_constexpr(self.elems_per_lane):
1121
- local_max[e] = cute.arch.fmax(local_max[e], score_vals[i, e])
1274
+ for e in cutlass.range_constexpr(elems_per_lane):
1275
+ for i in cutlass.range_constexpr(rows_per_warp):
1276
+ exp_score = cute.math.exp2(
1277
+ (score_vals[i, e] - local_max[e]) * Float32(rcp_ln2),
1278
+ fastmath=True,
1279
+ )
1280
+ local_sum[e] += exp_score
1281
+ local_product[e] += kv_vals[i, e] * exp_score
1282
+
1283
+ for e in cutlass.range_constexpr(elems_per_lane):
1284
+ s_max[warp_id, col_group, e] = local_max[e]
1285
+ s_sum[warp_id, col_group, e] = local_sum[e]
1286
+ s_product[warp_id, col_group, e] = local_product[e]
1287
+ cute.arch.sync_threads()
1288
+
1289
+ out_group = tid // num_warps
1290
+ final_lane = tid % num_warps
1291
+ final_groups_per_pass = const_expr(tb_size // num_warps)
1292
+ for pass_idx in cutlass.range_constexpr(
1293
+ head_tile // final_groups_per_pass
1294
+ ):
1295
+ out_idx = pass_idx * final_groups_per_pass + out_group
1296
+ out_lane = out_idx // elems_per_lane
1297
+ out_elem = out_idx % elems_per_lane
1298
+
1299
+ local_warp_max = s_max[final_lane, out_lane, out_elem]
1300
+ global_max = local_warp_max
1301
+ for step in cutlass.range_constexpr(final_reduce_steps):
1302
+ offset = const_expr(final_reduce_initial_offset >> step)
1303
+ global_max = cute.arch.fmax(
1304
+ global_max,
1305
+ cute.arch.shuffle_sync_bfly(
1306
+ global_max,
1307
+ offset=offset,
1308
+ mask_and_clamp=final_mask_and_clamp,
1309
+ ),
1310
+ )
1311
1123
- for e in cutlass.range_constexpr(self.elems_per_lane):
1124
- for i in cutlass.range_constexpr(self.rows_per_warp):
1125
- exp_score = cute.math.exp2(
1126
- (score_vals[i, e] - local_max[e]) * Float32(self.rcp_ln2),
1312
+ scale = cute.math.exp2(
1313
+ (local_warp_max - global_max) * Float32(rcp_ln2),
1314
fastmath=True,
1315
)
1129
- local_sum[e] += exp_score
1130
- local_product[e] += kv_vals[i, e] * exp_score
1131
-
1132
- for e in cutlass.range_constexpr(self.elems_per_lane):
1133
- s_max[warp_id, col_group, e] = local_max[e]
1134
- s_sum[warp_id, col_group, e] = local_sum[e]
1135
- s_product[warp_id, col_group, e] = local_product[e]
1136
- cute.arch.sync_threads()
1137
-
1138
- out_group = tid // self.num_warps
1139
- final_lane = tid % self.num_warps
1140
- final_groups_per_pass = const_expr(self.tb_size // self.num_warps)
1141
- for pass_idx in cutlass.range_constexpr(
1142
- self.head_tile // final_groups_per_pass
1143
- ):
1144
- out_idx = pass_idx * final_groups_per_pass + out_group
1145
- out_lane = out_idx // self.elems_per_lane
1146
- out_elem = out_idx % self.elems_per_lane
1147
-
1148
- local_warp_max = s_max[final_lane, out_lane, out_elem]
1149
- global_max = local_warp_max
1150
- for step in cutlass.range_constexpr(self.final_reduce_steps):
1151
- offset = const_expr(self.final_reduce_initial_offset >> step)
1152
- global_max = cute.arch.fmax(
1153
- global_max,
1154
- cute.arch.shuffle_sync_bfly(
1155
- global_max,
1316
+ global_sum = s_sum[final_lane, out_lane, out_elem] * scale
1317
+ global_product = s_product[final_lane, out_lane, out_elem] * scale
1318
+ for step in cutlass.range_constexpr(final_reduce_steps):
1319
+ offset = const_expr(final_reduce_initial_offset >> step)
1320
+ global_sum += cute.arch.shuffle_sync_bfly(
1321
+ global_sum,
1322
offset=offset,
1323
mask_and_clamp=final_mask_and_clamp,
1158
- ),
1159
- )
1160
-
1161
- scale = cute.math.exp2(
1162
- (local_warp_max - global_max) * Float32(self.rcp_ln2),
1163
- fastmath=True,
1164
- )
1165
- global_sum = s_sum[final_lane, out_lane, out_elem] * scale
1166
- global_product = s_product[final_lane, out_lane, out_elem] * scale
1167
- for step in cutlass.range_constexpr(self.final_reduce_steps):
1168
- offset = const_expr(self.final_reduce_initial_offset >> step)
1169
- global_sum += cute.arch.shuffle_sync_bfly(
1170
- global_sum,
1171
- offset=offset,
1172
- mask_and_clamp=final_mask_and_clamp,
1173
- )
1174
- global_product += cute.arch.shuffle_sync_bfly(
1175
- global_product,
1176
- offset=offset,
1177
- mask_and_clamp=final_mask_and_clamp,
1178
- )
1324
+ )
1325
+ global_product += cute.arch.shuffle_sync_bfly(
1326
+ global_product,
1327
+ offset=offset,
1328
+ mask_and_clamp=final_mask_and_clamp,
1329
+ )
1330
1180
- if final_lane == 0:
1181
- compressed_kv.iterator[
1182
- token_idx.to(Int64) * compressed_kv.stride[0]
1183
- + (split_idx * self.head_tile + out_idx).to(Int64)
1184
- ] = global_product / global_sum
1331
+ if final_lane == 0:
1332
+ compressed_kv.iterator[
1333
+ token_idx.to(Int64) * compressed_kv.stride[0]
1334
+ + (split_idx * head_tile + out_idx).to(Int64)
1335
+ ] = global_product / global_sum
1336
+
1337
+ @cute.jit
1338
+ def host_entrypoint(
1339
+ state_cache: cute.Tensor,
1340
+ token_to_req_indices: cute.Tensor,
1341
+ positions: cute.Tensor,
1342
+ slot_mapping: cute.Tensor,
1343
+ block_table: cute.Tensor,
1344
+ compressed_kv: cute.Tensor,
1345
+ stream: CUstream,
1346
+ ):
1347
+ grid = (slot_mapping.shape[0] * num_splits, 1, 1)
1348
+ device_kernel(
1349
+ state_cache,
1350
+ token_to_req_indices,
1351
+ positions,
1352
+ slot_mapping,
1353
+ block_table,
1354
+ compressed_kv,
1355
+ ).launch(grid=grid, block=(tb_size, 1, 1), stream=stream)
1356
+
1357
+ return host_entrypoint
1358
+
1359
+ def dispatch( # type: ignore[override]
1360
+ self, *, head_size: int, state_width: int
1361
+ ) -> CompileKey:
1362
+ return self.CompileKey(head_size=head_size, state_width=state_width)
1363
+
1364
+ def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]:
1365
+ hf_config = vllm_config.model_config.hf_config
1366
+ return self._trace_dispatch(self.dispatch)(
1367
+ head_size=hf_config.head_dim,
1368
+ state_width=hf_config.head_dim,
1369
+ enabled=128 in hf_config.compress_ratios,
1370
+ _when=lambda *, enabled: enabled,
1371
+ )
1372
1186
- @cache
1187
- @staticmethod
1188
- def compile(
1189
- head_size: int = 512,
1190
- state_width: int = 512,
1191
- ):
1373
+ def warmup_inputs(self, compile_key: CompileKey) -> tuple[Any, ...]:
1374
+ head_size = compile_key.head_size
1375
+ state_width = compile_key.state_width
1376
if head_size % SparseAttnCompressC128Block8Kernel.head_tile != 0:
1377
raise ValueError("head_size must be divisible by the 64-wide head tile.")
1378
num_positions = cute.sym_int()
@@ -1198,7 +1382,6 @@ class SparseAttnCompressC128Block8Kernel:
1382
state_cache_block_size = cute.sym_int()
1383
block_table_width = cute.sym_int()
1384
state_cache_width = state_width * 2
1201
-
1385
state_cache = cute.runtime.make_fake_tensor(
1386
Float32,
1387
(num_blocks, state_cache_block_size, state_cache_width),
@@ -1210,12 +1393,16 @@ class SparseAttnCompressC128Block8Kernel:
1393
assumed_align=16,
1394
)
1395
token_to_req_indices = make_fake_tensor(
1213
- Int32, (num_req_indices,), divisibility=4
1396
+ Int32,
1397
+ (num_req_indices,),
1398
+ divisibility=4,
1399
)
1400
positions = make_fake_tensor(Int64, (num_positions,), divisibility=8)
1401
slot_mapping = make_fake_tensor(Int64, (num_slots,), divisibility=8)
1402
block_table = make_fake_tensor(
1218
- Int32, (cute.sym_int(), block_table_width), divisibility=1
1403
+ Int32,
1404
+ (cute.sym_int(), block_table_width),
1405
+ divisibility=1,
1406
)
1407
compressed_kv = cute.runtime.make_fake_tensor(
1408
Float32,
@@ -1223,238 +1410,332 @@ class SparseAttnCompressC128Block8Kernel:
1410
stride=(head_size, 1),
1411
assumed_align=4,
1412
)
1226
-
1227
- kernel = SparseAttnCompressC128Block8Kernel(
1228
- head_size,
1229
- state_width,
1230
- )
1231
- stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
1232
- return cute.compile(
1233
- kernel,
1413
+ return (
1414
state_cache,
1415
token_to_req_indices,
1416
positions,
1417
slot_mapping,
1418
block_table,
1419
compressed_kv,
1240
- stream,
1241
- options="--enable-tvm-ffi",
1420
)
1421
1244
-
1245
-class SparseAttnNormRopeStoreKernel:
1246
- min_scale = 1.0e-4
1247
-
1248
- def __init__(
1249
- self,
1250
- head_size: int,
1251
- rope_head_dim: int,
1252
- fp8_max: float,
1253
- quant_block: int,
1254
- token_stride: int,
1255
- scale_dim: int,
1256
- compress_ratio: int,
1257
- static_kv_cache_block_size: int,
1258
- ):
1259
- self.head_dim = head_size
1260
- self.rope_dim = rope_head_dim
1261
- self.nope_dim = head_size - rope_head_dim
1262
- self.fp8_max = fp8_max
1263
- self.quant_block = quant_block
1264
- self.token_stride = token_stride
1265
- self.scale_dim = scale_dim
1266
- self.num_warps = head_size // quant_block
1267
- self.nope_blocks = self.nope_dim // quant_block
1268
- self.tb_size = head_size // 2
1269
- self.compress_ratio = compress_ratio
1270
- self.static_kv_cache_block_size = static_kv_cache_block_size
1271
-
1272
- @cute.jit
1422
+ @kernel_launcher
1423
def __call__(
1424
self,
1275
- compressed_kv: cute.Tensor,
1276
- positions: cute.Tensor,
1277
- slot_mapping: cute.Tensor,
1278
- rms_norm_weight: cute.Tensor,
1279
- rms_norm_eps: Float32,
1280
- cos_sin_cache: cute.Tensor,
1281
- k_cache: cute.Tensor,
1282
- kv_slot_mapping: cute.Tensor,
1283
- stream: CUstream,
1284
- ):
1285
- grid = (slot_mapping.shape[0], 1, 1)
1286
- self.kernel(
1287
- compressed_kv,
1425
+ *,
1426
+ state_cache: torch.Tensor,
1427
+ num_actual: int,
1428
+ token_to_req_indices: torch.Tensor,
1429
+ positions: torch.Tensor,
1430
+ slot_mapping: torch.Tensor,
1431
+ block_table: torch.Tensor,
1432
+ head_dim: int,
1433
+ compressed_kv: torch.Tensor | None = None,
1434
+ ) -> CuTeDSLLaunchSpec[SparseAttnCompressC128Block8Kernel.CompileKey]:
1435
+ compile_key = self.dispatch(head_size=head_dim, state_width=head_dim)
1436
+ if compressed_kv is None:
1437
+ compressed_kv = torch.empty(
1438
+ (num_actual, head_dim),
1439
+ dtype=torch.float32,
1440
+ device=state_cache.device,
1441
+ )
1442
+ launch_args = (
1443
+ state_cache,
1444
+ token_to_req_indices,
1445
positions,
1446
slot_mapping,
1290
- rms_norm_weight,
1291
- rms_norm_eps,
1292
- cos_sin_cache,
1293
- k_cache,
1294
- kv_slot_mapping,
1295
- ).launch(grid=grid, block=(self.tb_size, 1, 1), stream=stream)
1296
-
1297
- @cute.kernel
1298
- def kernel(
1299
- self,
1300
- compressed_kv: cute.Tensor,
1301
- positions: cute.Tensor,
1302
- slot_mapping: cute.Tensor,
1303
- rms_norm_weight: cute.Tensor,
1304
- rms_norm_eps: Float32,
1305
- cos_sin_cache: cute.Tensor,
1306
- k_cache: cute.Tensor,
1307
- kv_slot_mapping: cute.Tensor,
1308
- ):
1309
- token_idx, _, _ = cute.arch.block_idx()
1310
- tid, _, _ = cute.arch.thread_idx()
1311
- warp_id = cute.arch.make_warp_uniform(tid // 32)
1312
- lane_id = tid % 32
1313
- elem0 = tid * 2
1314
-
1315
- position = Int64(0)
1316
- kv_slot_idx = Int64(-1)
1317
- has_position = token_idx < positions.shape[0]
1318
- slot_id = Int64(-1)
1319
- if lane_id == 0:
1320
- slot_id = slot_mapping[token_idx]
1321
- if lane_id == 0 and has_position:
1322
- position = positions[token_idx]
1323
- has_kv_slot_idx = token_idx < kv_slot_mapping.shape[0]
1324
- if lane_id == 0 and has_kv_slot_idx:
1325
- kv_slot_idx = kv_slot_mapping[token_idx]
1326
- slot_id = cute.arch.shuffle_sync(slot_id, offset=0)
1327
- position = cute.arch.shuffle_sync(position, offset=0)
1328
- kv_slot_idx = cute.arch.shuffle_sync(kv_slot_idx, offset=0)
1329
- boundary = has_position and (
1330
- (position + Int64(1)) % Int64(self.compress_ratio) == Int64(0)
1447
+ block_table,
1448
+ compressed_kv,
1449
)
1332
- active = slot_id >= Int64(0) and boundary and kv_slot_idx >= Int64(0)
1333
-
1334
- if active:
1335
- k_cache_u16 = cute.recast_tensor(k_cache, Uint16)
1336
- k_cache_u32 = cute.recast_tensor(k_cache, Uint32)
1337
- static_block_size = Int64(self.static_kv_cache_block_size)
1338
- page = kv_slot_idx // static_block_size
1339
- kv_offset = kv_slot_idx - page * static_block_size
1340
- scale_row_offset = static_block_size * Int64(self.token_stride)
1341
- value_base = page * k_cache.stride[0] + kv_offset * Int64(self.token_stride)
1342
- scale_base = (
1343
- page * k_cache.stride[0]
1344
- + scale_row_offset
1345
- + kv_offset * Int64(self.scale_dim)
1346
- )
1347
- weight0 = rms_norm_weight[elem0].to(Float32)
1348
- weight1 = rms_norm_weight[elem0 + 1].to(Float32)
1349
-
1350
- base = token_idx.to(Int64) * compressed_kv.stride[0] + elem0.to(Int64)
1351
- x0 = compressed_kv.iterator[base]
1352
- x1 = compressed_kv.iterator[base + Int64(1)]
1353
-
1354
- local_sumsq = x0 * x0 + x1 * x1
1355
- warp_sum = local_sumsq
1356
- for step in cutlass.range_constexpr(5):
1357
- offset = const_expr(16 >> step)
1358
- warp_sum += cute.arch.shuffle_sync_bfly(warp_sum, offset)
1359
-
1360
- smem = cutlass.utils.SmemAllocator()
1361
- partial_sums = smem.allocate_tensor(
1362
- Float32, cute.make_layout((self.num_warps,)), byte_alignment=4
1363
- )
1450
+ return compile_key, launch_args, compressed_kv
1451
1365
- if lane_id == 0:
1366
- partial_sums[warp_id] = warp_sum
1367
- cute.arch.sync_threads()
1452
1369
- total = partial_sums[lane_id % self.num_warps]
1370
- sum_mask_and_clamp = const_expr(
1371
- (cute.arch.WARP_SIZE - self.num_warps) << 8 | (cute.arch.WARP_SIZE - 1)
1453
+class SparseAttnNormRopeStoreKernel(
1454
+ VllmCuTeDSLJitKernel["SparseAttnNormRopeStoreKernel.CompileKey"]
1455
+):
1456
+ min_scale = 1.0e-4
1457
+
1458
+ @dataclass(frozen=True)
1459
+ class CompileKey:
1460
+ head_size: int
1461
+ rope_head_dim: int
1462
+ fp8_max: float
1463
+ quant_block: int
1464
+ token_stride: int
1465
+ scale_dim: int
1466
+ kv_block_stride: int
1467
+ compress_ratio: int
1468
+ norm_weight_dtype: type[cutlass.Numeric]
1469
+ kv_cache_block_size: int
1470
+
1471
+ @staticmethod
1472
+ def kernel(compile_key: CompileKey) -> Any:
1473
+ min_scale = SparseAttnNormRopeStoreKernel.min_scale
1474
+ head_dim = compile_key.head_size
1475
+ rope_dim = compile_key.rope_head_dim
1476
+ nope_dim = head_dim - rope_dim
1477
+ fp8_max = compile_key.fp8_max
1478
+ quant_block = compile_key.quant_block
1479
+ token_stride = compile_key.token_stride
1480
+ scale_dim = compile_key.scale_dim
1481
+ num_warps = head_dim // quant_block
1482
+ nope_blocks = nope_dim // quant_block
1483
+ tb_size = head_dim // 2
1484
+ compress_ratio = compile_key.compress_ratio
1485
+ static_kv_cache_block_size = compile_key.kv_cache_block_size
1486
+
1487
+ @cute.kernel
1488
+ def device_kernel(
1489
+ compressed_kv: cute.Tensor,
1490
+ positions: cute.Tensor,
1491
+ slot_mapping: cute.Tensor,
1492
+ rms_norm_weight: cute.Tensor,
1493
+ rms_norm_eps: Float32,
1494
+ cos_sin_cache: cute.Tensor,
1495
+ k_cache: cute.Tensor,
1496
+ kv_slot_mapping: cute.Tensor,
1497
+ ):
1498
+ token_idx, _, _ = cute.arch.block_idx()
1499
+ tid, _, _ = cute.arch.thread_idx()
1500
+ warp_id = cute.arch.make_warp_uniform(tid // 32)
1501
+ lane_id = tid % 32
1502
+ elem0 = tid * 2
1503
+
1504
+ position = Int64(0)
1505
+ kv_slot_idx = Int64(-1)
1506
+ has_position = token_idx < positions.shape[0]
1507
+ slot_id = Int64(-1)
1508
+ if lane_id == 0:
1509
+ slot_id = slot_mapping[token_idx]
1510
+ if lane_id == 0 and has_position:
1511
+ position = positions[token_idx]
1512
+ has_kv_slot_idx = token_idx < kv_slot_mapping.shape[0]
1513
+ if lane_id == 0 and has_kv_slot_idx:
1514
+ kv_slot_idx = kv_slot_mapping[token_idx]
1515
+ slot_id = cute.arch.shuffle_sync(slot_id, offset=0)
1516
+ position = cute.arch.shuffle_sync(position, offset=0)
1517
+ kv_slot_idx = cute.arch.shuffle_sync(kv_slot_idx, offset=0)
1518
+ boundary = has_position and (
1519
+ (position + Int64(1)) % Int64(compress_ratio) == Int64(0)
1520
)
1373
- for step in cutlass.range_constexpr(3):
1374
- offset = const_expr(4 >> step)
1375
- total += cute.arch.shuffle_sync_bfly(
1376
- total,
1377
- offset,
1378
- mask_and_clamp=sum_mask_and_clamp,
1521
+ active = slot_id >= Int64(0) and boundary and kv_slot_idx >= Int64(0)
1522
+
1523
+ if active:
1524
+ k_cache_u16 = cute.recast_tensor(k_cache, Uint16)
1525
+ k_cache_u32 = cute.recast_tensor(k_cache, Uint32)
1526
+ static_block_size = Int64(static_kv_cache_block_size)
1527
+ page = kv_slot_idx // static_block_size
1528
+ kv_offset = kv_slot_idx - page * static_block_size
1529
+ scale_row_offset = static_block_size * Int64(token_stride)
1530
+ value_base = page * k_cache.stride[0] + kv_offset * Int64(token_stride)
1531
+ scale_base = (
1532
+ page * k_cache.stride[0]
1533
+ + scale_row_offset
1534
+ + kv_offset * Int64(scale_dim)
1535
)
1536
+ weight0 = rms_norm_weight[elem0].to(Float32)
1537
+ weight1 = rms_norm_weight[elem0 + 1].to(Float32)
1538
1381
- rrms = cute.math.rsqrt(
1382
- total / Float32(self.head_dim) + rms_norm_eps, fastmath=True
1383
- )
1384
- x0 = x0 * rrms * weight0
1385
- x1 = x1 * rrms * weight1
1386
-
1387
- if warp_id == self.nope_blocks:
1388
- pair_idx = lane_id
1389
- compressed_pos = position - Int64(self.compress_ratio - 1)
1390
- cs_base = compressed_pos * cos_sin_cache.stride[0] + pair_idx.to(Int64)
1391
- cos_v = cos_sin_cache.iterator[cs_base]
1392
- sin_v = cos_sin_cache.iterator[cs_base + Int64(self.rope_dim // 2)]
1393
- real = x0 * cos_v - x1 * sin_v
1394
- imag = x0 * sin_v + x1 * cos_v
1395
- packed = _fp32x2_to_bf16x2(real, imag)
1396
- out_base = value_base + Int64(self.nope_dim) + (lane_id * 4).to(Int64)
1397
- k_cache_u32.iterator[out_base // Int64(4)] = packed
1398
- else:
1399
- q_packed = _fp32x2_to_bf16x2(x0, x1)
1400
- q0, q1 = _bf16x2_to_fp32(q_packed)
1401
- abs0 = cute.math.absf(q0)
1402
- abs1 = cute.math.absf(q1)
1403
- local_absmax = cute.arch.fmax(abs0, abs1)
1404
- absmax = local_absmax
1539
+ base = token_idx.to(Int64) * compressed_kv.stride[0] + elem0.to(Int64)
1540
+ x0 = compressed_kv.iterator[base]
1541
+ x1 = compressed_kv.iterator[base + Int64(1)]
1542
+
1543
+ local_sumsq = x0 * x0 + x1 * x1
1544
+ warp_sum = local_sumsq
1545
for step in cutlass.range_constexpr(5):
1546
offset = const_expr(16 >> step)
1407
- absmax = cute.arch.fmax(
1408
- absmax, cute.arch.shuffle_sync_bfly(absmax, offset)
1409
- )
1410
- scale_raw = cute.arch.fmax(
1411
- Float32(self.min_scale),
1412
- absmax / Float32(self.fp8_max),
1413
- )
1414
- bits = _recast_val(scale_raw, Uint32)
1415
- ue8m0 = ((bits + Uint32(0x7FFFFF)) >> Uint32(23)) & Uint32(0xFF)
1416
- inv_scale = _recast_val((Uint32(254) - ue8m0) << Uint32(23), Float32)
1417
- y0 = cutlass.min(
1418
- cute.arch.fmax(q0 * inv_scale, Float32(-self.fp8_max)),
1419
- Float32(self.fp8_max),
1547
+ warp_sum += cute.arch.shuffle_sync_bfly(warp_sum, offset)
1548
+
1549
+ smem = cutlass.utils.SmemAllocator()
1550
+ partial_sums = smem.allocate_tensor(
1551
+ Float32, cute.make_layout((num_warps,)), byte_alignment=4
1552
)
1421
- y1 = cutlass.min(
1422
- cute.arch.fmax(q1 * inv_scale, Float32(-self.fp8_max)),
1423
- Float32(self.fp8_max),
1553
+
1554
+ if lane_id == 0:
1555
+ partial_sums[warp_id] = warp_sum
1556
+ cute.arch.sync_threads()
1557
+
1558
+ total = partial_sums[lane_id % num_warps]
1559
+ sum_mask_and_clamp = const_expr(
1560
+ (cute.arch.WARP_SIZE - num_warps) << 8 | (cute.arch.WARP_SIZE - 1)
1561
)
1425
- packed_fp8 = _fp32x2_to_fp8e4m3x2(y0, y1)
1426
- out_base = value_base + (warp_id * self.quant_block + lane_id * 2).to(
1427
- Int64
1562
+ for step in cutlass.range_constexpr(3):
1563
+ offset = const_expr(4 >> step)
1564
+ total += cute.arch.shuffle_sync_bfly(
1565
+ total,
1566
+ offset,
1567
+ mask_and_clamp=sum_mask_and_clamp,
1568
+ )
1569
+
1570
+ rrms = cute.math.rsqrt(
1571
+ total / Float32(head_dim) + rms_norm_eps, fastmath=True
1572
)
1429
- k_cache_u16.iterator[out_base // Int64(2)] = packed_fp8
1430
- if lane_id == 0:
1431
- k_cache.iterator[scale_base + warp_id.to(Int64)] = ue8m0.to(Uint8)
1432
- if warp_id == 0:
1433
- k_cache.iterator[scale_base + Int64(self.nope_blocks)] = Uint8(
1434
- 0
1573
+ x0 = x0 * rrms * weight0
1574
+ x1 = x1 * rrms * weight1
1575
+
1576
+ if warp_id == nope_blocks:
1577
+ pair_idx = lane_id
1578
+ compressed_pos = position - Int64(compress_ratio - 1)
1579
+ cs_base = compressed_pos * cos_sin_cache.stride[0] + pair_idx.to(
1580
+ Int64
1581
+ )
1582
+ cos_v = cos_sin_cache.iterator[cs_base]
1583
+ sin_v = cos_sin_cache.iterator[cs_base + Int64(rope_dim // 2)]
1584
+ real = x0 * cos_v - x1 * sin_v
1585
+ imag = x0 * sin_v + x1 * cos_v
1586
+ packed = _fp32x2_to_bf16x2(real, imag)
1587
+ out_base = value_base + Int64(nope_dim) + (lane_id * 4).to(Int64)
1588
+ k_cache_u32.iterator[out_base // Int64(4)] = packed
1589
+ else:
1590
+ q_packed = _fp32x2_to_bf16x2(x0, x1)
1591
+ q0, q1 = _bf16x2_to_fp32(q_packed)
1592
+ abs0 = cute.math.absf(q0)
1593
+ abs1 = cute.math.absf(q1)
1594
+ local_absmax = cute.arch.fmax(abs0, abs1)
1595
+ absmax = local_absmax
1596
+ for step in cutlass.range_constexpr(5):
1597
+ offset = const_expr(16 >> step)
1598
+ absmax = cute.arch.fmax(
1599
+ absmax, cute.arch.shuffle_sync_bfly(absmax, offset)
1600
)
1601
+ scale_raw = cute.arch.fmax(
1602
+ Float32(min_scale),
1603
+ absmax / Float32(fp8_max),
1604
+ )
1605
+ bits = _recast_val(scale_raw, Uint32)
1606
+ ue8m0 = ((bits + Uint32(0x7FFFFF)) >> Uint32(23)) & Uint32(0xFF)
1607
+ inv_scale = _recast_val(
1608
+ (Uint32(254) - ue8m0) << Uint32(23), Float32
1609
+ )
1610
+ y0 = cutlass.min(
1611
+ cute.arch.fmax(q0 * inv_scale, Float32(-fp8_max)),
1612
+ Float32(fp8_max),
1613
+ )
1614
+ y1 = cutlass.min(
1615
+ cute.arch.fmax(q1 * inv_scale, Float32(-fp8_max)),
1616
+ Float32(fp8_max),
1617
+ )
1618
+ packed_fp8 = _fp32x2_to_fp8e4m3x2(y0, y1)
1619
+ out_base = value_base + (warp_id * quant_block + lane_id * 2).to(
1620
+ Int64
1621
+ )
1622
+ k_cache_u16.iterator[out_base // Int64(2)] = packed_fp8
1623
+ if lane_id == 0:
1624
+ k_cache.iterator[scale_base + warp_id.to(Int64)] = ue8m0.to(
1625
+ Uint8
1626
+ )
1627
+ if warp_id == 0:
1628
+ k_cache.iterator[scale_base + Int64(nope_blocks)] = Uint8(0)
1629
+
1630
+ @cute.jit
1631
+ def host_entrypoint(
1632
+ compressed_kv: cute.Tensor,
1633
+ positions: cute.Tensor,
1634
+ slot_mapping: cute.Tensor,
1635
+ rms_norm_weight: cute.Tensor,
1636
+ rms_norm_eps: Float32,
1637
+ cos_sin_cache: cute.Tensor,
1638
+ k_cache: cute.Tensor,
1639
+ kv_slot_mapping: cute.Tensor,
1640
+ stream: CUstream,
1641
+ ):
1642
+ grid = (slot_mapping.shape[0], 1, 1)
1643
+ device_kernel(
1644
+ compressed_kv,
1645
+ positions,
1646
+ slot_mapping,
1647
+ rms_norm_weight,
1648
+ rms_norm_eps,
1649
+ cos_sin_cache,
1650
+ k_cache,
1651
+ kv_slot_mapping,
1652
+ ).launch(grid=grid, block=(tb_size, 1, 1), stream=stream)
1653
+
1654
+ return host_entrypoint
1655
+
1656
+ def dispatch( # type: ignore[override]
1657
+ self,
1658
+ *,
1659
+ compress_ratio: int,
1660
+ cache_block_size: int,
1661
+ cache_alignment: int,
1662
+ norm_weight_dtype: type[cutlass.Numeric],
1663
+ head_size: int,
1664
+ rope_head_dim: int,
1665
+ runtime_kv_block_stride: int | None = None,
1666
+ ) -> CompileKey:
1667
+ raw_kv_cache_block_size = cache_block_size // compress_ratio
1668
+ kv_cache_block_size = (
1669
+ raw_kv_cache_block_size if raw_kv_cache_block_size >= 1 else 1
1670
+ )
1671
+ token_stride = 576
1672
+ scale_dim = 8
1673
+ return self.CompileKey(
1674
+ head_size=head_size,
1675
+ rope_head_dim=rope_head_dim,
1676
+ fp8_max=448.0,
1677
+ quant_block=64,
1678
+ token_stride=token_stride,
1679
+ scale_dim=scale_dim,
1680
+ kv_block_stride=(
1681
+ runtime_kv_block_stride
1682
+ if runtime_kv_block_stride is not None
1683
+ else round_up(
1684
+ kv_cache_block_size * (token_stride + scale_dim),
1685
+ cache_alignment,
1686
+ )
1687
+ ),
1688
+ compress_ratio=compress_ratio,
1689
+ norm_weight_dtype=norm_weight_dtype,
1690
+ kv_cache_block_size=kv_cache_block_size,
1691
+ )
1692
1437
- @cache
1438
- @staticmethod
1439
- def compile(
1440
- head_size: int = 512,
1441
- rope_head_dim: int = 64,
1442
- fp8_max: float = 448.0,
1443
- quant_block: int = 64,
1444
- token_stride: int = 576,
1445
- scale_dim: int = 8,
1446
- kv_block_stride: int = 74752,
1447
- compress_ratio: int = 128,
1448
- norm_weight_dtype: type[cutlass.Numeric] = Float32,
1449
- static_kv_cache_block_size: int = 0,
1450
- ):
1693
+ def get_warmup_keys(
1694
+ self,
1695
+ vllm_config: Any,
1696
+ *,
1697
+ k_cache_prefix: str,
1698
+ compress_ratio: int,
1699
+ ) -> list[CompileKey]:
1700
+ hf_config = vllm_config.model_config.hf_config
1701
+ cache_config = vllm_config.cache_config
1702
+ k_cache_layer = vllm_config.compilation_config.static_forward_context[
1703
+ k_cache_prefix
1704
+ ]
1705
+ kv_cache = k_cache_layer.kv_cache
1706
+ return self._trace_dispatch(self.dispatch)(
1707
+ compress_ratio=compress_ratio,
1708
+ cache_block_size=cache_config.block_size,
1709
+ cache_alignment=576,
1710
+ norm_weight_dtype=torch_to_cute_dtype(vllm_config.model_config.dtype),
1711
+ head_size=hf_config.head_dim,
1712
+ rope_head_dim=hf_config.qk_rope_head_dim,
1713
+ runtime_kv_block_stride=kv_cache.stride(0),
1714
+ enabled=(
1715
+ cache_config.block_size > 0
1716
+ and compress_ratio == 128
1717
+ and kv_cache.ndim == 3
1718
+ and cache_config.cache_dtype == "fp8_ds_mla"
1719
+ ),
1720
+ _when=lambda *, enabled: enabled,
1721
+ )
1722
+
1723
+ def warmup_inputs(self, compile_key: CompileKey) -> tuple[Any, ...]:
1724
+ head_size = compile_key.head_size
1725
+ rope_head_dim = compile_key.rope_head_dim
1726
+ quant_block = compile_key.quant_block
1727
+ token_stride = compile_key.token_stride
1728
+ scale_dim = compile_key.scale_dim
1729
+ kv_block_stride = compile_key.kv_block_stride
1730
+ norm_weight_dtype = compile_key.norm_weight_dtype
1731
+ static_kv_cache_block_size = compile_key.kv_cache_block_size
1732
if quant_block != 64:
1733
raise ValueError(
1453
- "CuTe DSL sparse-attn store currently requires quant_block=64."
1734
+ "CuTe DSL sparse-attn store currently requires quant_block=64.",
1735
)
1736
if rope_head_dim != 64:
1737
raise ValueError(
1457
- "CuTe DSL sparse-attn store currently requires rope_head_dim=64."
1738
+ "CuTe DSL sparse-attn store currently requires rope_head_dim=64.",
1739
)
1740
if head_size % quant_block != 0:
1741
raise ValueError("head_size must be divisible by quant_block.")
@@ -1466,14 +1747,13 @@ class SparseAttnNormRopeStoreKernel:
1747
if static_kv_cache_block_size <= 0:
1748
raise ValueError(
1749
"CuTe DSL sparse-attn store requires a positive static "
1469
- "kv_cache_block_size."
1750
+ "kv_cache_block_size.",
1751
)
1752
num_positions = cute.sym_int()
1753
num_slots = cute.sym_int()
1754
num_kv_slots = cute.sym_int()
1755
max_pos = cute.sym_int()
1756
num_blocks = cute.sym_int()
1476
-
1757
compressed_kv = cute.runtime.make_fake_tensor(
1758
Float32,
1759
(num_slots, head_size),
@@ -1483,7 +1763,9 @@ class SparseAttnNormRopeStoreKernel:
1763
positions = make_fake_tensor(Int64, (num_positions,), divisibility=8)
1764
slot_mapping = make_fake_tensor(Int64, (num_slots,), divisibility=8)
1765
rms_norm_weight = make_fake_tensor(
1486
- norm_weight_dtype, (head_size,), divisibility=4
1766
+ norm_weight_dtype,
1767
+ (head_size,),
1768
+ divisibility=4,
1769
)
1770
cos_sin_cache = cute.runtime.make_fake_tensor(
1771
Float32,
@@ -1494,28 +1776,11 @@ class SparseAttnNormRopeStoreKernel:
1776
k_cache = cute.runtime.make_fake_tensor(
1777
Uint8,
1778
(num_blocks, cute.sym_int(), cute.sym_int()),
1497
- stride=(
1498
- kv_block_stride,
1499
- cute.sym_int64(divisibility=8),
1500
- 1,
1501
- ),
1779
+ stride=(kv_block_stride, cute.sym_int64(divisibility=8), 1),
1780
assumed_align=16,
1781
)
1782
kv_slot_mapping = make_fake_tensor(Int64, (num_kv_slots,), divisibility=8)
1505
-
1506
- kernel = SparseAttnNormRopeStoreKernel(
1507
- head_size,
1508
- rope_head_dim,
1509
- fp8_max,
1510
- quant_block,
1511
- token_stride,
1512
- scale_dim,
1513
- compress_ratio,
1514
- static_kv_cache_block_size,
1515
- )
1516
- stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
1517
- return cute.compile(
1518
- kernel,
1783
+ return (
1784
compressed_kv,
1785
positions,
1786
slot_mapping,
@@ -1524,211 +1789,275 @@ class SparseAttnNormRopeStoreKernel:
1789
cos_sin_cache,
1790
k_cache,
1791
kv_slot_mapping,
1527
- stream,
1528
- options="--enable-tvm-ffi",
1792
)
1793
1531
-
1532
-class SparseAttnNormRopeStoreFullKernel:
1533
- def __init__(
1534
- self,
1535
- head_size: int,
1536
- rope_head_dim: int,
1537
- fp8_max: float,
1538
- quant_block: int,
1539
- token_stride: int,
1540
- scale_dim: int,
1541
- compress_ratio: int,
1542
- store_full_fp8: bool = False,
1543
- ):
1544
- # Standalone (not inheriting the #44230-restructured legacy kernel):
1545
- # set attrs directly so the full-cache C128 path is decoupled.
1546
- self.head_dim = head_size
1547
- self.rope_dim = rope_head_dim
1548
- self.nope_dim = head_size - rope_head_dim
1549
- self.fp8_max = fp8_max
1550
- self.quant_block = quant_block
1551
- self.token_stride = token_stride
1552
- self.scale_dim = scale_dim
1553
- self.num_warps = head_size // quant_block
1554
- self.nope_blocks = self.nope_dim // quant_block
1555
- self.tb_size = head_size // 2
1556
- self.compress_ratio = compress_ratio
1557
- self.store_full_fp8 = store_full_fp8
1558
-
1559
- @cute.jit
1794
+ @kernel_launcher
1795
def __call__(
1796
self,
1562
- compressed_kv: cute.Tensor,
1563
- positions: cute.Tensor,
1564
- slot_mapping: cute.Tensor,
1565
- rms_norm_weight: cute.Tensor,
1566
- rms_norm_eps: Float32,
1567
- cos_sin_cache: cute.Tensor,
1568
- k_cache: cute.Tensor,
1569
- kv_slot_mapping: cute.Tensor,
1570
- kv_cache_block_size: Int64,
1571
- fp8_scale: cute.Tensor,
1572
- stream: CUstream,
1573
- ):
1574
- grid = (slot_mapping.shape[0], 1, 1)
1575
- self.kernel(
1797
+ *,
1798
+ compressed_kv: torch.Tensor,
1799
+ positions: torch.Tensor,
1800
+ slot_mapping: torch.Tensor,
1801
+ rms_norm_weight: torch.Tensor,
1802
+ rms_norm_eps: float,
1803
+ cos_sin_cache: torch.Tensor,
1804
+ kv_cache: torch.Tensor,
1805
+ kv_slot_mapping: torch.Tensor,
1806
+ compress_ratio: int,
1807
+ head_dim: int,
1808
+ rope_head_dim: int,
1809
+ ) -> CuTeDSLLaunchSpec[SparseAttnNormRopeStoreKernel.CompileKey]:
1810
+ compile_key = self.dispatch(
1811
+ compress_ratio=compress_ratio,
1812
+ cache_block_size=kv_cache.shape[1] * compress_ratio,
1813
+ cache_alignment=576,
1814
+ norm_weight_dtype=torch_to_cute_dtype(rms_norm_weight.dtype),
1815
+ head_size=head_dim,
1816
+ rope_head_dim=rope_head_dim,
1817
+ runtime_kv_block_stride=kv_cache.stride(0),
1818
+ )
1819
+ launch_args = (
1820
compressed_kv,
1821
positions,
1822
slot_mapping,
1823
rms_norm_weight,
1824
rms_norm_eps,
1825
cos_sin_cache,
1582
- k_cache,
1826
+ kv_cache,
1827
kv_slot_mapping,
1584
- kv_cache_block_size,
1585
- fp8_scale,
1586
- ).launch(grid=grid, block=(self.tb_size, 1, 1), stream=stream)
1587
-
1588
- @cute.kernel
1589
- def kernel(
1590
- self,
1591
- compressed_kv: cute.Tensor,
1592
- positions: cute.Tensor,
1593
- slot_mapping: cute.Tensor,
1594
- rms_norm_weight: cute.Tensor,
1595
- rms_norm_eps: Float32,
1596
- cos_sin_cache: cute.Tensor,
1597
- k_cache: cute.Tensor,
1598
- kv_slot_mapping: cute.Tensor,
1599
- kv_cache_block_size: Int64,
1600
- fp8_scale: cute.Tensor,
1601
- ):
1602
- token_idx, _, _ = cute.arch.block_idx()
1603
- tid, _, _ = cute.arch.thread_idx()
1604
- warp_id = cute.arch.make_warp_uniform(tid // 32)
1605
- lane_id = tid % 32
1606
- elem0 = tid * 2
1607
-
1608
- slot_id = slot_mapping[token_idx]
1609
- has_position = token_idx < positions.shape[0]
1610
- position = Int64(0)
1611
- if has_position:
1612
- position = positions[token_idx]
1613
- boundary = has_position and (
1614
- (position + Int64(1)) % Int64(self.compress_ratio) == Int64(0)
1828
)
1616
- has_kv_slot_idx = token_idx < kv_slot_mapping.shape[0]
1617
- kv_slot_idx = Int64(-1)
1618
- if has_kv_slot_idx:
1619
- kv_slot_idx = kv_slot_mapping[token_idx]
1620
- active = slot_id >= Int64(0) and boundary and kv_slot_idx >= Int64(0)
1621
-
1622
- if active:
1623
- base = token_idx.to(Int64) * compressed_kv.stride[0] + elem0.to(Int64)
1624
- x0 = compressed_kv.iterator[base]
1625
- x1 = compressed_kv.iterator[base + Int64(1)]
1626
-
1627
- local_sumsq = x0 * x0 + x1 * x1
1628
- warp_sum = local_sumsq
1629
- for step in cutlass.range_constexpr(5):
1630
- offset = const_expr(16 >> step)
1631
- warp_sum += cute.arch.shuffle_sync_bfly(warp_sum, offset)
1632
-
1633
- smem = cutlass.utils.SmemAllocator()
1634
- partial_sums = smem.allocate_tensor(
1635
- Float32, cute.make_layout((self.num_warps,)), byte_alignment=4
1636
- )
1637
- rrms_shared = smem.allocate_tensor(
1638
- Float32, cute.make_layout((1,)), byte_alignment=4
1639
- )
1829
+ return compile_key, launch_args
1830
1641
- if lane_id == 0:
1642
- partial_sums[warp_id] = warp_sum
1643
- cute.arch.sync_threads()
1644
- if tid == 0:
1645
- total = Float32(0.0)
1646
- for i in cutlass.range_constexpr(self.num_warps):
1647
- total += partial_sums[i]
1648
- rrms_shared[0] = cute.math.rsqrt(
1649
- total / Float32(self.head_dim) + rms_norm_eps, fastmath=True
1650
- )
1651
- cute.arch.sync_threads()
1831
1653
- rrms = rrms_shared[0]
1654
- x0 = x0 * rrms * rms_norm_weight[elem0].to(Float32)
1655
- x1 = x1 * rrms * rms_norm_weight[elem0 + 1].to(Float32)
1832
+class SparseAttnNormRopeStoreFullKernel(
1833
+ VllmCuTeDSLJitKernel["SparseAttnNormRopeStoreFullKernel.CompileKey"]
1834
+):
1835
+ @dataclass(frozen=True)
1836
+ class CompileKey:
1837
+ head_size: int
1838
+ rope_head_dim: int
1839
+ fp8_max: float
1840
+ quant_block: int
1841
+ compress_ratio: int
1842
+ store_full_fp8: bool
1843
+ norm_weight_dtype: type[cutlass.Numeric]
1844
1657
- page = kv_slot_idx // kv_cache_block_size
1658
- kv_offset = kv_slot_idx - page * kv_cache_block_size
1659
- value_base = page * k_cache.stride[0] + kv_offset * k_cache.stride[1]
1845
+ @staticmethod
1846
+ def kernel(compile_key: CompileKey) -> Any:
1847
+ head_dim = compile_key.head_size
1848
+ rope_dim = compile_key.rope_head_dim
1849
+ nope_dim = head_dim - rope_dim
1850
+ fp8_max = compile_key.fp8_max
1851
+ quant_block = compile_key.quant_block
1852
+ num_warps = head_dim // quant_block
1853
+ nope_blocks = nope_dim // quant_block
1854
+ tb_size = head_dim // 2
1855
+ compress_ratio = compile_key.compress_ratio
1856
+ store_full_fp8 = compile_key.store_full_fp8
1857
+
1858
+ @cute.kernel
1859
+ def device_kernel(
1860
+ compressed_kv: cute.Tensor,
1861
+ positions: cute.Tensor,
1862
+ slot_mapping: cute.Tensor,
1863
+ rms_norm_weight: cute.Tensor,
1864
+ rms_norm_eps: Float32,
1865
+ cos_sin_cache: cute.Tensor,
1866
+ k_cache: cute.Tensor,
1867
+ kv_slot_mapping: cute.Tensor,
1868
+ kv_cache_block_size: Int64,
1869
+ fp8_scale: cute.Tensor,
1870
+ ):
1871
+ token_idx, _, _ = cute.arch.block_idx()
1872
+ tid, _, _ = cute.arch.thread_idx()
1873
+ warp_id = cute.arch.make_warp_uniform(tid // 32)
1874
+ lane_id = tid % 32
1875
+ elem0 = tid * 2
1876
1661
- if const_expr(self.store_full_fp8):
1662
- k_cache_u16 = cute.recast_tensor(k_cache, Uint16)
1663
- inv_fp8 = Float32(1.0) / fp8_scale[0]
1664
- fp8_v0 = x0
1665
- fp8_v1 = x1
1666
- if warp_id == self.nope_blocks:
1667
- compressed_pos = (position // Int64(self.compress_ratio)) * Int64(
1668
- self.compress_ratio
1669
- )
1670
- pair_idx = lane_id
1671
- cs_base = compressed_pos * cos_sin_cache.stride[0] + pair_idx.to(
1672
- Int64
1673
- )
1674
- cos_v = cos_sin_cache.iterator[cs_base]
1675
- sin_v = cos_sin_cache.iterator[cs_base + Int64(self.rope_dim // 2)]
1676
- fp8_v0 = x0 * cos_v - x1 * sin_v
1677
- fp8_v1 = x0 * sin_v + x1 * cos_v
1678
- fp8_packed_bf16 = _fp32x2_to_bf16x2(fp8_v0, fp8_v1)
1679
- b0, b1 = _bf16x2_to_fp32(fp8_packed_bf16)
1680
- y0 = cutlass.min(
1681
- cutlass.max(b0 * inv_fp8, Float32(-self.fp8_max)),
1682
- Float32(self.fp8_max),
1877
+ slot_id = slot_mapping[token_idx]
1878
+ has_position = token_idx < positions.shape[0]
1879
+ position = Int64(0)
1880
+ if has_position:
1881
+ position = positions[token_idx]
1882
+ boundary = has_position and (
1883
+ (position + Int64(1)) % Int64(compress_ratio) == Int64(0)
1884
+ )
1885
+ has_kv_slot_idx = token_idx < kv_slot_mapping.shape[0]
1886
+ kv_slot_idx = Int64(-1)
1887
+ if has_kv_slot_idx:
1888
+ kv_slot_idx = kv_slot_mapping[token_idx]
1889
+ active = slot_id >= Int64(0) and boundary and kv_slot_idx >= Int64(0)
1890
+
1891
+ if active:
1892
+ base = token_idx.to(Int64) * compressed_kv.stride[0] + elem0.to(Int64)
1893
+ x0 = compressed_kv.iterator[base]
1894
+ x1 = compressed_kv.iterator[base + Int64(1)]
1895
+
1896
+ local_sumsq = x0 * x0 + x1 * x1
1897
+ warp_sum = local_sumsq
1898
+ for step in cutlass.range_constexpr(5):
1899
+ offset = const_expr(16 >> step)
1900
+ warp_sum += cute.arch.shuffle_sync_bfly(warp_sum, offset)
1901
+
1902
+ smem = cutlass.utils.SmemAllocator()
1903
+ partial_sums = smem.allocate_tensor(
1904
+ Float32, cute.make_layout((num_warps,)), byte_alignment=4
1905
)
1684
- y1 = cutlass.min(
1685
- cutlass.max(b1 * inv_fp8, Float32(-self.fp8_max)),
1686
- Float32(self.fp8_max),
1906
+ rrms_shared = smem.allocate_tensor(
1907
+ Float32, cute.make_layout((1,)), byte_alignment=4
1908
)
1688
- packed_fp8 = _fp32x2_to_fp8e4m3x2(y0, y1)
1689
- out_base = value_base + elem0.to(Int64)
1690
- k_cache_u16.iterator[out_base // Int64(2)] = packed_fp8
1691
- else:
1692
- k_cache_u32 = cute.recast_tensor(k_cache, Uint32)
1693
- bf16_v0 = x0
1694
- bf16_v1 = x1
1695
- if warp_id == self.nope_blocks:
1696
- compressed_pos = (position // Int64(self.compress_ratio)) * Int64(
1697
- self.compress_ratio
1909
+
1910
+ if lane_id == 0:
1911
+ partial_sums[warp_id] = warp_sum
1912
+ cute.arch.sync_threads()
1913
+ if tid == 0:
1914
+ total = Float32(0.0)
1915
+ for i in cutlass.range_constexpr(num_warps):
1916
+ total += partial_sums[i]
1917
+ rrms_shared[0] = cute.math.rsqrt(
1918
+ total / Float32(head_dim) + rms_norm_eps, fastmath=True
1919
)
1699
- pair_idx = lane_id
1700
- cs_base = compressed_pos * cos_sin_cache.stride[0] + pair_idx.to(
1701
- Int64
1920
+ cute.arch.sync_threads()
1921
+
1922
+ rrms = rrms_shared[0]
1923
+ x0 = x0 * rrms * rms_norm_weight[elem0].to(Float32)
1924
+ x1 = x1 * rrms * rms_norm_weight[elem0 + 1].to(Float32)
1925
+
1926
+ page = kv_slot_idx // kv_cache_block_size
1927
+ kv_offset = kv_slot_idx - page * kv_cache_block_size
1928
+ value_base = page * k_cache.stride[0] + kv_offset * k_cache.stride[1]
1929
+
1930
+ if const_expr(store_full_fp8):
1931
+ k_cache_u16 = cute.recast_tensor(k_cache, Uint16)
1932
+ inv_fp8 = Float32(1.0) / fp8_scale[0]
1933
+ fp8_v0 = x0
1934
+ fp8_v1 = x1
1935
+ if warp_id == nope_blocks:
1936
+ compressed_pos = (position // Int64(compress_ratio)) * Int64(
1937
+ compress_ratio
1938
+ )
1939
+ pair_idx = lane_id
1940
+ cs_base = compressed_pos * cos_sin_cache.stride[
1941
+ 0
1942
+ ] + pair_idx.to(Int64)
1943
+ cos_v = cos_sin_cache.iterator[cs_base]
1944
+ sin_v = cos_sin_cache.iterator[cs_base + Int64(rope_dim // 2)]
1945
+ fp8_v0 = x0 * cos_v - x1 * sin_v
1946
+ fp8_v1 = x0 * sin_v + x1 * cos_v
1947
+ fp8_packed_bf16 = _fp32x2_to_bf16x2(fp8_v0, fp8_v1)
1948
+ b0, b1 = _bf16x2_to_fp32(fp8_packed_bf16)
1949
+ y0 = cutlass.min(
1950
+ cutlass.max(b0 * inv_fp8, Float32(-fp8_max)),
1951
+ Float32(fp8_max),
1952
)
1703
- cos_v = cos_sin_cache.iterator[cs_base]
1704
- sin_v = cos_sin_cache.iterator[cs_base + Int64(self.rope_dim // 2)]
1705
- bf16_v0 = x0 * cos_v - x1 * sin_v
1706
- bf16_v1 = x0 * sin_v + x1 * cos_v
1707
- bf16_packed = _fp32x2_to_bf16x2(bf16_v0, bf16_v1)
1708
- out_base = value_base + (elem0 * 2).to(Int64)
1709
- k_cache_u32.iterator[out_base // Int64(4)] = bf16_packed
1710
-
1711
- @cache
1712
- @staticmethod
1713
- def compile(
1714
- head_size: int = 512,
1715
- rope_head_dim: int = 64,
1716
- fp8_max: float = 448.0,
1717
- quant_block: int = 64,
1718
- token_stride: int = 576,
1719
- scale_dim: int = 8,
1720
- kv_block_stride: int = 74752,
1721
- compress_ratio: int = 128,
1722
- store_full_fp8: bool = False,
1723
- norm_weight_dtype: type[cutlass.Numeric] = Float32,
1724
- ):
1953
+ y1 = cutlass.min(
1954
+ cutlass.max(b1 * inv_fp8, Float32(-fp8_max)),
1955
+ Float32(fp8_max),
1956
+ )
1957
+ packed_fp8 = _fp32x2_to_fp8e4m3x2(y0, y1)
1958
+ out_base = value_base + elem0.to(Int64)
1959
+ k_cache_u16.iterator[out_base // Int64(2)] = packed_fp8
1960
+ else:
1961
+ k_cache_u32 = cute.recast_tensor(k_cache, Uint32)
1962
+ bf16_v0 = x0
1963
+ bf16_v1 = x1
1964
+ if warp_id == nope_blocks:
1965
+ compressed_pos = (position // Int64(compress_ratio)) * Int64(
1966
+ compress_ratio
1967
+ )
1968
+ pair_idx = lane_id
1969
+ cs_base = compressed_pos * cos_sin_cache.stride[
1970
+ 0
1971
+ ] + pair_idx.to(Int64)
1972
+ cos_v = cos_sin_cache.iterator[cs_base]
1973
+ sin_v = cos_sin_cache.iterator[cs_base + Int64(rope_dim // 2)]
1974
+ bf16_v0 = x0 * cos_v - x1 * sin_v
1975
+ bf16_v1 = x0 * sin_v + x1 * cos_v
1976
+ bf16_packed = _fp32x2_to_bf16x2(bf16_v0, bf16_v1)
1977
+ out_base = value_base + (elem0 * 2).to(Int64)
1978
+ k_cache_u32.iterator[out_base // Int64(4)] = bf16_packed
1979
+
1980
+ @cute.jit
1981
+ def host_entrypoint(
1982
+ compressed_kv: cute.Tensor,
1983
+ positions: cute.Tensor,
1984
+ slot_mapping: cute.Tensor,
1985
+ rms_norm_weight: cute.Tensor,
1986
+ rms_norm_eps: Float32,
1987
+ cos_sin_cache: cute.Tensor,
1988
+ k_cache: cute.Tensor,
1989
+ kv_slot_mapping: cute.Tensor,
1990
+ kv_cache_block_size: Int64,
1991
+ fp8_scale: cute.Tensor,
1992
+ stream: CUstream,
1993
+ ):
1994
+ grid = (slot_mapping.shape[0], 1, 1)
1995
+ device_kernel(
1996
+ compressed_kv,
1997
+ positions,
1998
+ slot_mapping,
1999
+ rms_norm_weight,
2000
+ rms_norm_eps,
2001
+ cos_sin_cache,
2002
+ k_cache,
2003
+ kv_slot_mapping,
2004
+ kv_cache_block_size,
2005
+ fp8_scale,
2006
+ ).launch(grid=grid, block=(tb_size, 1, 1), stream=stream)
2007
+
2008
+ return host_entrypoint
2009
+
2010
+ def dispatch( # type: ignore[override]
2011
+ self,
2012
+ *,
2013
+ store_full_fp8: bool,
2014
+ norm_weight_dtype: type[cutlass.Numeric],
2015
+ head_size: int,
2016
+ rope_head_dim: int,
2017
+ compress_ratio: int,
2018
+ ) -> CompileKey:
2019
+ return self.CompileKey(
2020
+ head_size=head_size,
2021
+ rope_head_dim=rope_head_dim,
2022
+ compress_ratio=compress_ratio,
2023
+ fp8_max=448.0,
2024
+ quant_block=64,
2025
+ store_full_fp8=store_full_fp8,
2026
+ norm_weight_dtype=norm_weight_dtype,
2027
+ )
2028
+
2029
+ def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]:
2030
+ hf_config = vllm_config.model_config.hf_config
2031
+ cache_config = vllm_config.cache_config
2032
+ return self._trace_dispatch(self.dispatch)(
2033
+ compress_ratio=128,
2034
+ store_full_fp8=(
2035
+ cache_config.cache_dtype.startswith("fp8")
2036
+ and cache_config.cache_dtype != "fp8_ds_mla"
2037
+ ),
2038
+ norm_weight_dtype=torch_to_cute_dtype(vllm_config.model_config.dtype),
2039
+ head_size=hf_config.head_dim,
2040
+ rope_head_dim=hf_config.qk_rope_head_dim,
2041
+ enabled=(
2042
+ cache_config.block_size > 0
2043
+ and 128 in hf_config.compress_ratios
2044
+ and cache_config.cache_dtype != "fp8_ds_mla"
2045
+ ),
2046
+ _when=lambda *, enabled: enabled,
2047
+ )
2048
+
2049
+ def warmup_inputs(self, compile_key: CompileKey) -> tuple[Any, ...]:
2050
+ head_size = compile_key.head_size
2051
+ rope_head_dim = compile_key.rope_head_dim
2052
+ quant_block = compile_key.quant_block
2053
+ norm_weight_dtype = compile_key.norm_weight_dtype
2054
if quant_block != 64:
2055
raise ValueError(
1727
- "CuTe DSL sparse-attn store currently requires quant_block=64."
2056
+ "CuTe DSL sparse-attn store currently requires quant_block=64.",
2057
)
2058
if rope_head_dim != 64:
2059
raise ValueError(
1731
- "CuTe DSL sparse-attn store currently requires rope_head_dim=64."
2060
+ "CuTe DSL sparse-attn store currently requires rope_head_dim=64.",
2061
)
2062
if head_size % quant_block != 0:
2063
raise ValueError("head_size must be divisible by quant_block.")
@@ -1737,7 +2066,6 @@ class SparseAttnNormRopeStoreFullKernel:
2066
num_kv_slots = cute.sym_int()
2067
max_pos = cute.sym_int()
2068
num_blocks = cute.sym_int()
1740
-
2069
compressed_kv = cute.runtime.make_fake_tensor(
2070
Float32,
2071
(num_slots, head_size),
@@ -1747,7 +2075,9 @@ class SparseAttnNormRopeStoreFullKernel:
2075
positions = make_fake_tensor(Int64, (num_positions,), divisibility=8)
2076
slot_mapping = make_fake_tensor(Int64, (num_slots,), divisibility=8)
2077
rms_norm_weight = make_fake_tensor(
1750
- norm_weight_dtype, (head_size,), divisibility=4
2078
+ norm_weight_dtype,
2079
+ (head_size,),
2080
+ divisibility=4,
2081
)
2082
cos_sin_cache = cute.runtime.make_fake_tensor(
2083
Float32,
@@ -1758,29 +2088,12 @@ class SparseAttnNormRopeStoreFullKernel:
2088
k_cache = cute.runtime.make_fake_tensor(
2089
Uint8,
2090
(num_blocks, cute.sym_int(), cute.sym_int()),
1761
- stride=(
1762
- cute.sym_int64(divisibility=16),
1763
- cute.sym_int64(divisibility=8),
1764
- 1,
1765
- ),
2091
+ stride=(cute.sym_int64(divisibility=16), cute.sym_int64(divisibility=8), 1),
2092
assumed_align=16,
2093
)
2094
kv_slot_mapping = make_fake_tensor(Int64, (num_kv_slots,), divisibility=8)
2095
fp8_scale = make_fake_tensor(Float32, (1,), divisibility=1)
1770
-
1771
- kernel = SparseAttnNormRopeStoreFullKernel(
1772
- head_size,
1773
- rope_head_dim,
1774
- fp8_max,
1775
- quant_block,
1776
- token_stride,
1777
- scale_dim,
1778
- compress_ratio,
1779
- store_full_fp8,
1780
- )
1781
- stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
1782
- return cute.compile(
1783
- kernel,
2096
+ return (
2097
compressed_kv,
2098
positions,
2099
slot_mapping,
@@ -1791,83 +2104,57 @@ class SparseAttnNormRopeStoreFullKernel:
2104
kv_slot_mapping,
2105
Int64(0),
2106
fp8_scale,
1794
- stream,
1795
- options="--enable-tvm-ffi",
2107
)
2108
1798
-
1799
-def compile_split_sparse_attn_cutedsl(
1800
- head_size: int,
1801
- state_width: int,
1802
- block_size: int,
1803
- rope_head_dim: int,
1804
- fp8_max: float,
1805
- quant_block: int,
1806
- token_stride: int,
1807
- scale_dim: int,
1808
- kv_cache_block_size: int,
1809
- kv_block_stride: int,
1810
- compress_ratio: int,
1811
- overlap: bool,
1812
- rms_norm_weight_dtype: torch.dtype,
1813
- store_full_kv: bool = False,
1814
- store_full_fp8: bool = False,
1815
-):
1816
- if not (
1817
- head_size == 512
1818
- and state_width == head_size
1819
- and compress_ratio == 128
1820
- and not overlap
1821
- and block_size == 8
1822
- ):
1823
- raise ValueError(
1824
- "CuTe DSL split sparse-attn wrapper only supports the real "
1825
- "DeepSeek V4 C128 layout: head_size=512, state_width=512, "
1826
- "compress_ratio=128, overlap=False, block_size=8."
1827
- )
1828
- compress = SparseAttnCompressC128Block8Kernel.compile(
1829
- head_size=head_size,
1830
- state_width=state_width,
1831
- )
1832
- norm_weight_dtype = _TORCH_TO_CUTE[rms_norm_weight_dtype]
1833
- if store_full_kv:
1834
- # FlashInfer contiguous bf16/fp8 cache: standalone full-cache store.
1835
- store = SparseAttnNormRopeStoreFullKernel.compile(
1836
- head_size=head_size,
1837
- rope_head_dim=rope_head_dim,
1838
- fp8_max=fp8_max,
1839
- quant_block=quant_block,
1840
- token_stride=token_stride,
1841
- scale_dim=scale_dim,
1842
- kv_block_stride=kv_block_stride,
2109
+ @kernel_launcher
2110
+ def __call__(
2111
+ self,
2112
+ *,
2113
+ compressed_kv: torch.Tensor,
2114
+ positions: torch.Tensor,
2115
+ slot_mapping: torch.Tensor,
2116
+ rms_norm_weight: torch.Tensor,
2117
+ rms_norm_eps: float,
2118
+ cos_sin_cache: torch.Tensor,
2119
+ kv_cache: torch.Tensor,
2120
+ kv_slot_mapping: torch.Tensor,
2121
+ compress_ratio: int,
2122
+ store_full_fp8: bool,
2123
+ fp8_scale: torch.Tensor | None,
2124
+ head_dim: int,
2125
+ rope_head_dim: int,
2126
+ ) -> CuTeDSLLaunchSpec[SparseAttnNormRopeStoreFullKernel.CompileKey]:
2127
+ compile_key = self.dispatch(
2128
compress_ratio=compress_ratio,
2129
store_full_fp8=store_full_fp8,
1845
- norm_weight_dtype=norm_weight_dtype,
2130
+ norm_weight_dtype=torch_to_cute_dtype(rms_norm_weight.dtype),
2131
+ head_size=head_dim,
2132
+ rope_head_dim=rope_head_dim,
2133
)
1847
- else:
1848
- store = SparseAttnNormRopeStoreKernel.compile(
1849
- head_size,
1850
- rope_head_dim,
1851
- fp8_max,
1852
- quant_block,
1853
- token_stride,
1854
- scale_dim,
1855
- kv_block_stride,
1856
- compress_ratio,
1857
- norm_weight_dtype,
1858
- kv_cache_block_size,
2134
+ if fp8_scale is None:
2135
+ fp8_scale = torch.ones(1, dtype=torch.float32, device=kv_cache.device)
2136
+ launch_args = (
2137
+ compressed_kv,
2138
+ positions,
2139
+ slot_mapping,
2140
+ rms_norm_weight,
2141
+ rms_norm_eps,
2142
+ cos_sin_cache,
2143
+ kv_cache.view(torch.uint8),
2144
+ kv_slot_mapping,
2145
+ kv_cache.shape[1],
2146
+ fp8_scale,
2147
)
1860
- return compress, store
2148
+ return compile_key, launch_args
2149
2150
1863
-def split_kv_compress_norm_rope_insert_sparse_attn_cutedsl(
2151
+def fused_kv_compress_norm_rope_insert_sparse_attn_cutedsl(
2152
state_cache: torch.Tensor,
2153
token_to_req_indices: torch.Tensor,
2154
positions: torch.Tensor,
2155
slot_mapping: torch.Tensor,
2156
block_table: torch.Tensor,
2157
block_size: int,
1870
- compressed_kv: torch.Tensor,
2158
rms_norm_weight: torch.Tensor,
2159
rms_norm_eps: float,
2160
cos_sin_cache: torch.Tensor,
@@ -1876,102 +2163,92 @@ def split_kv_compress_norm_rope_insert_sparse_attn_cutedsl(
2163
kv_cache_block_size: int,
2164
kv_block_stride: int,
2165
head_size: int = 512,
1879
- state_width: int = 512,
2166
+ state_width: int = 1024,
2167
rope_head_dim: int = 64,
2168
fp8_max: float = 448.0,
2169
quant_block: int = 64,
2170
token_stride: int = 576,
2171
scale_dim: int = 8,
1885
- compress_ratio: int = 128,
1886
- overlap: bool = False,
2172
+ compress_ratio: int = 4,
2173
+ overlap: bool = True,
2174
store_full_kv: bool = False,
2175
store_full_fp8: bool = False,
2176
fp8_scale: torch.Tensor | None = None,
2177
) -> None:
2178
+ del fp8_max, quant_block, token_stride, scale_dim, kv_block_stride
2179
+ if not (
2180
+ head_size == 512
2181
+ and state_width == 2 * head_size
2182
+ and compress_ratio == 4
2183
+ and overlap
2184
+ ):
2185
+ raise ValueError(
2186
+ "CuTe DSL fused sparse-attn wrapper only supports the real "
2187
+ "DeepSeek V4 C4 layout: head_size=512, state_width=1024, "
2188
+ "compress_ratio=4, overlap=True."
2189
+ )
2190
if k_cache.ndim != 3:
2191
raise ValueError(
1893
- "CuTe DSL sparse-attn store expects the real DeepSeek V4 "
2192
+ "CuTe DSL sparse-attn fused store expects the real DeepSeek V4 "
2193
f"3D k_cache layout [num_blocks, block_size, 584], got ndim={k_cache.ndim}."
2194
)
1896
- if not store_full_kv and kv_cache_block_size != k_cache.shape[1]:
2195
+ if kv_cache_block_size != k_cache.shape[1]:
2196
raise ValueError(
1898
- "CuTe DSL split sparse-attn wrapper expected kv_cache_block_size "
2197
+ "CuTe DSL fused sparse-attn wrapper expected kv_cache_block_size "
2198
f"to match k_cache.shape[1], got {kv_cache_block_size} and "
2199
f"{k_cache.shape[1]}."
2200
)
2201
if positions.numel() == 0:
2202
return
1904
- if rms_norm_weight.dtype not in _TORCH_TO_CUTE:
1905
- raise ValueError(
1906
- "CuTe DSL sparse-attn store supports rms_norm_weight dtype "
1907
- f"bf16/fp32, got {rms_norm_weight.dtype}."
1908
- )
2203
if store_full_fp8 and not store_full_kv:
2204
raise ValueError("store_full_fp8 requires store_full_kv.")
1911
- compress, store = compile_split_sparse_attn_cutedsl(
1912
- head_size,
1913
- state_width,
1914
- block_size,
1915
- rope_head_dim,
1916
- fp8_max,
1917
- quant_block,
1918
- token_stride,
1919
- scale_dim,
1920
- kv_cache_block_size,
1921
- kv_block_stride,
1922
- compress_ratio,
1923
- overlap,
1924
- rms_norm_weight.dtype,
1925
- store_full_kv=store_full_kv,
1926
- store_full_fp8=store_full_fp8,
1927
- )
1928
- compress(
1929
- state_cache,
1930
- token_to_req_indices,
1931
- positions,
1932
- slot_mapping,
1933
- block_table,
1934
- compressed_kv,
1935
- )
2205
2206
if store_full_kv:
1938
- # Byte-addressed contiguous cache; block size + per-tensor scale are
1939
- # passed at call time (not baked into compile).
1940
- if fp8_scale is None:
1941
- fp8_scale = torch.ones(1, dtype=torch.float32, device=k_cache.device)
1942
- store(
1943
- compressed_kv,
1944
- positions,
1945
- slot_mapping,
1946
- rms_norm_weight,
1947
- rms_norm_eps,
1948
- cos_sin_cache,
1949
- k_cache.view(torch.uint8),
1950
- kv_slot_mapping,
1951
- kv_cache_block_size,
1952
- fp8_scale,
2207
+ _SPARSE_ATTN_COMPRESS_NORM_ROPE_STORE_FULL_C4_KERNEL(
2208
+ state_cache=state_cache,
2209
+ token_to_req_indices=token_to_req_indices,
2210
+ positions=positions,
2211
+ slot_mapping=slot_mapping,
2212
+ block_table=block_table,
2213
+ block_size=block_size,
2214
+ rms_norm_weight=rms_norm_weight,
2215
+ rms_norm_eps=rms_norm_eps,
2216
+ cos_sin_cache=cos_sin_cache,
2217
+ kv_cache=k_cache,
2218
+ kv_slot_mapping=kv_slot_mapping,
2219
+ compress_ratio=compress_ratio,
2220
+ store_full_fp8=store_full_fp8,
2221
+ fp8_scale=fp8_scale,
2222
+ head_dim=head_size,
2223
+ rope_head_dim=rope_head_dim,
2224
)
2225
return
1955
-
1956
- store(
1957
- compressed_kv,
1958
- positions,
1959
- slot_mapping,
1960
- rms_norm_weight,
1961
- rms_norm_eps,
1962
- cos_sin_cache,
1963
- k_cache,
1964
- kv_slot_mapping,
2226
+ _SPARSE_ATTN_COMPRESS_NORM_ROPE_STORE_C4_KERNEL(
2227
+ state_cache=state_cache,
2228
+ token_to_req_indices=token_to_req_indices,
2229
+ positions=positions,
2230
+ slot_mapping=slot_mapping,
2231
+ block_table=block_table,
2232
+ block_size=block_size,
2233
+ rms_norm_weight=rms_norm_weight,
2234
+ rms_norm_eps=rms_norm_eps,
2235
+ cos_sin_cache=cos_sin_cache,
2236
+ kv_cache=k_cache,
2237
+ kv_slot_mapping=kv_slot_mapping,
2238
+ compress_ratio=compress_ratio,
2239
+ head_dim=head_size,
2240
+ rope_head_dim=rope_head_dim,
2241
)
2242
2243
1968
-def fused_kv_compress_norm_rope_insert_sparse_attn_cutedsl(
2244
+def split_kv_compress_norm_rope_insert_sparse_attn_cutedsl(
2245
state_cache: torch.Tensor,
2246
token_to_req_indices: torch.Tensor,
2247
positions: torch.Tensor,
2248
slot_mapping: torch.Tensor,
2249
block_table: torch.Tensor,
2250
block_size: int,
2251
+ compressed_kv: torch.Tensor,
2252
rms_norm_weight: torch.Tensor,
2253
rms_norm_eps: float,
2254
cos_sin_cache: torch.Tensor,
@@ -1980,185 +2257,247 @@ def fused_kv_compress_norm_rope_insert_sparse_attn_cutedsl(
2257
kv_cache_block_size: int,
2258
kv_block_stride: int,
2259
head_size: int = 512,
1983
- state_width: int = 1024,
2260
+ state_width: int = 512,
2261
rope_head_dim: int = 64,
2262
fp8_max: float = 448.0,
2263
quant_block: int = 64,
2264
token_stride: int = 576,
2265
scale_dim: int = 8,
1989
- compress_ratio: int = 4,
1990
- overlap: bool = True,
2266
+ compress_ratio: int = 128,
2267
+ overlap: bool = False,
2268
store_full_kv: bool = False,
2269
store_full_fp8: bool = False,
2270
fp8_scale: torch.Tensor | None = None,
2271
) -> None:
1995
- if positions.numel() == 0:
1996
- return
1997
- norm_weight_dtype = _TORCH_TO_CUTE.get(rms_norm_weight.dtype)
1998
- if norm_weight_dtype is None:
2272
+ del fp8_max, quant_block, token_stride, scale_dim, kv_block_stride
2273
+ if not (
2274
+ head_size == 512
2275
+ and state_width == head_size
2276
+ and compress_ratio == 128
2277
+ and not overlap
2278
+ and block_size == 8
2279
+ ):
2280
raise ValueError(
2000
- "CuTe DSL sparse-attn fused store supports rms_norm_weight dtype "
2001
- f"bf16/fp32, got {rms_norm_weight.dtype}."
2281
+ "CuTe DSL split sparse-attn wrapper only supports the real "
2282
+ "DeepSeek V4 C128 layout: head_size=512, state_width=512, "
2283
+ "compress_ratio=128, overlap=False, block_size=8."
2284
)
2285
if k_cache.ndim != 3:
2286
raise ValueError(
2005
- "CuTe DSL sparse-attn fused store expects the real DeepSeek V4 "
2287
+ "CuTe DSL sparse-attn store expects the real DeepSeek V4 "
2288
f"3D k_cache layout [num_blocks, block_size, 584], got ndim={k_cache.ndim}."
2289
)
2290
+ if not store_full_kv and kv_cache_block_size != k_cache.shape[1]:
2291
+ raise ValueError(
2292
+ "CuTe DSL split sparse-attn wrapper expected kv_cache_block_size "
2293
+ f"to match k_cache.shape[1], got {kv_cache_block_size} and "
2294
+ f"{k_cache.shape[1]}."
2295
+ )
2296
+ if positions.numel() == 0:
2297
+ return
2298
if store_full_fp8 and not store_full_kv:
2299
raise ValueError("store_full_fp8 requires store_full_kv.")
2300
+
2301
+ _SPARSE_ATTN_COMPRESS_C128_BLOCK8_KERNEL(
2302
+ state_cache=state_cache,
2303
+ num_actual=compressed_kv.shape[0],
2304
+ token_to_req_indices=token_to_req_indices,
2305
+ positions=positions,
2306
+ slot_mapping=slot_mapping,
2307
+ block_table=block_table,
2308
+ head_dim=head_size,
2309
+ compressed_kv=compressed_kv,
2310
+ )
2311
if store_full_kv:
2011
- # FlashInfer contiguous bf16/fp8 cache: byte-addressed full-cache C4 store.
2012
- if fp8_scale is None:
2013
- fp8_scale = torch.ones(1, dtype=torch.float32, device=k_cache.device)
2014
- compiled = SparseAttnCompressNormRopeStoreFullC4Kernel.compile(
2015
- head_size=head_size,
2016
- state_width=state_width,
2017
- rope_head_dim=rope_head_dim,
2018
- fp8_max=fp8_max,
2019
- quant_block=quant_block,
2020
- token_stride=token_stride,
2021
- scale_dim=scale_dim,
2022
- kv_block_stride=kv_block_stride,
2312
+ _SPARSE_ATTN_NORM_ROPE_STORE_FULL_KERNEL(
2313
+ compressed_kv=compressed_kv,
2314
+ positions=positions,
2315
+ slot_mapping=slot_mapping,
2316
+ rms_norm_weight=rms_norm_weight,
2317
+ rms_norm_eps=rms_norm_eps,
2318
+ cos_sin_cache=cos_sin_cache,
2319
+ kv_cache=k_cache,
2320
+ kv_slot_mapping=kv_slot_mapping,
2321
compress_ratio=compress_ratio,
2024
- overlap=overlap,
2322
store_full_fp8=store_full_fp8,
2026
- norm_weight_dtype=norm_weight_dtype,
2027
- )
2028
- compiled(
2029
- state_cache,
2030
- token_to_req_indices,
2031
- positions,
2032
- slot_mapping,
2033
- block_table,
2034
- block_size,
2035
- rms_norm_weight,
2036
- rms_norm_eps,
2037
- cos_sin_cache,
2038
- k_cache.view(torch.uint8),
2039
- kv_slot_mapping,
2040
- kv_cache_block_size,
2041
- fp8_scale,
2323
+ fp8_scale=fp8_scale,
2324
+ head_dim=head_size,
2325
+ rope_head_dim=rope_head_dim,
2326
)
2327
return
2044
-
2045
- compiled = SparseAttnCompressNormRopeStoreC4Kernel.compile(
2046
- head_size=head_size,
2047
- state_width=state_width,
2048
- rope_head_dim=rope_head_dim,
2049
- fp8_max=fp8_max,
2050
- quant_block=quant_block,
2051
- token_stride=token_stride,
2052
- scale_dim=scale_dim,
2053
- kv_block_stride=kv_block_stride,
2328
+ _SPARSE_ATTN_NORM_ROPE_STORE_KERNEL(
2329
+ compressed_kv=compressed_kv,
2330
+ positions=positions,
2331
+ slot_mapping=slot_mapping,
2332
+ rms_norm_weight=rms_norm_weight,
2333
+ rms_norm_eps=rms_norm_eps,
2334
+ cos_sin_cache=cos_sin_cache,
2335
+ kv_cache=k_cache,
2336
+ kv_slot_mapping=kv_slot_mapping,
2337
compress_ratio=compress_ratio,
2055
- overlap=overlap,
2056
- norm_weight_dtype=norm_weight_dtype,
2057
- )
2058
- compiled(
2059
- state_cache,
2060
- token_to_req_indices,
2061
- positions,
2062
- slot_mapping,
2063
- block_table,
2064
- block_size,
2065
- rms_norm_weight,
2066
- rms_norm_eps,
2067
- cos_sin_cache,
2068
- k_cache,
2069
- kv_slot_mapping,
2070
- kv_cache_block_size,
2338
+ head_dim=head_size,
2339
+ rope_head_dim=rope_head_dim,
2340
)
2341
2342
2074
-def compress_norm_rope_store_cutedsl(
2075
- state_cache: torch.Tensor,
2076
- num_actual: int,
2077
- token_to_req_indices: torch.Tensor,
2078
- positions: torch.Tensor,
2079
- slot_mapping: torch.Tensor,
2080
- block_table: torch.Tensor,
2081
- block_size: int,
2082
- state_width: int,
2083
- cos_sin_cache: torch.Tensor,
2084
- kv_cache: torch.Tensor,
2085
- k_cache_metadata: Any,
2086
- pdl_kwargs: dict,
2087
- head_dim: int,
2088
- rope_head_dim: int,
2089
- compress_ratio: int,
2090
- overlap: bool,
2091
- use_fp4_cache: bool,
2092
- rms_norm_weight: torch.Tensor,
2093
- rms_norm_eps: float,
2094
- quant_block: int,
2095
- token_stride: int,
2096
- scale_dim: int,
2097
- store_full_kv: bool = False,
2098
- store_full_fp8: bool = False,
2099
- fp8_scale: torch.Tensor | None = None,
2100
-) -> None:
2101
- if compress_ratio == 4:
2102
- # For C4A, the single fused kernel is faster than the two-kernel version.
2103
- fused_kv_compress_norm_rope_insert_sparse_attn_cutedsl(
2104
- state_cache,
2105
- token_to_req_indices,
2106
- positions,
2107
- slot_mapping,
2108
- block_table,
2109
- block_size,
2110
- rms_norm_weight,
2111
- rms_norm_eps,
2112
- cos_sin_cache,
2113
- kv_cache,
2114
- k_cache_metadata.slot_mapping,
2115
- kv_cache.shape[1], # paged KV cache block size
2116
- kv_cache.stride(0),
2117
- head_size=head_dim,
2118
- state_width=state_width,
2119
- rope_head_dim=rope_head_dim,
2120
- fp8_max=448.0,
2121
- quant_block=quant_block,
2122
- token_stride=token_stride,
2123
- scale_dim=scale_dim,
2124
- compress_ratio=compress_ratio,
2125
- overlap=overlap,
2126
- store_full_kv=store_full_kv,
2127
- store_full_fp8=store_full_fp8,
2128
- fp8_scale=fp8_scale,
2129
- )
2130
- else:
2131
- # For C128, the two-kernel version is faster than the single fused kernel.
2132
- compressed_kv = torch.empty(
2133
- (num_actual, head_dim),
2134
- dtype=torch.float32,
2135
- device=state_cache.device,
2343
+class SparseAttnCompressorCuteDSL:
2344
+ def __init__(
2345
+ self,
2346
+ *,
2347
+ c4_store: SparseAttnCompressNormRopeStoreC4Kernel,
2348
+ c4_store_full: SparseAttnCompressNormRopeStoreFullC4Kernel,
2349
+ c128_compress: SparseAttnCompressC128Block8Kernel,
2350
+ c128_store: SparseAttnNormRopeStoreKernel,
2351
+ c128_store_full: SparseAttnNormRopeStoreFullKernel,
2352
+ ) -> None:
2353
+ self.c4_store = c4_store
2354
+ self.c4_store_full = c4_store_full
2355
+ self.c128_compress = c128_compress
2356
+ self.c128_store = c128_store
2357
+ self.c128_store_full = c128_store_full
2358
+
2359
+ def __call__(
2360
+ self,
2361
+ state_cache: torch.Tensor,
2362
+ num_actual: int,
2363
+ token_to_req_indices: torch.Tensor,
2364
+ positions: torch.Tensor,
2365
+ slot_mapping: torch.Tensor,
2366
+ block_table: torch.Tensor,
2367
+ block_size: int,
2368
+ state_width: int,
2369
+ cos_sin_cache: torch.Tensor,
2370
+ kv_cache: torch.Tensor,
2371
+ k_cache_metadata: Any,
2372
+ pdl_kwargs: dict,
2373
+ head_dim: int,
2374
+ rope_head_dim: int,
2375
+ compress_ratio: int,
2376
+ overlap: bool,
2377
+ use_fp4_cache: bool,
2378
+ rms_norm_weight: torch.Tensor,
2379
+ rms_norm_eps: float,
2380
+ quant_block: int,
2381
+ token_stride: int,
2382
+ scale_dim: int,
2383
+ store_full_kv: bool = False,
2384
+ store_full_fp8: bool = False,
2385
+ fp8_scale: torch.Tensor | None = None,
2386
+ ) -> None:
2387
+ del pdl_kwargs, state_width, overlap, use_fp4_cache
2388
+ del quant_block, token_stride, scale_dim
2389
+ if positions.numel() == 0:
2390
+ return
2391
+ if kv_cache.ndim != 3:
2392
+ raise ValueError(
2393
+ "CuTe DSL sparse-attn compressor expects the real DeepSeek V4 "
2394
+ f"3D kv_cache layout, got ndim={kv_cache.ndim}.",
2395
+ )
2396
+ if store_full_fp8 and (not store_full_kv):
2397
+ raise ValueError("store_full_fp8 requires store_full_kv.")
2398
+ if compress_ratio == 4:
2399
+ if store_full_kv:
2400
+ self.c4_store_full(
2401
+ state_cache=state_cache,
2402
+ token_to_req_indices=token_to_req_indices,
2403
+ positions=positions,
2404
+ slot_mapping=slot_mapping,
2405
+ block_table=block_table,
2406
+ block_size=block_size,
2407
+ rms_norm_weight=rms_norm_weight,
2408
+ rms_norm_eps=rms_norm_eps,
2409
+ cos_sin_cache=cos_sin_cache,
2410
+ kv_cache=kv_cache,
2411
+ kv_slot_mapping=k_cache_metadata.slot_mapping,
2412
+ compress_ratio=compress_ratio,
2413
+ store_full_fp8=store_full_fp8,
2414
+ fp8_scale=fp8_scale,
2415
+ head_dim=head_dim,
2416
+ rope_head_dim=rope_head_dim,
2417
+ )
2418
+ else:
2419
+ self.c4_store(
2420
+ state_cache=state_cache,
2421
+ token_to_req_indices=token_to_req_indices,
2422
+ positions=positions,
2423
+ slot_mapping=slot_mapping,
2424
+ block_table=block_table,
2425
+ block_size=block_size,
2426
+ rms_norm_weight=rms_norm_weight,
2427
+ rms_norm_eps=rms_norm_eps,
2428
+ cos_sin_cache=cos_sin_cache,
2429
+ kv_cache=kv_cache,
2430
+ kv_slot_mapping=k_cache_metadata.slot_mapping,
2431
+ compress_ratio=compress_ratio,
2432
+ head_dim=head_dim,
2433
+ rope_head_dim=rope_head_dim,
2434
+ )
2435
+ return
2436
+ if compress_ratio != 128:
2437
+ raise ValueError(
2438
+ "CuTe DSL sparse-attn compressor supports compress_ratio 4 or "
2439
+ f"128, got {compress_ratio}.",
2440
+ )
2441
+ if block_size != self.c128_compress.state_block_size:
2442
+ raise ValueError(
2443
+ "CuTe DSL C128 sparse-attn compressor requires state-cache "
2444
+ f"block_size={self.c128_compress.state_block_size}, got "
2445
+ f"{block_size}."
2446
+ )
2447
+ compressed_kv = self.c128_compress(
2448
+ state_cache=state_cache,
2449
+ num_actual=num_actual,
2450
+ token_to_req_indices=token_to_req_indices,
2451
+ positions=positions,
2452
+ slot_mapping=slot_mapping,
2453
+ block_table=block_table,
2454
+ head_dim=head_dim,
2455
)
2137
- split_kv_compress_norm_rope_insert_sparse_attn_cutedsl(
2138
- state_cache,
2139
- token_to_req_indices,
2140
- positions,
2141
- slot_mapping,
2142
- block_table,
2143
- block_size,
2144
- compressed_kv,
2145
- rms_norm_weight,
2146
- rms_norm_eps,
2147
- cos_sin_cache,
2148
- kv_cache,
2149
- k_cache_metadata.slot_mapping,
2150
- kv_cache.shape[1], # paged KV cache block size
2151
- kv_cache.stride(0),
2152
- head_size=head_dim,
2153
- state_width=state_width,
2154
- rope_head_dim=rope_head_dim,
2155
- fp8_max=448.0,
2156
- quant_block=quant_block,
2157
- token_stride=token_stride,
2158
- scale_dim=scale_dim,
2456
+ if store_full_kv:
2457
+ self.c128_store_full(
2458
+ compressed_kv=compressed_kv,
2459
+ positions=positions,
2460
+ slot_mapping=slot_mapping,
2461
+ rms_norm_weight=rms_norm_weight,
2462
+ rms_norm_eps=rms_norm_eps,
2463
+ cos_sin_cache=cos_sin_cache,
2464
+ kv_cache=kv_cache,
2465
+ kv_slot_mapping=k_cache_metadata.slot_mapping,
2466
+ compress_ratio=compress_ratio,
2467
+ store_full_fp8=store_full_fp8,
2468
+ fp8_scale=fp8_scale,
2469
+ head_dim=head_dim,
2470
+ rope_head_dim=rope_head_dim,
2471
+ )
2472
+ return
2473
+ self.c128_store(
2474
+ compressed_kv=compressed_kv,
2475
+ positions=positions,
2476
+ slot_mapping=slot_mapping,
2477
+ rms_norm_weight=rms_norm_weight,
2478
+ rms_norm_eps=rms_norm_eps,
2479
+ cos_sin_cache=cos_sin_cache,
2480
+ kv_cache=kv_cache,
2481
+ kv_slot_mapping=k_cache_metadata.slot_mapping,
2482
compress_ratio=compress_ratio,
2160
- overlap=overlap,
2161
- store_full_kv=store_full_kv,
2162
- store_full_fp8=store_full_fp8,
2163
- fp8_scale=fp8_scale,
2483
+ head_dim=head_dim,
2484
+ rope_head_dim=rope_head_dim,
2485
)
2486
+
2487
+
2488
+_SPARSE_ATTN_COMPRESS_NORM_ROPE_STORE_C4_KERNEL = (
2489
+ SparseAttnCompressNormRopeStoreC4Kernel()
2490
+)
2491
+_SPARSE_ATTN_COMPRESS_NORM_ROPE_STORE_FULL_C4_KERNEL = (
2492
+ SparseAttnCompressNormRopeStoreFullC4Kernel()
2493
+)
2494
+_SPARSE_ATTN_COMPRESS_C128_BLOCK8_KERNEL = SparseAttnCompressC128Block8Kernel()
2495
+_SPARSE_ATTN_NORM_ROPE_STORE_KERNEL = SparseAttnNormRopeStoreKernel()
2496
+_SPARSE_ATTN_NORM_ROPE_STORE_FULL_KERNEL = SparseAttnNormRopeStoreFullKernel()
2497
+_SPARSE_ATTN_COMPRESSOR_CUTEDSL_KERNEL = SparseAttnCompressorCuteDSL(
2498
+ c4_store=_SPARSE_ATTN_COMPRESS_NORM_ROPE_STORE_C4_KERNEL,
2499
+ c4_store_full=_SPARSE_ATTN_COMPRESS_NORM_ROPE_STORE_FULL_C4_KERNEL,
2500
+ c128_compress=_SPARSE_ATTN_COMPRESS_C128_BLOCK8_KERNEL,
2501
+ c128_store=_SPARSE_ATTN_NORM_ROPE_STORE_KERNEL,
2502
+ c128_store_full=_SPARSE_ATTN_NORM_ROPE_STORE_FULL_KERNEL,
2503
+)
vllm/models/deepseek_v4/sparse_mla.py
+1
-1
@@ -9,11 +9,11 @@ import torch
9
10
from vllm.config import VllmConfig
11
from vllm.config.cache import CacheDType
12
+from vllm.model_executor.warmup.jit_warmup import kernel_launcher
13
from vllm.model_executor.warmup.jit_warmup_triton_helper import (
14
LaunchSpec,
15
TritonWarmupTensor,
16
VllmTritonJitKernel,
16
- kernel_launcher,
17
)
18
from vllm.platforms import current_platform
19
from vllm.platforms.interface import DeviceCapability
vllm/models/deepseek_v4_1/common/ops/cache_utils.py
+9
-3
@@ -416,11 +416,17 @@ def dequantize_and_gather_k_cache(
416
if has_cutedsl():
417
# lazily import, otherwise some tests fail due to CUDA driver init failure.
418
from vllm.models.deepseek_v4.nvidia.ops.dequant_gather_k_cutedsl import (
419
- dequantize_and_gather_k_cache_cutedsl,
419
+ _DEQUANT_GATHER_K_CACHE_CUTEDSL_KERNEL,
420
)
421
422
- dequantize_and_gather_k_cache_cutedsl(
423
- out, k_cache, seq_lens, gather_lens, block_table, block_size, offset
422
+ _DEQUANT_GATHER_K_CACHE_CUTEDSL_KERNEL(
423
+ out=out,
424
+ k_cache=k_cache,
425
+ seq_lens=seq_lens,
426
+ gather_lens=gather_lens,
427
+ block_table=block_table,
428
+ block_size=block_size,
429
+ offset=offset,
430
)
431
return
432
vllm/v1/attention/backends/mla/compressor_utils.py
+1
-2
@@ -5,12 +5,11 @@ from typing import Any
5
6
import torch
7
8
-from vllm.model_executor.warmup.jit_warmup import zip_inputs
8
+from vllm.model_executor.warmup.jit_warmup import kernel_launcher, zip_inputs
9
from vllm.model_executor.warmup.jit_warmup_triton_helper import (
10
LaunchSpec,
11
TritonWarmupTensor,
12
VllmTritonJitKernel,
13
- kernel_launcher,
13
triton_scalar_specialization_rep,
14
)
15
from vllm.triton_utils import tl, triton
vllm/v1/attention/backends/mla/indexer.py
+1
-1
@@ -9,12 +9,12 @@ import vllm.envs as envs
9
from vllm.config import VllmConfig
10
from vllm.distributed import get_dcp_group, get_pcp_group
11
from vllm.logger import init_logger
12
+from vllm.model_executor.warmup.jit_warmup import kernel_launcher
13
from vllm.model_executor.warmup.jit_warmup_triton_helper import (
14
LaunchSpec,
15
TritonPointerInputVariant,
16
TritonWarmupTensor,
17
VllmTritonJitKernel,
17
- kernel_launcher,
18
triton_scalar_specialization_rep,
19
)
20
from vllm.platforms import current_platform
vllm/v1/attention/backends/mla/sparse_swa.py
+1
-1
@@ -9,13 +9,13 @@ from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config
9
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
10
from vllm.model_executor.warmup.jit_warmup import (
11
WarmupIntRange,
12
+ kernel_launcher,
13
zip_inputs,
14
)
15
from vllm.model_executor.warmup.jit_warmup_triton_helper import (
16
LaunchSpec,
17
TritonWarmupTensor,
18
VllmTritonJitKernel,
18
- kernel_launcher,
19
)
20
from vllm.platforms import current_platform
21
from vllm.triton_utils import tl, triton
vllm/v1/attention/backends/mla/sparse_utils.py
+1
-1
@@ -10,13 +10,13 @@ import torch
10
from vllm.config import VllmConfig
11
from vllm.distributed import get_dcp_group
12
from vllm.model_executor.warmup.jit_warmup import (
13
+ kernel_launcher,
14
zip_inputs,
15
)
16
from vllm.model_executor.warmup.jit_warmup_triton_helper import (
17
LaunchSpec,
18
TritonWarmupTensor,
19
VllmTritonJitKernel,
19
- kernel_launcher,
20
triton_scalar_specialization_rep,
21
)
22
from vllm.triton_utils import tl, triton
vllm/v1/attention/ops/common.py
+1
-1
@@ -5,11 +5,11 @@ from typing import Any
5
6
import torch
7
8
+from vllm.model_executor.warmup.jit_warmup import kernel_launcher
9
from vllm.model_executor.warmup.jit_warmup_triton_helper import (
10
LaunchSpec,
11
TritonWarmupTensor,
12
VllmTritonJitKernel,
12
- kernel_launcher,
13
)
14
from vllm.triton_utils import tl, triton
15
vllm/v1/attention/ops/dcp.py
+1
-1
@@ -18,12 +18,12 @@ from vllm.distributed import get_dcp_group
18
from vllm.logger import init_logger
19
from vllm.model_executor.warmup.jit_warmup import (
20
WarmupIntRange,
21
+ kernel_launcher,
22
)
23
from vllm.model_executor.warmup.jit_warmup_triton_helper import (
24
LaunchSpec,
25
TritonWarmupTensor,
26
VllmTritonJitKernel,
26
- kernel_launcher,
27
triton_scalar_specialization_rep,
28
)
29
from vllm.triton_utils import tl, triton
vllm/v1/worker/block_table.py
+1
-1
@@ -11,11 +11,11 @@ import torch
11
12
from vllm.distributed import get_dcp_group, get_pcp_group
13
from vllm.logger import init_logger
14
+from vllm.model_executor.warmup.jit_warmup import kernel_launcher
15
from vllm.model_executor.warmup.jit_warmup_triton_helper import (
16
LaunchSpec,
17
TritonWarmupTensor,
18
VllmTritonJitKernel,
18
- kernel_launcher,
19
triton_scalar_specialization_rep,
20
)
21
from vllm.triton_utils import tl, triton