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

[5/N][warmup][DSv4] Migrate NVIDIA CuTeDSL attention kernels (#53566)

Browse Source 
Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>
Co-authored-by: Codex <codex@openai.com>
This commit is contained in:
Roberto L. Castro
2026-09-12 00:05:59 -04:00
committed by GitHub
co-authored by Codex
parent 485421b1c3
commit 120ec4ebd2
28 changed files with 3124 additions and 2551 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/model_executor/test_jit_warmup_cutedsl_launcher.py
+38 -3
View File
@@ -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"]
tests/model_executor/test_jit_warmup_triton_launcher.py
+4 -4
View File
@@ -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
vllm/cute_utils/__init__.py
+8
View File
@@ -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",
vllm/model_executor/kernels/attention/dsa/dcp_indexer_cutedsl.py
+6 -12
View File
@@ -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()
vllm/model_executor/warmup/jit_warmup.py
+23
View File
@@ -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)
vllm/model_executor/warmup/jit_warmup_cutedsl_helper.py
+25 -26
View File
@@ -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
vllm/model_executor/warmup/jit_warmup_triton_helper.py
+14 -44
View File
@@ -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:
vllm/models/common/ops/fused_qk_rmsnorm.py
+1 -1
View File
@@ -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
vllm/models/deepseek_v4/attention.py
+23 -3
View File
@@ -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():
vllm/models/deepseek_v4/common/ops/cache_utils.py
+10 -4
View File
@@ -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
vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py
+1 -1
View File
@@ -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
vllm/models/deepseek_v4/common/ops/fused_indexer_q.py
+22 -22
View File
@@ -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(
vllm/models/deepseek_v4/common/ops/fused_inv_rope_fp8_quant.py
+1 -1
View File
@@ -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
vllm/models/deepseek_v4/common/ops/fused_mtp_input_rmsnorm.py
+1 -1
View File
@@ -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
vllm/models/deepseek_v4/common/ops/save_partial_states.py
+1 -1
View File
@@ -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
vllm/models/deepseek_v4/compressor.py
+31 -5
View File
@@ -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,
vllm/models/deepseek_v4/nvidia/ops/dequant_gather_k_cutedsl.py
+347 -280
View File
@@ -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()
vllm/models/deepseek_v4/nvidia/ops/fused_indexer_q_cutedsl.py
+576 -495
View File
File diff suppressed because it is too large Load Diff
vllm/models/deepseek_v4/nvidia/ops/sparse_attn_compress_cutedsl.py
+1975 -1636
View File
File diff suppressed because it is too large Load Diff
vllm/models/deepseek_v4/sparse_mla.py
+1 -1
View File
@@ -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
vllm/models/deepseek_v4_1/common/ops/cache_utils.py
+9 -3
View File
@@ -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
vllm/v1/attention/backends/mla/compressor_utils.py
+1 -2
View File
@@ -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
vllm/v1/attention/backends/mla/indexer.py
+1 -1
View File
@@ -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
vllm/v1/attention/backends/mla/sparse_swa.py
+1 -1
View File
@@ -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
vllm/v1/attention/backends/mla/sparse_utils.py
+1 -1
View File
@@ -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
vllm/v1/attention/ops/common.py
+1 -1
View File
@@ -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
vllm/v1/attention/ops/dcp.py
+1 -1
View File
@@ -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
vllm/v1/worker/block_table.py
+1 -1
View File
@@ -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