@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 196 lines 6.56 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """TokenSpeed CuTe DSL backend for MLA prefill."""
4
5 from typing import TYPE_CHECKING, ClassVar
6
7 import torch
8
9 from vllm.v1.attention.backends.mla.prefill.base import (
10 MLADimensions,
11 MLAPrefillBackend,
12 )
13
14 if TYPE_CHECKING:
15 from vllm.config import VllmConfig
16 from vllm.model_executor.layers.attention.mla_attention import (
17 MLACommonPrefillMetadata,
18 )
19 from vllm.platforms.interface import DeviceCapability
20
21
22 class TokenspeedMLAPrefillBackend(MLAPrefillBackend):
23 """TokenSpeed CuTe DSL backend for MLA prefill."""
24
25 supported_mla_dimensions: ClassVar[list[MLADimensions]] = [
26 MLADimensions(
27 qk_nope_head_dim=128,
28 qk_rope_head_dim=64,
29 v_head_dim=128,
30 ),
31 ]
32
33 @staticmethod
34 def get_name() -> str:
35 return "TOKENSPEED_MLA"
36
37 @classmethod
38 def supports_compute_capability(cls, device_capability: "DeviceCapability") -> bool:
39 return device_capability.major == 10
40
41 _INSTALL_HINT = (
42 "tokenspeed_mla package is not installed. "
43 "Install it with: `uv pip install tokenspeed-mla`"
44 )
45
46 @classmethod
47 def is_available(cls) -> bool:
48 try:
49 from tokenspeed_mla import (
50 tokenspeed_mla_prefill, # noqa: F401
51 )
52
53 return True
54 except ImportError:
55 return False
56
57 @classmethod
58 def validate_configuration(
59 cls,
60 device_capability,
61 selector_config,
62 ) -> list[str]:
63 # Replace the generic "required dependencies not available" message
64 # from the base class with a specific install hint so users know
65 # exactly which package to install when they explicitly select this
66 # backend without having tokenspeed_mla installed.
67 reasons = super().validate_configuration(device_capability, selector_config)
68 return [
69 cls._INSTALL_HINT if r == "required dependencies not available" else r
70 for r in reasons
71 ]
72
73 def __init__(
74 self,
75 num_heads: int,
76 scale: float,
77 kv_lora_rank: int,
78 qk_nope_head_dim: int,
79 qk_rope_head_dim: int,
80 v_head_dim: int,
81 vllm_config: "VllmConfig",
82 ) -> None:
83 super().__init__(
84 num_heads=num_heads,
85 scale=scale,
86 kv_lora_rank=kv_lora_rank,
87 qk_nope_head_dim=qk_nope_head_dim,
88 qk_rope_head_dim=qk_rope_head_dim,
89 v_head_dim=v_head_dim,
90 vllm_config=vllm_config,
91 )
92
93 # Pre-JIT BF16 and FP8 prefill kernels. Idempotent — also called from
94 # TokenspeedMLAImpl.__init__; second call is a no-op.
95 from tokenspeed_mla import warmup_compile_prefill
96
97 for q_dtype in (torch.bfloat16, torch.float8_e4m3fn):
98 warmup_compile_prefill(
99 q_dtype=q_dtype,
100 d_qk=qk_nope_head_dim + qk_rope_head_dim,
101 d_v=v_head_dim,
102 enable_pdl=False,
103 )
104
105 def prepare_metadata(
106 self,
107 prefill_metadata: "MLACommonPrefillMetadata",
108 ) -> None:
109 super().prepare_metadata(prefill_metadata)
110 # Kernel signature requires `seq_lens` but the implementation never reads
111 # it (per-batch lengths are derived from `cum_seq_lens` diffs); compute
112 # for parity with trtllm_ragged. cuda-graph padding in
113 # `query_start_loc` is saturated to `total_num_tokens`
114 # (gpu_model_runner.py:1905), so trailing diffs are 0 and padded batches
115 # are kernel no-ops — same reason trtllm passes the padded length as
116 # batch_size directly.
117 self._query_seq_lens = (
118 prefill_metadata.query_start_loc[1:] - prefill_metadata.query_start_loc[:-1]
119 )
120
121 def supports_out(self) -> bool:
122 # Output head dim is v_head_dim (unpadded); `out` supported since 0.1.8.
123 return True
124
125 def run_prefill_new_tokens(
126 self,
127 q: torch.Tensor,
128 k: torch.Tensor,
129 v: torch.Tensor,
130 return_softmax_lse: bool,
131 out: torch.Tensor | None = None,
132 output_scale: torch.Tensor | None = None,
133 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
134 from tokenspeed_mla import tokenspeed_mla_prefill
135
136 # `v` arrives as the second half of `kv_nope.split(...)` — a
137 # non-contiguous view of `kv_nope` along dim=-1. The kernel wraps inputs
138 # via `from_dlpack`, which preserves strides, but the FMHA kernel reads
139 # `v` assuming contiguous storage and silently produces wrong output on
140 # a strided view (verified on tokenspeed-mla 0.1.8). Force contiguity.
141 v = v.contiguous()
142
143 ret = tokenspeed_mla_prefill(
144 query=q,
145 key=k,
146 value=v,
147 seq_lens=self._query_seq_lens,
148 cum_seq_lens=self._prefill_metadata.query_start_loc,
149 max_seq_len=self._prefill_metadata.max_query_len,
150 batch_size=self._query_seq_lens.shape[0],
151 softmax_scale=self.scale,
152 is_causal=True,
153 return_lse=return_softmax_lse,
154 enable_pdl=False,
155 out=out,
156 )
157
158 if isinstance(ret, tuple):
159 # Convert from (q_len, num_heads) to (num_heads, q_len)
160 return ret[0], ret[1].transpose(0, 1)
161 return ret
162
163 def run_prefill_context_chunk(
164 self,
165 chunk: "MLACommonPrefillMetadata.ContextChunk",
166 q: torch.Tensor,
167 k: torch.Tensor,
168 v: torch.Tensor,
169 out: torch.Tensor | None = None,
170 ) -> tuple[torch.Tensor, torch.Tensor]:
171 from tokenspeed_mla import tokenspeed_mla_prefill
172
173 # See note in run_prefill_new_tokens — `v` is a split-view of `kv_nope`
174 # in the generic `_compute_prefill_context` and arrives non-contiguous.
175 # (K3's fused fp8 K/V pack already hands over a contiguous `v`.)
176 v = v.contiguous()
177
178 attn_out, lse = tokenspeed_mla_prefill(
179 query=q,
180 key=k,
181 value=v,
182 seq_lens=chunk.seq_lens,
183 cum_seq_lens=chunk.cu_seq_lens,
184 max_seq_len=chunk.max_seq_len,
185 batch_size=chunk.num_requests,
186 softmax_scale=self.scale,
187 is_causal=False,
188 return_lse=True,
189 cum_seq_lens_q=chunk.query_start_loc,
190 max_seq_len_q=chunk.max_query_len,
191 enable_pdl=False,
192 out=out,
193 )
194
195 # Convert from (q_len, num_heads) to (num_heads, q_len)
196 return attn_out, lse.transpose(0, 1)