| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """DeepSeek-V4 FlashMLA sparse backend, metadata, and metadata builder.""" |
| 4 | |
| 5 | from dataclasses import dataclass |
| 6 | from typing import Any, ClassVar |
| 7 | |
| 8 | import torch |
| 9 | |
| 10 | from vllm.config import VllmConfig |
| 11 | from vllm.config.cache import CacheDType |
| 12 | from vllm.model_executor.warmup.jit_warmup import kernel_launcher |
| 13 | from vllm.model_executor.warmup.jit_warmup_triton_helper import ( |
| 14 | LaunchSpec, |
| 15 | TritonWarmupTensor, |
| 16 | VllmTritonJitKernel, |
| 17 | ) |
| 18 | from vllm.platforms import current_platform |
| 19 | from vllm.platforms.interface import DeviceCapability |
| 20 | from vllm.triton_utils import tl, triton |
| 21 | from vllm.utils.math_utils import cdiv |
| 22 | from vllm.v1.attention.backend import ( |
| 23 | AttentionBackend, |
| 24 | AttentionCGSupport, |
| 25 | AttentionMetadata, |
| 26 | AttentionMetadataBuilder, |
| 27 | CommonAttentionMetadata, |
| 28 | MultipleOf, |
| 29 | ) |
| 30 | from vllm.v1.attention.backends.mla.compressor_utils import get_compressed_slot_mapping |
| 31 | from vllm.v1.attention.backends.utils import split_decodes_and_prefills |
| 32 | from vllm.v1.kv_cache_interface import AttentionSpec |
| 33 | |
| 34 | # Pad C128A topk width to this alignment. 128 covers both h_q=64 (B_TOPK=64) and |
| 35 | # h_q=128 (B_TOPK=128). FlashMLA decode asserts extra_topk % B_TOPK == 0; |
| 36 | # unaligned widths (e.g. 17 = ceil(2136/128)) crash the sm100 head64 kernel. |
| 37 | # Padded slots stay -1 and decode_lens caps them via topk_length, so the pad is a |
| 38 | # no-op at kernel level. Mirrors _SPARSE_PREFILL_TOPK_ALIGNMENT in cache_utils.py. |
| 39 | _C128A_TOPK_ALIGNMENT = 128 |
| 40 | |
| 41 | |
| 42 | class DeepseekV4SparseMLABackend(AttentionBackend): |
| 43 | """DeepSeek-V4 sparse-MLA backend base. |
| 44 | |
| 45 | Subclasses ``AttentionBackend`` directly (not the V3.2 |
| 46 | ``FlashMLASparseBackend``): DeepSeek-V4 runs its own attention layer |
| 47 | (``DeepseekV4Attention``), so it does not reuse the V3.2 builder or impl, and |
| 48 | only needs to declare its own metadata builder, KV-cache layout, and the |
| 49 | sparse-MLA capability flags. |
| 50 | """ |
| 51 | |
| 52 | supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16] |
| 53 | supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [ |
| 54 | "auto", |
| 55 | "fp8_ds_mla", |
| 56 | "fp8", # alias for fp8_ds_mla |
| 57 | ] |
| 58 | |
| 59 | @staticmethod |
| 60 | def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]: |
| 61 | return [256] |
| 62 | |
| 63 | @staticmethod |
| 64 | def get_builder_cls() -> type["DeepseekV4SparseMLAMetadataBuilder"]: |
| 65 | return DeepseekV4SparseMLAMetadataBuilder |
| 66 | |
| 67 | @staticmethod |
| 68 | def get_impl_cls() -> type[Any]: |
| 69 | # DeepSeek-V4 runs its attention through ``DeepseekV4Attention.forward``, |
| 70 | # not the generic ``Attention``/``MLAAttention`` layer, so the backend's |
| 71 | # impl class is never instantiated. |
| 72 | raise NotImplementedError( |
| 73 | "DeepseekV4SparseMLABackend has no separate impl class; DeepSeek-V4 " |
| 74 | "attention runs through DeepseekV4Attention." |
| 75 | ) |
| 76 | |
| 77 | @classmethod |
| 78 | def get_supported_head_sizes(cls) -> list[int]: |
| 79 | # DeepSeek V4 layout: 448 NoPE + 64 RoPE = 512. |
| 80 | return [512] |
| 81 | |
| 82 | @classmethod |
| 83 | def is_mla(cls) -> bool: |
| 84 | return True |
| 85 | |
| 86 | @classmethod |
| 87 | def is_sparse(cls) -> bool: |
| 88 | return True |
| 89 | |
| 90 | @classmethod |
| 91 | def supports_sink(cls) -> bool: |
| 92 | return True |
| 93 | |
| 94 | @classmethod |
| 95 | def supports_compute_capability(cls, capability: DeviceCapability) -> bool: |
| 96 | return capability.major in [9, 10] |
| 97 | |
| 98 | |
| 99 | @dataclass |
| 100 | class DeepseekV4FlashMLAMetadata(AttentionMetadata): |
| 101 | num_reqs: int |
| 102 | max_query_len: int |
| 103 | max_seq_len: int |
| 104 | |
| 105 | num_actual_tokens: int # Number of tokens excluding padding. |
| 106 | query_start_loc: torch.Tensor |
| 107 | slot_mapping: torch.Tensor |
| 108 | |
| 109 | block_table: torch.Tensor |
| 110 | req_id_per_token: torch.Tensor |
| 111 | block_size: int |
| 112 | topk_tokens: int |
| 113 | |
| 114 | # Pre-computed C128A metadata (compress_ratio == 128 only). |
| 115 | # Decode: global slot ids + valid-entry counts (fused from positions). |
| 116 | c128a_global_decode_topk_indices: torch.Tensor | None = None |
| 117 | c128a_decode_topk_lens: torch.Tensor | None = None |
| 118 | # Prefill: local topk indices (used by combine_topk_swa_indices). |
| 119 | c128a_prefill_topk_indices: torch.Tensor | None = None |
| 120 | |
| 121 | |
| 122 | class DeepseekV4SparseMLAMetadataBuilder( |
| 123 | AttentionMetadataBuilder[DeepseekV4FlashMLAMetadata] |
| 124 | ): |
| 125 | _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH |
| 126 | |
| 127 | def __init__( |
| 128 | self, |
| 129 | kv_cache_spec: AttentionSpec, |
| 130 | layer_names: list[str], |
| 131 | vllm_config: VllmConfig, |
| 132 | device: torch.device, |
| 133 | ) -> None: |
| 134 | super().__init__(kv_cache_spec, layer_names, vllm_config, device) |
| 135 | self.model_config = vllm_config.model_config |
| 136 | # Classify single-token queries (plus num_speculative_tokens via |
| 137 | # supports_spec_as_decode=True) as decodes; longer queries go to prefill. |
| 138 | self._init_reorder_batch_threshold(1, supports_spec_as_decode=True) |
| 139 | self.topk_tokens = self.model_config.hf_config.index_topk |
| 140 | |
| 141 | max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens |
| 142 | self.req_id_per_token_buffer = torch.empty( |
| 143 | (max_num_batched_tokens,), dtype=torch.int32, device=device |
| 144 | ) |
| 145 | |
| 146 | assert isinstance(self.kv_cache_spec.tokens_per_state, int) |
| 147 | self.compress_ratio = self.kv_cache_spec.tokens_per_state |
| 148 | |
| 149 | # Pre-allocate compressed slot mapping buffer for CUDA graph address |
| 150 | # stability when compress_ratio > 1. |
| 151 | if self.compress_ratio > 1: |
| 152 | self.compressed_slot_mapping_buffer = torch.empty( |
| 153 | max_num_batched_tokens, dtype=torch.int64, device=device |
| 154 | ) |
| 155 | |
| 156 | # Pre-allocate C128A topk buffers for CUDA graph address stability. |
| 157 | if self.compress_ratio == 128: |
| 158 | c128a_max_compressed = cdiv( |
| 159 | self.model_config.max_model_len, self.compress_ratio |
| 160 | ) |
| 161 | c128a_max_compressed = ( |
| 162 | cdiv(c128a_max_compressed, _C128A_TOPK_ALIGNMENT) |
| 163 | * _C128A_TOPK_ALIGNMENT |
| 164 | ) |
| 165 | # Stored so _build_c128a_metadata passes it as the kernel's |
| 166 | # max_compressed_tokens, matching the buffer stride. Otherwise the |
| 167 | # kernel's default 8192 iterates past row width and spills writes |
| 168 | # into adjacent rows (present in both decode and prefill branches of |
| 169 | # _build_c128a_topk_metadata_kernel). |
| 170 | self.c128a_max_compressed = c128a_max_compressed |
| 171 | self.c128a_global_decode_buffer = torch.empty( |
| 172 | (max_num_batched_tokens, c128a_max_compressed), |
| 173 | dtype=torch.int32, |
| 174 | device=device, |
| 175 | ) |
| 176 | self.c128a_decode_lens_buffer = torch.empty( |
| 177 | max_num_batched_tokens, dtype=torch.int32, device=device |
| 178 | ) |
| 179 | self.c128a_prefill_buffer = torch.empty( |
| 180 | (max_num_batched_tokens, c128a_max_compressed), |
| 181 | dtype=torch.int32, |
| 182 | device=device, |
| 183 | ) |
| 184 | |
| 185 | def build( |
| 186 | self, |
| 187 | common_prefix_len: int, |
| 188 | common_attn_metadata: CommonAttentionMetadata, |
| 189 | fast_build: bool = False, |
| 190 | ) -> DeepseekV4FlashMLAMetadata: |
| 191 | cm = common_attn_metadata |
| 192 | req_id_per_token = cm.token_to_req_indices(self.req_id_per_token_buffer) |
| 193 | |
| 194 | slot_mapping = cm.slot_mapping |
| 195 | if self.compress_ratio > 1: |
| 196 | slot_mapping = get_compressed_slot_mapping( |
| 197 | cm.num_actual_tokens, |
| 198 | cm.slot_mapping, |
| 199 | cm.query_start_loc, |
| 200 | cm.seq_lens, |
| 201 | cm.block_table_tensor.clamp_(min=0), |
| 202 | int(self.kv_cache_spec.num_states), |
| 203 | self.compress_ratio, |
| 204 | out=self.compressed_slot_mapping_buffer, |
| 205 | ) |
| 206 | |
| 207 | c128a_fields: dict[str, torch.Tensor | None] = {} |
| 208 | if self.compress_ratio == 128: |
| 209 | c128a_fields = self._build_c128a_metadata(cm, req_id_per_token) |
| 210 | |
| 211 | return DeepseekV4FlashMLAMetadata( |
| 212 | num_reqs=cm.num_reqs, |
| 213 | max_query_len=cm.max_query_len, |
| 214 | max_seq_len=cm.max_seq_len, |
| 215 | num_actual_tokens=cm.num_actual_tokens, |
| 216 | query_start_loc=cm.query_start_loc, |
| 217 | slot_mapping=slot_mapping, |
| 218 | block_table=cm.block_table_tensor, |
| 219 | req_id_per_token=req_id_per_token, |
| 220 | block_size=self.kv_cache_spec.block_size, |
| 221 | topk_tokens=self.topk_tokens, |
| 222 | c128a_global_decode_topk_indices=c128a_fields.get( |
| 223 | "c128a_global_decode_topk_indices" |
| 224 | ), |
| 225 | c128a_decode_topk_lens=c128a_fields.get("c128a_decode_topk_lens"), |
| 226 | c128a_prefill_topk_indices=c128a_fields.get("c128a_prefill_topk_indices"), |
| 227 | ) |
| 228 | |
| 229 | def _build_c128a_metadata( |
| 230 | self, |
| 231 | cm: CommonAttentionMetadata, |
| 232 | req_id_per_token: torch.Tensor, |
| 233 | ) -> dict[str, torch.Tensor | None]: |
| 234 | """Pre-compute C128A topk indices for DeepseekV4 (compress_ratio >= 128).""" |
| 235 | # Must match SWA's decode split (no `require_uniform=True`) so |
| 236 | # `c128a_global_decode_topk_indices.shape[0]` lines up with q in |
| 237 | # `_forward_decode`. The per-token C128A kernel handles non-uniform |
| 238 | # query lengths. |
| 239 | (num_decodes, _, num_decode_tokens, num_prefill_tokens) = ( |
| 240 | split_decodes_and_prefills( |
| 241 | cm, |
| 242 | decode_threshold=self.reorder_batch_threshold or 1, |
| 243 | ) |
| 244 | ) |
| 245 | |
| 246 | num_total = num_decode_tokens + num_prefill_tokens |
| 247 | if num_total == 0: |
| 248 | return {} |
| 249 | |
| 250 | assert cm.positions is not None, ( |
| 251 | "positions is required for C128A metadata build" |
| 252 | ) |
| 253 | active_topk_width = min( |
| 254 | max( |
| 255 | triton.next_power_of_2(max(cm.max_seq_len // self.compress_ratio, 1)), |
| 256 | _C128A_TOPK_ALIGNMENT, |
| 257 | ), |
| 258 | self.c128a_max_compressed, |
| 259 | ) |
| 260 | assert active_topk_width >= cm.max_seq_len // self.compress_ratio |
| 261 | assert active_topk_width % _C128A_TOPK_ALIGNMENT == 0 |
| 262 | block_size = self.kv_cache_spec.block_size // self.compress_ratio |
| 263 | global_decode, decode_lens, prefill_local = build_c128a_topk_metadata( |
| 264 | cm.positions[:num_total], |
| 265 | self.compress_ratio, |
| 266 | num_decode_tokens, |
| 267 | req_id_per_token, |
| 268 | cm.block_table_tensor[:num_decodes], |
| 269 | block_size, |
| 270 | cm.slot_mapping, |
| 271 | self.c128a_global_decode_buffer, |
| 272 | self.c128a_decode_lens_buffer, |
| 273 | self.c128a_prefill_buffer, |
| 274 | max_compressed_tokens=active_topk_width, |
| 275 | ) |
| 276 | |
| 277 | result: dict[str, torch.Tensor | None] = {} |
| 278 | if num_decode_tokens > 0: |
| 279 | result["c128a_global_decode_topk_indices"] = global_decode.view( |
| 280 | num_decode_tokens, 1, -1 |
| 281 | ) |
| 282 | result["c128a_decode_topk_lens"] = decode_lens |
| 283 | if num_prefill_tokens > 0: |
| 284 | result["c128a_prefill_topk_indices"] = prefill_local |
| 285 | return result |
| 286 | |
| 287 | |
| 288 | class DeepseekV4FlashMLAMetadataBuilder(DeepseekV4SparseMLAMetadataBuilder): |
| 289 | _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS |
| 290 | |
| 291 | |
| 292 | class DeepseekV4FlashMLABackend(DeepseekV4SparseMLABackend): |
| 293 | @staticmethod |
| 294 | def get_name() -> str: |
| 295 | return "FLASHMLA_SPARSE_DSV4" |
| 296 | |
| 297 | @staticmethod |
| 298 | def get_builder_cls() -> type[DeepseekV4FlashMLAMetadataBuilder]: |
| 299 | return DeepseekV4FlashMLAMetadataBuilder |
| 300 | |
| 301 | |
| 302 | def build_c128a_topk_metadata( |
| 303 | positions: torch.Tensor, |
| 304 | compress_ratio: int, |
| 305 | num_decode_tokens: int, |
| 306 | token_to_req_indices: torch.Tensor, |
| 307 | block_table: torch.Tensor, |
| 308 | block_size: int, |
| 309 | slot_mapping: torch.Tensor, |
| 310 | global_decode_buffer: torch.Tensor, |
| 311 | decode_lens_buffer: torch.Tensor, |
| 312 | prefill_buffer: torch.Tensor, |
| 313 | max_compressed_tokens: int = 8192, |
| 314 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| 315 | """Single kernel for all C128A tokens (decode + prefill). |
| 316 | |
| 317 | Decode tokens: position → block_table lookup → global slot ids + topk_lens. |
| 318 | Prefill tokens: position → local indices [0, ..., n-1, -1, ...]. |
| 319 | |
| 320 | Writes into pre-allocated buffers for CUDA graph address stability. |
| 321 | Returns views of the buffers. |
| 322 | """ |
| 323 | num_tokens = positions.shape[0] |
| 324 | num_prefill_tokens = num_tokens - num_decode_tokens |
| 325 | assert max_compressed_tokens % _C128A_TOPK_ALIGNMENT == 0 |
| 326 | assert ( |
| 327 | 0 |
| 328 | < max_compressed_tokens |
| 329 | <= min(global_decode_buffer.shape[1], prefill_buffer.shape[1]) |
| 330 | ) |
| 331 | assert global_decode_buffer.stride(-1) == prefill_buffer.stride(-1) == 1 |
| 332 | |
| 333 | # TODO: support adaptive-width decode on SM120 (needs the FlashInfer |
| 334 | # SM120 kernel to accept a real row stride for eidx). |
| 335 | capability = current_platform.get_device_capability() |
| 336 | if capability is not None and capability.major == 12: |
| 337 | global_decode = global_decode_buffer[:num_decode_tokens] |
| 338 | else: |
| 339 | global_decode = global_decode_buffer[:num_decode_tokens, :max_compressed_tokens] |
| 340 | decode_lens = decode_lens_buffer[:num_decode_tokens] |
| 341 | prefill_local = prefill_buffer[:num_prefill_tokens, :max_compressed_tokens] |
| 342 | assert global_decode.stride(0) == global_decode_buffer.stride(0) |
| 343 | assert prefill_local.stride(0) == prefill_buffer.stride(0) |
| 344 | |
| 345 | if num_tokens == 0: |
| 346 | return global_decode, decode_lens, prefill_local |
| 347 | |
| 348 | _BUILD_C128A_TOPK_METADATA_KERNEL( |
| 349 | global_decode_buffer, |
| 350 | decode_lens_buffer, |
| 351 | prefill_buffer, |
| 352 | positions, |
| 353 | compress_ratio, |
| 354 | max_compressed_tokens, |
| 355 | num_decode_tokens, |
| 356 | token_to_req_indices, |
| 357 | block_table, |
| 358 | block_size, |
| 359 | slot_mapping, |
| 360 | ) |
| 361 | return global_decode, decode_lens, prefill_local |
| 362 | |
| 363 | |
| 364 | class BuildC128ATopkMetadataKernel( |
| 365 | VllmTritonJitKernel["BuildC128ATopkMetadataKernel.CompileKey"] |
| 366 | ): |
| 367 | @dataclass(frozen=True) |
| 368 | class CompileKey: |
| 369 | compress_ratio: int |
| 370 | max_compressed_tokens: int |
| 371 | block_size: int |
| 372 | triton_block_size: int |
| 373 | |
| 374 | @staticmethod |
| 375 | @triton.jit( |
| 376 | do_not_specialize=[ |
| 377 | "global_decode_stride", |
| 378 | "prefill_local_stride", |
| 379 | "num_decode_tokens", |
| 380 | "block_table_stride", |
| 381 | ] |
| 382 | ) |
| 383 | def kernel( |
| 384 | # Decode outputs |
| 385 | global_decode_ptr, |
| 386 | global_decode_stride, |
| 387 | decode_lens_ptr, |
| 388 | # Prefill output |
| 389 | prefill_local_ptr, |
| 390 | prefill_local_stride, |
| 391 | # Inputs |
| 392 | positions_ptr, |
| 393 | compress_ratio, |
| 394 | max_compressed_tokens, |
| 395 | num_decode_tokens, |
| 396 | token_to_req_indices_ptr, |
| 397 | block_table_ptr, |
| 398 | block_table_stride, |
| 399 | block_size, |
| 400 | slot_mapping_ptr, |
| 401 | BLOCK_SIZE: tl.constexpr, |
| 402 | ): |
| 403 | token_idx = tl.program_id(0) |
| 404 | position = tl.load(positions_ptr + token_idx) |
| 405 | num_compressed = (position + 1) // compress_ratio |
| 406 | num_compressed = tl.minimum(num_compressed, max_compressed_tokens) |
| 407 | is_decode = token_idx < num_decode_tokens |
| 408 | |
| 409 | if is_decode: |
| 410 | # --- Decode: block-table lookup → global slot ids + count --- |
| 411 | is_valid_token = tl.load(slot_mapping_ptr + token_idx) >= 0 |
| 412 | req_idx = tl.load(token_to_req_indices_ptr + token_idx) |
| 413 | count = tl.zeros((), dtype=tl.int32) |
| 414 | for i in range(0, max_compressed_tokens, BLOCK_SIZE): |
| 415 | offset = i + tl.arange(0, BLOCK_SIZE) |
| 416 | mask = offset < max_compressed_tokens |
| 417 | is_valid = offset < num_compressed |
| 418 | |
| 419 | block_indices = offset // block_size |
| 420 | block_numbers = tl.load( |
| 421 | block_table_ptr + req_idx * block_table_stride + block_indices, |
| 422 | mask=mask & is_valid, |
| 423 | ) |
| 424 | block_offsets = offset % block_size |
| 425 | slot_ids = block_numbers * block_size + block_offsets |
| 426 | slot_ids = tl.where(is_valid, slot_ids, -1) |
| 427 | tl.store( |
| 428 | global_decode_ptr + token_idx * global_decode_stride + offset, |
| 429 | slot_ids, |
| 430 | mask=mask, |
| 431 | ) |
| 432 | count += tl.sum(is_valid.to(tl.int32), axis=0) |
| 433 | |
| 434 | tl.store( |
| 435 | decode_lens_ptr + token_idx, |
| 436 | tl.where(is_valid_token, count, 0), |
| 437 | ) |
| 438 | else: |
| 439 | # --- Prefill: write local indices --- |
| 440 | pfx_idx = token_idx - num_decode_tokens |
| 441 | for i in range(0, max_compressed_tokens, BLOCK_SIZE): |
| 442 | offset = i + tl.arange(0, BLOCK_SIZE) |
| 443 | mask = offset < max_compressed_tokens |
| 444 | tl.store( |
| 445 | prefill_local_ptr + pfx_idx * prefill_local_stride + offset, |
| 446 | tl.where(offset < num_compressed, offset, -1), |
| 447 | mask=mask, |
| 448 | ) |
| 449 | |
| 450 | def dispatch( # type: ignore[override] |
| 451 | self, |
| 452 | **compile_key_fields: int, |
| 453 | ) -> CompileKey: |
| 454 | return self.CompileKey(**compile_key_fields) |
| 455 | |
| 456 | def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]: |
| 457 | hf_config = vllm_config.model_config.hf_config |
| 458 | if getattr(hf_config, "model_type", None) != "deepseek_v4": |
| 459 | return [] |
| 460 | |
| 461 | compress_ratio = 128 |
| 462 | max_model_len = vllm_config.model_config.max_model_len |
| 463 | if max_model_len <= 0: |
| 464 | return [] |
| 465 | max_compressed_tokens = cdiv(max_model_len, compress_ratio) |
| 466 | max_compressed_tokens = ( |
| 467 | cdiv(max_compressed_tokens, _C128A_TOPK_ALIGNMENT) * _C128A_TOPK_ALIGNMENT |
| 468 | ) |
| 469 | return self._trace_dispatch(self.dispatch)( |
| 470 | compress_ratio=compress_ratio, |
| 471 | max_compressed_tokens=max_compressed_tokens, |
| 472 | # DeepSeek V4 sparse MLA uses 256-token KV pages; C128A metadata |
| 473 | # works in compressed-token units. |
| 474 | block_size=256 // compress_ratio, |
| 475 | triton_block_size=1024, |
| 476 | ) |
| 477 | |
| 478 | def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]: |
| 479 | int32_ptr = TritonWarmupTensor(torch.int32) |
| 480 | return dict( |
| 481 | global_decode_buffer=TritonWarmupTensor( |
| 482 | torch.int32, shape=(1, 1), strides=(1, 1) |
| 483 | ), |
| 484 | decode_lens_buffer=int32_ptr, |
| 485 | prefill_buffer=TritonWarmupTensor( |
| 486 | torch.int32, shape=(1, 1), strides=(1, 1) |
| 487 | ), |
| 488 | positions=TritonWarmupTensor(torch.int64), |
| 489 | compress_ratio=compile_key.compress_ratio, |
| 490 | max_compressed_tokens=compile_key.max_compressed_tokens, |
| 491 | num_decode_tokens=0, |
| 492 | token_to_req_indices=int32_ptr, |
| 493 | block_table=TritonWarmupTensor(torch.int32, shape=(1, 1), strides=(1, 1)), |
| 494 | block_size=compile_key.block_size, |
| 495 | slot_mapping=TritonWarmupTensor(torch.int64), |
| 496 | ) |
| 497 | |
| 498 | @kernel_launcher |
| 499 | def __call__( |
| 500 | self, |
| 501 | global_decode_buffer: torch.Tensor, |
| 502 | decode_lens_buffer: torch.Tensor, |
| 503 | prefill_buffer: torch.Tensor, |
| 504 | positions: torch.Tensor, |
| 505 | compress_ratio: int, |
| 506 | max_compressed_tokens: int, |
| 507 | num_decode_tokens: int, |
| 508 | token_to_req_indices: torch.Tensor, |
| 509 | block_table: torch.Tensor, |
| 510 | block_size: int, |
| 511 | slot_mapping: torch.Tensor, |
| 512 | ) -> LaunchSpec: |
| 513 | return (positions.shape[0],), dict( |
| 514 | global_decode_ptr=global_decode_buffer, |
| 515 | global_decode_stride=global_decode_buffer.stride(0), |
| 516 | decode_lens_ptr=decode_lens_buffer, |
| 517 | prefill_local_ptr=prefill_buffer, |
| 518 | prefill_local_stride=prefill_buffer.stride(0), |
| 519 | block_table_stride=block_table.stride(0), |
| 520 | BLOCK_SIZE=1024, |
| 521 | ) |
| 522 | |
| 523 | |
| 524 | _BUILD_C128A_TOPK_METADATA_KERNEL = BuildC128ATopkMetadataKernel() |