From 120ec4ebd2803c3cc63729e18fcfd42b5ba47447 Mon Sep 17 00:00:00 2001 From: "Roberto L. Castro" <38211239+LopezCastroRoberto@users.noreply.github.com> Date: Sat, 12 Sep 2026 06:05:59 +0200 Subject: [PATCH] [5/N][warmup][DSv4] Migrate NVIDIA CuTeDSL attention kernels (#53566) Signed-off-by: LopezCastroRoberto Co-authored-by: Codex --- .../test_jit_warmup_cutedsl_launcher.py | 41 +- .../test_jit_warmup_triton_launcher.py | 8 +- vllm/cute_utils/__init__.py | 8 + .../attention/dsa/dcp_indexer_cutedsl.py | 18 +- vllm/model_executor/warmup/jit_warmup.py | 23 + .../warmup/jit_warmup_cutedsl_helper.py | 51 +- .../warmup/jit_warmup_triton_helper.py | 58 +- vllm/models/common/ops/fused_qk_rmsnorm.py | 2 +- vllm/models/deepseek_v4/attention.py | 26 +- .../deepseek_v4/common/ops/cache_utils.py | 14 +- .../common/ops/fused_compress_quant_cache.py | 2 +- .../deepseek_v4/common/ops/fused_indexer_q.py | 44 +- .../common/ops/fused_inv_rope_fp8_quant.py | 2 +- .../common/ops/fused_mtp_input_rmsnorm.py | 2 +- .../common/ops/save_partial_states.py | 2 +- vllm/models/deepseek_v4/compressor.py | 36 +- .../nvidia/ops/dequant_gather_k_cutedsl.py | 627 +-- .../nvidia/ops/fused_indexer_q_cutedsl.py | 1071 ++--- .../ops/sparse_attn_compress_cutedsl.py | 3611 +++++++++-------- vllm/models/deepseek_v4/sparse_mla.py | 2 +- .../deepseek_v4_1/common/ops/cache_utils.py | 12 +- .../backends/mla/compressor_utils.py | 3 +- vllm/v1/attention/backends/mla/indexer.py | 2 +- vllm/v1/attention/backends/mla/sparse_swa.py | 2 +- .../v1/attention/backends/mla/sparse_utils.py | 2 +- vllm/v1/attention/ops/common.py | 2 +- vllm/v1/attention/ops/dcp.py | 2 +- vllm/v1/worker/block_table.py | 2 +- 28 files changed, 3124 insertions(+), 2551 deletions(-) diff --git a/tests/model_executor/test_jit_warmup_cutedsl_launcher.py b/tests/model_executor/test_jit_warmup_cutedsl_launcher.py index 2524c60d2e..76ec4d1200 100644 --- a/tests/model_executor/test_jit_warmup_cutedsl_launcher.py +++ b/tests/model_executor/test_jit_warmup_cutedsl_launcher.py @@ -5,10 +5,10 @@ from dataclasses import dataclass from typing import Any from vllm.model_executor.warmup import jit_warmup_cutedsl_helper +from vllm.model_executor.warmup.jit_warmup import kernel_launcher from vllm.model_executor.warmup.jit_warmup_cutedsl_helper import ( CuTeDSLLaunchSpec, VllmCuTeDSLJitKernel, - cutedsl_kernel_launcher, ) @@ -30,7 +30,7 @@ class _TestCuTeDSLKernel(VllmCuTeDSLJitKernel["_TestCuTeDSLKernel.CompileKey"]): def warmup_inputs(self, compile_key: CompileKey) -> tuple[Any, ...]: return (f"fake-{compile_key.variant}",) - @cutedsl_kernel_launcher + @kernel_launcher def __call__( self, payload: str, @@ -38,7 +38,7 @@ class _TestCuTeDSLKernel(VllmCuTeDSLJitKernel["_TestCuTeDSLKernel.CompileKey"]): variant: int, ) -> CuTeDSLLaunchSpec[CompileKey]: compile_key = self.dispatch(variant=variant) - return compile_key, (payload,), {"variant": variant} + return compile_key, (payload,) def test_cutedsl_launcher_reuses_compiled_executor(monkeypatch) -> None: @@ -74,3 +74,38 @@ def test_cutedsl_launcher_reuses_compiled_executor(monkeypatch) -> None: (("entry", 2), ("fake-2",)), ] assert launch_calls == [(1, ("runtime",)), (2, ("other",))] + + +def test_cutedsl_launcher_runs_post_launch_finalizer(monkeypatch) -> None: + events: list[str] = [] + + def compile_cutedsl(entry: tuple[str, int], *args: Any) -> Any: + def executor(*runtime_args: Any) -> None: + events.append("launch") + + return executor + + class FinalizedKernel(_TestCuTeDSLKernel): + @kernel_launcher + def __call__( + self, + payload: str, + *, + variant: int, + ) -> CuTeDSLLaunchSpec[_TestCuTeDSLKernel.CompileKey]: + compile_key = self.dispatch(variant=variant) + + def finalize() -> str: + events.append("finalize") + return "output" + + return compile_key, (payload,), "unused", finalize + + monkeypatch.setattr( + jit_warmup_cutedsl_helper, + "compile_cutedsl", + compile_cutedsl, + ) + + assert FinalizedKernel()("runtime", variant=1) == "output" + assert events == ["launch", "finalize"] diff --git a/tests/model_executor/test_jit_warmup_triton_launcher.py b/tests/model_executor/test_jit_warmup_triton_launcher.py index 2e624cc7f6..e76041532f 100644 --- a/tests/model_executor/test_jit_warmup_triton_launcher.py +++ b/tests/model_executor/test_jit_warmup_triton_launcher.py @@ -4,10 +4,10 @@ from dataclasses import dataclass from typing import Any +from vllm.model_executor.warmup.jit_warmup import kernel_launcher from vllm.model_executor.warmup.jit_warmup_triton_helper import ( LaunchSpec, VllmTritonJitKernel, - kernel_launcher, ) @@ -107,13 +107,13 @@ def test_compute_slot_mapping_uses_named_launcher_inputs(monkeypatch) -> None: ) launches: list[tuple[Any, ...]] = [] - def launch(grid: Any, inputs: Any, **kwargs: Any) -> None: - launches.append((grid, inputs, kwargs)) + def launch(launch_spec: Any, inputs: Any) -> None: + launches.append((launch_spec, inputs)) monkeypatch.setattr(owner, "launch", launch) owner.compile(compile_key) - grid, inputs, kwargs = launches[0] + (grid, kwargs), inputs = launches[0] assert grid == (2,) assert inputs["num_tokens"] == 2 assert inputs["block_table_stride"] == 128 diff --git a/vllm/cute_utils/__init__.py b/vllm/cute_utils/__init__.py index 23721aac44..033533a5c9 100644 --- a/vllm/cute_utils/__init__.py +++ b/vllm/cute_utils/__init__.py @@ -22,6 +22,14 @@ _TORCH_TO_CUTE_DTYPE = { torch.float32: Float32, } + +def torch_to_cute_dtype(dtype: torch.dtype): + try: + return _TORCH_TO_CUTE_DTYPE[dtype] + except KeyError as exc: + raise TypeError(f"Unsupported CuTe dtype {dtype}") from exc + + _CUTE_TO_PTX_DTYPE = { BFloat16: "bf16", Float16: "f16", diff --git a/vllm/model_executor/kernels/attention/dsa/dcp_indexer_cutedsl.py b/vllm/model_executor/kernels/attention/dsa/dcp_indexer_cutedsl.py index 29cdb2c589..34ce32856a 100644 --- a/vllm/model_executor/kernels/attention/dsa/dcp_indexer_cutedsl.py +++ b/vllm/model_executor/kernels/attention/dsa/dcp_indexer_cutedsl.py @@ -12,16 +12,15 @@ from cutlass import Float32, Int32, Uint32, Uint64 from quack.compile_utils import make_fake_tensor from vllm.cute_utils import recast_val +from vllm.model_executor.warmup.jit_warmup import kernel_launcher from vllm.model_executor.warmup.jit_warmup_cutedsl_helper import ( CuTeDSLLaunchSpec, VllmCuTeDSLJitKernel, - cutedsl_kernel_launcher, ) from vllm.model_executor.warmup.jit_warmup_triton_helper import ( LaunchSpec, TritonWarmupTensor, VllmTritonJitKernel, - kernel_launcher, ) from vllm.triton_utils import tl, triton @@ -569,7 +568,7 @@ class StableTopKFromGatheredCandidatesKernel( ) return gathered, out - @cutedsl_kernel_launcher + @kernel_launcher def __call__( self, gathered: torch.Tensor, @@ -577,16 +576,11 @@ class StableTopKFromGatheredCandidatesKernel( *, topk: int, ) -> CuTeDSLLaunchSpec[CompileKey]: - compile_key = self.dispatch(topk=topk, num_candidates=gathered.shape[1]) - return ( - compile_key, - (gathered, out), - { - "gathered_shape": tuple(gathered.shape), - "out_shape": tuple(out.shape), - "topk": topk, - }, + compile_key = self.dispatch( + topk=topk, + num_candidates=gathered.shape[1], ) + return compile_key, (gathered, out) _PACK_DCP_TOPK_CANDIDATES_KERNEL = PackDCPTopkCandidatesKernel() diff --git a/vllm/model_executor/warmup/jit_warmup.py b/vllm/model_executor/warmup/jit_warmup.py index 0be4bddea6..fc9d6fcdf8 100644 --- a/vllm/model_executor/warmup/jit_warmup.py +++ b/vllm/model_executor/warmup/jit_warmup.py @@ -15,6 +15,7 @@ from collections.abc import Callable, Iterable, Iterator, Mapping from contextlib import contextmanager from contextvars import ContextVar from dataclasses import dataclass, fields +from functools import wraps from typing import Any, Generic, TypeVar, cast __all__ = [ @@ -23,6 +24,7 @@ __all__ = [ "WarmupIntRange", "get_ast_full_name", "get_function_source_node", + "kernel_launcher", "zip_inputs", ] @@ -30,6 +32,26 @@ __all__ = [ CompileKeyT = TypeVar("CompileKeyT") +def kernel_launcher(call_fn: Callable[..., Any]) -> Callable[..., Any]: + """Delegate a declarative launch specification to its backend owner.""" + signature = inspect.signature(call_fn) + + @wraps(call_fn) + def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any: + launch_spec = call_fn(self, *args, **kwargs) + if not self.bind_launch_inputs: + return self.launch(launch_spec, {}) + + bound = signature.bind(self, *args, **kwargs) + bound.apply_defaults() + inputs = { + name: value for name, value in bound.arguments.items() if name != "self" + } + return self.launch(launch_spec, inputs) + + return wrapper + + @dataclass(frozen=True) class WarmupIntRange: """Expand integers with range semantics or a custom monotonic progression.""" @@ -549,6 +571,7 @@ class VllmJitKernel(Generic[CompileKeyT], ABC): """Kernel wrapper that owns dispatch, warmup keys, and compilation.""" CompileKey: type[CompileKeyT] + bind_launch_inputs = True def __init__(self) -> None: self._dispatch_trace = _trace_compile_key_dispatch(self.dispatch) diff --git a/vllm/model_executor/warmup/jit_warmup_cutedsl_helper.py b/vllm/model_executor/warmup/jit_warmup_cutedsl_helper.py index 7311cc18a4..879c8c6998 100644 --- a/vllm/model_executor/warmup/jit_warmup_cutedsl_helper.py +++ b/vllm/model_executor/warmup/jit_warmup_cutedsl_helper.py @@ -3,19 +3,25 @@ from abc import abstractmethod from collections.abc import Callable, Mapping -from functools import wraps from typing import Any, ClassVar, Generic, TypeAlias, TypeVar -from vllm.model_executor.warmup.jit_warmup import VllmJitKernel +from vllm.model_executor.warmup.jit_warmup import ( + VllmJitKernel, +) DEFAULT_CUTEDSL_COMPILE_OPTIONS = "--enable-tvm-ffi" CompileKeyT = TypeVar("CompileKeyT") -CuTeDSLLaunchSpec: TypeAlias = tuple[ - CompileKeyT, - tuple[Any, ...], - Mapping[str, Any] | None, -] +CuTeDSLLaunchSpec: TypeAlias = ( + tuple[CompileKeyT, tuple[Any, ...]] + | tuple[CompileKeyT, tuple[Any, ...], Any] + | tuple[ + CompileKeyT, + tuple[Any, ...], + Any, + Callable[[], Any], + ] +) def cutedsl_fake_stream(*, use_tvm_ffi_env_stream: bool = True) -> Any: @@ -44,6 +50,7 @@ class VllmCuTeDSLJitKernel(VllmJitKernel[CompileKeyT], Generic[CompileKeyT]): """CuTeDSL owner whose compiled executor is shared by warmup and runtime.""" kernel: ClassVar[Any] + bind_launch_inputs = False @abstractmethod def warmup_inputs(self, compile_key: CompileKeyT) -> tuple[Any, ...]: @@ -58,23 +65,15 @@ class VllmCuTeDSLJitKernel(VllmJitKernel[CompileKeyT], Generic[CompileKeyT]): *self.warmup_inputs(compile_key), ) - -def cutedsl_kernel_launcher( - call_fn: Callable[..., CuTeDSLLaunchSpec[CompileKeyT]], -) -> Callable[..., Any]: - """Invoke a cached CuTeDSL executor from a declarative ``__call__``.""" - - @wraps(call_fn) - def wrapper( - self: VllmCuTeDSLJitKernel[CompileKeyT], - *args: Any, - **kwargs: Any, + def launch( + self, + launch_spec: CuTeDSLLaunchSpec[CompileKeyT], + _inputs: Mapping[str, Any], ) -> Any: - compile_key, launch_args, runtime_context = call_fn(self, *args, **kwargs) - executor = self._get_or_compile( - compile_key, - runtime_context=runtime_context, - ) - return executor(*launch_args) - - return wrapper + # (compile_key, args), optionally followed by output and epilogue. + compile_key, launch_args = launch_spec[:2] + executor = self._get_or_compile(compile_key) + result = executor(*launch_args) + if len(launch_spec) == 4: + return launch_spec[3]() + return launch_spec[2] if len(launch_spec) == 3 else result diff --git a/vllm/model_executor/warmup/jit_warmup_triton_helper.py b/vllm/model_executor/warmup/jit_warmup_triton_helper.py index 7fffa12ec9..b019404cc9 100644 --- a/vllm/model_executor/warmup/jit_warmup_triton_helper.py +++ b/vllm/model_executor/warmup/jit_warmup_triton_helper.py @@ -5,7 +5,7 @@ import inspect from abc import abstractmethod from collections.abc import Callable, Mapping from dataclasses import dataclass -from functools import cached_property, wraps +from functools import cached_property from typing import Any, ClassVar, Generic, TypeVar, cast from vllm.model_executor.warmup.jit_warmup import ( @@ -128,14 +128,16 @@ class VllmTritonJitKernel(VllmJitKernel[CompileKeyT], Generic[CompileKeyT]): def launch( self, - grid: tuple[int, ...] | None, + launch_spec: LaunchSpec, inputs: Mapping[str, Any], - /, - *, - kernel: Any = None, - **kwargs: Any, ) -> Any: - kernel = kernel if kernel is not None else self.kernel + if len(launch_spec) == 3: + grid, launch_kwargs, outputs = launch_spec + else: + grid, launch_kwargs = launch_spec + outputs = None + kwargs = dict(launch_kwargs) + kernel = kwargs.pop("kernel", None) or self.kernel runtime_launcher = kwargs.pop("_runtime_launcher", None) runtime_launcher_arg_count = kwargs.pop("_runtime_launcher_arg_count", 0) for name, value in inputs.items(): @@ -151,52 +153,20 @@ class VllmTritonJitKernel(VllmJitKernel[CompileKeyT], Generic[CompileKeyT]): kwargs["launch_pdl"] = self._warming_compile_key.launch_pdl warmup = getattr(kernel, "warmup", None) assert warmup is not None - return warmup(grid=(1,), **kwargs) + warmup(grid=(1,), **kwargs) + return outputs if grid is None: - return None + return outputs if runtime_launcher is not None: regular_args = [ kwargs.pop(name) for name in self._kernel_arg_names[:runtime_launcher_arg_count] ] - return runtime_launcher(kernel, grid, *regular_args, **kwargs) - return kernel[grid](**kwargs) - - -def kernel_launcher( - call_fn: Callable[..., LaunchSpec], -) -> Callable[..., Any]: - """Launch a Triton kernel from a declarative ``__call__`` specification. - - ``call_fn`` returns either ``(grid, launch_kwargs)`` or, when it allocates - its own outputs, ``(grid, launch_kwargs, outputs)``. The declared outputs are - returned to the caller. A ``grid`` of ``None`` skips the launch (e.g. an - empty-token batch) while still returning the outputs. - """ - signature = inspect.signature(call_fn) - - @wraps(call_fn) - def wrapper( - self: VllmTritonJitKernel[Any], - *args: Any, - **kwargs: Any, - ) -> Any: - spec = call_fn(self, *args, **kwargs) - if len(spec) == 3: - grid, launch_kwargs, outputs = spec + runtime_launcher(kernel, grid, *regular_args, **kwargs) else: - grid, launch_kwargs = spec - outputs = None - bound = signature.bind(self, *args, **kwargs) - bound.apply_defaults() - inputs = { - name: value for name, value in bound.arguments.items() if name != "self" - } - self.launch(grid, inputs, **launch_kwargs) + kernel[grid](**kwargs) return outputs - return wrapper - @dataclass(frozen=True) class TritonPointerInputVariant: diff --git a/vllm/models/common/ops/fused_qk_rmsnorm.py b/vllm/models/common/ops/fused_qk_rmsnorm.py index 2ee6cc2854..e64e818210 100644 --- a/vllm/models/common/ops/fused_qk_rmsnorm.py +++ b/vllm/models/common/ops/fused_qk_rmsnorm.py @@ -5,11 +5,11 @@ from typing import Any import torch +from vllm.model_executor.warmup.jit_warmup import kernel_launcher from vllm.model_executor.warmup.jit_warmup_triton_helper import ( LaunchSpec, TritonWarmupTensor, VllmTritonJitKernel, - kernel_launcher, ) from vllm.platforms import current_platform from vllm.triton_utils import tl, triton diff --git a/vllm/models/deepseek_v4/attention.py b/vllm/models/deepseek_v4/attention.py index 12205bc043..d32a48cbe4 100644 --- a/vllm/models/deepseek_v4/attention.py +++ b/vllm/models/deepseek_v4/attention.py @@ -50,6 +50,7 @@ from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.model_executor.models.utils import extract_layer_index from vllm.models.deepseek_v4.common.rope import build_deepseek_v4_rope from vllm.models.deepseek_v4.compressor import DeepseekCompressor +from vllm.platforms import current_platform from vllm.triton_utils import tl, triton from vllm.utils.multi_stream_utils import ( execute_in_parallel, @@ -420,7 +421,6 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): _COMBINE_TOPK_SWA_INDICES_KERNEL.register_warmup() - from vllm.platforms import current_platform from vllm.utils.import_utils import has_cutedsl _FUSED_Q_KV_RMSNORM_KERNEL.register_warmup() @@ -447,7 +447,13 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): ) _COMPUTE_GLOBAL_TOPK_INDICES_AND_LENS_KERNEL.register_warmup() - if not has_cutedsl(): + if has_cutedsl(): + from vllm.models.deepseek_v4.nvidia.ops.dequant_gather_k_cutedsl import ( # noqa: E501 + _DEQUANT_GATHER_K_CACHE_CUTEDSL_KERNEL, + ) + + _DEQUANT_GATHER_K_CACHE_CUTEDSL_KERNEL.register_warmup() + else: _DEQUANTIZE_AND_GATHER_K_CACHE_KERNEL.register_warmup() elif backend_name == "FLASHINFER_MLA_SPARSE_DSV4": from vllm.models.deepseek_v4.common.ops.cache_utils import ( @@ -934,6 +940,21 @@ class DeepseekV4Indexer(nn.Module): "Using %s indexer cache for Lightning Indexer.", "MXFP4" if self.use_fp4_kv else "FP8", ) + if vllm_config.kernel_config.enable_jit_warmup: + from vllm.utils.import_utils import has_cutedsl + + if current_platform.is_cuda() and has_cutedsl(): + from vllm.models.deepseek_v4.nvidia.ops.fused_indexer_q_cutedsl import ( # noqa: E501 + _INDEXER_Q_FP8_KERNEL, + _INDEXER_Q_MXFP4_KERNEL, + ) + + indexer_q_kernel = ( + _INDEXER_Q_MXFP4_KERNEL + if self.use_fp4_kv + else _INDEXER_Q_FP8_KERNEL + ) + indexer_q_kernel.register_warmup() # no tensor parallel, just replicated self.wq_b = ReplicatedLinear( @@ -1016,7 +1037,6 @@ class DeepseekV4Indexer(nn.Module): ] if vllm_config.kernel_config.enable_jit_warmup: - from vllm.platforms import current_platform from vllm.utils.import_utils import has_cutedsl if not has_cutedsl() and not current_platform.is_xpu(): diff --git a/vllm/models/deepseek_v4/common/ops/cache_utils.py b/vllm/models/deepseek_v4/common/ops/cache_utils.py index 1bd0706f0f..d98d6a297a 100644 --- a/vllm/models/deepseek_v4/common/ops/cache_utils.py +++ b/vllm/models/deepseek_v4/common/ops/cache_utils.py @@ -24,6 +24,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( ) from vllm.model_executor.warmup.jit_warmup import ( WarmupIntRange, + kernel_launcher, zip_inputs, ) from vllm.model_executor.warmup.jit_warmup_triton_helper import ( @@ -31,7 +32,6 @@ from vllm.model_executor.warmup.jit_warmup_triton_helper import ( TritonPointerInputVariant, TritonWarmupTensor, VllmTritonJitKernel, - kernel_launcher, triton_scalar_specialization_rep, ) from vllm.platforms import current_platform @@ -521,11 +521,17 @@ def dequantize_and_gather_k_cache( if has_cutedsl(): # lazily import, otherwise some tests fail due to CUDA driver init failure. from vllm.models.deepseek_v4.nvidia.ops.dequant_gather_k_cutedsl import ( - dequantize_and_gather_k_cache_cutedsl, + _DEQUANT_GATHER_K_CACHE_CUTEDSL_KERNEL, ) - dequantize_and_gather_k_cache_cutedsl( - out, k_cache, seq_lens, gather_lens, block_table, block_size, offset + _DEQUANT_GATHER_K_CACHE_CUTEDSL_KERNEL( + out=out, + k_cache=k_cache, + seq_lens=seq_lens, + gather_lens=gather_lens, + block_table=block_table, + block_size=block_size, + offset=offset, ) return diff --git a/vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py b/vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py index e0bc2a0428..326a94f4f4 100644 --- a/vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py +++ b/vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py @@ -26,11 +26,11 @@ from typing import Any import torch +from vllm.model_executor.warmup.jit_warmup import kernel_launcher from vllm.model_executor.warmup.jit_warmup_triton_helper import ( LaunchSpec, TritonWarmupTensor, VllmTritonJitKernel, - kernel_launcher, ) from vllm.platforms import current_platform from vllm.triton_utils import tl, triton diff --git a/vllm/models/deepseek_v4/common/ops/fused_indexer_q.py b/vllm/models/deepseek_v4/common/ops/fused_indexer_q.py index ae014750bf..cf6e7235f9 100644 --- a/vllm/models/deepseek_v4/common/ops/fused_indexer_q.py +++ b/vllm/models/deepseek_v4/common/ops/fused_indexer_q.py @@ -6,11 +6,11 @@ from typing import Any import torch +from vllm.model_executor.warmup.jit_warmup import kernel_launcher from vllm.model_executor.warmup.jit_warmup_triton_helper import ( LaunchSpec, TritonWarmupTensor, VllmTritonJitKernel, - kernel_launcher, ) from vllm.platforms import current_platform from vllm.triton_utils import tl, triton @@ -598,19 +598,19 @@ def fused_indexer_q_rope_quant( if has_cutedsl(): # lazily import, otherwise some tests fail due to CUDA driver init failure. from vllm.models.deepseek_v4.nvidia.ops.fused_indexer_q_cutedsl import ( - fused_indexer_q_rope_quant_mxfp4_cutedsl, + _INDEXER_Q_MXFP4_KERNEL, ) - fused_indexer_q_rope_quant_mxfp4_cutedsl( - positions, - index_q, - index_q_cos_sin_cache, - index_weights, - index_weights_softmax_scale, - index_weights_head_scale, - index_q_packed, - index_q_scale, - index_weights_out, + _INDEXER_Q_MXFP4_KERNEL( + positions=positions, + q=index_q, + cos_sin_cache=index_q_cos_sin_cache, + weights=index_weights, + weights_softmax_scale=index_weights_softmax_scale, + weights_head_scale=index_weights_head_scale, + q_packed=index_q_packed, + q_scale=index_q_scale, + weights_out=index_weights_out, ) elif current_platform.is_xpu(): torch.ops.vllm.xpu_deepseek_fused_indexer_q_rope_mxfp4( @@ -654,18 +654,18 @@ def fused_indexer_q_rope_quant( if has_cutedsl(): # lazily import, otherwise some tests fail due to CUDA driver init failure. from vllm.models.deepseek_v4.nvidia.ops.fused_indexer_q_cutedsl import ( - fused_indexer_q_rope_quant_fp8_cutedsl, + _INDEXER_Q_FP8_KERNEL, ) - fused_indexer_q_rope_quant_fp8_cutedsl( - positions, - index_q, - index_q_cos_sin_cache, - index_weights, - index_weights_softmax_scale, - index_weights_head_scale, - index_q_fp8, - index_weights_out, + _INDEXER_Q_FP8_KERNEL( + positions=positions, + q=index_q, + cos_sin_cache=index_q_cos_sin_cache, + weights=index_weights, + weights_softmax_scale=index_weights_softmax_scale, + weights_head_scale=index_weights_head_scale, + q_fp8=index_q_fp8.view(torch.uint8), + weights_out=index_weights_out, ) elif current_platform.is_xpu(): torch.ops.vllm.xpu_deepseek_fused_indexer_q_rope_fp8( diff --git a/vllm/models/deepseek_v4/common/ops/fused_inv_rope_fp8_quant.py b/vllm/models/deepseek_v4/common/ops/fused_inv_rope_fp8_quant.py index 948d84be5c..77b1aa2236 100644 --- a/vllm/models/deepseek_v4/common/ops/fused_inv_rope_fp8_quant.py +++ b/vllm/models/deepseek_v4/common/ops/fused_inv_rope_fp8_quant.py @@ -12,11 +12,11 @@ from typing import Any import torch +from vllm.model_executor.warmup.jit_warmup import kernel_launcher from vllm.model_executor.warmup.jit_warmup_triton_helper import ( LaunchSpec, TritonWarmupTensor, VllmTritonJitKernel, - kernel_launcher, ) from vllm.platforms import current_platform from vllm.triton_utils import tl, triton diff --git a/vllm/models/deepseek_v4/common/ops/fused_mtp_input_rmsnorm.py b/vllm/models/deepseek_v4/common/ops/fused_mtp_input_rmsnorm.py index c71043f1ae..59ac4d27c6 100644 --- a/vllm/models/deepseek_v4/common/ops/fused_mtp_input_rmsnorm.py +++ b/vllm/models/deepseek_v4/common/ops/fused_mtp_input_rmsnorm.py @@ -22,11 +22,11 @@ from typing import Any import torch +from vllm.model_executor.warmup.jit_warmup import kernel_launcher from vllm.model_executor.warmup.jit_warmup_triton_helper import ( LaunchSpec, TritonWarmupTensor, VllmTritonJitKernel, - kernel_launcher, ) from vllm.triton_utils import tl, triton from vllm.utils.math_utils import next_power_of_2 diff --git a/vllm/models/deepseek_v4/common/ops/save_partial_states.py b/vllm/models/deepseek_v4/common/ops/save_partial_states.py index 0537feab29..6a554b9b4e 100644 --- a/vllm/models/deepseek_v4/common/ops/save_partial_states.py +++ b/vllm/models/deepseek_v4/common/ops/save_partial_states.py @@ -6,11 +6,11 @@ from typing import Any import torch +from vllm.model_executor.warmup.jit_warmup import kernel_launcher from vllm.model_executor.warmup.jit_warmup_triton_helper import ( LaunchSpec, TritonWarmupTensor, VllmTritonJitKernel, - kernel_launcher, ) from vllm.triton_utils import tl, triton from vllm.utils.math_utils import next_power_of_2 diff --git a/vllm/models/deepseek_v4/compressor.py b/vllm/models/deepseek_v4/compressor.py index 4da343b2c9..41af0c7660 100644 --- a/vllm/models/deepseek_v4/compressor.py +++ b/vllm/models/deepseek_v4/compressor.py @@ -203,8 +203,8 @@ class DeepseekCompressor(nn.Module): prologue (kv/score split, save_partial_states launch). The compress → norm → RoPE → store step is dispatched to a triton kernel (``compress_norm_rope_store_triton``) by default, except for the NVIDIA - head_dim=128 indexer path which uses the cutedsl kernel - (``compress_norm_rope_store_cutedsl``) for better performance. + head_dim=512 path which uses the CuTeDSL compressor kernels for better + performance. """ def __init__( @@ -317,7 +317,33 @@ class DeepseekCompressor(nn.Module): head_dim=self.head_dim, compress_ratio=self.compress_ratio, ) - if self.head_dim != 512: + if current_platform.is_cuda() and self.head_dim == 512: + from vllm.models.deepseek_v4.nvidia.ops.sparse_attn_compress_cutedsl import ( # noqa: E501 + _SPARSE_ATTN_COMPRESS_C128_BLOCK8_KERNEL, + _SPARSE_ATTN_COMPRESS_NORM_ROPE_STORE_C4_KERNEL, + _SPARSE_ATTN_COMPRESS_NORM_ROPE_STORE_FULL_C4_KERNEL, + _SPARSE_ATTN_NORM_ROPE_STORE_FULL_KERNEL, + _SPARSE_ATTN_NORM_ROPE_STORE_KERNEL, + ) + + store_full_kv = vllm_config.cache_config.cache_dtype != "fp8_ds_mla" + if self.compress_ratio == 4: + ( + _SPARSE_ATTN_COMPRESS_NORM_ROPE_STORE_FULL_C4_KERNEL + if store_full_kv + else _SPARSE_ATTN_COMPRESS_NORM_ROPE_STORE_C4_KERNEL + ).register_warmup() + else: + _SPARSE_ATTN_COMPRESS_C128_BLOCK8_KERNEL.register_warmup() + if store_full_kv: + _SPARSE_ATTN_NORM_ROPE_STORE_FULL_KERNEL.register_warmup() + else: + _SPARSE_ATTN_NORM_ROPE_STORE_KERNEL.register_warmup( + vllm_config, + k_cache_prefix=self.k_cache_prefix, + compress_ratio=self.compress_ratio, + ) + else: from vllm.models.deepseek_v4.common.ops.fused_compress_quant_cache import ( # noqa: E501 _FUSED_KV_COMPRESS_NORM_ROPE_INSERT_INDEXER_TRITON_KERNEL, ) @@ -419,13 +445,13 @@ class DeepseekCompressor(nn.Module): compress_norm_rope_store_fn: Any if current_platform.is_cuda() and self.head_dim == 512: from .nvidia.ops.sparse_attn_compress_cutedsl import ( - compress_norm_rope_store_cutedsl, + _SPARSE_ATTN_COMPRESSOR_CUTEDSL_KERNEL, ) # head=512 on CUDA always uses cutedsl, for both the fp8_ds_mla # layout and the plain full-cache layout. The full-cache flags # are consumed only here. - compress_norm_rope_store_fn = compress_norm_rope_store_cutedsl + compress_norm_rope_store_fn = _SPARSE_ATTN_COMPRESSOR_CUTEDSL_KERNEL extra_kwargs: dict[str, Any] = dict( store_full_kv=store_full_kv, store_full_fp8=store_full_fp8, diff --git a/vllm/models/deepseek_v4/nvidia/ops/dequant_gather_k_cutedsl.py b/vllm/models/deepseek_v4/nvidia/ops/dequant_gather_k_cutedsl.py index a4514293c9..fcb12740f2 100644 --- a/vllm/models/deepseek_v4/nvidia/ops/dequant_gather_k_cutedsl.py +++ b/vllm/models/deepseek_v4/nvidia/ops/dequant_gather_k_cutedsl.py @@ -1,7 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from functools import cache +from dataclasses import dataclass +from typing import Any import cutlass import cutlass.cute as cute @@ -12,302 +13,340 @@ from cutlass.cute.nvgpu import cpasync from quack.compile_utils import make_fake_tensor from vllm.cute_utils import _bf16x2_mul, cvt +from vllm.model_executor.warmup.jit_warmup import kernel_launcher, zip_inputs +from vllm.model_executor.warmup.jit_warmup_cutedsl_helper import ( + CuTeDSLLaunchSpec, + VllmCuTeDSLJitKernel, +) -def dequantize_and_gather_k_cache_cutedsl( - out: torch.Tensor, - k_cache: torch.Tensor, - seq_lens: torch.Tensor, - gather_lens: torch.Tensor | None, - block_table: torch.Tensor, - block_size: int, - offset: int, -) -> None: - DequantGatherKCacheKernel.compile( - block_size=block_size, - has_gather_lens=gather_lens is not None, - )(out, k_cache, seq_lens, gather_lens, block_table, offset) - - -class DequantGatherKCacheKernel: +class DequantGatherKCacheKernel( + VllmCuTeDSLJitKernel["DequantGatherKCacheKernel.CompileKey"] +): # Hard-coded for DSv4. head_dim = 512 + fp8_dim = 448 group_size = 64 # 1 scale per 64 elems + num_warps = 4 + tb_size = num_warps * 32 + num_stages = 4 - def __init__(self, fp8_dim: int = 448, block_size: int = 64): - self.fp8_dim = fp8_dim - self.bf16_dim = self.head_dim - fp8_dim - self.data_dim = fp8_dim + self.bf16_dim * 2 - self.block_size = block_size + @dataclass(frozen=True) + class CompileKey: + block_size: int + has_gather_lens: bool - self.num_warps = 4 - self.tb_size = self.num_warps * 32 - self.num_stages = 4 + @staticmethod + def kernel(compile_key: CompileKey) -> Any: + head_dim = DequantGatherKCacheKernel.head_dim + fp8_dim = DequantGatherKCacheKernel.fp8_dim + group_size = DequantGatherKCacheKernel.group_size + num_warps = DequantGatherKCacheKernel.num_warps + tb_size = DequantGatherKCacheKernel.tb_size + num_stages = DequantGatherKCacheKernel.num_stages + bf16_dim = head_dim - fp8_dim + data_dim = fp8_dim + bf16_dim * 2 + block_size = compile_key.block_size - @cute.jit - def __call__( - self, - out: cute.Tensor, - k_cache: cute.Tensor, - seq_lens: cute.Tensor, - gather_lens: cute.Tensor | None, - block_table: cute.Tensor, - offset: Int32, - stream: CUstream, - ): - # Split k_cache into k_data and k_scale. Each [block_size, head_bytes] - # block is actually a concat of - # [block_size, fp8_dim + bf16_dim * 2] and [block_size, 8]. - k_data = cute.make_tensor( - k_cache.iterator, - layout=cute.make_layout( - (k_cache.shape[0], self.block_size, self.data_dim), - stride=(k_cache.stride[0], self.data_dim, 1), - ), - ) - k_scale = cute.make_tensor( - k_cache.iterator + (self.block_size * self.data_dim), - layout=cute.make_layout( - (k_cache.shape[0], self.block_size, 8), - stride=(k_cache.stride[0], 8, 1), - ), - ) + @cute.jit + def load_g2s( + k_data_slice: cute.Tensor, + k_scale: cute.Tensor, + block_table: cute.Tensor, + s_kdata_slice: cute.Tensor, + s_kscale: cute.Tensor, + req_id, + pos, + lane_id, + stage_id, + ): + op = cpasync.CopyG2SOp(cute.nvgpu.LoadCacheMode.GLOBAL) + cp16_atom = cute.make_copy_atom(op, Uint32, num_bits_per_copy=128) + cp8_atom = cute.make_copy_atom( + cpasync.CopyG2SOp(), Uint8, num_bits_per_copy=64 + ) + page_id = block_table[req_id, pos // block_size] + block_offset = pos % block_size - num_reqs = out.shape[0] - max_rows = out.shape[1] - offset - num_workers = cutlass.max( - 1, cutlass.min(cute.ceil_div(max_rows, 4), cute.ceil_div(8192, num_reqs)) - ) - grid = (num_reqs, num_workers, 1) - self.kernel( - out, - k_data, - k_scale, - seq_lens, - gather_lens, - block_table, - offset, - ).launch(grid=grid, block=(self.tb_size, 1, 1), stream=stream) - - @cute.jit - def load_g2s( - self, - k_data_slice: cute.Tensor, - k_scale: cute.Tensor, - block_table: cute.Tensor, - s_kdata_slice: cute.Tensor, - s_kscale: cute.Tensor, - req_id, - pos, - lane_id, - stage_id, - ): - # k_data_slice: [num_blocks, block_size, (16, data_dim/16)] - # s_kdata_slice: [(4, data_dim/16), num_stages] - - op = cpasync.CopyG2SOp(cute.nvgpu.LoadCacheMode.GLOBAL) - cp16_atom = cute.make_copy_atom(op, Uint32, num_bits_per_copy=128) - cp8_atom = cute.make_copy_atom(cpasync.CopyG2SOp(), Uint8, num_bits_per_copy=64) - page_id = block_table[req_id, pos // self.block_size] - block_offset = pos % self.block_size - - # Load the first 512 bytes (32x16B). - idx = lane_id - src = k_data_slice[page_id, block_offset, (None, idx)] - cute.copy( - cp16_atom, - cute.recast_tensor(src, Uint32), - s_kdata_slice[(None, idx), stage_id], - ) - - # Load the tail 64 bytes. - idx += 32 - if idx < cutlass.const_expr(self.data_dim // 16): + # Load the first 512 bytes (32x16B). + idx = lane_id src = k_data_slice[page_id, block_offset, (None, idx)] cute.copy( cp16_atom, cute.recast_tensor(src, Uint32), s_kdata_slice[(None, idx), stage_id], ) - elif idx == cutlass.const_expr(self.data_dim // 16): - cute.copy( - cp8_atom, - k_scale[page_id, block_offset, None], - s_kscale[None, stage_id], - ) - @cute.kernel - def kernel( - self, - out: cute.Tensor, - k_data: cute.Tensor, - k_scale: cute.Tensor, - seq_lens: cute.Tensor, - gather_lens: cute.Tensor | None, - block_table: cute.Tensor, - offset: Int32, - ): - req_id, worker_id, _ = cute.arch.block_idx() - tid, _, _ = cute.arch.thread_idx() - warp_id = cute.arch.make_warp_uniform(tid // 32) - lane_id = tid % 32 - - _, num_workers, _ = cute.arch.grid_dim() - - # Prepare smem. - smem = cutlass.utils.SmemAllocator() - s_kdata = smem.allocate_tensor( - Uint32, - cute.make_layout((self.data_dim // 4, self.num_warps, self.num_stages)), - byte_alignment=16, - )[None, warp_id, None] - s_kscale = smem.allocate_tensor( - Uint8, - cute.make_layout((8, self.num_warps, self.num_stages)), - byte_alignment=8, - )[None, warp_id, None] - - # Prepare for 16B cp.async, also for BF16 smem loads later. - k_data_slice = cute.logical_divide(k_data, (None, None, 16)) - s_kdata_16B_slice = cute.logical_divide(s_kdata, (4, None)) - - # Load FP8 elems in 8B units, so once dequantized, they are 16B units. - s_kdata_8B_slice = cute.logical_divide(s_kdata, (2, None)) - - # 16B st.global. - out_slice = cute.logical_divide(out, (None, None, 8)) - - cp_op = cute.nvgpu.CopyUniversalOp() - cp8_atom = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=64) - cp16_atom = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=128) - - seq_len = seq_lens[req_id] - gather_len = seq_len - if cutlass.const_expr(gather_lens is not None): - gather_len = gather_lens[req_id] # type: ignore[index] - start_pos = seq_len - gather_len - - # Start prefetch. - for i in cutlass.range_constexpr(self.num_stages - 1): - next_pos = ( - start_pos - + worker_id * self.num_warps - + warp_id - + i * num_workers * self.num_warps - ) - if next_pos < seq_len: - self.load_g2s( - k_data_slice, - k_scale, - block_table, - s_kdata_16B_slice, - s_kscale, - req_id, - next_pos, - lane_id, - i, - ) - cute.arch.cp_async_commit_group() - prefetch_stage = self.num_stages - 1 - compute_stage = 0 - - # Main loop. - for i in range( - worker_id * self.num_warps + warp_id, - gather_len, - num_workers * self.num_warps, - ): - pos = start_pos + i - - # Prefetch next stage. - next_pos = pos + num_workers * self.num_warps * (self.num_stages - 1) - if next_pos < seq_len: - self.load_g2s( - k_data_slice, - k_scale, - block_table, - s_kdata_16B_slice, - s_kscale, - req_id, - next_pos, - lane_id, - prefetch_stage, - ) - prefetch_stage = (prefetch_stage + 1) % self.num_stages - cute.arch.cp_async_commit_group() - - # Wait for gmem->smem to finish. - cute.arch.cp_async_wait_group(self.num_stages - 1) - cute.arch.sync_warp() - - # There are 512 elems per token. As a warp, data0 holds the first - # 256 elems and data1 holds the second 256 elems, i.e. each thread - # holds 8 FP8 elems. This keeps the dequantized 8 BF16 elems as - # contiguous 16B global stores. On Blackwell, this might not be - # necessary as we have 32B global stores, but doing it this way - # does not seem to be slower. - data0 = cute.make_rmem_tensor((2,), Uint32) - data1 = cute.make_rmem_tensor((2,), Uint32) - cute.copy(cp8_atom, s_kdata_8B_slice[(None, lane_id), compute_stage], data0) - cute.copy( - cp8_atom, - s_kdata_8B_slice[(None, lane_id + 32), compute_stage], - data1, - ) - - # Convert to bf16x2 via bit manipulation. FP8 scales are per 64 - # elements. An 8-element chunk advances the scale index by - # chunk_id * 8 // group_size. - scale0_u32 = Uint32(s_kscale[lane_id * 8 // self.group_size, compute_stage]) - scale0_bf16x2 = (scale0_u32 << Uint32(23)) | (scale0_u32 << Uint32(7)) - scale1_u32 = Uint32( - s_kscale[(lane_id + 32) * 8 // self.group_size, compute_stage] - ) - scale1_bf16x2 = (scale1_u32 << Uint32(23)) | (scale1_u32 << Uint32(7)) - - # cvt.rn.scaled::n2::ue8m0.bf16x2.e4m3x2 requires PTX 9.2 - # (CUDA 13.2). - dequant0 = cute.make_rmem_tensor(4, Uint32) - dequant1 = cute.make_rmem_tensor(4, Uint32) - for j in cutlass.range_constexpr(2): - tmp0 = cvt.fp8x4_to_bf16x4(data0[j]) - tmp1 = cvt.fp8x4_to_bf16x4(data1[j]) - - # BF16 multiply is safe because the scales are exact powers of 2. - dequant0[j * 2] = _bf16x2_mul(tmp0[0], scale0_bf16x2) - dequant1[j * 2] = _bf16x2_mul(tmp1[0], scale1_bf16x2) - dequant0[j * 2 + 1] = _bf16x2_mul(tmp0[1], scale0_bf16x2) - dequant1[j * 2 + 1] = _bf16x2_mul(tmp1[1], scale1_bf16x2) - - # Last 64 elems are BF16 tail, corresponds to dequant1 of last - # 8 threads. We have 448 FP8 + 64 BF16 -> 28x 16B for FP8 + - # 8x 16B for BF16. - if lane_id + 32 >= self.fp8_dim // 8: - idx = self.fp8_dim // 16 + (lane_id + 32) - self.fp8_dim // 8 + # Load the tail 64 bytes. + idx += 32 + if idx < cutlass.const_expr(data_dim // 16): + src = k_data_slice[page_id, block_offset, (None, idx)] cute.copy( cp16_atom, - s_kdata_16B_slice[(None, idx), compute_stage], - dequant1, + cute.recast_tensor(src, Uint32), + s_kdata_slice[(None, idx), stage_id], + ) + elif idx == cutlass.const_expr(data_dim // 16): + cute.copy( + cp8_atom, + k_scale[page_id, block_offset, None], + s_kscale[None, stage_id], ) - # Store two 16B BF16 chunks per lane: first half, then second half. - dst = out_slice[req_id, offset + i, (None, lane_id)] - cute.copy(cp16_atom, dequant0, cute.recast_tensor(dst, Uint32)) + @cute.kernel + def device_kernel( + out: cute.Tensor, + k_data: cute.Tensor, + k_scale: cute.Tensor, + seq_lens: cute.Tensor, + gather_lens: cute.Tensor | None, + block_table: cute.Tensor, + offset: Int32, + ): + req_id, worker_id, _ = cute.arch.block_idx() + tid, _, _ = cute.arch.thread_idx() + warp_id = cute.arch.make_warp_uniform(tid // 32) + lane_id = tid % 32 - dst = out_slice[req_id, offset + i, (None, lane_id + 32)] - cute.copy(cp16_atom, dequant1, cute.recast_tensor(dst, Uint32)) + _, num_workers, _ = cute.arch.grid_dim() - compute_stage = (compute_stage + 1) % self.num_stages + # Prepare smem. + smem = cutlass.utils.SmemAllocator() + s_kdata = smem.allocate_tensor( + Uint32, + cute.make_layout((data_dim // 4, num_warps, num_stages)), + byte_alignment=16, + )[None, warp_id, None] + s_kscale = smem.allocate_tensor( + Uint8, + cute.make_layout((8, num_warps, num_stages)), + byte_alignment=8, + )[None, warp_id, None] - @cache - @staticmethod - def compile( - fp8_dim: int = 448, - block_size: int = 64, - has_gather_lens: bool = True, - ): + # Prepare for 16B cp.async, also for BF16 smem loads later. + k_data_slice = cute.logical_divide(k_data, (None, None, 16)) + s_kdata_16B_slice = cute.logical_divide(s_kdata, (4, None)) + + # Load FP8 elems in 8B units, so once dequantized, they are 16B units. + s_kdata_8B_slice = cute.logical_divide(s_kdata, (2, None)) + + # 16B st.global. + out_slice = cute.logical_divide(out, (None, None, 8)) + + cp_op = cute.nvgpu.CopyUniversalOp() + cp8_atom = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=64) + cp16_atom = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=128) + + seq_len = seq_lens[req_id] + gather_len = seq_len + if cutlass.const_expr(gather_lens is not None): + gather_len = gather_lens[req_id] # type: ignore[index] + start_pos = seq_len - gather_len + + # Start prefetch. + for i in cutlass.range_constexpr(num_stages - 1): + next_pos = ( + start_pos + + worker_id * num_warps + + warp_id + + i * num_workers * num_warps + ) + if next_pos < seq_len: + load_g2s( + k_data_slice, + k_scale, + block_table, + s_kdata_16B_slice, + s_kscale, + req_id, + next_pos, + lane_id, + i, + ) + cute.arch.cp_async_commit_group() + prefetch_stage = num_stages - 1 + compute_stage = 0 + + # Main loop. + for i in range( + worker_id * num_warps + warp_id, + gather_len, + num_workers * num_warps, + ): + pos = start_pos + i + + # Prefetch next stage. + next_pos = pos + num_workers * num_warps * (num_stages - 1) + if next_pos < seq_len: + load_g2s( + k_data_slice, + k_scale, + block_table, + s_kdata_16B_slice, + s_kscale, + req_id, + next_pos, + lane_id, + prefetch_stage, + ) + prefetch_stage = (prefetch_stage + 1) % num_stages + cute.arch.cp_async_commit_group() + + # Wait for gmem->smem to finish. + cute.arch.cp_async_wait_group(num_stages - 1) + cute.arch.sync_warp() + + # There are 512 elems per token. As a warp, data0 holds the first + # 256 elems and data1 holds the second 256 elems, i.e. each thread + # holds 8 FP8 elems. This keeps the dequantized 8 BF16 elems as + # contiguous 16B global stores. On Blackwell, this might not be + # necessary as we have 32B global stores, but doing it this way + # does not seem to be slower. + data0 = cute.make_rmem_tensor((2,), Uint32) + data1 = cute.make_rmem_tensor((2,), Uint32) + cute.copy( + cp8_atom, s_kdata_8B_slice[(None, lane_id), compute_stage], data0 + ) + cute.copy( + cp8_atom, + s_kdata_8B_slice[(None, lane_id + 32), compute_stage], + data1, + ) + + # Convert to bf16x2 via bit manipulation. FP8 scales are per 64 + # elements. An 8-element chunk advances the scale index by + # chunk_id * 8 // group_size. + scale0_u32 = Uint32(s_kscale[lane_id * 8 // group_size, compute_stage]) + scale0_bf16x2 = (scale0_u32 << Uint32(23)) | (scale0_u32 << Uint32(7)) + scale1_u32 = Uint32( + s_kscale[(lane_id + 32) * 8 // group_size, compute_stage] + ) + scale1_bf16x2 = (scale1_u32 << Uint32(23)) | (scale1_u32 << Uint32(7)) + + # cvt.rn.scaled::n2::ue8m0.bf16x2.e4m3x2 requires PTX 9.2 + # (CUDA 13.2). + dequant0 = cute.make_rmem_tensor(4, Uint32) + dequant1 = cute.make_rmem_tensor(4, Uint32) + for j in cutlass.range_constexpr(2): + tmp0 = cvt.fp8x4_to_bf16x4(data0[j]) + tmp1 = cvt.fp8x4_to_bf16x4(data1[j]) + + # BF16 multiply is safe because the scales are exact powers of 2. + dequant0[j * 2] = _bf16x2_mul(tmp0[0], scale0_bf16x2) + dequant1[j * 2] = _bf16x2_mul(tmp1[0], scale1_bf16x2) + dequant0[j * 2 + 1] = _bf16x2_mul(tmp0[1], scale0_bf16x2) + dequant1[j * 2 + 1] = _bf16x2_mul(tmp1[1], scale1_bf16x2) + + # Last 64 elems are BF16 tail, corresponds to dequant1 of last + # 8 threads. We have 448 FP8 + 64 BF16 -> 28x 16B for FP8 + + # 8x 16B for BF16. + if lane_id + 32 >= fp8_dim // 8: + idx = fp8_dim // 16 + (lane_id + 32) - fp8_dim // 8 + cute.copy( + cp16_atom, + s_kdata_16B_slice[(None, idx), compute_stage], + dequant1, + ) + + # Store two 16B BF16 chunks per lane: first half, then second half. + dst = out_slice[req_id, offset + i, (None, lane_id)] + cute.copy(cp16_atom, dequant0, cute.recast_tensor(dst, Uint32)) + + dst = out_slice[req_id, offset + i, (None, lane_id + 32)] + cute.copy(cp16_atom, dequant1, cute.recast_tensor(dst, Uint32)) + + compute_stage = (compute_stage + 1) % num_stages + + @cute.jit + def host_entrypoint( + out: cute.Tensor, + k_cache: cute.Tensor, + seq_lens: cute.Tensor, + gather_lens: cute.Tensor | None, + block_table: cute.Tensor, + offset: Int32, + stream: CUstream, + ): + # Split k_cache into k_data and k_scale. Each [block_size, head_bytes] + # block is actually a concat of + # [block_size, fp8_dim + bf16_dim * 2] and [block_size, 8]. + k_data = cute.make_tensor( + k_cache.iterator, + layout=cute.make_layout( + (k_cache.shape[0], block_size, data_dim), + stride=(k_cache.stride[0], data_dim, 1), + ), + ) + # k_data_slice: [num_blocks, block_size, (16, data_dim/16)] + # s_kdata_slice: [(4, data_dim/16), num_stages] + k_scale = cute.make_tensor( + k_cache.iterator + (block_size * data_dim), + layout=cute.make_layout( + (k_cache.shape[0], block_size, 8), + stride=(k_cache.stride[0], 8, 1), + ), + ) + + num_reqs = out.shape[0] + max_rows = out.shape[1] - offset + num_workers = cutlass.max( + 1, + cutlass.min(cute.ceil_div(max_rows, 4), cute.ceil_div(8192, num_reqs)), + ) + grid = (num_reqs, num_workers, 1) + device_kernel( + out, + k_data, + k_scale, + seq_lens, + gather_lens, + block_table, + offset, + ).launch(grid=grid, block=(tb_size, 1, 1), stream=stream) + + return host_entrypoint + + def dispatch( # type: ignore[override] + self, + *, + block_size: int, + has_gather_lens: bool, + ) -> CompileKey: + return self.CompileKey( + block_size=block_size, + has_gather_lens=has_gather_lens, + ) + + def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]: + block_size = vllm_config.cache_config.block_size + compress_ratios = vllm_config.model_config.hf_config.compress_ratios + return self._trace_dispatch(self.dispatch)( + zip_inputs( + dict( + block_size=block_size, + has_gather_lens=True, + enabled=True, + ), + dict( + block_size=max(1, block_size // 4), + has_gather_lens=False, + enabled=4 in compress_ratios, + ), + dict( + block_size=max(1, block_size // 128), + has_gather_lens=False, + enabled=128 in compress_ratios, + ), + ), + valid_block_size=block_size > 0, + _when=lambda *, enabled, valid_block_size: enabled and valid_block_size, + ) + + def warmup_inputs(self, compile_key: CompileKey) -> tuple[Any, ...]: num_reqs = cute.sym_int() - head_dim = DequantGatherKCacheKernel.head_dim + head_dim = self.head_dim + fp8_dim = self.fp8_dim + block_size = compile_key.block_size head_bytes = fp8_dim + (head_dim - fp8_dim) * 2 + 8 out = make_fake_tensor(BFloat16, (num_reqs, cute.sym_int(), head_dim), 16) @@ -318,19 +357,47 @@ class DequantGatherKCacheKernel: assumed_align=32, ) seq_lens = make_fake_tensor(Int32, (num_reqs,)) - gather_lens = make_fake_tensor(Int32, (num_reqs,)) if has_gather_lens else None + gather_lens = ( + make_fake_tensor(Int32, (num_reqs,)) + if compile_key.has_gather_lens + else None + ) block_table = make_fake_tensor(Int32, (num_reqs, cute.sym_int())) - kernel = DequantGatherKCacheKernel(fp8_dim, block_size) - stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) - return cute.compile( - kernel, + return ( out, k_cache, seq_lens, gather_lens, block_table, Int32(0), - stream, - options="--enable-tvm-ffi", ) + + @kernel_launcher + def __call__( + self, + *, + out: torch.Tensor, + k_cache: torch.Tensor, + seq_lens: torch.Tensor, + gather_lens: torch.Tensor | None, + block_table: torch.Tensor, + block_size: int, + offset: int, + ) -> CuTeDSLLaunchSpec["DequantGatherKCacheKernel.CompileKey"]: + compile_key = self.dispatch( + block_size=block_size, + has_gather_lens=gather_lens is not None, + ) + launch_args = ( + out, + k_cache, + seq_lens, + gather_lens, + block_table, + offset, + ) + return compile_key, launch_args + + +_DEQUANT_GATHER_K_CACHE_CUTEDSL_KERNEL = DequantGatherKCacheKernel() diff --git a/vllm/models/deepseek_v4/nvidia/ops/fused_indexer_q_cutedsl.py b/vllm/models/deepseek_v4/nvidia/ops/fused_indexer_q_cutedsl.py index 16e98abf23..6684e37f37 100644 --- a/vllm/models/deepseek_v4/nvidia/ops/fused_indexer_q_cutedsl.py +++ b/vllm/models/deepseek_v4/nvidia/ops/fused_indexer_q_cutedsl.py @@ -1,6 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from functools import cache +# ruff: noqa: E501 + +from dataclasses import dataclass +from typing import Any import cutlass import cutlass.cute as cute @@ -14,594 +17,638 @@ from vllm.cute_utils import ( _bf16x2_max, cvt, recast_val, + torch_to_cute_dtype, +) +from vllm.model_executor.warmup.jit_warmup import kernel_launcher +from vllm.model_executor.warmup.jit_warmup_cutedsl_helper import ( + CuTeDSLLaunchSpec, + VllmCuTeDSLJitKernel, ) from vllm.vllm_flash_attn.cute import utils as cute_utils # MXFP4: 32 elements per block, packed 2 nibbles per byte, ue8m0 block scale. MXFP4_BLOCK_SIZE = 32 -_TORCH_TO_CUTE = { - torch.bfloat16: BFloat16, - torch.float32: Float32, -} +@cute.jit +def _load_q_and_rope( + positions: cute.Tensor, + q: cute.Tensor, + cos_sin_cache: cute.Tensor, + num_heads: cutlass.Constexpr, + coarsen: cutlass.Constexpr, + tb_size: cutlass.Constexpr, + subwarp_size: cutlass.Constexpr, + threads_per_token: cutlass.Constexpr, + nope_dim: cutlass.Constexpr, + rope_dim: cutlass.Constexpr, + cos_sin_dtype: cutlass.Constexpr, +): + """Compute thread indices, load Q (BF16), and apply interleaved RoPE. -def fused_indexer_q_rope_quant_mxfp4_cutedsl( - positions: torch.Tensor, - index_q: torch.Tensor, - index_q_cos_sin_cache: torch.Tensor, - index_weights: torch.Tensor, - index_weights_softmax_scale: float, - index_weights_head_scale: float, - index_q_packed: torch.Tensor, - index_q_scale: torch.Tensor, - index_weights_out: torch.Tensor, -) -> None: - num_tokens, num_heads, head_dim = index_q.shape - rope_dim = index_q_cos_sin_cache.shape[-1] - rope_type = _TORCH_TO_CUTE[index_q_cos_sin_cache.dtype] - - # compile all variants at first invocation - for coarsen in (1, 4): - IndexerQMxFp4Kernel.compile(head_dim, rope_dim, num_heads, rope_type, coarsen) - - # heuristic - coarsen = 1 if num_tokens < 512 else 4 - compiled = IndexerQMxFp4Kernel.compile( - head_dim, rope_dim, num_heads, rope_type, coarsen - ) - scale = float(index_weights_softmax_scale * index_weights_head_scale) - compiled( - positions, - index_q, - index_q_cos_sin_cache, - index_weights, - index_q_packed, - index_q_scale, - index_weights_out, - scale, - ) - - -def fused_indexer_q_rope_quant_fp8_cutedsl( - positions: torch.Tensor, - index_q: torch.Tensor, - index_q_cos_sin_cache: torch.Tensor, - index_weights: torch.Tensor, - index_weights_softmax_scale: float, - index_weights_head_scale: float, - index_q_fp8: torch.Tensor, - index_weights_out: torch.Tensor, -) -> None: - num_tokens, num_heads, head_dim = index_q.shape - rope_dim = index_q_cos_sin_cache.shape[-1] - rope_type = _TORCH_TO_CUTE[index_q_cos_sin_cache.dtype] - - for coarsen in (1, 4): - IndexerQFp8Kernel.compile(head_dim, rope_dim, num_heads, rope_type, coarsen) - - coarsen = 1 if num_tokens < 512 else 4 - compiled = IndexerQFp8Kernel.compile( - head_dim, rope_dim, num_heads, rope_type, coarsen - ) - scale = float(index_weights_softmax_scale * index_weights_head_scale) - # The cute kernel treats the FP8 buffer as raw bytes (Uint8). - compiled( - positions, - index_q, - index_q_cos_sin_cache, - index_weights, - index_q_fp8.view(torch.uint8), - index_weights_out, - scale, - ) - - -class IndexerQRopeQuantKernel: - """Shared infrastructure for indexer-Q RoPE+quant fused kernels. - - Subclasses implement ``kernel`` for a particular Q quantization scheme - (MXFP4, FP8 e4m3, …). The base class owns the launch geometry and the - common preamble: thread/token addressing, the BF16 Q load, and the - interleaved-RoPE pass over the trailing ``rope_dim`` lanes. + Returns a tuple + (q_bf16x2, tid, global_tid, sublane, token_id, head_tile_id, + head_start, in_bounds, num_token_heads) + where ``q_bf16x2`` is a (coarsen, 8) rmem tile of Uint32 packed + bf16x2 pairs covering the 16 BF16 lanes owned by this thread for + each of ``coarsen`` heads. RoPE is applied in place to the + trailing ``rope_dim`` lanes; the leading nope lanes pass through. """ + block_id, _, _ = cute.arch.block_idx() + tid, _, _ = cute.arch.thread_idx() - def __init__( - self, - head_dim: int = 128, - rope_dim: int = 64, - num_heads: int = 64, - cos_sin_dtype: type[cutlass.Numeric] = Float32, - coarsen: int = 4, - ): - self.head_dim = head_dim - self.rope_dim = rope_dim - self.nope_dim = head_dim - rope_dim - self.num_heads = num_heads - self.cos_sin_dtype = cos_sin_dtype + num_tokens = q.shape[0] + num_token_heads = num_tokens * num_heads + global_tid = block_id * tb_size + tid + global_subwarp_id = global_tid // subwarp_size + sublane = tid % subwarp_size + + token_id = global_subwarp_id // (num_heads // coarsen) + head_tile_id = global_subwarp_id % (num_heads // coarsen) + head_start = head_tile_id * coarsen + + # NOTE: token_id may exceed bounds, hence we need to add load/store guards + # we can't do early exit because CuteDSL doesn't support it. and we also need + # all threads in a warp to be active since we utilize warp shuffle later. + # must_in_bounds is constexpr, True when a token spans a whole number of + # threadblocks. Then num_tokens * threads_per_token is always a multiple of + # tb_size, so the trailing block never straddles num_tokens and the compiler + # removes the bounds check. Otherwise the guard has to stay: e.g. 32 heads + # with coarsen=4 needs 64 threads per token, so an odd num_tokens leaves half + # of the last block addressing token num_tokens. + must_in_bounds = cutlass.const_expr(threads_per_token % tb_size == 0) + in_bounds = must_in_bounds or (token_id < num_tokens) + + cp_op = cute.nvgpu.CopyUniversalOp() + + _layout = cute.make_layout((coarsen, 8), stride=(8, 1)) + q_bf16x2 = cute.make_rmem_tensor(_layout, Uint32) + + if in_bounds: + # we can't do cute.copy() on the whole 2D tile directly because + # cute.copy() wants the 1st mode to be covered by the copy atom, + # and other modes as for loop. there is no fast way to + # "transpose" the tensor view. + q_tile = cute.local_tile( + q[token_id, None, None], + tiler=(coarsen, 16), + coord=(head_tile_id, sublane), + ) + cp_u32x8 = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=256) + for i in cutlass.range_constexpr(coarsen): + src = cute.recast_tensor(q_tile[i, None], Uint32) + cute.copy(cp_u32x8, src, q_bf16x2[i, None]) + + # RoPE applies only to the trailing rope_dim values. We keep the rounded + # BF16 result in q_bits so the later amax and quantization see BF16. + # cos_sin_cache layout: [max_pos, rope_dim] + if in_bounds and sublane * 16 >= nope_dim: + cos_vals = cute.make_rmem_tensor((8,), Float32) + sin_vals = cute.make_rmem_tensor((8,), Float32) + + pos = positions[token_id] + + # select 8 elems from cos and sin + cos_id = sublane - nope_dim // 16 + sin_id = cos_id + rope_dim // 16 + cos_src = cute.local_tile(cos_sin_cache[pos, None], tiler=(8,), coord=(cos_id,)) + sin_src = cute.local_tile(cos_sin_cache[pos, None], tiler=(8,), coord=(sin_id,)) + + cp_f32x8 = cute.make_copy_atom(cp_op, Float32, num_bits_per_copy=256) + cp_u32x4 = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=128) + + if const_expr(cos_sin_dtype is Float32): + cute.copy(cp_f32x8, cos_src, cos_vals) + cute.copy(cp_f32x8, sin_src, sin_vals) + else: + cos_bf16x2 = cute.make_rmem_tensor((4,), Uint32) + sin_bf16x2 = cute.make_rmem_tensor((4,), Uint32) + cute.copy(cp_u32x4, cute.recast_tensor(cos_src, Uint32), cos_bf16x2) + cute.copy(cp_u32x4, cute.recast_tensor(sin_src, Uint32), sin_bf16x2) + + for i in cutlass.range_constexpr(4): + cos0, cos1 = cvt.bf16x2_to_fp32x2(cos_bf16x2[i]) + sin0, sin1 = cvt.bf16x2_to_fp32x2(sin_bf16x2[i]) + cos_vals[i * 2] = cos0 + cos_vals[i * 2 + 1] = cos1 + sin_vals[i * 2] = sin0 + sin_vals[i * 2 + 1] = sin1 + + for i in cutlass.range_constexpr(coarsen): + for j in cutlass.range_constexpr(8): + q0, q1 = cvt.bf16x2_to_fp32x2(q_bf16x2[i, j]) + rot0 = q0 * cos_vals[j] - q1 * sin_vals[j] + rot1 = q0 * sin_vals[j] + q1 * cos_vals[j] + # convert back to BF16 to match numerics + q_bf16x2[i, j] = cvt.fp32x2_to_bf16x2(rot0, rot1) + + return ( + q_bf16x2, + tid, + global_tid, + sublane, + token_id, + head_tile_id, + head_start, + in_bounds, + num_token_heads, + ) + + +class IndexerQMxFp4Kernel(VllmCuTeDSLJitKernel["IndexerQMxFp4Kernel.CompileKey"]): + tb_size = 128 + + @dataclass(frozen=True) + class CompileKey: + head_dim: int + rope_dim: int + num_heads: int + cos_sin_dtype: type[cutlass.Numeric] + coarsen: int + + @staticmethod + def kernel(compile_key: CompileKey) -> Any: + head_dim = compile_key.head_dim + rope_dim = compile_key.rope_dim + nope_dim = head_dim - rope_dim + num_heads = compile_key.num_heads + cos_sin_dtype = compile_key.cos_sin_dtype # process multiple heads at the same time to armotize RoPE load costs + coarsen = compile_key.coarsen + tb_size = IndexerQMxFp4Kernel.tb_size assert num_heads % coarsen == 0 - self.coarsen = coarsen # later we will use 32B load = 16 BF16 elems # thus, head_dim=128 requires 8 threads to handle. # let's call subwarp = 8 threads. - self.subwarp_size = head_dim // 16 - self.tb_size = 128 - self.threads_per_token = (self.num_heads // self.coarsen) * self.subwarp_size + subwarp_size = head_dim // 16 + threads_per_token = (num_heads // coarsen) * subwarp_size - @cute.jit - def _load_q_and_rope( - self, - positions: cute.Tensor, - q: cute.Tensor, - cos_sin_cache: cute.Tensor, - ): - """Compute thread indices, load Q (BF16), and apply interleaved RoPE. + @cute.kernel + def device_kernel( + positions: cute.Tensor, + q: cute.Tensor, + cos_sin_cache: cute.Tensor, + weights: cute.Tensor, + q_quant: cute.Tensor, + q_scale: cute.Tensor, + weights_out: cute.Tensor, + scale: Float32, + ): + ( + q_bf16x2, + tid, + global_tid, + sublane, + token_id, + head_tile_id, + head_start, + in_bounds, + num_token_heads, + ) = _load_q_and_rope( + positions, + q, + cos_sin_cache, + num_heads, + coarsen, + tb_size, + subwarp_size, + threads_per_token, + nope_dim, + rope_dim, + cos_sin_dtype, + ) - Returns a tuple - (q_bf16x2, tid, global_tid, sublane, token_id, head_tile_id, - head_start, in_bounds, num_token_heads) - where ``q_bf16x2`` is a (coarsen, 8) rmem tile of Uint32 packed - bf16x2 pairs covering the 16 BF16 lanes owned by this thread for - each of ``coarsen`` heads. RoPE is applied in place to the - trailing ``rope_dim`` lanes; the leading nope lanes pass through. - """ - block_id, _, _ = cute.arch.block_idx() - tid, _, _ = cute.arch.thread_idx() + cp_op = cute.nvgpu.CopyUniversalOp() - num_tokens = q.shape[0] - num_token_heads = num_tokens * self.num_heads - global_tid = block_id * self.tb_size + tid - - global_subwarp_id = global_tid // self.subwarp_size - sublane = tid % self.subwarp_size - - token_id = global_subwarp_id // (self.num_heads // self.coarsen) - head_tile_id = global_subwarp_id % (self.num_heads // self.coarsen) - head_start = head_tile_id * self.coarsen - - # NOTE: token_id may exceed bounds, hence we need to add load/store guards - # we can't do early exit because CuteDSL doesn't support it. and we also need - # all threads in a warp to be active since we utilize warp shuffle later. - # must_in_bounds is constexpr, True when a token spans a whole number of - # threadblocks. Then num_tokens * threads_per_token is always a multiple of - # tb_size, so the trailing block never straddles num_tokens and the compiler - # removes the bounds check. Otherwise the guard has to stay: e.g. 32 heads - # with coarsen=4 needs 64 threads per token, so an odd num_tokens leaves half - # of the last block addressing token num_tokens. - must_in_bounds = cutlass.const_expr(self.threads_per_token % self.tb_size == 0) - in_bounds = must_in_bounds or (token_id < num_tokens) - - cp_op = cute.nvgpu.CopyUniversalOp() - - _layout = cute.make_layout((self.coarsen, 8), stride=(8, 1)) - q_bf16x2 = cute.make_rmem_tensor(_layout, Uint32) - - if in_bounds: - # we can't do cute.copy() on the whole 2D tile directly because - # cute.copy() wants the 1st mode to be covered by the copy atom, - # and other modes as for loop. there is no fast way to - # "transpose" the tensor view. - q_tile = cute.local_tile( - q[token_id, None, None], - tiler=(self.coarsen, 16), + # layout: [coarsen, 8] + q_fp4_tile = cute.local_tile( + q_quant[token_id, None, None], + tiler=(coarsen, 8), coord=(head_tile_id, sublane), ) - cp_u32x8 = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=256) - for i in cutlass.range_constexpr(self.coarsen): - src = cute.recast_tensor(q_tile[i, None], Uint32) - cute.copy(cp_u32x8, src, q_bf16x2[i, None]) - # RoPE applies only to the trailing rope_dim values. We keep the rounded - # BF16 result in q_bits so the later amax and quantization see BF16. - # cos_sin_cache layout: [max_pos, rope_dim] - if in_bounds and sublane * 16 >= self.nope_dim: - cos_vals = cute.make_rmem_tensor((8,), Float32) - sin_vals = cute.make_rmem_tensor((8,), Float32) + for i in cutlass.range_constexpr(coarsen): + # compute amax in packed bf16x2 to save instructions + # Each thread holds 16 elems. Two adjacent threads form one 32-elem + # MXFP4 block, so a width-2 shuffle gives the block amax. + amax_bf16x2 = _bf16x2_abs(q_bf16x2[i, 0]) + for j in cutlass.range_constexpr(1, 8): + amax_bf16x2 = _bf16x2_max(amax_bf16x2, _bf16x2_abs(q_bf16x2[i, j])) + amax_bf16x2 = cute_utils.warp_reduce( + amax_bf16x2, + _bf16x2_max, + width=MXFP4_BLOCK_SIZE // 16, + ) + amax_pair = cvt.bf16x2_to_fp32x2(amax_bf16x2) + amax = cute_utils.fmax(amax_pair[0], amax_pair[1]) - pos = positions[token_id] + if in_bounds: + # compute block scale with bit manipulation + # UE8M0 stores ceil(log2(fp4_scale)) + 127. Adding the mantissa mask + # increments the exponent whenever fp4_scale is not exactly a power of 2 + eps = cutlass.const_expr(float.fromhex("0x6p-126")) + fp4_scale = cute_utils.fmax(amax, eps) * Float32(1.0 / 6.0) + bits = recast_val(fp4_scale, Uint32) + ue8m0 = cute_utils.shr_u32( + bits + Uint32(0x7FFFFF), Uint32(23) + ) & Uint32(0xFF) - # select 8 elems from cos and sin - cos_id = sublane - self.nope_dim // 16 - sin_id = cos_id + self.rope_dim // 16 - cos_src = cute.local_tile( - cos_sin_cache[pos, None], tiler=(8,), coord=(cos_id,) - ) - sin_src = cute.local_tile( - cos_sin_cache[pos, None], tiler=(8,), coord=(sin_id,) - ) + # Only one of the two threads in an MXFP4 block writes the shared scale. + if tid % 2 == 0: + mx_block = sublane // 2 + q_scale[token_id, head_start + i, mx_block] = Uint8(ue8m0) - cp_f32x8 = cute.make_copy_atom(cp_op, Float32, num_bits_per_copy=256) - cp_u32x4 = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=128) + # If scale = 2^A and ue8m0 = A + 127, then inverse scale has exponent + # -A + 127 = 254 - ue8m0. + inv_scale_bits = (Uint32(254) - ue8m0) << Uint32(23) + inv_fp4_scale = recast_val(inv_scale_bits, Float32) - if const_expr(self.cos_sin_dtype is Float32): - cute.copy(cp_f32x8, cos_src, cos_vals) - cute.copy(cp_f32x8, sin_src, sin_vals) - else: - cos_bf16x2 = cute.make_rmem_tensor((4,), Uint32) - sin_bf16x2 = cute.make_rmem_tensor((4,), Uint32) - cute.copy(cp_u32x4, cute.recast_tensor(cos_src, Uint32), cos_bf16x2) - cute.copy(cp_u32x4, cute.recast_tensor(sin_src, Uint32), sin_bf16x2) + vals = cute.make_rmem_tensor(16, Float32) + for j in cutlass.range_constexpr(8): + q0, q1 = cvt.bf16x2_to_fp32x2(q_bf16x2[i, j]) + vals[j * 2] = q0 * inv_fp4_scale + vals[j * 2 + 1] = q1 * inv_fp4_scale - for i in cutlass.range_constexpr(4): - cos0, cos1 = cvt.bf16x2_to_fp32x2(cos_bf16x2[i]) - sin0, sin1 = cvt.bf16x2_to_fp32x2(sin_bf16x2[i]) - cos_vals[i * 2] = cos0 - cos_vals[i * 2 + 1] = cos1 - sin_vals[i * 2] = sin0 - sin_vals[i * 2 + 1] = sin1 + # pack to FP4 + packed = cute.make_rmem_tensor((2,), Uint32) + packed[0] = cvt.fp32x8_to_fp4x8(vals, 0) + packed[1] = cvt.fp32x8_to_fp4x8(vals, 8) - for i in cutlass.range_constexpr(self.coarsen): - for j in cutlass.range_constexpr(8): - q0, q1 = cvt.bf16x2_to_fp32x2(q_bf16x2[i, j]) - rot0 = q0 * cos_vals[j] - q1 * sin_vals[j] - rot1 = q0 * sin_vals[j] + q1 * cos_vals[j] - # convert back to BF16 to match numerics - q_bf16x2[i, j] = cvt.fp32x2_to_bf16x2(rot0, rot1) + dst = q_fp4_tile[i, None] + cp_u32x2 = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=64) + cute.copy(cp_u32x2, packed, cute.recast_tensor(dst, Uint32)) - return ( - q_bf16x2, - tid, - global_tid, - sublane, - token_id, - head_tile_id, - head_start, - in_bounds, - num_token_heads, + # Weight scaling is independent of the Q subwarp work. The first + # num_tokens * num_heads logical threads cover one weight each. + if global_tid < num_token_heads: + weight_token_id = global_tid // num_heads + weight_head_id = global_tid % num_heads + weights_out[weight_token_id, weight_head_id] = ( + weights[weight_token_id, weight_head_id].to(Float32) * scale + ) + + @cute.jit + def host_entrypoint( + positions: cute.Tensor, + q: cute.Tensor, + cos_sin_cache: cute.Tensor, + weights: cute.Tensor, + q_quant: cute.Tensor, + q_scale: cute.Tensor, + weights_out: cute.Tensor, + scale: Float32, + stream: CUstream, + ): + total_threads = q.shape[0] * threads_per_token + grid = (cute.ceil_div(total_threads, tb_size), 1, 1) + device_kernel( + positions, + q, + cos_sin_cache, + weights, + q_quant, + q_scale, + weights_out, + scale, + ).launch(grid=grid, block=(tb_size, 1, 1), stream=stream) + + return host_entrypoint + + def dispatch( # type: ignore[override] + self, + *, + cos_sin_dtype: type[cutlass.Numeric], + **compile_key_fields: int, + ) -> CompileKey: + return self.CompileKey( + **compile_key_fields, + cos_sin_dtype=cos_sin_dtype, ) + def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]: + from vllm.v1.attention.backends.mla.indexer import dsa_indexer_uses_fp4 -class IndexerQMxFp4Kernel(IndexerQRopeQuantKernel): - """Eight-thread subwarps process one ``(token, head)`` row.""" + hf_config = vllm_config.model_config.hf_config + if hf_config is None or not dsa_indexer_uses_fp4(vllm_config): + return [] - @cute.jit - def __call__( - self, - positions: cute.Tensor, - q: cute.Tensor, - cos_sin_cache: cute.Tensor, - weights: cute.Tensor, - q_quant: cute.Tensor, - q_scale: cute.Tensor, - weights_out: cute.Tensor, - scale: Float32, - stream: CUstream, - ): - total_threads = q.shape[0] * self.threads_per_token - grid = (cute.ceil_div(total_threads, self.tb_size), 1, 1) - self.kernel( - positions, - q, - cos_sin_cache, - weights, - q_quant, - q_scale, - weights_out, - scale, - ).launch(grid=grid, block=(self.tb_size, 1, 1), stream=stream) + num_heads = int(getattr(hf_config, "index_n_heads", 0) or 0) + head_dim = int(getattr(hf_config, "index_head_dim", 0) or 0) + rope_dim = int(getattr(hf_config, "qk_rope_head_dim", 0) or 0) + if num_heads <= 0 or head_dim <= 0 or rope_dim <= 0: + return [] - @cute.kernel - def kernel( - self, - positions: cute.Tensor, - q: cute.Tensor, - cos_sin_cache: cute.Tensor, - weights: cute.Tensor, - q_quant: cute.Tensor, - q_scale: cute.Tensor, - weights_out: cute.Tensor, - scale: Float32, - ): - ( - q_bf16x2, - tid, - global_tid, - sublane, - token_id, - head_tile_id, - head_start, - in_bounds, - num_token_heads, - ) = self._load_q_and_rope(positions, q, cos_sin_cache) - - cp_op = cute.nvgpu.CopyUniversalOp() - - # layout: [coarsen, 8] - q_fp4_tile = cute.local_tile( - q_quant[token_id, None, None], - tiler=(self.coarsen, 8), - coord=(head_tile_id, sublane), + return self._trace_dispatch(self.dispatch)( + head_dim=head_dim, + rope_dim=rope_dim, + num_heads=num_heads, + cos_sin_dtype=Float32, + coarsen=(1, 4), ) - for i in cutlass.range_constexpr(self.coarsen): - # compute amax in packed bf16x2 to save instructions - # Each thread holds 16 elems. Two adjacent threads form one 32-elem - # MXFP4 block, so a width-2 shuffle gives the block amax. - amax_bf16x2 = _bf16x2_abs(q_bf16x2[i, 0]) - for j in cutlass.range_constexpr(1, 8): - amax_bf16x2 = _bf16x2_max(amax_bf16x2, _bf16x2_abs(q_bf16x2[i, j])) - amax_bf16x2 = cute_utils.warp_reduce( - amax_bf16x2, - _bf16x2_max, - width=MXFP4_BLOCK_SIZE // 16, - ) - amax_pair = cvt.bf16x2_to_fp32x2(amax_bf16x2) - amax = cute_utils.fmax(amax_pair[0], amax_pair[1]) - - if in_bounds: - # compute block scale with bit manipulation - # UE8M0 stores ceil(log2(fp4_scale)) + 127. Adding the mantissa mask - # increments the exponent whenever fp4_scale is not exactly a power of 2 - eps = cutlass.const_expr(float.fromhex("0x6p-126")) - fp4_scale = cute_utils.fmax(amax, eps) * Float32(1.0 / 6.0) - bits = recast_val(fp4_scale, Uint32) - ue8m0 = cute_utils.shr_u32( - bits + Uint32(0x7FFFFF), Uint32(23) - ) & Uint32(0xFF) - - # Only one of the two threads in an MXFP4 block writes the shared scale. - if tid % 2 == 0: - mx_block = sublane // 2 - q_scale[token_id, head_start + i, mx_block] = Uint8(ue8m0) - - # If scale = 2^A and ue8m0 = A + 127, then inverse scale has exponent - # -A + 127 = 254 - ue8m0. - inv_scale_bits = (Uint32(254) - ue8m0) << Uint32(23) - inv_fp4_scale = recast_val(inv_scale_bits, Float32) - - vals = cute.make_rmem_tensor(16, Float32) - for j in cutlass.range_constexpr(8): - q0, q1 = cvt.bf16x2_to_fp32x2(q_bf16x2[i, j]) - vals[j * 2] = q0 * inv_fp4_scale - vals[j * 2 + 1] = q1 * inv_fp4_scale - - # pack to FP4 - packed = cute.make_rmem_tensor((2,), Uint32) - packed[0] = cvt.fp32x8_to_fp4x8(vals, 0) - packed[1] = cvt.fp32x8_to_fp4x8(vals, 8) - - dst = q_fp4_tile[i, None] - cp_u32x2 = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=64) - cute.copy(cp_u32x2, packed, cute.recast_tensor(dst, Uint32)) - - # Weight scaling is independent of the Q subwarp work. The first - # num_tokens * num_heads logical threads cover one weight each. - if global_tid < num_token_heads: - weight_token_id = global_tid // self.num_heads - weight_head_id = global_tid % self.num_heads - weights_out[weight_token_id, weight_head_id] = ( - weights[weight_token_id, weight_head_id].to(Float32) * scale - ) - - @cache - @staticmethod - def compile( - head_dim: int = 128, - rope_dim: int = 64, - num_heads: int = 64, - cos_sin_dtype: type[cutlass.Numeric] = Float32, - coarsen: int = 4, - ): + def warmup_inputs(self, compile_key: CompileKey) -> tuple[Any, ...]: num_tokens = cute.sym_int() max_pos = cute.sym_int() - q = make_fake_tensor( - BFloat16, (num_tokens, num_heads, head_dim), divisibility=16 + BFloat16, + (num_tokens, compile_key.num_heads, compile_key.head_dim), + divisibility=16, ) positions = make_fake_tensor(Int64, (num_tokens,), divisibility=1) cos_sin_cache = make_fake_tensor( - cos_sin_dtype, - (max_pos, rope_dim), + compile_key.cos_sin_dtype, + (max_pos, compile_key.rope_dim), divisibility=8, ) - weights = make_fake_tensor(BFloat16, (num_tokens, num_heads), divisibility=8) - q_fp4 = make_fake_tensor( + weights = make_fake_tensor( + BFloat16, (num_tokens, compile_key.num_heads), divisibility=8 + ) + q_packed = make_fake_tensor( Uint8, - (num_tokens, num_heads, head_dim // 2), + (num_tokens, compile_key.num_heads, compile_key.head_dim // 2), divisibility=16, ) q_scale = make_fake_tensor( Uint8, - (num_tokens, num_heads, head_dim // MXFP4_BLOCK_SIZE), + ( + num_tokens, + compile_key.num_heads, + compile_key.head_dim // MXFP4_BLOCK_SIZE, + ), divisibility=4, ) - weights_out = make_fake_tensor(Float32, (num_tokens, num_heads), divisibility=4) - - kernel = IndexerQMxFp4Kernel( - head_dim, rope_dim, num_heads, cos_sin_dtype, coarsen + weights_out = make_fake_tensor( + Float32, (num_tokens, compile_key.num_heads), divisibility=4 ) - stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) - return cute.compile( - kernel, + return ( positions, q, cos_sin_cache, weights, - q_fp4, + q_packed, q_scale, weights_out, Float32(0.0), - stream, - options="--enable-tvm-ffi", ) - -class IndexerQFp8Kernel(IndexerQRopeQuantKernel): - """Eight-thread subwarps process one ``(token, head)`` row and emit - float8 e4m3fn with a single per-(token, head) scalar scale folded - into the per-token weight (mirrors ``_fused_indexer_q_rope_quant_kernel``). - """ - - def __init__( - self, - head_dim: int = 128, - rope_dim: int = 64, - num_heads: int = 64, - cos_sin_dtype: type[cutlass.Numeric] = Float32, - coarsen: int = 4, - ): - super().__init__(head_dim, rope_dim, num_heads, cos_sin_dtype, coarsen) - # Each subwarp owns `coarsen` heads; we use the first `coarsen` - # threads of the subwarp to write the per-head weights using the - # fp8 scale computed in the matching loop iteration. - assert self.coarsen <= self.subwarp_size, ( - f"FP8 kernel requires coarsen ({self.coarsen}) <= " - f"subwarp_size ({self.subwarp_size}) for the weight-fold step" - ) - - @cute.jit + @kernel_launcher def __call__( self, - positions: cute.Tensor, - q: cute.Tensor, - cos_sin_cache: cute.Tensor, - weights: cute.Tensor, - q_fp8: cute.Tensor, - weights_out: cute.Tensor, - scale: Float32, - stream: CUstream, - ): - total_threads = q.shape[0] * self.threads_per_token - grid = (cute.ceil_div(total_threads, self.tb_size), 1, 1) - self.kernel( + *, + positions: torch.Tensor, + q: torch.Tensor, + cos_sin_cache: torch.Tensor, + weights: torch.Tensor, + weights_softmax_scale: float, + weights_head_scale: float, + q_packed: torch.Tensor, + q_scale: torch.Tensor, + weights_out: torch.Tensor, + ) -> CuTeDSLLaunchSpec["IndexerQMxFp4Kernel.CompileKey"]: + num_tokens, num_heads, head_dim = q.shape + compile_key = self.dispatch( + head_dim=head_dim, + rope_dim=cos_sin_cache.shape[-1], + num_heads=num_heads, + cos_sin_dtype=torch_to_cute_dtype(cos_sin_cache.dtype), + coarsen=1 if num_tokens < 512 else 4, + ) + launch_args = ( positions, q, cos_sin_cache, weights, - q_fp8, + q_packed, + q_scale, weights_out, - scale, - ).launch(grid=grid, block=(self.tb_size, 1, 1), stream=stream) + float(weights_softmax_scale * weights_head_scale), + ) + return compile_key, launch_args - @cute.kernel - def kernel( - self, - positions: cute.Tensor, - q: cute.Tensor, - cos_sin_cache: cute.Tensor, - weights: cute.Tensor, - q_fp8: cute.Tensor, - weights_out: cute.Tensor, - scale: Float32, - ): - ( - q_bf16x2, - _tid, - _global_tid, - sublane, - token_id, - head_tile_id, - head_start, - in_bounds, - _num_token_heads, - ) = self._load_q_and_rope(positions, q, cos_sin_cache) - cp_op = cute.nvgpu.CopyUniversalOp() +class IndexerQFp8Kernel(VllmCuTeDSLJitKernel["IndexerQFp8Kernel.CompileKey"]): + tb_size = 128 - # layout: [coarsen, 16] bytes (one e4m3fn per element). - q_fp8_tile = cute.local_tile( - q_fp8[token_id, None, None], - tiler=(self.coarsen, 16), - coord=(head_tile_id, sublane), + @dataclass(frozen=True) + class CompileKey: + head_dim: int + rope_dim: int + num_heads: int + cos_sin_dtype: type[cutlass.Numeric] + coarsen: int + + @staticmethod + def kernel(compile_key: CompileKey) -> Any: + head_dim = compile_key.head_dim + rope_dim = compile_key.rope_dim + nope_dim = head_dim - rope_dim + num_heads = compile_key.num_heads + cos_sin_dtype = compile_key.cos_sin_dtype + # process multiple heads at the same time to armotize RoPE load costs + coarsen = compile_key.coarsen + tb_size = IndexerQFp8Kernel.tb_size + assert num_heads % coarsen == 0 + + # later we will use 32B load = 16 BF16 elems + # thus, head_dim=128 requires 8 threads to handle. + # let's call subwarp = 8 threads. + subwarp_size = head_dim // 16 + threads_per_token = (num_heads // coarsen) * subwarp_size + # Each subwarp owns `coarsen` heads; we use the first `coarsen` + # threads of the subwarp to write the per-head weights using the + # fp8 scale computed in the matching loop iteration. + assert coarsen <= subwarp_size, ( + f"FP8 kernel requires coarsen ({coarsen}) <= " + f"subwarp_size ({subwarp_size}) for the weight-fold step" ) - for i in cutlass.range_constexpr(self.coarsen): - # Reduce amax across the full head_dim: each thread already holds - # the max over its 16 lanes; a width=subwarp_size warp shuffle - # spreads the head-wide max to every lane in the subwarp. - amax_bf16x2 = _bf16x2_abs(q_bf16x2[i, 0]) - for j in cutlass.range_constexpr(1, 8): - amax_bf16x2 = _bf16x2_max(amax_bf16x2, _bf16x2_abs(q_bf16x2[i, j])) - amax_bf16x2 = cute_utils.warp_reduce( - amax_bf16x2, - _bf16x2_max, - width=self.subwarp_size, + @cute.kernel + def device_kernel( + positions: cute.Tensor, + q: cute.Tensor, + cos_sin_cache: cute.Tensor, + weights: cute.Tensor, + q_fp8: cute.Tensor, + weights_out: cute.Tensor, + scale: Float32, + ): + ( + q_bf16x2, + _tid, + _global_tid, + sublane, + token_id, + head_tile_id, + head_start, + in_bounds, + _num_token_heads, + ) = _load_q_and_rope( + positions, + q, + cos_sin_cache, + num_heads, + coarsen, + tb_size, + subwarp_size, + threads_per_token, + nope_dim, + rope_dim, + cos_sin_dtype, ) - amax_pair = cvt.bf16x2_to_fp32x2(amax_bf16x2) - amax = cute_utils.fmax(amax_pair[0], amax_pair[1]) - # scale = max(amax, eps) / fp8_max, then rounded UP to the next - # power of two. Adding the mantissa mask before shifting out the - # mantissa bumps the exponent whenever s isn't a pure pow2. - fp32_scale = cute_utils.fmax(amax, Float32(1e-4)) * Float32(1.0 / 448.0) - bits = recast_val(fp32_scale, Uint32) - scale_exp = cute_utils.shr_u32( - bits + Uint32(0x7FFFFF), Uint32(23) - ) & Uint32(0xFF) + cp_op = cute.nvgpu.CopyUniversalOp() - # rounded scale = 2^(scale_exp - 127); bit pattern is scale_exp << 23 - fp8_scale_bits = scale_exp << Uint32(23) - fp8_scale = recast_val(fp8_scale_bits, Float32) - # inverse = 2^-(scale_exp - 127); bit pattern is (254 - scale_exp) << 23 - inv_scale_bits = (Uint32(254) - scale_exp) << Uint32(23) - inv_fp8_scale = recast_val(inv_scale_bits, Float32) + # layout: [coarsen, 16] bytes (one e4m3fn per element). + q_fp8_tile = cute.local_tile( + q_fp8[token_id, None, None], + tiler=(coarsen, 16), + coord=(head_tile_id, sublane), + ) - # Weight fold: weights_out = weights * q_scale * scale_combined. - # All threads in the subwarp share the same fp8_scale after the - # warp_reduce above, so we let thread `sublane == i` write the - # weight for head `head_start + i`. - if in_bounds and sublane == i: - head_id = head_start + i - weights_out[token_id, head_id] = ( - weights[token_id, head_id].to(Float32) * scale * fp8_scale + for i in cutlass.range_constexpr(coarsen): + # Reduce amax across the full head_dim: each thread already holds + # the max over its 16 lanes; a width=subwarp_size warp shuffle + # spreads the head-wide max to every lane in the subwarp. + amax_bf16x2 = _bf16x2_abs(q_bf16x2[i, 0]) + for j in cutlass.range_constexpr(1, 8): + amax_bf16x2 = _bf16x2_max(amax_bf16x2, _bf16x2_abs(q_bf16x2[i, j])) + amax_bf16x2 = cute_utils.warp_reduce( + amax_bf16x2, + _bf16x2_max, + width=subwarp_size, ) + amax_pair = cvt.bf16x2_to_fp32x2(amax_bf16x2) + amax = cute_utils.fmax(amax_pair[0], amax_pair[1]) - if in_bounds: - # 16 BF16 → 16 e4m3 bytes per thread, packed into 4 b32s - # (one cp.async-shaped 128-bit store per row). - packed = cute.make_rmem_tensor((4,), Uint32) - for j in cutlass.range_constexpr(4): - q0, q1 = cvt.bf16x2_to_fp32x2(q_bf16x2[i, j * 2]) - q2, q3 = cvt.bf16x2_to_fp32x2(q_bf16x2[i, j * 2 + 1]) - packed[j] = cvt.fp32x4_to_fp8x4( - q0 * inv_fp8_scale, - q1 * inv_fp8_scale, - q2 * inv_fp8_scale, - q3 * inv_fp8_scale, + # scale = max(amax, eps) / fp8_max, then rounded UP to the next + # power of two. Adding the mantissa mask before shifting out the + # mantissa bumps the exponent whenever s isn't a pure pow2. + fp32_scale = cute_utils.fmax(amax, Float32(1e-4)) * Float32(1.0 / 448.0) + bits = recast_val(fp32_scale, Uint32) + scale_exp = cute_utils.shr_u32( + bits + Uint32(0x7FFFFF), Uint32(23) + ) & Uint32(0xFF) + + # rounded scale = 2^(scale_exp - 127); bit pattern is scale_exp << 23 + fp8_scale_bits = scale_exp << Uint32(23) + fp8_scale = recast_val(fp8_scale_bits, Float32) + # inverse = 2^-(scale_exp - 127); bit pattern is (254 - scale_exp) << 23 + inv_scale_bits = (Uint32(254) - scale_exp) << Uint32(23) + inv_fp8_scale = recast_val(inv_scale_bits, Float32) + + # Weight fold: weights_out = weights * q_scale * scale_combined. + # All threads in the subwarp share the same fp8_scale after the + # warp_reduce above, so we let thread `sublane == i` write the + # weight for head `head_start + i`. + if in_bounds and sublane == i: + head_id = head_start + i + weights_out[token_id, head_id] = ( + weights[token_id, head_id].to(Float32) * scale * fp8_scale ) - dst = q_fp8_tile[i, None] - cp_u32x4 = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=128) - cute.copy(cp_u32x4, packed, cute.recast_tensor(dst, Uint32)) + if in_bounds: + # 16 BF16 → 16 e4m3 bytes per thread, packed into 4 b32s + # (one cp.async-shaped 128-bit store per row). + packed = cute.make_rmem_tensor((4,), Uint32) + for j in cutlass.range_constexpr(4): + q0, q1 = cvt.bf16x2_to_fp32x2(q_bf16x2[i, j * 2]) + q2, q3 = cvt.bf16x2_to_fp32x2(q_bf16x2[i, j * 2 + 1]) + packed[j] = cvt.fp32x4_to_fp8x4( + q0 * inv_fp8_scale, + q1 * inv_fp8_scale, + q2 * inv_fp8_scale, + q3 * inv_fp8_scale, + ) - @cache - @staticmethod - def compile( - head_dim: int = 128, - rope_dim: int = 64, - num_heads: int = 64, - cos_sin_dtype: type[cutlass.Numeric] = Float32, - coarsen: int = 4, - ): + dst = q_fp8_tile[i, None] + cp_u32x4 = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=128) + cute.copy(cp_u32x4, packed, cute.recast_tensor(dst, Uint32)) + + @cute.jit + def host_entrypoint( + positions: cute.Tensor, + q: cute.Tensor, + cos_sin_cache: cute.Tensor, + weights: cute.Tensor, + q_fp8: cute.Tensor, + weights_out: cute.Tensor, + scale: Float32, + stream: CUstream, + ): + total_threads = q.shape[0] * threads_per_token + grid = (cute.ceil_div(total_threads, tb_size), 1, 1) + device_kernel( + positions, + q, + cos_sin_cache, + weights, + q_fp8, + weights_out, + scale, + ).launch(grid=grid, block=(tb_size, 1, 1), stream=stream) + + return host_entrypoint + + def dispatch( # type: ignore[override] + self, + *, + cos_sin_dtype: type[cutlass.Numeric], + **compile_key_fields: int, + ) -> CompileKey: + return self.CompileKey( + **compile_key_fields, + cos_sin_dtype=cos_sin_dtype, + ) + + def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]: + from vllm.v1.attention.backends.mla.indexer import dsa_indexer_uses_fp4 + + hf_config = vllm_config.model_config.hf_config + if hf_config is None or dsa_indexer_uses_fp4(vllm_config): + return [] + + num_heads = int(getattr(hf_config, "index_n_heads", 0) or 0) + head_dim = int(getattr(hf_config, "index_head_dim", 0) or 0) + rope_dim = int(getattr(hf_config, "qk_rope_head_dim", 0) or 0) + if num_heads <= 0 or head_dim <= 0 or rope_dim <= 0: + return [] + + return self._trace_dispatch(self.dispatch)( + head_dim=head_dim, + rope_dim=rope_dim, + num_heads=num_heads, + cos_sin_dtype=Float32, + coarsen=(1, 4), + ) + + def warmup_inputs(self, compile_key: CompileKey) -> tuple[Any, ...]: num_tokens = cute.sym_int() max_pos = cute.sym_int() - q = make_fake_tensor( - BFloat16, (num_tokens, num_heads, head_dim), divisibility=16 + BFloat16, + (num_tokens, compile_key.num_heads, compile_key.head_dim), + divisibility=16, ) positions = make_fake_tensor(Int64, (num_tokens,), divisibility=1) cos_sin_cache = make_fake_tensor( - cos_sin_dtype, - (max_pos, rope_dim), + compile_key.cos_sin_dtype, + (max_pos, compile_key.rope_dim), divisibility=8, ) - weights = make_fake_tensor(BFloat16, (num_tokens, num_heads), divisibility=8) + weights = make_fake_tensor( + BFloat16, (num_tokens, compile_key.num_heads), divisibility=8 + ) q_fp8 = make_fake_tensor( Uint8, - (num_tokens, num_heads, head_dim), + (num_tokens, compile_key.num_heads, compile_key.head_dim), divisibility=16, ) - weights_out = make_fake_tensor(Float32, (num_tokens, num_heads), divisibility=4) - - kernel = IndexerQFp8Kernel( - head_dim, rope_dim, num_heads, cos_sin_dtype, coarsen + weights_out = make_fake_tensor( + Float32, (num_tokens, compile_key.num_heads), divisibility=4 ) - stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) - return cute.compile( - kernel, + return ( positions, q, cos_sin_cache, @@ -609,6 +656,40 @@ class IndexerQFp8Kernel(IndexerQRopeQuantKernel): q_fp8, weights_out, Float32(0.0), - stream, - options="--enable-tvm-ffi", ) + + @kernel_launcher + def __call__( + self, + *, + positions: torch.Tensor, + q: torch.Tensor, + cos_sin_cache: torch.Tensor, + weights: torch.Tensor, + weights_softmax_scale: float, + weights_head_scale: float, + q_fp8: torch.Tensor, + weights_out: torch.Tensor, + ) -> CuTeDSLLaunchSpec["IndexerQFp8Kernel.CompileKey"]: + num_tokens, num_heads, head_dim = q.shape + compile_key = self.dispatch( + head_dim=head_dim, + rope_dim=cos_sin_cache.shape[-1], + num_heads=num_heads, + cos_sin_dtype=torch_to_cute_dtype(cos_sin_cache.dtype), + coarsen=1 if num_tokens < 512 else 4, + ) + launch_args = ( + positions, + q, + cos_sin_cache, + weights, + q_fp8, + weights_out, + float(weights_softmax_scale * weights_head_scale), + ) + return compile_key, launch_args + + +_INDEXER_Q_MXFP4_KERNEL = IndexerQMxFp4Kernel() +_INDEXER_Q_FP8_KERNEL = IndexerQFp8Kernel() diff --git a/vllm/models/deepseek_v4/nvidia/ops/sparse_attn_compress_cutedsl.py b/vllm/models/deepseek_v4/nvidia/ops/sparse_attn_compress_cutedsl.py index 4ff4b232d1..74866b0960 100644 --- a/vllm/models/deepseek_v4/nvidia/ops/sparse_attn_compress_cutedsl.py +++ b/vllm/models/deepseek_v4/nvidia/ops/sparse_attn_compress_cutedsl.py @@ -7,22 +7,25 @@ The public wrappers provide the C4 fused and C128 split kernels. from __future__ import annotations -from functools import cache +from dataclasses import dataclass from typing import Any import cutlass import cutlass.cute as cute import torch from cuda.bindings.driver import CUstream -from cutlass import BFloat16, Float32, Int32, Int64, Uint8, Uint16, Uint32, const_expr +from cutlass import Float32, Int32, Int64, Uint8, Uint16, Uint32, const_expr from cutlass._mlir.dialects import llvm from cutlass.cutlass_dsl import T, dsl_user_op from quack.compile_utils import make_fake_tensor -_TORCH_TO_CUTE = { - torch.bfloat16: BFloat16, - torch.float32: Float32, -} +from vllm.cute_utils import torch_to_cute_dtype +from vllm.model_executor.warmup.jit_warmup import kernel_launcher +from vllm.model_executor.warmup.jit_warmup_cutedsl_helper import ( + CuTeDSLLaunchSpec, + VllmCuTeDSLJitKernel, +) +from vllm.utils.math_utils import round_up @dsl_user_op @@ -72,361 +75,406 @@ def _fp32x2_to_fp8e4m3x2(a: Float32, b: Float32, *, loc=None, ip=None) -> Uint16 return Uint16(out) -class SparseAttnCompressNormRopeStoreC4Kernel: +class SparseAttnCompressNormRopeStoreC4Kernel( + VllmCuTeDSLJitKernel["SparseAttnCompressNormRopeStoreC4Kernel.CompileKey"] +): min_scale = 1.0e-4 rcp_ln2 = 1.4426950408889634 - def __init__( - self, - head_size: int, - state_width: int, - rope_head_dim: int, - fp8_max: float, - quant_block: int, - token_stride: int, - scale_dim: int, - compress_ratio: int, - overlap: bool, - ): - self.head_dim = head_size - self.state_width = state_width - self.rope_dim = rope_head_dim - self.nope_dim = head_size - rope_head_dim - self.fp8_max = fp8_max - self.quant_block = quant_block - self.token_stride = token_stride - self.scale_dim = scale_dim - self.elems_per_lane = 8 - self.copy_elems = 4 - self.copy_chunks = self.elems_per_lane // self.copy_elems - self.lanes_per_group = quant_block // self.elems_per_lane - self.groups_per_warp = 32 // self.lanes_per_group - self.scale_reduce_steps = self.lanes_per_group.bit_length() - 1 - self.scale_reduce_offset = self.lanes_per_group // 2 - self.num_warps = (head_size // quant_block) // self.groups_per_warp - self.nope_blocks = self.nope_dim // quant_block - self.tb_size = self.num_warps * 32 - self.compress_ratio = compress_ratio - self.overlap = overlap - self.window = (1 + int(overlap)) * compress_ratio + @dataclass(frozen=True) + class CompileKey: + head_size: int + state_width: int + rope_head_dim: int + fp8_max: float + quant_block: int + token_stride: int + scale_dim: int + compress_ratio: int + overlap: bool + norm_weight_dtype: type[cutlass.Numeric] - @cute.jit - def __call__( - self, - state_cache: cute.Tensor, - token_to_req_indices: cute.Tensor, - positions: cute.Tensor, - slot_mapping: cute.Tensor, - block_table: cute.Tensor, - block_size: Int64, - rms_norm_weight: cute.Tensor, - rms_norm_eps: Float32, - cos_sin_cache: cute.Tensor, - k_cache: cute.Tensor, - kv_slot_mapping: cute.Tensor, - kv_cache_block_size: Int64, - stream: CUstream, - ): - grid = (slot_mapping.shape[0], 1, 1) - self.kernel( - state_cache, - token_to_req_indices, - positions, - slot_mapping, - block_table, - block_size, - rms_norm_weight, - rms_norm_eps, - cos_sin_cache, - k_cache, - kv_slot_mapping, - kv_cache_block_size, - ).launch(grid=grid, block=(self.tb_size, 1, 1), stream=stream) + @staticmethod + def kernel(compile_key: CompileKey) -> Any: + min_scale = SparseAttnCompressNormRopeStoreC4Kernel.min_scale + rcp_ln2 = SparseAttnCompressNormRopeStoreC4Kernel.rcp_ln2 + head_dim = compile_key.head_size + state_width = compile_key.state_width + rope_dim = compile_key.rope_head_dim + nope_dim = head_dim - rope_dim + fp8_max = compile_key.fp8_max + quant_block = compile_key.quant_block + token_stride = compile_key.token_stride + scale_dim = compile_key.scale_dim + elems_per_lane = 8 + copy_elems = 4 + copy_chunks = elems_per_lane // copy_elems + lanes_per_group = quant_block // elems_per_lane + groups_per_warp = 32 // lanes_per_group + scale_reduce_steps = lanes_per_group.bit_length() - 1 + scale_reduce_offset = lanes_per_group // 2 + num_warps = (head_dim // quant_block) // groups_per_warp + nope_blocks = nope_dim // quant_block + tb_size = num_warps * 32 + compress_ratio = compile_key.compress_ratio + overlap = compile_key.overlap + window = (1 + int(overlap)) * compress_ratio - @cute.kernel - def kernel( - self, - state_cache: cute.Tensor, - token_to_req_indices: cute.Tensor, - positions: cute.Tensor, - slot_mapping: cute.Tensor, - block_table: cute.Tensor, - block_size: Int64, - rms_norm_weight: cute.Tensor, - rms_norm_eps: Float32, - cos_sin_cache: cute.Tensor, - k_cache: cute.Tensor, - kv_slot_mapping: cute.Tensor, - kv_cache_block_size: Int64, - ): - token_idx, _, _ = cute.arch.block_idx() - tid, _, _ = cute.arch.thread_idx() - warp_id = cute.arch.make_warp_uniform(tid // 32) - lane_id = tid % 32 - group_lane = lane_id % self.lanes_per_group - group_idx = warp_id * self.groups_per_warp + lane_id // self.lanes_per_group - elem_base = group_idx * self.quant_block + group_lane * self.elems_per_lane + @cute.kernel + def device_kernel( + state_cache: cute.Tensor, + token_to_req_indices: cute.Tensor, + positions: cute.Tensor, + slot_mapping: cute.Tensor, + block_table: cute.Tensor, + block_size: Int64, + rms_norm_weight: cute.Tensor, + rms_norm_eps: Float32, + cos_sin_cache: cute.Tensor, + k_cache: cute.Tensor, + kv_slot_mapping: cute.Tensor, + kv_cache_block_size: Int64, + ): + token_idx, _, _ = cute.arch.block_idx() + tid, _, _ = cute.arch.thread_idx() + warp_id = cute.arch.make_warp_uniform(tid // 32) + lane_id = tid % 32 + group_lane = lane_id % lanes_per_group + group_idx = warp_id * groups_per_warp + lane_id // lanes_per_group + elem_base = group_idx * quant_block + group_lane * elems_per_lane - slot_id = slot_mapping[token_idx] - has_position = token_idx < positions.shape[0] - position = Int64(0) - if has_position: - position = positions[token_idx] - boundary = has_position and ( - (position + Int64(1)) % Int64(self.compress_ratio) == Int64(0) - ) - has_req_idx = token_idx < token_to_req_indices.shape[0] - has_kv_slot_idx = token_idx < kv_slot_mapping.shape[0] - kv_slot_idx = Int64(-1) - if has_kv_slot_idx: - kv_slot_idx = kv_slot_mapping[token_idx] - active = ( - slot_id >= Int64(0) and has_req_idx and boundary and kv_slot_idx >= Int64(0) - ) - - if active: - req_idx = token_to_req_indices[token_idx] - start = position - Int64(self.window - 1) - - smem = cutlass.utils.SmemAllocator() - s_block_numbers = smem.allocate_tensor( - Int32, cute.make_layout((self.window,)), byte_alignment=4 + slot_id = slot_mapping[token_idx] + has_position = token_idx < positions.shape[0] + position = Int64(0) + if has_position: + position = positions[token_idx] + boundary = has_position and ( + (position + Int64(1)) % Int64(compress_ratio) == Int64(0) ) - partial_sums = smem.allocate_tensor( - Float32, cute.make_layout((self.num_warps,)), byte_alignment=4 - ) - rrms_shared = smem.allocate_tensor( - Float32, cute.make_layout((1,)), byte_alignment=4 + has_req_idx = token_idx < token_to_req_indices.shape[0] + has_kv_slot_idx = token_idx < kv_slot_mapping.shape[0] + kv_slot_idx = Int64(-1) + if has_kv_slot_idx: + kv_slot_idx = kv_slot_mapping[token_idx] + active = ( + slot_id >= Int64(0) + and has_req_idx + and boundary + and kv_slot_idx >= Int64(0) ) - for row in cutlass.range_constexpr(self.window): - pos = start + Int64(row) - if tid == row: - block_number_i32 = Int32(0) + if active: + req_idx = token_to_req_indices[token_idx] + start = position - Int64(window - 1) + + smem = cutlass.utils.SmemAllocator() + s_block_numbers = smem.allocate_tensor( + Int32, cute.make_layout((window,)), byte_alignment=4 + ) + partial_sums = smem.allocate_tensor( + Float32, cute.make_layout((num_warps,)), byte_alignment=4 + ) + rrms_shared = smem.allocate_tensor( + Float32, cute.make_layout((1,)), byte_alignment=4 + ) + + for row in cutlass.range_constexpr(window): + pos = start + Int64(row) + if tid == row: + block_number_i32 = Int32(0) + if pos >= Int64(0): + block_index = pos // block_size + block_number_i32 = block_table[req_idx, block_index] + s_block_numbers[row] = block_number_i32 + cute.arch.sync_threads() + + local_max = cute.make_rmem_tensor((elems_per_lane,), Float32) + local_sum = cute.make_rmem_tensor((elems_per_lane,), Float32) + local_product = cute.make_rmem_tensor((elems_per_lane,), Float32) + + for e in cutlass.range_constexpr(elems_per_lane): + local_max[e] = -Float32.inf + local_sum[e] = Float32(0.0) + local_product[e] = Float32(0.0) + + cp_f32x4 = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), Float32, num_bits_per_copy=128 + ) + copy_layout = cute.make_layout( + (copy_chunks, copy_elems), + stride=(copy_elems, 1), + ) + kv_vals = cute.make_rmem_tensor(copy_layout, Float32) + score_vals = cute.make_rmem_tensor(copy_layout, Float32) + + for row in cutlass.range_constexpr(window): + pos = start + Int64(row) if pos >= Int64(0): block_index = pos // block_size - block_number_i32 = block_table[req_idx, block_index] - s_block_numbers[row] = block_number_i32 - cute.arch.sync_threads() + block_offset = pos - block_index * block_size + block_number = s_block_numbers[row].to(Int64) + head_offset = Int64((row // compress_ratio) * head_dim) + row_tensor = state_cache[block_number, block_offset, None] + for chunk in cutlass.range_constexpr(copy_chunks): + copy_elem = const_expr(chunk * copy_elems) + col_tile = ( + head_offset + (elem_base + Int32(copy_elem)).to(Int64) + ) // Int64(copy_elems) + kv_src = cute.local_tile( + row_tensor, + tiler=(copy_elems,), + coord=(col_tile,), + ) + score_src = cute.local_tile( + row_tensor, + tiler=(copy_elems,), + coord=(col_tile + Int64(state_width // copy_elems),), + ) + cute.copy(cp_f32x4, kv_src, kv_vals[chunk, None]) + cute.copy(cp_f32x4, score_src, score_vals[chunk, None]) - local_max = cute.make_rmem_tensor((self.elems_per_lane,), Float32) - local_sum = cute.make_rmem_tensor((self.elems_per_lane,), Float32) - local_product = cute.make_rmem_tensor((self.elems_per_lane,), Float32) + for e in cutlass.range_constexpr(elems_per_lane): + chunk = const_expr(e // copy_elems) + copy_elem = const_expr(e % copy_elems) + score = score_vals[chunk, copy_elem] + kv = kv_vals[chunk, copy_elem] + new_max = cute.arch.fmax(local_max[e], score) + old_scale = cute.math.exp2( + (local_max[e] - new_max) * Float32(rcp_ln2), + fastmath=True, + ) + new_scale = cute.math.exp2( + (score - new_max) * Float32(rcp_ln2), + fastmath=True, + ) + local_sum[e] = local_sum[e] * old_scale + new_scale + local_product[e] = ( + local_product[e] * old_scale + kv * new_scale + ) + local_max[e] = new_max - for e in cutlass.range_constexpr(self.elems_per_lane): - local_max[e] = -Float32.inf - local_sum[e] = Float32(0.0) - local_product[e] = Float32(0.0) + x = cute.make_rmem_tensor((elems_per_lane,), Float32) + local_sumsq = Float32(0.0) + for e in cutlass.range_constexpr(elems_per_lane): + x[e] = local_product[e] / local_sum[e] + local_sumsq += x[e] * x[e] - cp_f32x4 = cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), Float32, num_bits_per_copy=128 - ) - copy_layout = cute.make_layout( - (self.copy_chunks, self.copy_elems), - stride=(self.copy_elems, 1), - ) - kv_vals = cute.make_rmem_tensor(copy_layout, Float32) - score_vals = cute.make_rmem_tensor(copy_layout, Float32) + warp_sum = local_sumsq + for step in cutlass.range_constexpr(5): + offset = const_expr(16 >> step) + warp_sum += cute.arch.shuffle_sync_bfly(warp_sum, offset) - for row in cutlass.range_constexpr(self.window): - pos = start + Int64(row) - if pos >= Int64(0): - block_index = pos // block_size - block_offset = pos - block_index * block_size - block_number = s_block_numbers[row].to(Int64) - head_offset = Int64((row // self.compress_ratio) * self.head_dim) - row_tensor = state_cache[block_number, block_offset, None] - for chunk in cutlass.range_constexpr(self.copy_chunks): - copy_elem = const_expr(chunk * self.copy_elems) - col_tile = ( - head_offset + (elem_base + Int32(copy_elem)).to(Int64) - ) // Int64(self.copy_elems) - kv_src = cute.local_tile( - row_tensor, - tiler=(self.copy_elems,), - coord=(col_tile,), + if lane_id == 0: + partial_sums[warp_id] = warp_sum + cute.arch.sync_threads() + if tid == 0: + total = Float32(0.0) + for i in cutlass.range_constexpr(num_warps): + total += partial_sums[i] + rrms_shared[0] = cute.math.rsqrt( + total / Float32(head_dim) + rms_norm_eps, fastmath=True + ) + cute.arch.sync_threads() + + rrms = rrms_shared[0] + for e in cutlass.range_constexpr(elems_per_lane): + elem = elem_base + e + x[e] = x[e] * rrms * rms_norm_weight[elem].to(Float32) + + k_cache_u16 = cute.recast_tensor(k_cache, Uint16) + k_cache_u32 = cute.recast_tensor(k_cache, Uint32) + page = kv_slot_idx // kv_cache_block_size + kv_offset = kv_slot_idx - page * kv_cache_block_size + value_base = page * k_cache.stride[0] + kv_offset * Int64(token_stride) + scale_base = ( + page * k_cache.stride[0] + + kv_cache_block_size * Int64(token_stride) + + kv_offset * Int64(scale_dim) + ) + + if group_idx == nope_blocks: + compressed_pos = (position // Int64(compress_ratio)) * Int64( + compress_ratio + ) + for pair in cutlass.range_constexpr(elems_per_lane // 2): + elem = const_expr(pair * 2) + pair_idx = (elem_base - nope_dim) // 2 + Int32(pair) + cos_v = cos_sin_cache[compressed_pos, pair_idx] + sin_v = cos_sin_cache[ + compressed_pos, pair_idx + Int32(rope_dim // 2) + ] + real = x[elem] * cos_v - x[elem + 1] * sin_v + imag = x[elem] * sin_v + x[elem + 1] * cos_v + packed = _fp32x2_to_bf16x2(real, imag) + out_base = ( + value_base + + Int64(nope_dim) + + ((elem_base - nope_dim + Int32(elem)) * 2).to(Int64) ) - score_src = cute.local_tile( - row_tensor, - tiler=(self.copy_elems,), - coord=( - col_tile + Int64(self.state_width // self.copy_elems), + k_cache_u32.iterator[out_base // Int64(4)] = packed + else: + q = cute.make_rmem_tensor((elems_per_lane,), Float32) + local_absmax = Float32(0.0) + for pair in cutlass.range_constexpr(elems_per_lane // 2): + elem = const_expr(pair * 2) + q_packed = _fp32x2_to_bf16x2(x[elem], x[elem + 1]) + q0, q1 = _bf16x2_to_fp32(q_packed) + q[elem] = q0 + q[elem + 1] = q1 + local_absmax = cute.arch.fmax( + local_absmax, + cute.arch.fmax(cute.math.absf(q0), cute.math.absf(q1)), + ) + absmax = local_absmax + group_mask_and_clamp = const_expr( + (cute.arch.WARP_SIZE - lanes_per_group) << 8 + | (cute.arch.WARP_SIZE - 1) + ) + for step in cutlass.range_constexpr(scale_reduce_steps): + offset = const_expr(scale_reduce_offset >> step) + absmax = cute.arch.fmax( + absmax, + cute.arch.shuffle_sync_bfly( + absmax, + offset=offset, + mask_and_clamp=group_mask_and_clamp, ), ) - cute.copy(cp_f32x4, kv_src, kv_vals[chunk, None]) - cute.copy(cp_f32x4, score_src, score_vals[chunk, None]) - - for e in cutlass.range_constexpr(self.elems_per_lane): - chunk = const_expr(e // self.copy_elems) - copy_elem = const_expr(e % self.copy_elems) - score = score_vals[chunk, copy_elem] - kv = kv_vals[chunk, copy_elem] - new_max = cute.arch.fmax(local_max[e], score) - old_scale = cute.math.exp2( - (local_max[e] - new_max) * Float32(self.rcp_ln2), - fastmath=True, + scale_raw = cute.arch.fmax( + Float32(min_scale), + absmax / Float32(fp8_max), + ) + bits = _recast_val(scale_raw, Uint32) + ue8m0 = ((bits + Uint32(0x7FFFFF)) >> Uint32(23)) & Uint32(0xFF) + inv_scale = _recast_val( + (Uint32(254) - ue8m0) << Uint32(23), Float32 + ) + for pair in cutlass.range_constexpr(elems_per_lane // 2): + elem = const_expr(pair * 2) + y0 = cutlass.min( + cute.arch.fmax(q[elem] * inv_scale, Float32(-fp8_max)), + Float32(fp8_max), ) - new_scale = cute.math.exp2( - (score - new_max) * Float32(self.rcp_ln2), - fastmath=True, + y1 = cutlass.min( + cute.arch.fmax(q[elem + 1] * inv_scale, Float32(-fp8_max)), + Float32(fp8_max), ) - local_sum[e] = local_sum[e] * old_scale + new_scale - local_product[e] = local_product[e] * old_scale + kv * new_scale - local_max[e] = new_max - - x = cute.make_rmem_tensor((self.elems_per_lane,), Float32) - local_sumsq = Float32(0.0) - for e in cutlass.range_constexpr(self.elems_per_lane): - x[e] = local_product[e] / local_sum[e] - local_sumsq += x[e] * x[e] - - warp_sum = local_sumsq - for step in cutlass.range_constexpr(5): - offset = const_expr(16 >> step) - warp_sum += cute.arch.shuffle_sync_bfly(warp_sum, offset) - - if lane_id == 0: - partial_sums[warp_id] = warp_sum - cute.arch.sync_threads() - if tid == 0: - total = Float32(0.0) - for i in cutlass.range_constexpr(self.num_warps): - total += partial_sums[i] - rrms_shared[0] = cute.math.rsqrt( - total / Float32(self.head_dim) + rms_norm_eps, fastmath=True - ) - cute.arch.sync_threads() - - rrms = rrms_shared[0] - for e in cutlass.range_constexpr(self.elems_per_lane): - elem = elem_base + e - x[e] = x[e] * rrms * rms_norm_weight[elem].to(Float32) - - k_cache_u16 = cute.recast_tensor(k_cache, Uint16) - k_cache_u32 = cute.recast_tensor(k_cache, Uint32) - page = kv_slot_idx // kv_cache_block_size - kv_offset = kv_slot_idx - page * kv_cache_block_size - value_base = page * k_cache.stride[0] + kv_offset * Int64(self.token_stride) - scale_base = ( - page * k_cache.stride[0] - + kv_cache_block_size * Int64(self.token_stride) - + kv_offset * Int64(self.scale_dim) - ) - - if group_idx == self.nope_blocks: - compressed_pos = (position // Int64(self.compress_ratio)) * Int64( - self.compress_ratio - ) - for pair in cutlass.range_constexpr(self.elems_per_lane // 2): - elem = const_expr(pair * 2) - pair_idx = (elem_base - self.nope_dim) // 2 + Int32(pair) - cos_v = cos_sin_cache[compressed_pos, pair_idx] - sin_v = cos_sin_cache[ - compressed_pos, pair_idx + Int32(self.rope_dim // 2) - ] - real = x[elem] * cos_v - x[elem + 1] * sin_v - imag = x[elem] * sin_v + x[elem + 1] * cos_v - packed = _fp32x2_to_bf16x2(real, imag) - out_base = ( - value_base - + Int64(self.nope_dim) - + ((elem_base - self.nope_dim + Int32(elem)) * 2).to(Int64) - ) - k_cache_u32.iterator[out_base // Int64(4)] = packed - else: - q = cute.make_rmem_tensor((self.elems_per_lane,), Float32) - local_absmax = Float32(0.0) - for pair in cutlass.range_constexpr(self.elems_per_lane // 2): - elem = const_expr(pair * 2) - q_packed = _fp32x2_to_bf16x2(x[elem], x[elem + 1]) - q0, q1 = _bf16x2_to_fp32(q_packed) - q[elem] = q0 - q[elem + 1] = q1 - local_absmax = cute.arch.fmax( - local_absmax, - cute.arch.fmax(cute.math.absf(q0), cute.math.absf(q1)), - ) - absmax = local_absmax - group_mask_and_clamp = const_expr( - (cute.arch.WARP_SIZE - self.lanes_per_group) << 8 - | (cute.arch.WARP_SIZE - 1) - ) - for step in cutlass.range_constexpr(self.scale_reduce_steps): - offset = const_expr(self.scale_reduce_offset >> step) - absmax = cute.arch.fmax( - absmax, - cute.arch.shuffle_sync_bfly( - absmax, - offset=offset, - mask_and_clamp=group_mask_and_clamp, - ), - ) - scale_raw = cute.arch.fmax( - Float32(self.min_scale), - absmax / Float32(self.fp8_max), - ) - bits = _recast_val(scale_raw, Uint32) - ue8m0 = ((bits + Uint32(0x7FFFFF)) >> Uint32(23)) & Uint32(0xFF) - inv_scale = _recast_val((Uint32(254) - ue8m0) << Uint32(23), Float32) - for pair in cutlass.range_constexpr(self.elems_per_lane // 2): - elem = const_expr(pair * 2) - y0 = cutlass.min( - cute.arch.fmax(q[elem] * inv_scale, Float32(-self.fp8_max)), - Float32(self.fp8_max), - ) - y1 = cutlass.min( - cute.arch.fmax(q[elem + 1] * inv_scale, Float32(-self.fp8_max)), - Float32(self.fp8_max), - ) - packed_fp8 = _fp32x2_to_fp8e4m3x2(y0, y1) - out_base = value_base + (elem_base + Int32(elem)).to(Int64) - k_cache_u16.iterator[out_base // Int64(2)] = packed_fp8 - if group_lane == 0: - k_cache.iterator[scale_base + group_idx.to(Int64)] = ue8m0.to(Uint8) - if group_idx == 0: - k_cache.iterator[scale_base + Int64(self.nope_blocks)] = Uint8( - 0 + packed_fp8 = _fp32x2_to_fp8e4m3x2(y0, y1) + out_base = value_base + (elem_base + Int32(elem)).to(Int64) + k_cache_u16.iterator[out_base // Int64(2)] = packed_fp8 + if group_lane == 0: + k_cache.iterator[scale_base + group_idx.to(Int64)] = ue8m0.to( + Uint8 ) + if group_idx == 0: + k_cache.iterator[scale_base + Int64(nope_blocks)] = Uint8(0) - @cache - @staticmethod - def compile( - head_size: int = 512, - state_width: int = 1024, - rope_head_dim: int = 64, - fp8_max: float = 448.0, - quant_block: int = 64, - token_stride: int = 576, - scale_dim: int = 8, - kv_block_stride: int = 74752, - compress_ratio: int = 4, - overlap: bool = True, - norm_weight_dtype: type[cutlass.Numeric] = Float32, - ): + @cute.jit + def host_entrypoint( + state_cache: cute.Tensor, + token_to_req_indices: cute.Tensor, + positions: cute.Tensor, + slot_mapping: cute.Tensor, + block_table: cute.Tensor, + block_size: Int64, + rms_norm_weight: cute.Tensor, + rms_norm_eps: Float32, + cos_sin_cache: cute.Tensor, + k_cache: cute.Tensor, + kv_slot_mapping: cute.Tensor, + kv_cache_block_size: Int64, + stream: CUstream, + ): + grid = (slot_mapping.shape[0], 1, 1) + device_kernel( + state_cache, + token_to_req_indices, + positions, + slot_mapping, + block_table, + block_size, + rms_norm_weight, + rms_norm_eps, + cos_sin_cache, + k_cache, + kv_slot_mapping, + kv_cache_block_size, + ).launch(grid=grid, block=(tb_size, 1, 1), stream=stream) + + return host_entrypoint + + def dispatch( # type: ignore[override] + self, + *, + compress_ratio: int, + norm_weight_dtype: type[cutlass.Numeric], + head_size: int, + rope_head_dim: int, + ) -> CompileKey: + overlap = compress_ratio == 4 + state_width = head_size * (2 if overlap else 1) + return self.CompileKey( + head_size=head_size, + state_width=state_width, + rope_head_dim=rope_head_dim, + fp8_max=448.0, + quant_block=64, + token_stride=576, + scale_dim=8, + compress_ratio=compress_ratio, + overlap=overlap, + norm_weight_dtype=norm_weight_dtype, + ) + + def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]: + hf_config = vllm_config.model_config.hf_config + cache_config = vllm_config.cache_config + return self._trace_dispatch(self.dispatch)( + compress_ratio=4, + norm_weight_dtype=torch_to_cute_dtype(vllm_config.model_config.dtype), + head_size=hf_config.head_dim, + rope_head_dim=hf_config.qk_rope_head_dim, + enabled=( + cache_config.block_size > 0 + and 4 in hf_config.compress_ratios + and cache_config.cache_dtype == "fp8_ds_mla" + ), + _when=lambda *, enabled: enabled, + ) + + def warmup_inputs(self, compile_key: CompileKey) -> tuple[Any, ...]: + head_size = compile_key.head_size + state_width = compile_key.state_width + rope_head_dim = compile_key.rope_head_dim + quant_block = compile_key.quant_block + token_stride = compile_key.token_stride + scale_dim = compile_key.scale_dim + compress_ratio = compile_key.compress_ratio + overlap = compile_key.overlap + norm_weight_dtype = compile_key.norm_weight_dtype if compress_ratio != 4 or not overlap: raise ValueError("CuTe DSL C4 fused sparse-attn requires C4 overlap.") if head_size != 512: raise ValueError( - "CuTe DSL C4 fused sparse-attn currently requires head_size=512." + "CuTe DSL C4 fused sparse-attn currently requires head_size=512.", ) if state_width != 2 * head_size: raise ValueError( - "CuTe DSL C4 fused sparse-attn requires state_width=2*head_size." + "CuTe DSL C4 fused sparse-attn requires state_width=2*head_size.", ) if quant_block != 64: raise ValueError( - "CuTe DSL C4 fused sparse-attn currently requires quant_block=64." + "CuTe DSL C4 fused sparse-attn currently requires quant_block=64.", ) if rope_head_dim != 64: raise ValueError( - "CuTe DSL C4 fused sparse-attn currently requires rope_head_dim=64." + "CuTe DSL C4 fused sparse-attn currently requires rope_head_dim=64.", ) if token_stride < head_size + rope_head_dim: raise ValueError("token_stride is too small for the packed FP8/BF16 row.") expected_scale_dim = (head_size - rope_head_dim) // quant_block + 1 if scale_dim < expected_scale_dim: raise ValueError("scale_dim is too small for the UE8M0 scale row.") - num_positions = cute.sym_int() num_slots = cute.sym_int() num_req_indices = cute.sym_int() @@ -437,7 +485,6 @@ class SparseAttnCompressNormRopeStoreC4Kernel: block_table_width = cute.sym_int() max_pos = cute.sym_int() state_cache_width = state_width * 2 - state_cache = cute.runtime.make_fake_tensor( Float32, (num_state_blocks, state_cache_block_size, state_cache_width), @@ -449,15 +496,21 @@ class SparseAttnCompressNormRopeStoreC4Kernel: assumed_align=16, ) token_to_req_indices = make_fake_tensor( - Int32, (num_req_indices,), divisibility=4 + Int32, + (num_req_indices,), + divisibility=4, ) positions = make_fake_tensor(Int64, (num_positions,), divisibility=8) slot_mapping = make_fake_tensor(Int64, (num_slots,), divisibility=8) block_table = make_fake_tensor( - Int32, (cute.sym_int(), block_table_width), divisibility=1 + Int32, + (cute.sym_int(), block_table_width), + divisibility=1, ) rms_norm_weight = make_fake_tensor( - norm_weight_dtype, (head_size,), divisibility=4 + norm_weight_dtype, + (head_size,), + divisibility=4, ) cos_sin_cache = cute.runtime.make_fake_tensor( Float32, @@ -468,29 +521,12 @@ class SparseAttnCompressNormRopeStoreC4Kernel: k_cache = cute.runtime.make_fake_tensor( Uint8, (num_kv_blocks, cute.sym_int(), cute.sym_int()), - stride=( - cute.sym_int64(divisibility=16), - cute.sym_int64(divisibility=8), - 1, - ), + stride=(cute.sym_int64(divisibility=16), cute.sym_int64(divisibility=8), 1), assumed_align=16, ) kv_slot_mapping = make_fake_tensor(Int64, (num_kv_slots,), divisibility=8) - kernel = SparseAttnCompressNormRopeStoreC4Kernel( - head_size, - state_width, - rope_head_dim, - fp8_max, - quant_block, - token_stride, - scale_dim, - compress_ratio, - overlap, - ) - stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) - return cute.compile( - kernel, + return ( state_cache, token_to_req_indices, positions, @@ -503,60 +539,34 @@ class SparseAttnCompressNormRopeStoreC4Kernel: k_cache, kv_slot_mapping, Int64(0), - stream, - options="--enable-tvm-ffi", ) - -class SparseAttnCompressNormRopeStoreFullC4Kernel( - SparseAttnCompressNormRopeStoreC4Kernel -): - def __init__( - self, - head_size: int, - state_width: int, - rope_head_dim: int, - fp8_max: float, - quant_block: int, - token_stride: int, - scale_dim: int, - compress_ratio: int, - overlap: bool, - store_full_fp8: bool = False, - ): - super().__init__( - head_size, - state_width, - rope_head_dim, - fp8_max, - quant_block, - token_stride, - scale_dim, - compress_ratio, - overlap, - ) - self.store_full_fp8 = store_full_fp8 - - @cute.jit + @kernel_launcher def __call__( self, - state_cache: cute.Tensor, - token_to_req_indices: cute.Tensor, - positions: cute.Tensor, - slot_mapping: cute.Tensor, - block_table: cute.Tensor, - block_size: Int64, - rms_norm_weight: cute.Tensor, - rms_norm_eps: Float32, - cos_sin_cache: cute.Tensor, - k_cache: cute.Tensor, - kv_slot_mapping: cute.Tensor, - kv_cache_block_size: Int64, - fp8_scale: cute.Tensor, - stream: CUstream, - ): - grid = (slot_mapping.shape[0], 1, 1) - self.kernel( + *, + state_cache: torch.Tensor, + token_to_req_indices: torch.Tensor, + positions: torch.Tensor, + slot_mapping: torch.Tensor, + block_table: torch.Tensor, + block_size: int, + rms_norm_weight: torch.Tensor, + rms_norm_eps: float, + cos_sin_cache: torch.Tensor, + kv_cache: torch.Tensor, + kv_slot_mapping: torch.Tensor, + compress_ratio: int, + head_dim: int, + rope_head_dim: int, + ) -> CuTeDSLLaunchSpec[SparseAttnCompressNormRopeStoreC4Kernel.CompileKey]: + compile_key = self.dispatch( + compress_ratio=compress_ratio, + norm_weight_dtype=torch_to_cute_dtype(rms_norm_weight.dtype), + head_size=head_dim, + rope_head_dim=rope_head_dim, + ) + launch_args = ( state_cache, token_to_req_indices, positions, @@ -566,283 +576,395 @@ class SparseAttnCompressNormRopeStoreFullC4Kernel( rms_norm_weight, rms_norm_eps, cos_sin_cache, - k_cache, + kv_cache, kv_slot_mapping, - kv_cache_block_size, - fp8_scale, - ).launch(grid=grid, block=(self.tb_size, 1, 1), stream=stream) - - @cute.kernel - def kernel( - self, - state_cache: cute.Tensor, - token_to_req_indices: cute.Tensor, - positions: cute.Tensor, - slot_mapping: cute.Tensor, - block_table: cute.Tensor, - block_size: Int64, - rms_norm_weight: cute.Tensor, - rms_norm_eps: Float32, - cos_sin_cache: cute.Tensor, - k_cache: cute.Tensor, - kv_slot_mapping: cute.Tensor, - kv_cache_block_size: Int64, - fp8_scale: cute.Tensor, - ): - token_idx, _, _ = cute.arch.block_idx() - tid, _, _ = cute.arch.thread_idx() - warp_id = cute.arch.make_warp_uniform(tid // 32) - lane_id = tid % 32 - group_lane = lane_id % self.lanes_per_group - group_idx = warp_id * self.groups_per_warp + lane_id // self.lanes_per_group - elem_base = group_idx * self.quant_block + group_lane * self.elems_per_lane - - slot_id = slot_mapping[token_idx] - has_position = token_idx < positions.shape[0] - position = Int64(0) - if has_position: - position = positions[token_idx] - boundary = has_position and ( - (position + Int64(1)) % Int64(self.compress_ratio) == Int64(0) - ) - has_req_idx = token_idx < token_to_req_indices.shape[0] - has_kv_slot_idx = token_idx < kv_slot_mapping.shape[0] - kv_slot_idx = Int64(-1) - if has_kv_slot_idx: - kv_slot_idx = kv_slot_mapping[token_idx] - active = ( - slot_id >= Int64(0) and has_req_idx and boundary and kv_slot_idx >= Int64(0) + kv_cache.shape[1], ) + return compile_key, launch_args - if active: - req_idx = token_to_req_indices[token_idx] - start = position - Int64(self.window - 1) - smem = cutlass.utils.SmemAllocator() - s_block_numbers = smem.allocate_tensor( - Int32, cute.make_layout((self.window,)), byte_alignment=4 +class SparseAttnCompressNormRopeStoreFullC4Kernel( + VllmCuTeDSLJitKernel["SparseAttnCompressNormRopeStoreFullC4Kernel.CompileKey"] +): + rcp_ln2 = 1.4426950408889634 + + @dataclass(frozen=True) + class CompileKey: + head_size: int + state_width: int + rope_head_dim: int + fp8_max: float + quant_block: int + compress_ratio: int + overlap: bool + store_full_fp8: bool + norm_weight_dtype: type[cutlass.Numeric] + + @staticmethod + def kernel(compile_key: CompileKey) -> Any: + rcp_ln2 = SparseAttnCompressNormRopeStoreFullC4Kernel.rcp_ln2 + head_dim = compile_key.head_size + state_width = compile_key.state_width + rope_dim = compile_key.rope_head_dim + nope_dim = head_dim - rope_dim + fp8_max = compile_key.fp8_max + quant_block = compile_key.quant_block + elems_per_lane = 8 + copy_elems = 4 + copy_chunks = elems_per_lane // copy_elems + lanes_per_group = quant_block // elems_per_lane + groups_per_warp = 32 // lanes_per_group + num_warps = (head_dim // quant_block) // groups_per_warp + nope_blocks = nope_dim // quant_block + tb_size = num_warps * 32 + compress_ratio = compile_key.compress_ratio + overlap = compile_key.overlap + window = (1 + int(overlap)) * compress_ratio + store_full_fp8 = compile_key.store_full_fp8 + + @cute.kernel + def device_kernel( + state_cache: cute.Tensor, + token_to_req_indices: cute.Tensor, + positions: cute.Tensor, + slot_mapping: cute.Tensor, + block_table: cute.Tensor, + block_size: Int64, + rms_norm_weight: cute.Tensor, + rms_norm_eps: Float32, + cos_sin_cache: cute.Tensor, + k_cache: cute.Tensor, + kv_slot_mapping: cute.Tensor, + kv_cache_block_size: Int64, + fp8_scale: cute.Tensor, + ): + token_idx, _, _ = cute.arch.block_idx() + tid, _, _ = cute.arch.thread_idx() + warp_id = cute.arch.make_warp_uniform(tid // 32) + lane_id = tid % 32 + group_lane = lane_id % lanes_per_group + group_idx = warp_id * groups_per_warp + lane_id // lanes_per_group + elem_base = group_idx * quant_block + group_lane * elems_per_lane + + slot_id = slot_mapping[token_idx] + has_position = token_idx < positions.shape[0] + position = Int64(0) + if has_position: + position = positions[token_idx] + boundary = has_position and ( + (position + Int64(1)) % Int64(compress_ratio) == Int64(0) ) - partial_sums = smem.allocate_tensor( - Float32, cute.make_layout((self.num_warps,)), byte_alignment=4 - ) - rrms_shared = smem.allocate_tensor( - Float32, cute.make_layout((1,)), byte_alignment=4 + has_req_idx = token_idx < token_to_req_indices.shape[0] + has_kv_slot_idx = token_idx < kv_slot_mapping.shape[0] + kv_slot_idx = Int64(-1) + if has_kv_slot_idx: + kv_slot_idx = kv_slot_mapping[token_idx] + active = ( + slot_id >= Int64(0) + and has_req_idx + and boundary + and kv_slot_idx >= Int64(0) ) - for row in cutlass.range_constexpr(self.window): - pos = start + Int64(row) - if tid == row: - block_number_i32 = Int32(0) + if active: + req_idx = token_to_req_indices[token_idx] + start = position - Int64(window - 1) + + smem = cutlass.utils.SmemAllocator() + s_block_numbers = smem.allocate_tensor( + Int32, cute.make_layout((window,)), byte_alignment=4 + ) + partial_sums = smem.allocate_tensor( + Float32, cute.make_layout((num_warps,)), byte_alignment=4 + ) + rrms_shared = smem.allocate_tensor( + Float32, cute.make_layout((1,)), byte_alignment=4 + ) + + for row in cutlass.range_constexpr(window): + pos = start + Int64(row) + if tid == row: + block_number_i32 = Int32(0) + if pos >= Int64(0): + block_index = pos // block_size + block_number_i32 = block_table[req_idx, block_index] + s_block_numbers[row] = block_number_i32 + cute.arch.sync_threads() + + local_max = cute.make_rmem_tensor((elems_per_lane,), Float32) + local_sum = cute.make_rmem_tensor((elems_per_lane,), Float32) + local_product = cute.make_rmem_tensor((elems_per_lane,), Float32) + + for e in cutlass.range_constexpr(elems_per_lane): + local_max[e] = -Float32.inf + local_sum[e] = Float32(0.0) + local_product[e] = Float32(0.0) + + cp_f32x4 = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), Float32, num_bits_per_copy=128 + ) + copy_layout = cute.make_layout( + (copy_chunks, copy_elems), + stride=(copy_elems, 1), + ) + kv_vals = cute.make_rmem_tensor(copy_layout, Float32) + score_vals = cute.make_rmem_tensor(copy_layout, Float32) + + for row in cutlass.range_constexpr(window): + pos = start + Int64(row) if pos >= Int64(0): block_index = pos // block_size - block_number_i32 = block_table[req_idx, block_index] - s_block_numbers[row] = block_number_i32 - cute.arch.sync_threads() + block_offset = pos - block_index * block_size + block_number = s_block_numbers[row].to(Int64) + head_offset = Int64((row // compress_ratio) * head_dim) + row_tensor = state_cache[block_number, block_offset, None] + for chunk in cutlass.range_constexpr(copy_chunks): + copy_elem = const_expr(chunk * copy_elems) + col_tile = ( + head_offset + (elem_base + Int32(copy_elem)).to(Int64) + ) // Int64(copy_elems) + kv_src = cute.local_tile( + row_tensor, + tiler=(copy_elems,), + coord=(col_tile,), + ) + score_src = cute.local_tile( + row_tensor, + tiler=(copy_elems,), + coord=(col_tile + Int64(state_width // copy_elems),), + ) + cute.copy(cp_f32x4, kv_src, kv_vals[chunk, None]) + cute.copy(cp_f32x4, score_src, score_vals[chunk, None]) - local_max = cute.make_rmem_tensor((self.elems_per_lane,), Float32) - local_sum = cute.make_rmem_tensor((self.elems_per_lane,), Float32) - local_product = cute.make_rmem_tensor((self.elems_per_lane,), Float32) + for e in cutlass.range_constexpr(elems_per_lane): + chunk = const_expr(e // copy_elems) + copy_elem = const_expr(e % copy_elems) + score = score_vals[chunk, copy_elem] + kv = kv_vals[chunk, copy_elem] + new_max = cute.arch.fmax(local_max[e], score) + old_scale = cute.math.exp2( + (local_max[e] - new_max) * Float32(rcp_ln2), + fastmath=True, + ) + new_scale = cute.math.exp2( + (score - new_max) * Float32(rcp_ln2), + fastmath=True, + ) + local_sum[e] = local_sum[e] * old_scale + new_scale + local_product[e] = ( + local_product[e] * old_scale + kv * new_scale + ) + local_max[e] = new_max - for e in cutlass.range_constexpr(self.elems_per_lane): - local_max[e] = -Float32.inf - local_sum[e] = Float32(0.0) - local_product[e] = Float32(0.0) + x = cute.make_rmem_tensor((elems_per_lane,), Float32) + local_sumsq = Float32(0.0) + for e in cutlass.range_constexpr(elems_per_lane): + x[e] = local_product[e] / local_sum[e] + local_sumsq += x[e] * x[e] - cp_f32x4 = cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), Float32, num_bits_per_copy=128 - ) - copy_layout = cute.make_layout( - (self.copy_chunks, self.copy_elems), - stride=(self.copy_elems, 1), - ) - kv_vals = cute.make_rmem_tensor(copy_layout, Float32) - score_vals = cute.make_rmem_tensor(copy_layout, Float32) + warp_sum = local_sumsq + for step in cutlass.range_constexpr(5): + offset = const_expr(16 >> step) + warp_sum += cute.arch.shuffle_sync_bfly(warp_sum, offset) - for row in cutlass.range_constexpr(self.window): - pos = start + Int64(row) - if pos >= Int64(0): - block_index = pos // block_size - block_offset = pos - block_index * block_size - block_number = s_block_numbers[row].to(Int64) - head_offset = Int64((row // self.compress_ratio) * self.head_dim) - row_tensor = state_cache[block_number, block_offset, None] - for chunk in cutlass.range_constexpr(self.copy_chunks): - copy_elem = const_expr(chunk * self.copy_elems) - col_tile = ( - head_offset + (elem_base + Int32(copy_elem)).to(Int64) - ) // Int64(self.copy_elems) - kv_src = cute.local_tile( - row_tensor, - tiler=(self.copy_elems,), - coord=(col_tile,), - ) - score_src = cute.local_tile( - row_tensor, - tiler=(self.copy_elems,), - coord=( - col_tile + Int64(self.state_width // self.copy_elems), - ), - ) - cute.copy(cp_f32x4, kv_src, kv_vals[chunk, None]) - cute.copy(cp_f32x4, score_src, score_vals[chunk, None]) - - for e in cutlass.range_constexpr(self.elems_per_lane): - chunk = const_expr(e // self.copy_elems) - copy_elem = const_expr(e % self.copy_elems) - score = score_vals[chunk, copy_elem] - kv = kv_vals[chunk, copy_elem] - new_max = cute.arch.fmax(local_max[e], score) - old_scale = cute.math.exp2( - (local_max[e] - new_max) * Float32(self.rcp_ln2), - fastmath=True, - ) - new_scale = cute.math.exp2( - (score - new_max) * Float32(self.rcp_ln2), - fastmath=True, - ) - local_sum[e] = local_sum[e] * old_scale + new_scale - local_product[e] = local_product[e] * old_scale + kv * new_scale - local_max[e] = new_max - - x = cute.make_rmem_tensor((self.elems_per_lane,), Float32) - local_sumsq = Float32(0.0) - for e in cutlass.range_constexpr(self.elems_per_lane): - x[e] = local_product[e] / local_sum[e] - local_sumsq += x[e] * x[e] - - warp_sum = local_sumsq - for step in cutlass.range_constexpr(5): - offset = const_expr(16 >> step) - warp_sum += cute.arch.shuffle_sync_bfly(warp_sum, offset) - - if lane_id == 0: - partial_sums[warp_id] = warp_sum - cute.arch.sync_threads() - if tid == 0: - total = Float32(0.0) - for i in cutlass.range_constexpr(self.num_warps): - total += partial_sums[i] - rrms_shared[0] = cute.math.rsqrt( - total / Float32(self.head_dim) + rms_norm_eps, fastmath=True - ) - cute.arch.sync_threads() - - rrms = rrms_shared[0] - for e in cutlass.range_constexpr(self.elems_per_lane): - elem = elem_base + e - x[e] = x[e] * rrms * rms_norm_weight[elem].to(Float32) - - page = kv_slot_idx // kv_cache_block_size - kv_offset = kv_slot_idx - page * kv_cache_block_size - value_base = page * k_cache.stride[0] + kv_offset * k_cache.stride[1] - - if const_expr(self.store_full_fp8): - k_cache_u16 = cute.recast_tensor(k_cache, Uint16) - inv_fp8 = Float32(1.0) / fp8_scale[0] - if group_idx == self.nope_blocks: - compressed_pos = (position // Int64(self.compress_ratio)) * Int64( - self.compress_ratio + if lane_id == 0: + partial_sums[warp_id] = warp_sum + cute.arch.sync_threads() + if tid == 0: + total = Float32(0.0) + for i in cutlass.range_constexpr(num_warps): + total += partial_sums[i] + rrms_shared[0] = cute.math.rsqrt( + total / Float32(head_dim) + rms_norm_eps, fastmath=True ) - for pair in cutlass.range_constexpr(self.elems_per_lane // 2): - elem = const_expr(pair * 2) - pair_idx = (elem_base - self.nope_dim) // 2 + Int32(pair) - cos_v = cos_sin_cache[compressed_pos, pair_idx] - sin_v = cos_sin_cache[ - compressed_pos, pair_idx + Int32(self.rope_dim // 2) - ] - real = x[elem] * cos_v - x[elem + 1] * sin_v - imag = x[elem] * sin_v + x[elem + 1] * cos_v - packed_bf16 = _fp32x2_to_bf16x2(real, imag) - b0, b1 = _bf16x2_to_fp32(packed_bf16) - y0 = cutlass.min( - cutlass.max(b0 * inv_fp8, Float32(-self.fp8_max)), - Float32(self.fp8_max), - ) - y1 = cutlass.min( - cutlass.max(b1 * inv_fp8, Float32(-self.fp8_max)), - Float32(self.fp8_max), - ) - packed_fp8 = _fp32x2_to_fp8e4m3x2(y0, y1) - out_base = value_base + (elem_base + Int32(elem)).to(Int64) - k_cache_u16.iterator[out_base // Int64(2)] = packed_fp8 - else: - for pair in cutlass.range_constexpr(self.elems_per_lane // 2): - elem = const_expr(pair * 2) - packed_bf16 = _fp32x2_to_bf16x2(x[elem], x[elem + 1]) - b0, b1 = _bf16x2_to_fp32(packed_bf16) - y0 = cutlass.min( - cutlass.max(b0 * inv_fp8, Float32(-self.fp8_max)), - Float32(self.fp8_max), - ) - y1 = cutlass.min( - cutlass.max(b1 * inv_fp8, Float32(-self.fp8_max)), - Float32(self.fp8_max), - ) - packed_fp8 = _fp32x2_to_fp8e4m3x2(y0, y1) - out_base = value_base + (elem_base + Int32(elem)).to(Int64) - k_cache_u16.iterator[out_base // Int64(2)] = packed_fp8 - else: - k_cache_u32 = cute.recast_tensor(k_cache, Uint32) - if group_idx == self.nope_blocks: - compressed_pos = (position // Int64(self.compress_ratio)) * Int64( - self.compress_ratio - ) - for pair in cutlass.range_constexpr(self.elems_per_lane // 2): - elem = const_expr(pair * 2) - pair_idx = (elem_base - self.nope_dim) // 2 + Int32(pair) - cos_v = cos_sin_cache[compressed_pos, pair_idx] - sin_v = cos_sin_cache[ - compressed_pos, pair_idx + Int32(self.rope_dim // 2) - ] - real = x[elem] * cos_v - x[elem + 1] * sin_v - imag = x[elem] * sin_v + x[elem + 1] * cos_v - packed_bf16 = _fp32x2_to_bf16x2(real, imag) - out_base = value_base + ((elem_base + Int32(elem)) * 2).to( - Int64 - ) - k_cache_u32.iterator[out_base // Int64(4)] = packed_bf16 - else: - for pair in cutlass.range_constexpr(self.elems_per_lane // 2): - elem = const_expr(pair * 2) - packed_bf16 = _fp32x2_to_bf16x2(x[elem], x[elem + 1]) - out_base = value_base + ((elem_base + Int32(elem)) * 2).to( - Int64 - ) - k_cache_u32.iterator[out_base // Int64(4)] = packed_bf16 + cute.arch.sync_threads() - @cache - @staticmethod - def compile( - head_size: int = 512, - state_width: int = 1024, - rope_head_dim: int = 64, - fp8_max: float = 448.0, - quant_block: int = 64, - token_stride: int = 576, - scale_dim: int = 8, - kv_block_stride: int = 74752, - compress_ratio: int = 4, - overlap: bool = True, - store_full_fp8: bool = False, - norm_weight_dtype: type[cutlass.Numeric] = Float32, - ): + rrms = rrms_shared[0] + for e in cutlass.range_constexpr(elems_per_lane): + elem = elem_base + e + x[e] = x[e] * rrms * rms_norm_weight[elem].to(Float32) + + page = kv_slot_idx // kv_cache_block_size + kv_offset = kv_slot_idx - page * kv_cache_block_size + value_base = page * k_cache.stride[0] + kv_offset * k_cache.stride[1] + + if const_expr(store_full_fp8): + k_cache_u16 = cute.recast_tensor(k_cache, Uint16) + inv_fp8 = Float32(1.0) / fp8_scale[0] + if group_idx == nope_blocks: + compressed_pos = (position // Int64(compress_ratio)) * Int64( + compress_ratio + ) + for pair in cutlass.range_constexpr(elems_per_lane // 2): + elem = const_expr(pair * 2) + pair_idx = (elem_base - nope_dim) // 2 + Int32(pair) + cos_v = cos_sin_cache[compressed_pos, pair_idx] + sin_v = cos_sin_cache[ + compressed_pos, pair_idx + Int32(rope_dim // 2) + ] + real = x[elem] * cos_v - x[elem + 1] * sin_v + imag = x[elem] * sin_v + x[elem + 1] * cos_v + packed_bf16 = _fp32x2_to_bf16x2(real, imag) + b0, b1 = _bf16x2_to_fp32(packed_bf16) + y0 = cutlass.min( + cutlass.max(b0 * inv_fp8, Float32(-fp8_max)), + Float32(fp8_max), + ) + y1 = cutlass.min( + cutlass.max(b1 * inv_fp8, Float32(-fp8_max)), + Float32(fp8_max), + ) + packed_fp8 = _fp32x2_to_fp8e4m3x2(y0, y1) + out_base = value_base + (elem_base + Int32(elem)).to(Int64) + k_cache_u16.iterator[out_base // Int64(2)] = packed_fp8 + else: + for pair in cutlass.range_constexpr(elems_per_lane // 2): + elem = const_expr(pair * 2) + packed_bf16 = _fp32x2_to_bf16x2(x[elem], x[elem + 1]) + b0, b1 = _bf16x2_to_fp32(packed_bf16) + y0 = cutlass.min( + cutlass.max(b0 * inv_fp8, Float32(-fp8_max)), + Float32(fp8_max), + ) + y1 = cutlass.min( + cutlass.max(b1 * inv_fp8, Float32(-fp8_max)), + Float32(fp8_max), + ) + packed_fp8 = _fp32x2_to_fp8e4m3x2(y0, y1) + out_base = value_base + (elem_base + Int32(elem)).to(Int64) + k_cache_u16.iterator[out_base // Int64(2)] = packed_fp8 + else: + k_cache_u32 = cute.recast_tensor(k_cache, Uint32) + if group_idx == nope_blocks: + compressed_pos = (position // Int64(compress_ratio)) * Int64( + compress_ratio + ) + for pair in cutlass.range_constexpr(elems_per_lane // 2): + elem = const_expr(pair * 2) + pair_idx = (elem_base - nope_dim) // 2 + Int32(pair) + cos_v = cos_sin_cache[compressed_pos, pair_idx] + sin_v = cos_sin_cache[ + compressed_pos, pair_idx + Int32(rope_dim // 2) + ] + real = x[elem] * cos_v - x[elem + 1] * sin_v + imag = x[elem] * sin_v + x[elem + 1] * cos_v + packed_bf16 = _fp32x2_to_bf16x2(real, imag) + out_base = value_base + ((elem_base + Int32(elem)) * 2).to( + Int64 + ) + k_cache_u32.iterator[out_base // Int64(4)] = packed_bf16 + else: + for pair in cutlass.range_constexpr(elems_per_lane // 2): + elem = const_expr(pair * 2) + packed_bf16 = _fp32x2_to_bf16x2(x[elem], x[elem + 1]) + out_base = value_base + ((elem_base + Int32(elem)) * 2).to( + Int64 + ) + k_cache_u32.iterator[out_base // Int64(4)] = packed_bf16 + + @cute.jit + def host_entrypoint( + state_cache: cute.Tensor, + token_to_req_indices: cute.Tensor, + positions: cute.Tensor, + slot_mapping: cute.Tensor, + block_table: cute.Tensor, + block_size: Int64, + rms_norm_weight: cute.Tensor, + rms_norm_eps: Float32, + cos_sin_cache: cute.Tensor, + k_cache: cute.Tensor, + kv_slot_mapping: cute.Tensor, + kv_cache_block_size: Int64, + fp8_scale: cute.Tensor, + stream: CUstream, + ): + grid = (slot_mapping.shape[0], 1, 1) + device_kernel( + state_cache, + token_to_req_indices, + positions, + slot_mapping, + block_table, + block_size, + rms_norm_weight, + rms_norm_eps, + cos_sin_cache, + k_cache, + kv_slot_mapping, + kv_cache_block_size, + fp8_scale, + ).launch(grid=grid, block=(tb_size, 1, 1), stream=stream) + + return host_entrypoint + + def dispatch( # type: ignore[override] + self, + *, + compress_ratio: int, + store_full_fp8: bool, + norm_weight_dtype: type[cutlass.Numeric], + head_size: int, + rope_head_dim: int, + ) -> CompileKey: + overlap = compress_ratio == 4 + return self.CompileKey( + head_size=head_size, + state_width=head_size * (2 if overlap else 1), + rope_head_dim=rope_head_dim, + fp8_max=448.0, + quant_block=64, + compress_ratio=compress_ratio, + overlap=overlap, + store_full_fp8=store_full_fp8, + norm_weight_dtype=norm_weight_dtype, + ) + + def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]: + hf_config = vllm_config.model_config.hf_config + cache_config = vllm_config.cache_config + return self._trace_dispatch(self.dispatch)( + compress_ratio=4, + store_full_fp8=( + cache_config.cache_dtype.startswith("fp8") + and cache_config.cache_dtype != "fp8_ds_mla" + ), + norm_weight_dtype=torch_to_cute_dtype(vllm_config.model_config.dtype), + head_size=hf_config.head_dim, + rope_head_dim=hf_config.qk_rope_head_dim, + enabled=( + cache_config.block_size > 0 + and 4 in hf_config.compress_ratios + and cache_config.cache_dtype != "fp8_ds_mla" + ), + _when=lambda *, enabled: enabled, + ) + + def warmup_inputs(self, compile_key: CompileKey) -> tuple[Any, ...]: + head_size = compile_key.head_size + state_width = compile_key.state_width + rope_head_dim = compile_key.rope_head_dim + quant_block = compile_key.quant_block + compress_ratio = compile_key.compress_ratio + overlap = compile_key.overlap + norm_weight_dtype = compile_key.norm_weight_dtype if compress_ratio != 4 or not overlap: raise ValueError("CuTe DSL C4 fused sparse-attn requires C4 overlap.") if head_size != 512: raise ValueError( - "CuTe DSL C4 fused sparse-attn currently requires head_size=512." + "CuTe DSL C4 fused sparse-attn currently requires head_size=512.", ) if state_width != 2 * head_size: raise ValueError( - "CuTe DSL C4 fused sparse-attn requires state_width=2*head_size." + "CuTe DSL C4 fused sparse-attn requires state_width=2*head_size.", ) if quant_block != 64: raise ValueError( - "CuTe DSL C4 fused sparse-attn currently requires quant_block=64." + "CuTe DSL C4 fused sparse-attn currently requires quant_block=64.", ) if rope_head_dim != 64: raise ValueError( - "CuTe DSL C4 fused sparse-attn currently requires rope_head_dim=64." + "CuTe DSL C4 fused sparse-attn currently requires rope_head_dim=64.", ) num_positions = cute.sym_int() num_slots = cute.sym_int() @@ -854,7 +976,6 @@ class SparseAttnCompressNormRopeStoreFullC4Kernel( block_table_width = cute.sym_int() max_pos = cute.sym_int() state_cache_width = state_width * 2 - state_cache = cute.runtime.make_fake_tensor( Float32, (num_state_blocks, state_cache_block_size, state_cache_width), @@ -866,15 +987,21 @@ class SparseAttnCompressNormRopeStoreFullC4Kernel( assumed_align=16, ) token_to_req_indices = make_fake_tensor( - Int32, (num_req_indices,), divisibility=4 + Int32, + (num_req_indices,), + divisibility=4, ) positions = make_fake_tensor(Int64, (num_positions,), divisibility=8) slot_mapping = make_fake_tensor(Int64, (num_slots,), divisibility=8) block_table = make_fake_tensor( - Int32, (cute.sym_int(), block_table_width), divisibility=1 + Int32, + (cute.sym_int(), block_table_width), + divisibility=1, ) rms_norm_weight = make_fake_tensor( - norm_weight_dtype, (head_size,), divisibility=4 + norm_weight_dtype, + (head_size,), + divisibility=4, ) cos_sin_cache = cute.runtime.make_fake_tensor( Float32, @@ -885,31 +1012,12 @@ class SparseAttnCompressNormRopeStoreFullC4Kernel( k_cache = cute.runtime.make_fake_tensor( Uint8, (num_kv_blocks, cute.sym_int(), cute.sym_int()), - stride=( - cute.sym_int64(divisibility=16), - cute.sym_int64(divisibility=8), - 1, - ), + stride=(cute.sym_int64(divisibility=16), cute.sym_int64(divisibility=8), 1), assumed_align=16, ) kv_slot_mapping = make_fake_tensor(Int64, (num_kv_slots,), divisibility=8) fp8_scale = make_fake_tensor(Float32, (1,), divisibility=1) - - kernel = SparseAttnCompressNormRopeStoreFullC4Kernel( - head_size, - state_width, - rope_head_dim, - fp8_max, - quant_block, - token_stride, - scale_dim, - compress_ratio, - overlap, - store_full_fp8, - ) - stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) - return cute.compile( - kernel, + return ( state_cache, token_to_req_indices, positions, @@ -923,12 +1031,59 @@ class SparseAttnCompressNormRopeStoreFullC4Kernel( kv_slot_mapping, Int64(0), fp8_scale, - stream, - options="--enable-tvm-ffi", ) + @kernel_launcher + def __call__( + self, + *, + state_cache: torch.Tensor, + token_to_req_indices: torch.Tensor, + positions: torch.Tensor, + slot_mapping: torch.Tensor, + block_table: torch.Tensor, + block_size: int, + rms_norm_weight: torch.Tensor, + rms_norm_eps: float, + cos_sin_cache: torch.Tensor, + kv_cache: torch.Tensor, + kv_slot_mapping: torch.Tensor, + compress_ratio: int, + store_full_fp8: bool, + fp8_scale: torch.Tensor | None, + head_dim: int, + rope_head_dim: int, + ) -> CuTeDSLLaunchSpec[SparseAttnCompressNormRopeStoreFullC4Kernel.CompileKey]: + compile_key = self.dispatch( + compress_ratio=compress_ratio, + store_full_fp8=store_full_fp8, + norm_weight_dtype=torch_to_cute_dtype(rms_norm_weight.dtype), + head_size=head_dim, + rope_head_dim=rope_head_dim, + ) + if fp8_scale is None: + fp8_scale = torch.ones(1, dtype=torch.float32, device=kv_cache.device) + launch_args = ( + state_cache, + token_to_req_indices, + positions, + slot_mapping, + block_table, + block_size, + rms_norm_weight, + rms_norm_eps, + cos_sin_cache, + kv_cache.view(torch.uint8), + kv_slot_mapping, + kv_cache.shape[1], + fp8_scale, + ) + return compile_key, launch_args -class SparseAttnCompressC128Block8Kernel: + +class SparseAttnCompressC128Block8Kernel( + VllmCuTeDSLJitKernel["SparseAttnCompressC128Block8Kernel.CompileKey"] +): head_tile = 64 rows_per_warp = 16 elems_per_lane = 2 @@ -942,253 +1097,282 @@ class SparseAttnCompressC128Block8Kernel: state_block_size = 8 rcp_ln2 = 1.4426950408889634 - def __init__( - self, - head_size: int, - state_width: int, - ): - self.head_dim = head_size - self.num_splits = head_size // self.head_tile - self.state_width = state_width + @dataclass(frozen=True) + class CompileKey: + head_size: int + state_width: int - @cute.jit - def __call__( - self, - state_cache: cute.Tensor, - token_to_req_indices: cute.Tensor, - positions: cute.Tensor, - slot_mapping: cute.Tensor, - block_table: cute.Tensor, - compressed_kv: cute.Tensor, - stream: CUstream, - ): - grid = (slot_mapping.shape[0] * self.num_splits, 1, 1) - self.kernel( - state_cache, - token_to_req_indices, - positions, - slot_mapping, - block_table, - compressed_kv, - ).launch(grid=grid, block=(self.tb_size, 1, 1), stream=stream) - - @cute.kernel - def kernel( - self, - state_cache: cute.Tensor, - token_to_req_indices: cute.Tensor, - positions: cute.Tensor, - slot_mapping: cute.Tensor, - block_table: cute.Tensor, - compressed_kv: cute.Tensor, - ): - block_id, _, _ = cute.arch.block_idx() - tid, _, _ = cute.arch.thread_idx() - warp_id = cute.arch.make_warp_uniform(tid // 32) - lane_id = tid % 32 - col_group = lane_id % self.lanes_per_row - - token_idx = block_id // self.num_splits - split_idx = block_id - token_idx * self.num_splits - col_base = split_idx * self.head_tile + col_group * self.elems_per_lane - - position = Int64(0) - req_idx = Int32(0) - slot_id = Int64(-1) - has_position = token_idx < positions.shape[0] - has_req_idx = token_idx < token_to_req_indices.shape[0] - if lane_id == 0: - slot_id = slot_mapping[token_idx] - if lane_id == 0 and has_position: - position = positions[token_idx] - if lane_id == 0 and has_req_idx: - req_idx = token_to_req_indices[token_idx] - slot_id = cute.arch.shuffle_sync(slot_id, offset=0) - position = cute.arch.shuffle_sync(position, offset=0) - req_idx = cute.arch.shuffle_sync(req_idx, offset=0) - boundary = has_position and ( - (position + Int64(1)) % Int64(self.compress_ratio) == Int64(0) + @staticmethod + def kernel(compile_key: CompileKey) -> Any: + head_tile = SparseAttnCompressC128Block8Kernel.head_tile + rows_per_warp = SparseAttnCompressC128Block8Kernel.rows_per_warp + elems_per_lane = SparseAttnCompressC128Block8Kernel.elems_per_lane + lanes_per_row = SparseAttnCompressC128Block8Kernel.lanes_per_row + num_warps = SparseAttnCompressC128Block8Kernel.num_warps + stats_lane_stride = SparseAttnCompressC128Block8Kernel.stats_lane_stride + final_reduce_steps = SparseAttnCompressC128Block8Kernel.final_reduce_steps + final_reduce_initial_offset = ( + SparseAttnCompressC128Block8Kernel.final_reduce_initial_offset ) - start = position - Int64(self.compress_ratio - 1) - active = slot_id >= Int64(0) and has_req_idx and boundary + tb_size = SparseAttnCompressC128Block8Kernel.tb_size + compress_ratio = SparseAttnCompressC128Block8Kernel.compress_ratio + state_block_size = SparseAttnCompressC128Block8Kernel.state_block_size + rcp_ln2 = SparseAttnCompressC128Block8Kernel.rcp_ln2 + head_dim = compile_key.head_size + num_splits = head_dim // head_tile + state_width = compile_key.state_width - if active: - smem = cutlass.utils.SmemAllocator() - s_max = smem.allocate_tensor( - Float32, - cute.make_layout( - ( - self.num_warps, - self.lanes_per_row, - self.elems_per_lane, - ), - stride=( - self.stats_lane_stride * self.elems_per_lane, - self.elems_per_lane, - 1, - ), - ), - byte_alignment=4, - ) - s_sum = smem.allocate_tensor( - Float32, - cute.make_layout( - ( - self.num_warps, - self.lanes_per_row, - self.elems_per_lane, - ), - stride=( - self.stats_lane_stride * self.elems_per_lane, - self.elems_per_lane, - 1, - ), - ), - byte_alignment=4, - ) - s_product = smem.allocate_tensor( - Float32, - cute.make_layout( - ( - self.num_warps, - self.lanes_per_row, - self.elems_per_lane, - ), - stride=( - self.stats_lane_stride * self.elems_per_lane, - self.elems_per_lane, - 1, - ), - ), - byte_alignment=4, - ) + @cute.kernel + def device_kernel( + state_cache: cute.Tensor, + token_to_req_indices: cute.Tensor, + positions: cute.Tensor, + slot_mapping: cute.Tensor, + block_table: cute.Tensor, + compressed_kv: cute.Tensor, + ): + block_id, _, _ = cute.arch.block_idx() + tid, _, _ = cute.arch.thread_idx() + warp_id = cute.arch.make_warp_uniform(tid // 32) + lane_id = tid % 32 + col_group = lane_id % lanes_per_row - row_layout = cute.make_layout( - (self.rows_per_warp, self.elems_per_lane), - stride=(self.elems_per_lane, 1), - ) - kv_vals = cute.make_rmem_tensor(row_layout, Float32) - score_vals = cute.make_rmem_tensor(row_layout, Float32) - local_max = cute.make_rmem_tensor((self.elems_per_lane,), Float32) - local_sum = cute.make_rmem_tensor((self.elems_per_lane,), Float32) - local_product = cute.make_rmem_tensor((self.elems_per_lane,), Float32) + token_idx = block_id // num_splits + split_idx = block_id - token_idx * num_splits + col_base = split_idx * head_tile + col_group * elems_per_lane - for e in cutlass.range_constexpr(self.elems_per_lane): - local_max[e] = -Float32.inf - local_sum[e] = Float32(0.0) - local_product[e] = Float32(0.0) - - first_block_index = start // Int64(self.state_block_size) - warp_block_index = first_block_index + (warp_id * 2).to(Int64) - block0_i32 = Int32(0) - block1_i32 = Int32(0) + position = Int64(0) + req_idx = Int32(0) + slot_id = Int64(-1) + has_position = token_idx < positions.shape[0] + has_req_idx = token_idx < token_to_req_indices.shape[0] if lane_id == 0: - block0_i32 = block_table[req_idx, warp_block_index] - block1_i32 = block_table[req_idx, warp_block_index + Int64(1)] - block0_i32 = cute.arch.shuffle_sync(block0_i32, offset=0) - block1_i32 = cute.arch.shuffle_sync(block1_i32, offset=0) - - cp_f32x2 = cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), Float32, num_bits_per_copy=64 + slot_id = slot_mapping[token_idx] + if lane_id == 0 and has_position: + position = positions[token_idx] + if lane_id == 0 and has_req_idx: + req_idx = token_to_req_indices[token_idx] + slot_id = cute.arch.shuffle_sync(slot_id, offset=0) + position = cute.arch.shuffle_sync(position, offset=0) + req_idx = cute.arch.shuffle_sync(req_idx, offset=0) + boundary = has_position and ( + (position + Int64(1)) % Int64(compress_ratio) == Int64(0) ) - final_mask_and_clamp = const_expr( - (cute.arch.WARP_SIZE - self.num_warps) << 8 | (cute.arch.WARP_SIZE - 1) - ) - col_tile = col_base.to(Int64) // Int64(self.elems_per_lane) - score_col_tile = col_tile + Int64(self.state_width // self.elems_per_lane) + start = position - Int64(compress_ratio - 1) + active = slot_id >= Int64(0) and has_req_idx and boundary - for i in cutlass.range_constexpr(self.rows_per_warp): - block_number_i32 = block0_i32 - block_offset = Int64(i) - if const_expr(i >= self.state_block_size): - block_number_i32 = block1_i32 - block_offset = Int64(i - self.state_block_size) - row_tensor = state_cache[block_number_i32.to(Int64), block_offset, None] - kv_src = cute.local_tile( - row_tensor, - tiler=(self.elems_per_lane,), - coord=(col_tile,), + if active: + smem = cutlass.utils.SmemAllocator() + s_max = smem.allocate_tensor( + Float32, + cute.make_layout( + ( + num_warps, + lanes_per_row, + elems_per_lane, + ), + stride=( + stats_lane_stride * elems_per_lane, + elems_per_lane, + 1, + ), + ), + byte_alignment=4, ) - score_src = cute.local_tile( - row_tensor, - tiler=(self.elems_per_lane,), - coord=(score_col_tile,), + s_sum = smem.allocate_tensor( + Float32, + cute.make_layout( + ( + num_warps, + lanes_per_row, + elems_per_lane, + ), + stride=( + stats_lane_stride * elems_per_lane, + elems_per_lane, + 1, + ), + ), + byte_alignment=4, + ) + s_product = smem.allocate_tensor( + Float32, + cute.make_layout( + ( + num_warps, + lanes_per_row, + elems_per_lane, + ), + stride=( + stats_lane_stride * elems_per_lane, + elems_per_lane, + 1, + ), + ), + byte_alignment=4, ) - cute.copy(cp_f32x2, kv_src, kv_vals[i, None]) - cute.copy(cp_f32x2, score_src, score_vals[i, None]) - for e in cutlass.range_constexpr(self.elems_per_lane): - local_max[e] = cute.arch.fmax(local_max[e], score_vals[i, e]) + row_layout = cute.make_layout( + (rows_per_warp, elems_per_lane), + stride=(elems_per_lane, 1), + ) + kv_vals = cute.make_rmem_tensor(row_layout, Float32) + score_vals = cute.make_rmem_tensor(row_layout, Float32) + local_max = cute.make_rmem_tensor((elems_per_lane,), Float32) + local_sum = cute.make_rmem_tensor((elems_per_lane,), Float32) + local_product = cute.make_rmem_tensor((elems_per_lane,), Float32) - for e in cutlass.range_constexpr(self.elems_per_lane): - for i in cutlass.range_constexpr(self.rows_per_warp): - exp_score = cute.math.exp2( - (score_vals[i, e] - local_max[e]) * Float32(self.rcp_ln2), + for e in cutlass.range_constexpr(elems_per_lane): + local_max[e] = -Float32.inf + local_sum[e] = Float32(0.0) + local_product[e] = Float32(0.0) + + first_block_index = start // Int64(state_block_size) + warp_block_index = first_block_index + (warp_id * 2).to(Int64) + block0_i32 = Int32(0) + block1_i32 = Int32(0) + if lane_id == 0: + block0_i32 = block_table[req_idx, warp_block_index] + block1_i32 = block_table[req_idx, warp_block_index + Int64(1)] + block0_i32 = cute.arch.shuffle_sync(block0_i32, offset=0) + block1_i32 = cute.arch.shuffle_sync(block1_i32, offset=0) + + cp_f32x2 = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), Float32, num_bits_per_copy=64 + ) + final_mask_and_clamp = const_expr( + (cute.arch.WARP_SIZE - num_warps) << 8 | (cute.arch.WARP_SIZE - 1) + ) + col_tile = col_base.to(Int64) // Int64(elems_per_lane) + score_col_tile = col_tile + Int64(state_width // elems_per_lane) + + for i in cutlass.range_constexpr(rows_per_warp): + block_number_i32 = block0_i32 + block_offset = Int64(i) + if const_expr(i >= state_block_size): + block_number_i32 = block1_i32 + block_offset = Int64(i - state_block_size) + row_tensor = state_cache[ + block_number_i32.to(Int64), block_offset, None + ] + kv_src = cute.local_tile( + row_tensor, + tiler=(elems_per_lane,), + coord=(col_tile,), + ) + score_src = cute.local_tile( + row_tensor, + tiler=(elems_per_lane,), + coord=(score_col_tile,), + ) + cute.copy(cp_f32x2, kv_src, kv_vals[i, None]) + cute.copy(cp_f32x2, score_src, score_vals[i, None]) + + for e in cutlass.range_constexpr(elems_per_lane): + local_max[e] = cute.arch.fmax(local_max[e], score_vals[i, e]) + + for e in cutlass.range_constexpr(elems_per_lane): + for i in cutlass.range_constexpr(rows_per_warp): + exp_score = cute.math.exp2( + (score_vals[i, e] - local_max[e]) * Float32(rcp_ln2), + fastmath=True, + ) + local_sum[e] += exp_score + local_product[e] += kv_vals[i, e] * exp_score + + for e in cutlass.range_constexpr(elems_per_lane): + s_max[warp_id, col_group, e] = local_max[e] + s_sum[warp_id, col_group, e] = local_sum[e] + s_product[warp_id, col_group, e] = local_product[e] + cute.arch.sync_threads() + + out_group = tid // num_warps + final_lane = tid % num_warps + final_groups_per_pass = const_expr(tb_size // num_warps) + for pass_idx in cutlass.range_constexpr( + head_tile // final_groups_per_pass + ): + out_idx = pass_idx * final_groups_per_pass + out_group + out_lane = out_idx // elems_per_lane + out_elem = out_idx % elems_per_lane + + local_warp_max = s_max[final_lane, out_lane, out_elem] + global_max = local_warp_max + for step in cutlass.range_constexpr(final_reduce_steps): + offset = const_expr(final_reduce_initial_offset >> step) + global_max = cute.arch.fmax( + global_max, + cute.arch.shuffle_sync_bfly( + global_max, + offset=offset, + mask_and_clamp=final_mask_and_clamp, + ), + ) + + scale = cute.math.exp2( + (local_warp_max - global_max) * Float32(rcp_ln2), fastmath=True, ) - local_sum[e] += exp_score - local_product[e] += kv_vals[i, e] * exp_score - - for e in cutlass.range_constexpr(self.elems_per_lane): - s_max[warp_id, col_group, e] = local_max[e] - s_sum[warp_id, col_group, e] = local_sum[e] - s_product[warp_id, col_group, e] = local_product[e] - cute.arch.sync_threads() - - out_group = tid // self.num_warps - final_lane = tid % self.num_warps - final_groups_per_pass = const_expr(self.tb_size // self.num_warps) - for pass_idx in cutlass.range_constexpr( - self.head_tile // final_groups_per_pass - ): - out_idx = pass_idx * final_groups_per_pass + out_group - out_lane = out_idx // self.elems_per_lane - out_elem = out_idx % self.elems_per_lane - - local_warp_max = s_max[final_lane, out_lane, out_elem] - global_max = local_warp_max - for step in cutlass.range_constexpr(self.final_reduce_steps): - offset = const_expr(self.final_reduce_initial_offset >> step) - global_max = cute.arch.fmax( - global_max, - cute.arch.shuffle_sync_bfly( - global_max, + global_sum = s_sum[final_lane, out_lane, out_elem] * scale + global_product = s_product[final_lane, out_lane, out_elem] * scale + for step in cutlass.range_constexpr(final_reduce_steps): + offset = const_expr(final_reduce_initial_offset >> step) + global_sum += cute.arch.shuffle_sync_bfly( + global_sum, offset=offset, mask_and_clamp=final_mask_and_clamp, - ), - ) + ) + global_product += cute.arch.shuffle_sync_bfly( + global_product, + offset=offset, + mask_and_clamp=final_mask_and_clamp, + ) - scale = cute.math.exp2( - (local_warp_max - global_max) * Float32(self.rcp_ln2), - fastmath=True, - ) - global_sum = s_sum[final_lane, out_lane, out_elem] * scale - global_product = s_product[final_lane, out_lane, out_elem] * scale - for step in cutlass.range_constexpr(self.final_reduce_steps): - offset = const_expr(self.final_reduce_initial_offset >> step) - global_sum += cute.arch.shuffle_sync_bfly( - global_sum, - offset=offset, - mask_and_clamp=final_mask_and_clamp, - ) - global_product += cute.arch.shuffle_sync_bfly( - global_product, - offset=offset, - mask_and_clamp=final_mask_and_clamp, - ) + if final_lane == 0: + compressed_kv.iterator[ + token_idx.to(Int64) * compressed_kv.stride[0] + + (split_idx * head_tile + out_idx).to(Int64) + ] = global_product / global_sum - if final_lane == 0: - compressed_kv.iterator[ - token_idx.to(Int64) * compressed_kv.stride[0] - + (split_idx * self.head_tile + out_idx).to(Int64) - ] = global_product / global_sum + @cute.jit + def host_entrypoint( + state_cache: cute.Tensor, + token_to_req_indices: cute.Tensor, + positions: cute.Tensor, + slot_mapping: cute.Tensor, + block_table: cute.Tensor, + compressed_kv: cute.Tensor, + stream: CUstream, + ): + grid = (slot_mapping.shape[0] * num_splits, 1, 1) + device_kernel( + state_cache, + token_to_req_indices, + positions, + slot_mapping, + block_table, + compressed_kv, + ).launch(grid=grid, block=(tb_size, 1, 1), stream=stream) - @cache - @staticmethod - def compile( - head_size: int = 512, - state_width: int = 512, - ): + return host_entrypoint + + def dispatch( # type: ignore[override] + self, *, head_size: int, state_width: int + ) -> CompileKey: + return self.CompileKey(head_size=head_size, state_width=state_width) + + def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]: + hf_config = vllm_config.model_config.hf_config + return self._trace_dispatch(self.dispatch)( + head_size=hf_config.head_dim, + state_width=hf_config.head_dim, + enabled=128 in hf_config.compress_ratios, + _when=lambda *, enabled: enabled, + ) + + def warmup_inputs(self, compile_key: CompileKey) -> tuple[Any, ...]: + head_size = compile_key.head_size + state_width = compile_key.state_width if head_size % SparseAttnCompressC128Block8Kernel.head_tile != 0: raise ValueError("head_size must be divisible by the 64-wide head tile.") num_positions = cute.sym_int() @@ -1198,7 +1382,6 @@ class SparseAttnCompressC128Block8Kernel: state_cache_block_size = cute.sym_int() block_table_width = cute.sym_int() state_cache_width = state_width * 2 - state_cache = cute.runtime.make_fake_tensor( Float32, (num_blocks, state_cache_block_size, state_cache_width), @@ -1210,12 +1393,16 @@ class SparseAttnCompressC128Block8Kernel: assumed_align=16, ) token_to_req_indices = make_fake_tensor( - Int32, (num_req_indices,), divisibility=4 + Int32, + (num_req_indices,), + divisibility=4, ) positions = make_fake_tensor(Int64, (num_positions,), divisibility=8) slot_mapping = make_fake_tensor(Int64, (num_slots,), divisibility=8) block_table = make_fake_tensor( - Int32, (cute.sym_int(), block_table_width), divisibility=1 + Int32, + (cute.sym_int(), block_table_width), + divisibility=1, ) compressed_kv = cute.runtime.make_fake_tensor( Float32, @@ -1223,238 +1410,332 @@ class SparseAttnCompressC128Block8Kernel: stride=(head_size, 1), assumed_align=4, ) - - kernel = SparseAttnCompressC128Block8Kernel( - head_size, - state_width, - ) - stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) - return cute.compile( - kernel, + return ( state_cache, token_to_req_indices, positions, slot_mapping, block_table, compressed_kv, - stream, - options="--enable-tvm-ffi", ) - -class SparseAttnNormRopeStoreKernel: - min_scale = 1.0e-4 - - def __init__( - self, - head_size: int, - rope_head_dim: int, - fp8_max: float, - quant_block: int, - token_stride: int, - scale_dim: int, - compress_ratio: int, - static_kv_cache_block_size: int, - ): - self.head_dim = head_size - self.rope_dim = rope_head_dim - self.nope_dim = head_size - rope_head_dim - self.fp8_max = fp8_max - self.quant_block = quant_block - self.token_stride = token_stride - self.scale_dim = scale_dim - self.num_warps = head_size // quant_block - self.nope_blocks = self.nope_dim // quant_block - self.tb_size = head_size // 2 - self.compress_ratio = compress_ratio - self.static_kv_cache_block_size = static_kv_cache_block_size - - @cute.jit + @kernel_launcher def __call__( self, - compressed_kv: cute.Tensor, - positions: cute.Tensor, - slot_mapping: cute.Tensor, - rms_norm_weight: cute.Tensor, - rms_norm_eps: Float32, - cos_sin_cache: cute.Tensor, - k_cache: cute.Tensor, - kv_slot_mapping: cute.Tensor, - stream: CUstream, - ): - grid = (slot_mapping.shape[0], 1, 1) - self.kernel( - compressed_kv, + *, + state_cache: torch.Tensor, + num_actual: int, + token_to_req_indices: torch.Tensor, + positions: torch.Tensor, + slot_mapping: torch.Tensor, + block_table: torch.Tensor, + head_dim: int, + compressed_kv: torch.Tensor | None = None, + ) -> CuTeDSLLaunchSpec[SparseAttnCompressC128Block8Kernel.CompileKey]: + compile_key = self.dispatch(head_size=head_dim, state_width=head_dim) + if compressed_kv is None: + compressed_kv = torch.empty( + (num_actual, head_dim), + dtype=torch.float32, + device=state_cache.device, + ) + launch_args = ( + state_cache, + token_to_req_indices, positions, slot_mapping, - rms_norm_weight, - rms_norm_eps, - cos_sin_cache, - k_cache, - kv_slot_mapping, - ).launch(grid=grid, block=(self.tb_size, 1, 1), stream=stream) - - @cute.kernel - def kernel( - self, - compressed_kv: cute.Tensor, - positions: cute.Tensor, - slot_mapping: cute.Tensor, - rms_norm_weight: cute.Tensor, - rms_norm_eps: Float32, - cos_sin_cache: cute.Tensor, - k_cache: cute.Tensor, - kv_slot_mapping: cute.Tensor, - ): - token_idx, _, _ = cute.arch.block_idx() - tid, _, _ = cute.arch.thread_idx() - warp_id = cute.arch.make_warp_uniform(tid // 32) - lane_id = tid % 32 - elem0 = tid * 2 - - position = Int64(0) - kv_slot_idx = Int64(-1) - has_position = token_idx < positions.shape[0] - slot_id = Int64(-1) - if lane_id == 0: - slot_id = slot_mapping[token_idx] - if lane_id == 0 and has_position: - position = positions[token_idx] - has_kv_slot_idx = token_idx < kv_slot_mapping.shape[0] - if lane_id == 0 and has_kv_slot_idx: - kv_slot_idx = kv_slot_mapping[token_idx] - slot_id = cute.arch.shuffle_sync(slot_id, offset=0) - position = cute.arch.shuffle_sync(position, offset=0) - kv_slot_idx = cute.arch.shuffle_sync(kv_slot_idx, offset=0) - boundary = has_position and ( - (position + Int64(1)) % Int64(self.compress_ratio) == Int64(0) + block_table, + compressed_kv, ) - active = slot_id >= Int64(0) and boundary and kv_slot_idx >= Int64(0) + return compile_key, launch_args, compressed_kv - if active: - k_cache_u16 = cute.recast_tensor(k_cache, Uint16) - k_cache_u32 = cute.recast_tensor(k_cache, Uint32) - static_block_size = Int64(self.static_kv_cache_block_size) - page = kv_slot_idx // static_block_size - kv_offset = kv_slot_idx - page * static_block_size - scale_row_offset = static_block_size * Int64(self.token_stride) - value_base = page * k_cache.stride[0] + kv_offset * Int64(self.token_stride) - scale_base = ( - page * k_cache.stride[0] - + scale_row_offset - + kv_offset * Int64(self.scale_dim) - ) - weight0 = rms_norm_weight[elem0].to(Float32) - weight1 = rms_norm_weight[elem0 + 1].to(Float32) - base = token_idx.to(Int64) * compressed_kv.stride[0] + elem0.to(Int64) - x0 = compressed_kv.iterator[base] - x1 = compressed_kv.iterator[base + Int64(1)] +class SparseAttnNormRopeStoreKernel( + VllmCuTeDSLJitKernel["SparseAttnNormRopeStoreKernel.CompileKey"] +): + min_scale = 1.0e-4 - local_sumsq = x0 * x0 + x1 * x1 - warp_sum = local_sumsq - for step in cutlass.range_constexpr(5): - offset = const_expr(16 >> step) - warp_sum += cute.arch.shuffle_sync_bfly(warp_sum, offset) + @dataclass(frozen=True) + class CompileKey: + head_size: int + rope_head_dim: int + fp8_max: float + quant_block: int + token_stride: int + scale_dim: int + kv_block_stride: int + compress_ratio: int + norm_weight_dtype: type[cutlass.Numeric] + kv_cache_block_size: int - smem = cutlass.utils.SmemAllocator() - partial_sums = smem.allocate_tensor( - Float32, cute.make_layout((self.num_warps,)), byte_alignment=4 - ) + @staticmethod + def kernel(compile_key: CompileKey) -> Any: + min_scale = SparseAttnNormRopeStoreKernel.min_scale + head_dim = compile_key.head_size + rope_dim = compile_key.rope_head_dim + nope_dim = head_dim - rope_dim + fp8_max = compile_key.fp8_max + quant_block = compile_key.quant_block + token_stride = compile_key.token_stride + scale_dim = compile_key.scale_dim + num_warps = head_dim // quant_block + nope_blocks = nope_dim // quant_block + tb_size = head_dim // 2 + compress_ratio = compile_key.compress_ratio + static_kv_cache_block_size = compile_key.kv_cache_block_size + @cute.kernel + def device_kernel( + compressed_kv: cute.Tensor, + positions: cute.Tensor, + slot_mapping: cute.Tensor, + rms_norm_weight: cute.Tensor, + rms_norm_eps: Float32, + cos_sin_cache: cute.Tensor, + k_cache: cute.Tensor, + kv_slot_mapping: cute.Tensor, + ): + token_idx, _, _ = cute.arch.block_idx() + tid, _, _ = cute.arch.thread_idx() + warp_id = cute.arch.make_warp_uniform(tid // 32) + lane_id = tid % 32 + elem0 = tid * 2 + + position = Int64(0) + kv_slot_idx = Int64(-1) + has_position = token_idx < positions.shape[0] + slot_id = Int64(-1) if lane_id == 0: - partial_sums[warp_id] = warp_sum - cute.arch.sync_threads() - - total = partial_sums[lane_id % self.num_warps] - sum_mask_and_clamp = const_expr( - (cute.arch.WARP_SIZE - self.num_warps) << 8 | (cute.arch.WARP_SIZE - 1) + slot_id = slot_mapping[token_idx] + if lane_id == 0 and has_position: + position = positions[token_idx] + has_kv_slot_idx = token_idx < kv_slot_mapping.shape[0] + if lane_id == 0 and has_kv_slot_idx: + kv_slot_idx = kv_slot_mapping[token_idx] + slot_id = cute.arch.shuffle_sync(slot_id, offset=0) + position = cute.arch.shuffle_sync(position, offset=0) + kv_slot_idx = cute.arch.shuffle_sync(kv_slot_idx, offset=0) + boundary = has_position and ( + (position + Int64(1)) % Int64(compress_ratio) == Int64(0) ) - for step in cutlass.range_constexpr(3): - offset = const_expr(4 >> step) - total += cute.arch.shuffle_sync_bfly( - total, - offset, - mask_and_clamp=sum_mask_and_clamp, + active = slot_id >= Int64(0) and boundary and kv_slot_idx >= Int64(0) + + if active: + k_cache_u16 = cute.recast_tensor(k_cache, Uint16) + k_cache_u32 = cute.recast_tensor(k_cache, Uint32) + static_block_size = Int64(static_kv_cache_block_size) + page = kv_slot_idx // static_block_size + kv_offset = kv_slot_idx - page * static_block_size + scale_row_offset = static_block_size * Int64(token_stride) + value_base = page * k_cache.stride[0] + kv_offset * Int64(token_stride) + scale_base = ( + page * k_cache.stride[0] + + scale_row_offset + + kv_offset * Int64(scale_dim) ) + weight0 = rms_norm_weight[elem0].to(Float32) + weight1 = rms_norm_weight[elem0 + 1].to(Float32) - rrms = cute.math.rsqrt( - total / Float32(self.head_dim) + rms_norm_eps, fastmath=True - ) - x0 = x0 * rrms * weight0 - x1 = x1 * rrms * weight1 + base = token_idx.to(Int64) * compressed_kv.stride[0] + elem0.to(Int64) + x0 = compressed_kv.iterator[base] + x1 = compressed_kv.iterator[base + Int64(1)] - if warp_id == self.nope_blocks: - pair_idx = lane_id - compressed_pos = position - Int64(self.compress_ratio - 1) - cs_base = compressed_pos * cos_sin_cache.stride[0] + pair_idx.to(Int64) - cos_v = cos_sin_cache.iterator[cs_base] - sin_v = cos_sin_cache.iterator[cs_base + Int64(self.rope_dim // 2)] - real = x0 * cos_v - x1 * sin_v - imag = x0 * sin_v + x1 * cos_v - packed = _fp32x2_to_bf16x2(real, imag) - out_base = value_base + Int64(self.nope_dim) + (lane_id * 4).to(Int64) - k_cache_u32.iterator[out_base // Int64(4)] = packed - else: - q_packed = _fp32x2_to_bf16x2(x0, x1) - q0, q1 = _bf16x2_to_fp32(q_packed) - abs0 = cute.math.absf(q0) - abs1 = cute.math.absf(q1) - local_absmax = cute.arch.fmax(abs0, abs1) - absmax = local_absmax + local_sumsq = x0 * x0 + x1 * x1 + warp_sum = local_sumsq for step in cutlass.range_constexpr(5): offset = const_expr(16 >> step) - absmax = cute.arch.fmax( - absmax, cute.arch.shuffle_sync_bfly(absmax, offset) - ) - scale_raw = cute.arch.fmax( - Float32(self.min_scale), - absmax / Float32(self.fp8_max), - ) - bits = _recast_val(scale_raw, Uint32) - ue8m0 = ((bits + Uint32(0x7FFFFF)) >> Uint32(23)) & Uint32(0xFF) - inv_scale = _recast_val((Uint32(254) - ue8m0) << Uint32(23), Float32) - y0 = cutlass.min( - cute.arch.fmax(q0 * inv_scale, Float32(-self.fp8_max)), - Float32(self.fp8_max), - ) - y1 = cutlass.min( - cute.arch.fmax(q1 * inv_scale, Float32(-self.fp8_max)), - Float32(self.fp8_max), - ) - packed_fp8 = _fp32x2_to_fp8e4m3x2(y0, y1) - out_base = value_base + (warp_id * self.quant_block + lane_id * 2).to( - Int64 - ) - k_cache_u16.iterator[out_base // Int64(2)] = packed_fp8 - if lane_id == 0: - k_cache.iterator[scale_base + warp_id.to(Int64)] = ue8m0.to(Uint8) - if warp_id == 0: - k_cache.iterator[scale_base + Int64(self.nope_blocks)] = Uint8( - 0 - ) + warp_sum += cute.arch.shuffle_sync_bfly(warp_sum, offset) - @cache - @staticmethod - def compile( - head_size: int = 512, - rope_head_dim: int = 64, - fp8_max: float = 448.0, - quant_block: int = 64, - token_stride: int = 576, - scale_dim: int = 8, - kv_block_stride: int = 74752, - compress_ratio: int = 128, - norm_weight_dtype: type[cutlass.Numeric] = Float32, - static_kv_cache_block_size: int = 0, - ): + smem = cutlass.utils.SmemAllocator() + partial_sums = smem.allocate_tensor( + Float32, cute.make_layout((num_warps,)), byte_alignment=4 + ) + + if lane_id == 0: + partial_sums[warp_id] = warp_sum + cute.arch.sync_threads() + + total = partial_sums[lane_id % num_warps] + sum_mask_and_clamp = const_expr( + (cute.arch.WARP_SIZE - num_warps) << 8 | (cute.arch.WARP_SIZE - 1) + ) + for step in cutlass.range_constexpr(3): + offset = const_expr(4 >> step) + total += cute.arch.shuffle_sync_bfly( + total, + offset, + mask_and_clamp=sum_mask_and_clamp, + ) + + rrms = cute.math.rsqrt( + total / Float32(head_dim) + rms_norm_eps, fastmath=True + ) + x0 = x0 * rrms * weight0 + x1 = x1 * rrms * weight1 + + if warp_id == nope_blocks: + pair_idx = lane_id + compressed_pos = position - Int64(compress_ratio - 1) + cs_base = compressed_pos * cos_sin_cache.stride[0] + pair_idx.to( + Int64 + ) + cos_v = cos_sin_cache.iterator[cs_base] + sin_v = cos_sin_cache.iterator[cs_base + Int64(rope_dim // 2)] + real = x0 * cos_v - x1 * sin_v + imag = x0 * sin_v + x1 * cos_v + packed = _fp32x2_to_bf16x2(real, imag) + out_base = value_base + Int64(nope_dim) + (lane_id * 4).to(Int64) + k_cache_u32.iterator[out_base // Int64(4)] = packed + else: + q_packed = _fp32x2_to_bf16x2(x0, x1) + q0, q1 = _bf16x2_to_fp32(q_packed) + abs0 = cute.math.absf(q0) + abs1 = cute.math.absf(q1) + local_absmax = cute.arch.fmax(abs0, abs1) + absmax = local_absmax + for step in cutlass.range_constexpr(5): + offset = const_expr(16 >> step) + absmax = cute.arch.fmax( + absmax, cute.arch.shuffle_sync_bfly(absmax, offset) + ) + scale_raw = cute.arch.fmax( + Float32(min_scale), + absmax / Float32(fp8_max), + ) + bits = _recast_val(scale_raw, Uint32) + ue8m0 = ((bits + Uint32(0x7FFFFF)) >> Uint32(23)) & Uint32(0xFF) + inv_scale = _recast_val( + (Uint32(254) - ue8m0) << Uint32(23), Float32 + ) + y0 = cutlass.min( + cute.arch.fmax(q0 * inv_scale, Float32(-fp8_max)), + Float32(fp8_max), + ) + y1 = cutlass.min( + cute.arch.fmax(q1 * inv_scale, Float32(-fp8_max)), + Float32(fp8_max), + ) + packed_fp8 = _fp32x2_to_fp8e4m3x2(y0, y1) + out_base = value_base + (warp_id * quant_block + lane_id * 2).to( + Int64 + ) + k_cache_u16.iterator[out_base // Int64(2)] = packed_fp8 + if lane_id == 0: + k_cache.iterator[scale_base + warp_id.to(Int64)] = ue8m0.to( + Uint8 + ) + if warp_id == 0: + k_cache.iterator[scale_base + Int64(nope_blocks)] = Uint8(0) + + @cute.jit + def host_entrypoint( + compressed_kv: cute.Tensor, + positions: cute.Tensor, + slot_mapping: cute.Tensor, + rms_norm_weight: cute.Tensor, + rms_norm_eps: Float32, + cos_sin_cache: cute.Tensor, + k_cache: cute.Tensor, + kv_slot_mapping: cute.Tensor, + stream: CUstream, + ): + grid = (slot_mapping.shape[0], 1, 1) + device_kernel( + compressed_kv, + positions, + slot_mapping, + rms_norm_weight, + rms_norm_eps, + cos_sin_cache, + k_cache, + kv_slot_mapping, + ).launch(grid=grid, block=(tb_size, 1, 1), stream=stream) + + return host_entrypoint + + def dispatch( # type: ignore[override] + self, + *, + compress_ratio: int, + cache_block_size: int, + cache_alignment: int, + norm_weight_dtype: type[cutlass.Numeric], + head_size: int, + rope_head_dim: int, + runtime_kv_block_stride: int | None = None, + ) -> CompileKey: + raw_kv_cache_block_size = cache_block_size // compress_ratio + kv_cache_block_size = ( + raw_kv_cache_block_size if raw_kv_cache_block_size >= 1 else 1 + ) + token_stride = 576 + scale_dim = 8 + return self.CompileKey( + head_size=head_size, + rope_head_dim=rope_head_dim, + fp8_max=448.0, + quant_block=64, + token_stride=token_stride, + scale_dim=scale_dim, + kv_block_stride=( + runtime_kv_block_stride + if runtime_kv_block_stride is not None + else round_up( + kv_cache_block_size * (token_stride + scale_dim), + cache_alignment, + ) + ), + compress_ratio=compress_ratio, + norm_weight_dtype=norm_weight_dtype, + kv_cache_block_size=kv_cache_block_size, + ) + + def get_warmup_keys( + self, + vllm_config: Any, + *, + k_cache_prefix: str, + compress_ratio: int, + ) -> list[CompileKey]: + hf_config = vllm_config.model_config.hf_config + cache_config = vllm_config.cache_config + k_cache_layer = vllm_config.compilation_config.static_forward_context[ + k_cache_prefix + ] + kv_cache = k_cache_layer.kv_cache + return self._trace_dispatch(self.dispatch)( + compress_ratio=compress_ratio, + cache_block_size=cache_config.block_size, + cache_alignment=576, + norm_weight_dtype=torch_to_cute_dtype(vllm_config.model_config.dtype), + head_size=hf_config.head_dim, + rope_head_dim=hf_config.qk_rope_head_dim, + runtime_kv_block_stride=kv_cache.stride(0), + enabled=( + cache_config.block_size > 0 + and compress_ratio == 128 + and kv_cache.ndim == 3 + and cache_config.cache_dtype == "fp8_ds_mla" + ), + _when=lambda *, enabled: enabled, + ) + + def warmup_inputs(self, compile_key: CompileKey) -> tuple[Any, ...]: + head_size = compile_key.head_size + rope_head_dim = compile_key.rope_head_dim + quant_block = compile_key.quant_block + token_stride = compile_key.token_stride + scale_dim = compile_key.scale_dim + kv_block_stride = compile_key.kv_block_stride + norm_weight_dtype = compile_key.norm_weight_dtype + static_kv_cache_block_size = compile_key.kv_cache_block_size if quant_block != 64: raise ValueError( - "CuTe DSL sparse-attn store currently requires quant_block=64." + "CuTe DSL sparse-attn store currently requires quant_block=64.", ) if rope_head_dim != 64: raise ValueError( - "CuTe DSL sparse-attn store currently requires rope_head_dim=64." + "CuTe DSL sparse-attn store currently requires rope_head_dim=64.", ) if head_size % quant_block != 0: raise ValueError("head_size must be divisible by quant_block.") @@ -1466,14 +1747,13 @@ class SparseAttnNormRopeStoreKernel: if static_kv_cache_block_size <= 0: raise ValueError( "CuTe DSL sparse-attn store requires a positive static " - "kv_cache_block_size." + "kv_cache_block_size.", ) num_positions = cute.sym_int() num_slots = cute.sym_int() num_kv_slots = cute.sym_int() max_pos = cute.sym_int() num_blocks = cute.sym_int() - compressed_kv = cute.runtime.make_fake_tensor( Float32, (num_slots, head_size), @@ -1483,7 +1763,9 @@ class SparseAttnNormRopeStoreKernel: positions = make_fake_tensor(Int64, (num_positions,), divisibility=8) slot_mapping = make_fake_tensor(Int64, (num_slots,), divisibility=8) rms_norm_weight = make_fake_tensor( - norm_weight_dtype, (head_size,), divisibility=4 + norm_weight_dtype, + (head_size,), + divisibility=4, ) cos_sin_cache = cute.runtime.make_fake_tensor( Float32, @@ -1494,28 +1776,11 @@ class SparseAttnNormRopeStoreKernel: k_cache = cute.runtime.make_fake_tensor( Uint8, (num_blocks, cute.sym_int(), cute.sym_int()), - stride=( - kv_block_stride, - cute.sym_int64(divisibility=8), - 1, - ), + stride=(kv_block_stride, cute.sym_int64(divisibility=8), 1), assumed_align=16, ) kv_slot_mapping = make_fake_tensor(Int64, (num_kv_slots,), divisibility=8) - - kernel = SparseAttnNormRopeStoreKernel( - head_size, - rope_head_dim, - fp8_max, - quant_block, - token_stride, - scale_dim, - compress_ratio, - static_kv_cache_block_size, - ) - stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) - return cute.compile( - kernel, + return ( compressed_kv, positions, slot_mapping, @@ -1524,211 +1789,275 @@ class SparseAttnNormRopeStoreKernel: cos_sin_cache, k_cache, kv_slot_mapping, - stream, - options="--enable-tvm-ffi", ) - -class SparseAttnNormRopeStoreFullKernel: - def __init__( - self, - head_size: int, - rope_head_dim: int, - fp8_max: float, - quant_block: int, - token_stride: int, - scale_dim: int, - compress_ratio: int, - store_full_fp8: bool = False, - ): - # Standalone (not inheriting the #44230-restructured legacy kernel): - # set attrs directly so the full-cache C128 path is decoupled. - self.head_dim = head_size - self.rope_dim = rope_head_dim - self.nope_dim = head_size - rope_head_dim - self.fp8_max = fp8_max - self.quant_block = quant_block - self.token_stride = token_stride - self.scale_dim = scale_dim - self.num_warps = head_size // quant_block - self.nope_blocks = self.nope_dim // quant_block - self.tb_size = head_size // 2 - self.compress_ratio = compress_ratio - self.store_full_fp8 = store_full_fp8 - - @cute.jit + @kernel_launcher def __call__( self, - compressed_kv: cute.Tensor, - positions: cute.Tensor, - slot_mapping: cute.Tensor, - rms_norm_weight: cute.Tensor, - rms_norm_eps: Float32, - cos_sin_cache: cute.Tensor, - k_cache: cute.Tensor, - kv_slot_mapping: cute.Tensor, - kv_cache_block_size: Int64, - fp8_scale: cute.Tensor, - stream: CUstream, - ): - grid = (slot_mapping.shape[0], 1, 1) - self.kernel( + *, + compressed_kv: torch.Tensor, + positions: torch.Tensor, + slot_mapping: torch.Tensor, + rms_norm_weight: torch.Tensor, + rms_norm_eps: float, + cos_sin_cache: torch.Tensor, + kv_cache: torch.Tensor, + kv_slot_mapping: torch.Tensor, + compress_ratio: int, + head_dim: int, + rope_head_dim: int, + ) -> CuTeDSLLaunchSpec[SparseAttnNormRopeStoreKernel.CompileKey]: + compile_key = self.dispatch( + compress_ratio=compress_ratio, + cache_block_size=kv_cache.shape[1] * compress_ratio, + cache_alignment=576, + norm_weight_dtype=torch_to_cute_dtype(rms_norm_weight.dtype), + head_size=head_dim, + rope_head_dim=rope_head_dim, + runtime_kv_block_stride=kv_cache.stride(0), + ) + launch_args = ( compressed_kv, positions, slot_mapping, rms_norm_weight, rms_norm_eps, cos_sin_cache, - k_cache, + kv_cache, kv_slot_mapping, - kv_cache_block_size, - fp8_scale, - ).launch(grid=grid, block=(self.tb_size, 1, 1), stream=stream) - - @cute.kernel - def kernel( - self, - compressed_kv: cute.Tensor, - positions: cute.Tensor, - slot_mapping: cute.Tensor, - rms_norm_weight: cute.Tensor, - rms_norm_eps: Float32, - cos_sin_cache: cute.Tensor, - k_cache: cute.Tensor, - kv_slot_mapping: cute.Tensor, - kv_cache_block_size: Int64, - fp8_scale: cute.Tensor, - ): - token_idx, _, _ = cute.arch.block_idx() - tid, _, _ = cute.arch.thread_idx() - warp_id = cute.arch.make_warp_uniform(tid // 32) - lane_id = tid % 32 - elem0 = tid * 2 - - slot_id = slot_mapping[token_idx] - has_position = token_idx < positions.shape[0] - position = Int64(0) - if has_position: - position = positions[token_idx] - boundary = has_position and ( - (position + Int64(1)) % Int64(self.compress_ratio) == Int64(0) ) - has_kv_slot_idx = token_idx < kv_slot_mapping.shape[0] - kv_slot_idx = Int64(-1) - if has_kv_slot_idx: - kv_slot_idx = kv_slot_mapping[token_idx] - active = slot_id >= Int64(0) and boundary and kv_slot_idx >= Int64(0) + return compile_key, launch_args - if active: - base = token_idx.to(Int64) * compressed_kv.stride[0] + elem0.to(Int64) - x0 = compressed_kv.iterator[base] - x1 = compressed_kv.iterator[base + Int64(1)] - local_sumsq = x0 * x0 + x1 * x1 - warp_sum = local_sumsq - for step in cutlass.range_constexpr(5): - offset = const_expr(16 >> step) - warp_sum += cute.arch.shuffle_sync_bfly(warp_sum, offset) +class SparseAttnNormRopeStoreFullKernel( + VllmCuTeDSLJitKernel["SparseAttnNormRopeStoreFullKernel.CompileKey"] +): + @dataclass(frozen=True) + class CompileKey: + head_size: int + rope_head_dim: int + fp8_max: float + quant_block: int + compress_ratio: int + store_full_fp8: bool + norm_weight_dtype: type[cutlass.Numeric] - smem = cutlass.utils.SmemAllocator() - partial_sums = smem.allocate_tensor( - Float32, cute.make_layout((self.num_warps,)), byte_alignment=4 - ) - rrms_shared = smem.allocate_tensor( - Float32, cute.make_layout((1,)), byte_alignment=4 - ) - - if lane_id == 0: - partial_sums[warp_id] = warp_sum - cute.arch.sync_threads() - if tid == 0: - total = Float32(0.0) - for i in cutlass.range_constexpr(self.num_warps): - total += partial_sums[i] - rrms_shared[0] = cute.math.rsqrt( - total / Float32(self.head_dim) + rms_norm_eps, fastmath=True - ) - cute.arch.sync_threads() - - rrms = rrms_shared[0] - x0 = x0 * rrms * rms_norm_weight[elem0].to(Float32) - x1 = x1 * rrms * rms_norm_weight[elem0 + 1].to(Float32) - - page = kv_slot_idx // kv_cache_block_size - kv_offset = kv_slot_idx - page * kv_cache_block_size - value_base = page * k_cache.stride[0] + kv_offset * k_cache.stride[1] - - if const_expr(self.store_full_fp8): - k_cache_u16 = cute.recast_tensor(k_cache, Uint16) - inv_fp8 = Float32(1.0) / fp8_scale[0] - fp8_v0 = x0 - fp8_v1 = x1 - if warp_id == self.nope_blocks: - compressed_pos = (position // Int64(self.compress_ratio)) * Int64( - self.compress_ratio - ) - pair_idx = lane_id - cs_base = compressed_pos * cos_sin_cache.stride[0] + pair_idx.to( - Int64 - ) - cos_v = cos_sin_cache.iterator[cs_base] - sin_v = cos_sin_cache.iterator[cs_base + Int64(self.rope_dim // 2)] - fp8_v0 = x0 * cos_v - x1 * sin_v - fp8_v1 = x0 * sin_v + x1 * cos_v - fp8_packed_bf16 = _fp32x2_to_bf16x2(fp8_v0, fp8_v1) - b0, b1 = _bf16x2_to_fp32(fp8_packed_bf16) - y0 = cutlass.min( - cutlass.max(b0 * inv_fp8, Float32(-self.fp8_max)), - Float32(self.fp8_max), - ) - y1 = cutlass.min( - cutlass.max(b1 * inv_fp8, Float32(-self.fp8_max)), - Float32(self.fp8_max), - ) - packed_fp8 = _fp32x2_to_fp8e4m3x2(y0, y1) - out_base = value_base + elem0.to(Int64) - k_cache_u16.iterator[out_base // Int64(2)] = packed_fp8 - else: - k_cache_u32 = cute.recast_tensor(k_cache, Uint32) - bf16_v0 = x0 - bf16_v1 = x1 - if warp_id == self.nope_blocks: - compressed_pos = (position // Int64(self.compress_ratio)) * Int64( - self.compress_ratio - ) - pair_idx = lane_id - cs_base = compressed_pos * cos_sin_cache.stride[0] + pair_idx.to( - Int64 - ) - cos_v = cos_sin_cache.iterator[cs_base] - sin_v = cos_sin_cache.iterator[cs_base + Int64(self.rope_dim // 2)] - bf16_v0 = x0 * cos_v - x1 * sin_v - bf16_v1 = x0 * sin_v + x1 * cos_v - bf16_packed = _fp32x2_to_bf16x2(bf16_v0, bf16_v1) - out_base = value_base + (elem0 * 2).to(Int64) - k_cache_u32.iterator[out_base // Int64(4)] = bf16_packed - - @cache @staticmethod - def compile( - head_size: int = 512, - rope_head_dim: int = 64, - fp8_max: float = 448.0, - quant_block: int = 64, - token_stride: int = 576, - scale_dim: int = 8, - kv_block_stride: int = 74752, - compress_ratio: int = 128, - store_full_fp8: bool = False, - norm_weight_dtype: type[cutlass.Numeric] = Float32, - ): + def kernel(compile_key: CompileKey) -> Any: + head_dim = compile_key.head_size + rope_dim = compile_key.rope_head_dim + nope_dim = head_dim - rope_dim + fp8_max = compile_key.fp8_max + quant_block = compile_key.quant_block + num_warps = head_dim // quant_block + nope_blocks = nope_dim // quant_block + tb_size = head_dim // 2 + compress_ratio = compile_key.compress_ratio + store_full_fp8 = compile_key.store_full_fp8 + + @cute.kernel + def device_kernel( + compressed_kv: cute.Tensor, + positions: cute.Tensor, + slot_mapping: cute.Tensor, + rms_norm_weight: cute.Tensor, + rms_norm_eps: Float32, + cos_sin_cache: cute.Tensor, + k_cache: cute.Tensor, + kv_slot_mapping: cute.Tensor, + kv_cache_block_size: Int64, + fp8_scale: cute.Tensor, + ): + token_idx, _, _ = cute.arch.block_idx() + tid, _, _ = cute.arch.thread_idx() + warp_id = cute.arch.make_warp_uniform(tid // 32) + lane_id = tid % 32 + elem0 = tid * 2 + + slot_id = slot_mapping[token_idx] + has_position = token_idx < positions.shape[0] + position = Int64(0) + if has_position: + position = positions[token_idx] + boundary = has_position and ( + (position + Int64(1)) % Int64(compress_ratio) == Int64(0) + ) + has_kv_slot_idx = token_idx < kv_slot_mapping.shape[0] + kv_slot_idx = Int64(-1) + if has_kv_slot_idx: + kv_slot_idx = kv_slot_mapping[token_idx] + active = slot_id >= Int64(0) and boundary and kv_slot_idx >= Int64(0) + + if active: + base = token_idx.to(Int64) * compressed_kv.stride[0] + elem0.to(Int64) + x0 = compressed_kv.iterator[base] + x1 = compressed_kv.iterator[base + Int64(1)] + + local_sumsq = x0 * x0 + x1 * x1 + warp_sum = local_sumsq + for step in cutlass.range_constexpr(5): + offset = const_expr(16 >> step) + warp_sum += cute.arch.shuffle_sync_bfly(warp_sum, offset) + + smem = cutlass.utils.SmemAllocator() + partial_sums = smem.allocate_tensor( + Float32, cute.make_layout((num_warps,)), byte_alignment=4 + ) + rrms_shared = smem.allocate_tensor( + Float32, cute.make_layout((1,)), byte_alignment=4 + ) + + if lane_id == 0: + partial_sums[warp_id] = warp_sum + cute.arch.sync_threads() + if tid == 0: + total = Float32(0.0) + for i in cutlass.range_constexpr(num_warps): + total += partial_sums[i] + rrms_shared[0] = cute.math.rsqrt( + total / Float32(head_dim) + rms_norm_eps, fastmath=True + ) + cute.arch.sync_threads() + + rrms = rrms_shared[0] + x0 = x0 * rrms * rms_norm_weight[elem0].to(Float32) + x1 = x1 * rrms * rms_norm_weight[elem0 + 1].to(Float32) + + page = kv_slot_idx // kv_cache_block_size + kv_offset = kv_slot_idx - page * kv_cache_block_size + value_base = page * k_cache.stride[0] + kv_offset * k_cache.stride[1] + + if const_expr(store_full_fp8): + k_cache_u16 = cute.recast_tensor(k_cache, Uint16) + inv_fp8 = Float32(1.0) / fp8_scale[0] + fp8_v0 = x0 + fp8_v1 = x1 + if warp_id == nope_blocks: + compressed_pos = (position // Int64(compress_ratio)) * Int64( + compress_ratio + ) + pair_idx = lane_id + cs_base = compressed_pos * cos_sin_cache.stride[ + 0 + ] + pair_idx.to(Int64) + cos_v = cos_sin_cache.iterator[cs_base] + sin_v = cos_sin_cache.iterator[cs_base + Int64(rope_dim // 2)] + fp8_v0 = x0 * cos_v - x1 * sin_v + fp8_v1 = x0 * sin_v + x1 * cos_v + fp8_packed_bf16 = _fp32x2_to_bf16x2(fp8_v0, fp8_v1) + b0, b1 = _bf16x2_to_fp32(fp8_packed_bf16) + y0 = cutlass.min( + cutlass.max(b0 * inv_fp8, Float32(-fp8_max)), + Float32(fp8_max), + ) + y1 = cutlass.min( + cutlass.max(b1 * inv_fp8, Float32(-fp8_max)), + Float32(fp8_max), + ) + packed_fp8 = _fp32x2_to_fp8e4m3x2(y0, y1) + out_base = value_base + elem0.to(Int64) + k_cache_u16.iterator[out_base // Int64(2)] = packed_fp8 + else: + k_cache_u32 = cute.recast_tensor(k_cache, Uint32) + bf16_v0 = x0 + bf16_v1 = x1 + if warp_id == nope_blocks: + compressed_pos = (position // Int64(compress_ratio)) * Int64( + compress_ratio + ) + pair_idx = lane_id + cs_base = compressed_pos * cos_sin_cache.stride[ + 0 + ] + pair_idx.to(Int64) + cos_v = cos_sin_cache.iterator[cs_base] + sin_v = cos_sin_cache.iterator[cs_base + Int64(rope_dim // 2)] + bf16_v0 = x0 * cos_v - x1 * sin_v + bf16_v1 = x0 * sin_v + x1 * cos_v + bf16_packed = _fp32x2_to_bf16x2(bf16_v0, bf16_v1) + out_base = value_base + (elem0 * 2).to(Int64) + k_cache_u32.iterator[out_base // Int64(4)] = bf16_packed + + @cute.jit + def host_entrypoint( + compressed_kv: cute.Tensor, + positions: cute.Tensor, + slot_mapping: cute.Tensor, + rms_norm_weight: cute.Tensor, + rms_norm_eps: Float32, + cos_sin_cache: cute.Tensor, + k_cache: cute.Tensor, + kv_slot_mapping: cute.Tensor, + kv_cache_block_size: Int64, + fp8_scale: cute.Tensor, + stream: CUstream, + ): + grid = (slot_mapping.shape[0], 1, 1) + device_kernel( + compressed_kv, + positions, + slot_mapping, + rms_norm_weight, + rms_norm_eps, + cos_sin_cache, + k_cache, + kv_slot_mapping, + kv_cache_block_size, + fp8_scale, + ).launch(grid=grid, block=(tb_size, 1, 1), stream=stream) + + return host_entrypoint + + def dispatch( # type: ignore[override] + self, + *, + store_full_fp8: bool, + norm_weight_dtype: type[cutlass.Numeric], + head_size: int, + rope_head_dim: int, + compress_ratio: int, + ) -> CompileKey: + return self.CompileKey( + head_size=head_size, + rope_head_dim=rope_head_dim, + compress_ratio=compress_ratio, + fp8_max=448.0, + quant_block=64, + store_full_fp8=store_full_fp8, + norm_weight_dtype=norm_weight_dtype, + ) + + def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]: + hf_config = vllm_config.model_config.hf_config + cache_config = vllm_config.cache_config + return self._trace_dispatch(self.dispatch)( + compress_ratio=128, + store_full_fp8=( + cache_config.cache_dtype.startswith("fp8") + and cache_config.cache_dtype != "fp8_ds_mla" + ), + norm_weight_dtype=torch_to_cute_dtype(vllm_config.model_config.dtype), + head_size=hf_config.head_dim, + rope_head_dim=hf_config.qk_rope_head_dim, + enabled=( + cache_config.block_size > 0 + and 128 in hf_config.compress_ratios + and cache_config.cache_dtype != "fp8_ds_mla" + ), + _when=lambda *, enabled: enabled, + ) + + def warmup_inputs(self, compile_key: CompileKey) -> tuple[Any, ...]: + head_size = compile_key.head_size + rope_head_dim = compile_key.rope_head_dim + quant_block = compile_key.quant_block + norm_weight_dtype = compile_key.norm_weight_dtype if quant_block != 64: raise ValueError( - "CuTe DSL sparse-attn store currently requires quant_block=64." + "CuTe DSL sparse-attn store currently requires quant_block=64.", ) if rope_head_dim != 64: raise ValueError( - "CuTe DSL sparse-attn store currently requires rope_head_dim=64." + "CuTe DSL sparse-attn store currently requires rope_head_dim=64.", ) if head_size % quant_block != 0: raise ValueError("head_size must be divisible by quant_block.") @@ -1737,7 +2066,6 @@ class SparseAttnNormRopeStoreFullKernel: num_kv_slots = cute.sym_int() max_pos = cute.sym_int() num_blocks = cute.sym_int() - compressed_kv = cute.runtime.make_fake_tensor( Float32, (num_slots, head_size), @@ -1747,7 +2075,9 @@ class SparseAttnNormRopeStoreFullKernel: positions = make_fake_tensor(Int64, (num_positions,), divisibility=8) slot_mapping = make_fake_tensor(Int64, (num_slots,), divisibility=8) rms_norm_weight = make_fake_tensor( - norm_weight_dtype, (head_size,), divisibility=4 + norm_weight_dtype, + (head_size,), + divisibility=4, ) cos_sin_cache = cute.runtime.make_fake_tensor( Float32, @@ -1758,29 +2088,12 @@ class SparseAttnNormRopeStoreFullKernel: k_cache = cute.runtime.make_fake_tensor( Uint8, (num_blocks, cute.sym_int(), cute.sym_int()), - stride=( - cute.sym_int64(divisibility=16), - cute.sym_int64(divisibility=8), - 1, - ), + stride=(cute.sym_int64(divisibility=16), cute.sym_int64(divisibility=8), 1), assumed_align=16, ) kv_slot_mapping = make_fake_tensor(Int64, (num_kv_slots,), divisibility=8) fp8_scale = make_fake_tensor(Float32, (1,), divisibility=1) - - kernel = SparseAttnNormRopeStoreFullKernel( - head_size, - rope_head_dim, - fp8_max, - quant_block, - token_stride, - scale_dim, - compress_ratio, - store_full_fp8, - ) - stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) - return cute.compile( - kernel, + return ( compressed_kv, positions, slot_mapping, @@ -1791,73 +2104,141 @@ class SparseAttnNormRopeStoreFullKernel: kv_slot_mapping, Int64(0), fp8_scale, - stream, - options="--enable-tvm-ffi", ) - -def compile_split_sparse_attn_cutedsl( - head_size: int, - state_width: int, - block_size: int, - rope_head_dim: int, - fp8_max: float, - quant_block: int, - token_stride: int, - scale_dim: int, - kv_cache_block_size: int, - kv_block_stride: int, - compress_ratio: int, - overlap: bool, - rms_norm_weight_dtype: torch.dtype, - store_full_kv: bool = False, - store_full_fp8: bool = False, -): - if not ( - head_size == 512 - and state_width == head_size - and compress_ratio == 128 - and not overlap - and block_size == 8 - ): - raise ValueError( - "CuTe DSL split sparse-attn wrapper only supports the real " - "DeepSeek V4 C128 layout: head_size=512, state_width=512, " - "compress_ratio=128, overlap=False, block_size=8." - ) - compress = SparseAttnCompressC128Block8Kernel.compile( - head_size=head_size, - state_width=state_width, - ) - norm_weight_dtype = _TORCH_TO_CUTE[rms_norm_weight_dtype] - if store_full_kv: - # FlashInfer contiguous bf16/fp8 cache: standalone full-cache store. - store = SparseAttnNormRopeStoreFullKernel.compile( - head_size=head_size, - rope_head_dim=rope_head_dim, - fp8_max=fp8_max, - quant_block=quant_block, - token_stride=token_stride, - scale_dim=scale_dim, - kv_block_stride=kv_block_stride, + @kernel_launcher + def __call__( + self, + *, + compressed_kv: torch.Tensor, + positions: torch.Tensor, + slot_mapping: torch.Tensor, + rms_norm_weight: torch.Tensor, + rms_norm_eps: float, + cos_sin_cache: torch.Tensor, + kv_cache: torch.Tensor, + kv_slot_mapping: torch.Tensor, + compress_ratio: int, + store_full_fp8: bool, + fp8_scale: torch.Tensor | None, + head_dim: int, + rope_head_dim: int, + ) -> CuTeDSLLaunchSpec[SparseAttnNormRopeStoreFullKernel.CompileKey]: + compile_key = self.dispatch( compress_ratio=compress_ratio, store_full_fp8=store_full_fp8, - norm_weight_dtype=norm_weight_dtype, + norm_weight_dtype=torch_to_cute_dtype(rms_norm_weight.dtype), + head_size=head_dim, + rope_head_dim=rope_head_dim, ) - else: - store = SparseAttnNormRopeStoreKernel.compile( - head_size, - rope_head_dim, - fp8_max, - quant_block, - token_stride, - scale_dim, - kv_block_stride, - compress_ratio, - norm_weight_dtype, - kv_cache_block_size, + if fp8_scale is None: + fp8_scale = torch.ones(1, dtype=torch.float32, device=kv_cache.device) + launch_args = ( + compressed_kv, + positions, + slot_mapping, + rms_norm_weight, + rms_norm_eps, + cos_sin_cache, + kv_cache.view(torch.uint8), + kv_slot_mapping, + kv_cache.shape[1], + fp8_scale, ) - return compress, store + return compile_key, launch_args + + +def fused_kv_compress_norm_rope_insert_sparse_attn_cutedsl( + state_cache: torch.Tensor, + token_to_req_indices: torch.Tensor, + positions: torch.Tensor, + slot_mapping: torch.Tensor, + block_table: torch.Tensor, + block_size: int, + rms_norm_weight: torch.Tensor, + rms_norm_eps: float, + cos_sin_cache: torch.Tensor, + k_cache: torch.Tensor, + kv_slot_mapping: torch.Tensor, + kv_cache_block_size: int, + kv_block_stride: int, + head_size: int = 512, + state_width: int = 1024, + rope_head_dim: int = 64, + fp8_max: float = 448.0, + quant_block: int = 64, + token_stride: int = 576, + scale_dim: int = 8, + compress_ratio: int = 4, + overlap: bool = True, + store_full_kv: bool = False, + store_full_fp8: bool = False, + fp8_scale: torch.Tensor | None = None, +) -> None: + del fp8_max, quant_block, token_stride, scale_dim, kv_block_stride + if not ( + head_size == 512 + and state_width == 2 * head_size + and compress_ratio == 4 + and overlap + ): + raise ValueError( + "CuTe DSL fused sparse-attn wrapper only supports the real " + "DeepSeek V4 C4 layout: head_size=512, state_width=1024, " + "compress_ratio=4, overlap=True." + ) + if k_cache.ndim != 3: + raise ValueError( + "CuTe DSL sparse-attn fused store expects the real DeepSeek V4 " + f"3D k_cache layout [num_blocks, block_size, 584], got ndim={k_cache.ndim}." + ) + if kv_cache_block_size != k_cache.shape[1]: + raise ValueError( + "CuTe DSL fused sparse-attn wrapper expected kv_cache_block_size " + f"to match k_cache.shape[1], got {kv_cache_block_size} and " + f"{k_cache.shape[1]}." + ) + if positions.numel() == 0: + return + if store_full_fp8 and not store_full_kv: + raise ValueError("store_full_fp8 requires store_full_kv.") + + if store_full_kv: + _SPARSE_ATTN_COMPRESS_NORM_ROPE_STORE_FULL_C4_KERNEL( + state_cache=state_cache, + token_to_req_indices=token_to_req_indices, + positions=positions, + slot_mapping=slot_mapping, + block_table=block_table, + block_size=block_size, + rms_norm_weight=rms_norm_weight, + rms_norm_eps=rms_norm_eps, + cos_sin_cache=cos_sin_cache, + kv_cache=k_cache, + kv_slot_mapping=kv_slot_mapping, + compress_ratio=compress_ratio, + store_full_fp8=store_full_fp8, + fp8_scale=fp8_scale, + head_dim=head_size, + rope_head_dim=rope_head_dim, + ) + return + _SPARSE_ATTN_COMPRESS_NORM_ROPE_STORE_C4_KERNEL( + state_cache=state_cache, + token_to_req_indices=token_to_req_indices, + positions=positions, + slot_mapping=slot_mapping, + block_table=block_table, + block_size=block_size, + rms_norm_weight=rms_norm_weight, + rms_norm_eps=rms_norm_eps, + cos_sin_cache=cos_sin_cache, + kv_cache=k_cache, + kv_slot_mapping=kv_slot_mapping, + compress_ratio=compress_ratio, + head_dim=head_size, + rope_head_dim=rope_head_dim, + ) def split_kv_compress_norm_rope_insert_sparse_attn_cutedsl( @@ -1888,6 +2269,19 @@ def split_kv_compress_norm_rope_insert_sparse_attn_cutedsl( store_full_fp8: bool = False, fp8_scale: torch.Tensor | None = None, ) -> None: + del fp8_max, quant_block, token_stride, scale_dim, kv_block_stride + if not ( + head_size == 512 + and state_width == head_size + and compress_ratio == 128 + and not overlap + and block_size == 8 + ): + raise ValueError( + "CuTe DSL split sparse-attn wrapper only supports the real " + "DeepSeek V4 C128 layout: head_size=512, state_width=512, " + "compress_ratio=128, overlap=False, block_size=8." + ) if k_cache.ndim != 3: raise ValueError( "CuTe DSL sparse-attn store expects the real DeepSeek V4 " @@ -1901,264 +2295,209 @@ def split_kv_compress_norm_rope_insert_sparse_attn_cutedsl( ) if positions.numel() == 0: return - if rms_norm_weight.dtype not in _TORCH_TO_CUTE: - raise ValueError( - "CuTe DSL sparse-attn store supports rms_norm_weight dtype " - f"bf16/fp32, got {rms_norm_weight.dtype}." - ) if store_full_fp8 and not store_full_kv: raise ValueError("store_full_fp8 requires store_full_kv.") - compress, store = compile_split_sparse_attn_cutedsl( - head_size, - state_width, - block_size, - rope_head_dim, - fp8_max, - quant_block, - token_stride, - scale_dim, - kv_cache_block_size, - kv_block_stride, - compress_ratio, - overlap, - rms_norm_weight.dtype, - store_full_kv=store_full_kv, - store_full_fp8=store_full_fp8, - ) - compress( - state_cache, - token_to_req_indices, - positions, - slot_mapping, - block_table, - compressed_kv, - ) + _SPARSE_ATTN_COMPRESS_C128_BLOCK8_KERNEL( + state_cache=state_cache, + num_actual=compressed_kv.shape[0], + token_to_req_indices=token_to_req_indices, + positions=positions, + slot_mapping=slot_mapping, + block_table=block_table, + head_dim=head_size, + compressed_kv=compressed_kv, + ) if store_full_kv: - # Byte-addressed contiguous cache; block size + per-tensor scale are - # passed at call time (not baked into compile). - if fp8_scale is None: - fp8_scale = torch.ones(1, dtype=torch.float32, device=k_cache.device) - store( - compressed_kv, - positions, - slot_mapping, - rms_norm_weight, - rms_norm_eps, - cos_sin_cache, - k_cache.view(torch.uint8), - kv_slot_mapping, - kv_cache_block_size, - fp8_scale, - ) - return - - store( - compressed_kv, - positions, - slot_mapping, - rms_norm_weight, - rms_norm_eps, - cos_sin_cache, - k_cache, - kv_slot_mapping, - ) - - -def fused_kv_compress_norm_rope_insert_sparse_attn_cutedsl( - state_cache: torch.Tensor, - token_to_req_indices: torch.Tensor, - positions: torch.Tensor, - slot_mapping: torch.Tensor, - block_table: torch.Tensor, - block_size: int, - rms_norm_weight: torch.Tensor, - rms_norm_eps: float, - cos_sin_cache: torch.Tensor, - k_cache: torch.Tensor, - kv_slot_mapping: torch.Tensor, - kv_cache_block_size: int, - kv_block_stride: int, - head_size: int = 512, - state_width: int = 1024, - rope_head_dim: int = 64, - fp8_max: float = 448.0, - quant_block: int = 64, - token_stride: int = 576, - scale_dim: int = 8, - compress_ratio: int = 4, - overlap: bool = True, - store_full_kv: bool = False, - store_full_fp8: bool = False, - fp8_scale: torch.Tensor | None = None, -) -> None: - if positions.numel() == 0: - return - norm_weight_dtype = _TORCH_TO_CUTE.get(rms_norm_weight.dtype) - if norm_weight_dtype is None: - raise ValueError( - "CuTe DSL sparse-attn fused store supports rms_norm_weight dtype " - f"bf16/fp32, got {rms_norm_weight.dtype}." - ) - if k_cache.ndim != 3: - raise ValueError( - "CuTe DSL sparse-attn fused store expects the real DeepSeek V4 " - f"3D k_cache layout [num_blocks, block_size, 584], got ndim={k_cache.ndim}." - ) - if store_full_fp8 and not store_full_kv: - raise ValueError("store_full_fp8 requires store_full_kv.") - if store_full_kv: - # FlashInfer contiguous bf16/fp8 cache: byte-addressed full-cache C4 store. - if fp8_scale is None: - fp8_scale = torch.ones(1, dtype=torch.float32, device=k_cache.device) - compiled = SparseAttnCompressNormRopeStoreFullC4Kernel.compile( - head_size=head_size, - state_width=state_width, - rope_head_dim=rope_head_dim, - fp8_max=fp8_max, - quant_block=quant_block, - token_stride=token_stride, - scale_dim=scale_dim, - kv_block_stride=kv_block_stride, + _SPARSE_ATTN_NORM_ROPE_STORE_FULL_KERNEL( + compressed_kv=compressed_kv, + positions=positions, + slot_mapping=slot_mapping, + rms_norm_weight=rms_norm_weight, + rms_norm_eps=rms_norm_eps, + cos_sin_cache=cos_sin_cache, + kv_cache=k_cache, + kv_slot_mapping=kv_slot_mapping, compress_ratio=compress_ratio, - overlap=overlap, store_full_fp8=store_full_fp8, - norm_weight_dtype=norm_weight_dtype, - ) - compiled( - state_cache, - token_to_req_indices, - positions, - slot_mapping, - block_table, - block_size, - rms_norm_weight, - rms_norm_eps, - cos_sin_cache, - k_cache.view(torch.uint8), - kv_slot_mapping, - kv_cache_block_size, - fp8_scale, + fp8_scale=fp8_scale, + head_dim=head_size, + rope_head_dim=rope_head_dim, ) return - - compiled = SparseAttnCompressNormRopeStoreC4Kernel.compile( - head_size=head_size, - state_width=state_width, - rope_head_dim=rope_head_dim, - fp8_max=fp8_max, - quant_block=quant_block, - token_stride=token_stride, - scale_dim=scale_dim, - kv_block_stride=kv_block_stride, + _SPARSE_ATTN_NORM_ROPE_STORE_KERNEL( + compressed_kv=compressed_kv, + positions=positions, + slot_mapping=slot_mapping, + rms_norm_weight=rms_norm_weight, + rms_norm_eps=rms_norm_eps, + cos_sin_cache=cos_sin_cache, + kv_cache=k_cache, + kv_slot_mapping=kv_slot_mapping, compress_ratio=compress_ratio, - overlap=overlap, - norm_weight_dtype=norm_weight_dtype, - ) - compiled( - state_cache, - token_to_req_indices, - positions, - slot_mapping, - block_table, - block_size, - rms_norm_weight, - rms_norm_eps, - cos_sin_cache, - k_cache, - kv_slot_mapping, - kv_cache_block_size, + head_dim=head_size, + rope_head_dim=rope_head_dim, ) -def compress_norm_rope_store_cutedsl( - state_cache: torch.Tensor, - num_actual: int, - token_to_req_indices: torch.Tensor, - positions: torch.Tensor, - slot_mapping: torch.Tensor, - block_table: torch.Tensor, - block_size: int, - state_width: int, - cos_sin_cache: torch.Tensor, - kv_cache: torch.Tensor, - k_cache_metadata: Any, - pdl_kwargs: dict, - head_dim: int, - rope_head_dim: int, - compress_ratio: int, - overlap: bool, - use_fp4_cache: bool, - rms_norm_weight: torch.Tensor, - rms_norm_eps: float, - quant_block: int, - token_stride: int, - scale_dim: int, - store_full_kv: bool = False, - store_full_fp8: bool = False, - fp8_scale: torch.Tensor | None = None, -) -> None: - if compress_ratio == 4: - # For C4A, the single fused kernel is faster than the two-kernel version. - fused_kv_compress_norm_rope_insert_sparse_attn_cutedsl( - state_cache, - token_to_req_indices, - positions, - slot_mapping, - block_table, - block_size, - rms_norm_weight, - rms_norm_eps, - cos_sin_cache, - kv_cache, - k_cache_metadata.slot_mapping, - kv_cache.shape[1], # paged KV cache block size - kv_cache.stride(0), - head_size=head_dim, - state_width=state_width, - rope_head_dim=rope_head_dim, - fp8_max=448.0, - quant_block=quant_block, - token_stride=token_stride, - scale_dim=scale_dim, +class SparseAttnCompressorCuteDSL: + def __init__( + self, + *, + c4_store: SparseAttnCompressNormRopeStoreC4Kernel, + c4_store_full: SparseAttnCompressNormRopeStoreFullC4Kernel, + c128_compress: SparseAttnCompressC128Block8Kernel, + c128_store: SparseAttnNormRopeStoreKernel, + c128_store_full: SparseAttnNormRopeStoreFullKernel, + ) -> None: + self.c4_store = c4_store + self.c4_store_full = c4_store_full + self.c128_compress = c128_compress + self.c128_store = c128_store + self.c128_store_full = c128_store_full + + def __call__( + self, + state_cache: torch.Tensor, + num_actual: int, + token_to_req_indices: torch.Tensor, + positions: torch.Tensor, + slot_mapping: torch.Tensor, + block_table: torch.Tensor, + block_size: int, + state_width: int, + cos_sin_cache: torch.Tensor, + kv_cache: torch.Tensor, + k_cache_metadata: Any, + pdl_kwargs: dict, + head_dim: int, + rope_head_dim: int, + compress_ratio: int, + overlap: bool, + use_fp4_cache: bool, + rms_norm_weight: torch.Tensor, + rms_norm_eps: float, + quant_block: int, + token_stride: int, + scale_dim: int, + store_full_kv: bool = False, + store_full_fp8: bool = False, + fp8_scale: torch.Tensor | None = None, + ) -> None: + del pdl_kwargs, state_width, overlap, use_fp4_cache + del quant_block, token_stride, scale_dim + if positions.numel() == 0: + return + if kv_cache.ndim != 3: + raise ValueError( + "CuTe DSL sparse-attn compressor expects the real DeepSeek V4 " + f"3D kv_cache layout, got ndim={kv_cache.ndim}.", + ) + if store_full_fp8 and (not store_full_kv): + raise ValueError("store_full_fp8 requires store_full_kv.") + if compress_ratio == 4: + if store_full_kv: + self.c4_store_full( + state_cache=state_cache, + token_to_req_indices=token_to_req_indices, + positions=positions, + slot_mapping=slot_mapping, + block_table=block_table, + block_size=block_size, + rms_norm_weight=rms_norm_weight, + rms_norm_eps=rms_norm_eps, + cos_sin_cache=cos_sin_cache, + kv_cache=kv_cache, + kv_slot_mapping=k_cache_metadata.slot_mapping, + compress_ratio=compress_ratio, + store_full_fp8=store_full_fp8, + fp8_scale=fp8_scale, + head_dim=head_dim, + rope_head_dim=rope_head_dim, + ) + else: + self.c4_store( + state_cache=state_cache, + token_to_req_indices=token_to_req_indices, + positions=positions, + slot_mapping=slot_mapping, + block_table=block_table, + block_size=block_size, + rms_norm_weight=rms_norm_weight, + rms_norm_eps=rms_norm_eps, + cos_sin_cache=cos_sin_cache, + kv_cache=kv_cache, + kv_slot_mapping=k_cache_metadata.slot_mapping, + compress_ratio=compress_ratio, + head_dim=head_dim, + rope_head_dim=rope_head_dim, + ) + return + if compress_ratio != 128: + raise ValueError( + "CuTe DSL sparse-attn compressor supports compress_ratio 4 or " + f"128, got {compress_ratio}.", + ) + if block_size != self.c128_compress.state_block_size: + raise ValueError( + "CuTe DSL C128 sparse-attn compressor requires state-cache " + f"block_size={self.c128_compress.state_block_size}, got " + f"{block_size}." + ) + compressed_kv = self.c128_compress( + state_cache=state_cache, + num_actual=num_actual, + token_to_req_indices=token_to_req_indices, + positions=positions, + slot_mapping=slot_mapping, + block_table=block_table, + head_dim=head_dim, + ) + if store_full_kv: + self.c128_store_full( + compressed_kv=compressed_kv, + positions=positions, + slot_mapping=slot_mapping, + rms_norm_weight=rms_norm_weight, + rms_norm_eps=rms_norm_eps, + cos_sin_cache=cos_sin_cache, + kv_cache=kv_cache, + kv_slot_mapping=k_cache_metadata.slot_mapping, + compress_ratio=compress_ratio, + store_full_fp8=store_full_fp8, + fp8_scale=fp8_scale, + head_dim=head_dim, + rope_head_dim=rope_head_dim, + ) + return + self.c128_store( + compressed_kv=compressed_kv, + positions=positions, + slot_mapping=slot_mapping, + rms_norm_weight=rms_norm_weight, + rms_norm_eps=rms_norm_eps, + cos_sin_cache=cos_sin_cache, + kv_cache=kv_cache, + kv_slot_mapping=k_cache_metadata.slot_mapping, compress_ratio=compress_ratio, - overlap=overlap, - store_full_kv=store_full_kv, - store_full_fp8=store_full_fp8, - fp8_scale=fp8_scale, - ) - else: - # For C128, the two-kernel version is faster than the single fused kernel. - compressed_kv = torch.empty( - (num_actual, head_dim), - dtype=torch.float32, - device=state_cache.device, - ) - split_kv_compress_norm_rope_insert_sparse_attn_cutedsl( - state_cache, - token_to_req_indices, - positions, - slot_mapping, - block_table, - block_size, - compressed_kv, - rms_norm_weight, - rms_norm_eps, - cos_sin_cache, - kv_cache, - k_cache_metadata.slot_mapping, - kv_cache.shape[1], # paged KV cache block size - kv_cache.stride(0), - head_size=head_dim, - state_width=state_width, + head_dim=head_dim, rope_head_dim=rope_head_dim, - fp8_max=448.0, - quant_block=quant_block, - token_stride=token_stride, - scale_dim=scale_dim, - compress_ratio=compress_ratio, - overlap=overlap, - store_full_kv=store_full_kv, - store_full_fp8=store_full_fp8, - fp8_scale=fp8_scale, ) + + +_SPARSE_ATTN_COMPRESS_NORM_ROPE_STORE_C4_KERNEL = ( + SparseAttnCompressNormRopeStoreC4Kernel() +) +_SPARSE_ATTN_COMPRESS_NORM_ROPE_STORE_FULL_C4_KERNEL = ( + SparseAttnCompressNormRopeStoreFullC4Kernel() +) +_SPARSE_ATTN_COMPRESS_C128_BLOCK8_KERNEL = SparseAttnCompressC128Block8Kernel() +_SPARSE_ATTN_NORM_ROPE_STORE_KERNEL = SparseAttnNormRopeStoreKernel() +_SPARSE_ATTN_NORM_ROPE_STORE_FULL_KERNEL = SparseAttnNormRopeStoreFullKernel() +_SPARSE_ATTN_COMPRESSOR_CUTEDSL_KERNEL = SparseAttnCompressorCuteDSL( + c4_store=_SPARSE_ATTN_COMPRESS_NORM_ROPE_STORE_C4_KERNEL, + c4_store_full=_SPARSE_ATTN_COMPRESS_NORM_ROPE_STORE_FULL_C4_KERNEL, + c128_compress=_SPARSE_ATTN_COMPRESS_C128_BLOCK8_KERNEL, + c128_store=_SPARSE_ATTN_NORM_ROPE_STORE_KERNEL, + c128_store_full=_SPARSE_ATTN_NORM_ROPE_STORE_FULL_KERNEL, +) diff --git a/vllm/models/deepseek_v4/sparse_mla.py b/vllm/models/deepseek_v4/sparse_mla.py index 81659f342a..4953bdc925 100644 --- a/vllm/models/deepseek_v4/sparse_mla.py +++ b/vllm/models/deepseek_v4/sparse_mla.py @@ -9,11 +9,11 @@ import torch from vllm.config import VllmConfig from vllm.config.cache import CacheDType +from vllm.model_executor.warmup.jit_warmup import kernel_launcher from vllm.model_executor.warmup.jit_warmup_triton_helper import ( LaunchSpec, TritonWarmupTensor, VllmTritonJitKernel, - kernel_launcher, ) from vllm.platforms import current_platform from vllm.platforms.interface import DeviceCapability diff --git a/vllm/models/deepseek_v4_1/common/ops/cache_utils.py b/vllm/models/deepseek_v4_1/common/ops/cache_utils.py index e2da6f5e2b..64462eb73b 100644 --- a/vllm/models/deepseek_v4_1/common/ops/cache_utils.py +++ b/vllm/models/deepseek_v4_1/common/ops/cache_utils.py @@ -416,11 +416,17 @@ def dequantize_and_gather_k_cache( if has_cutedsl(): # lazily import, otherwise some tests fail due to CUDA driver init failure. from vllm.models.deepseek_v4.nvidia.ops.dequant_gather_k_cutedsl import ( - dequantize_and_gather_k_cache_cutedsl, + _DEQUANT_GATHER_K_CACHE_CUTEDSL_KERNEL, ) - dequantize_and_gather_k_cache_cutedsl( - out, k_cache, seq_lens, gather_lens, block_table, block_size, offset + _DEQUANT_GATHER_K_CACHE_CUTEDSL_KERNEL( + out=out, + k_cache=k_cache, + seq_lens=seq_lens, + gather_lens=gather_lens, + block_table=block_table, + block_size=block_size, + offset=offset, ) return diff --git a/vllm/v1/attention/backends/mla/compressor_utils.py b/vllm/v1/attention/backends/mla/compressor_utils.py index 78c30c0f5f..b32af7234b 100644 --- a/vllm/v1/attention/backends/mla/compressor_utils.py +++ b/vllm/v1/attention/backends/mla/compressor_utils.py @@ -5,12 +5,11 @@ from typing import Any import torch -from vllm.model_executor.warmup.jit_warmup import zip_inputs +from vllm.model_executor.warmup.jit_warmup import kernel_launcher, zip_inputs from vllm.model_executor.warmup.jit_warmup_triton_helper import ( LaunchSpec, TritonWarmupTensor, VllmTritonJitKernel, - kernel_launcher, triton_scalar_specialization_rep, ) from vllm.triton_utils import tl, triton diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 0eb642bd76..59fb109dac 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -9,12 +9,12 @@ import vllm.envs as envs from vllm.config import VllmConfig from vllm.distributed import get_dcp_group, get_pcp_group from vllm.logger import init_logger +from vllm.model_executor.warmup.jit_warmup import kernel_launcher from vllm.model_executor.warmup.jit_warmup_triton_helper import ( LaunchSpec, TritonPointerInputVariant, TritonWarmupTensor, VllmTritonJitKernel, - kernel_launcher, triton_scalar_specialization_rep, ) from vllm.platforms import current_platform diff --git a/vllm/v1/attention/backends/mla/sparse_swa.py b/vllm/v1/attention/backends/mla/sparse_swa.py index 02e6336a38..f47554d81f 100644 --- a/vllm/v1/attention/backends/mla/sparse_swa.py +++ b/vllm/v1/attention/backends/mla/sparse_swa.py @@ -9,13 +9,13 @@ from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.model_executor.warmup.jit_warmup import ( WarmupIntRange, + kernel_launcher, zip_inputs, ) from vllm.model_executor.warmup.jit_warmup_triton_helper import ( LaunchSpec, TritonWarmupTensor, VllmTritonJitKernel, - kernel_launcher, ) from vllm.platforms import current_platform from vllm.triton_utils import tl, triton diff --git a/vllm/v1/attention/backends/mla/sparse_utils.py b/vllm/v1/attention/backends/mla/sparse_utils.py index 329d48be21..8d196f9630 100644 --- a/vllm/v1/attention/backends/mla/sparse_utils.py +++ b/vllm/v1/attention/backends/mla/sparse_utils.py @@ -10,13 +10,13 @@ import torch from vllm.config import VllmConfig from vllm.distributed import get_dcp_group from vllm.model_executor.warmup.jit_warmup import ( + kernel_launcher, zip_inputs, ) from vllm.model_executor.warmup.jit_warmup_triton_helper import ( LaunchSpec, TritonWarmupTensor, VllmTritonJitKernel, - kernel_launcher, triton_scalar_specialization_rep, ) from vllm.triton_utils import tl, triton diff --git a/vllm/v1/attention/ops/common.py b/vllm/v1/attention/ops/common.py index 88ce371226..b561817055 100644 --- a/vllm/v1/attention/ops/common.py +++ b/vllm/v1/attention/ops/common.py @@ -5,11 +5,11 @@ from typing import Any import torch +from vllm.model_executor.warmup.jit_warmup import kernel_launcher from vllm.model_executor.warmup.jit_warmup_triton_helper import ( LaunchSpec, TritonWarmupTensor, VllmTritonJitKernel, - kernel_launcher, ) from vllm.triton_utils import tl, triton diff --git a/vllm/v1/attention/ops/dcp.py b/vllm/v1/attention/ops/dcp.py index 66a6bace46..bd20c9020d 100644 --- a/vllm/v1/attention/ops/dcp.py +++ b/vllm/v1/attention/ops/dcp.py @@ -18,12 +18,12 @@ from vllm.distributed import get_dcp_group from vllm.logger import init_logger from vllm.model_executor.warmup.jit_warmup import ( WarmupIntRange, + kernel_launcher, ) from vllm.model_executor.warmup.jit_warmup_triton_helper import ( LaunchSpec, TritonWarmupTensor, VllmTritonJitKernel, - kernel_launcher, triton_scalar_specialization_rep, ) from vllm.triton_utils import tl, triton diff --git a/vllm/v1/worker/block_table.py b/vllm/v1/worker/block_table.py index e8e6f80853..33b52652a1 100644 --- a/vllm/v1/worker/block_table.py +++ b/vllm/v1/worker/block_table.py @@ -11,11 +11,11 @@ import torch from vllm.distributed import get_dcp_group, get_pcp_group from vllm.logger import init_logger +from vllm.model_executor.warmup.jit_warmup import kernel_launcher from vllm.model_executor.warmup.jit_warmup_triton_helper import ( LaunchSpec, TritonWarmupTensor, VllmTritonJitKernel, - kernel_launcher, triton_scalar_specialization_rep, ) from vllm.triton_utils import tl, triton