| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from dataclasses import dataclass |
| 5 | from typing import TYPE_CHECKING, Any, ClassVar |
| 6 | |
| 7 | import torch |
| 8 | from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla |
| 9 | from flashinfer.utils import ( |
| 10 | get_device_sm_count, |
| 11 | get_trtllm_gen_multi_ctas_kv_counter_bytes, |
| 12 | ) |
| 13 | |
| 14 | from vllm.config import get_current_vllm_config |
| 15 | from vllm.config.cache import CacheDType |
| 16 | from vllm.logger import init_logger |
| 17 | from vllm.model_executor.layers.attention.mla_attention import ( |
| 18 | MLACommonBackend, |
| 19 | MLACommonDecodeMetadata, |
| 20 | MLACommonImpl, |
| 21 | MLACommonMetadata, |
| 22 | MLACommonMetadataBuilder, |
| 23 | QueryLenSupport, |
| 24 | ) |
| 25 | from vllm.platforms.interface import DeviceCapability |
| 26 | from vllm.utils.torch_utils import is_quantized_kv_cache |
| 27 | from vllm.v1.attention.backend import ( |
| 28 | AttentionCGSupport, |
| 29 | AttentionLayer, |
| 30 | AttentionType, |
| 31 | MultipleOf, |
| 32 | ) |
| 33 | |
| 34 | if TYPE_CHECKING: |
| 35 | from vllm.config import VllmConfig |
| 36 | from vllm.v1.kv_cache_interface import AttentionSpec |
| 37 | |
| 38 | logger = init_logger(__name__) |
| 39 | |
| 40 | |
| 41 | def _trtllm_gen_mla_decode_supports_num_heads(num_heads: int) -> bool: |
| 42 | """True if trtllm-gen's MLA decode kernel supports this query head count. |
| 43 | |
| 44 | The kernel groups Q heads into CTAs of ``min(num_heads, tileSizeQ)`` and |
| 45 | requires ``num_heads`` divisible by that, else raises "The |
| 46 | numHeadsQ/numHeadsKv is not supported" (flashinfer fmhaKernels.cuh). |
| 47 | ``tileSizeQ`` is 8/16 for ``num_heads <= 8``/``<= 32`` (SwapsMmaAb) else 64; |
| 48 | treat ``> 32`` as tile 64 (safe bound). E.g. 96/24 -> False, 48/64/128 -> True. |
| 49 | """ |
| 50 | if num_heads <= 8: |
| 51 | tile = 8 |
| 52 | elif num_heads <= 32: |
| 53 | tile = 16 |
| 54 | else: |
| 55 | tile = 64 |
| 56 | return num_heads % min(num_heads, tile) == 0 |
| 57 | |
| 58 | |
| 59 | def _select_mla_decode_backend(num_heads: int) -> str | None: |
| 60 | """cute-dsl for head counts trtllm-gen cannot tile, else None (=> auto).""" |
| 61 | if not _trtllm_gen_mla_decode_supports_num_heads(num_heads): |
| 62 | logger.warning_once( |
| 63 | "trtllm-gen MLA decode does not support num_heads=%d " |
| 64 | "(query/kv head ratio); falling back to the cute-dsl backend.", |
| 65 | num_heads, |
| 66 | ) |
| 67 | return "cute-dsl" |
| 68 | return None |
| 69 | |
| 70 | |
| 71 | FLASHINFER_MLA_WORKSPACE_BUFFER_SIZE = 128 * 1024 * 1024 |
| 72 | FLASHINFER_MLA_LSE_WORKSPACE_BUFFER_SIZE = 256 * 1024 * 1024 |
| 73 | |
| 74 | _fi_workspace: torch.Tensor | None = None |
| 75 | |
| 76 | |
| 77 | def _get_workspace_buffer(return_lse: bool) -> torch.Tensor: |
| 78 | global _fi_workspace |
| 79 | |
| 80 | buffer_size = ( |
| 81 | FLASHINFER_MLA_LSE_WORKSPACE_BUFFER_SIZE |
| 82 | if return_lse |
| 83 | else FLASHINFER_MLA_WORKSPACE_BUFFER_SIZE |
| 84 | ) |
| 85 | if _fi_workspace is None or _fi_workspace.numel() < buffer_size: |
| 86 | # FlashInfer's CuteDSL MLA-decode tactic requires an int8 workspace; |
| 87 | # the trtllm-gen path views it as uint8, so int8 is safe for all backends. |
| 88 | _fi_workspace = torch.zeros(buffer_size, dtype=torch.int8, device="cuda") |
| 89 | return _fi_workspace |
| 90 | |
| 91 | |
| 92 | _fi_multi_ctas_kv_counter: torch.Tensor | None = None |
| 93 | |
| 94 | |
| 95 | def _get_multi_ctas_kv_counter_buffer( |
| 96 | min_bytes: int, device: torch.device |
| 97 | ) -> torch.Tensor: |
| 98 | """Persistent, zero-initialized trtllm-gen multi-CTA-KV counter buffer. |
| 99 | |
| 100 | trtllm-gen's multi-CTA-KV MLA decode kernel resets these semaphores to zero |
| 101 | at the end of every launch, so the buffer only needs zeroing once. The |
| 102 | public ``trtllm_batch_decode_with_kv_cache_mla`` entry point builds a fresh |
| 103 | runner per call, so without a caller-owned buffer it re-allocates and |
| 104 | re-zeros this counter on every decode step (a tiny ``FillFunctor<uint8>`` |
| 105 | launch right before the FMHA). Owning it here and passing it in removes that |
| 106 | per-step launch. ``min_bytes`` is sized to the worst-case batch so the |
| 107 | buffer is allocated once and never reallocated after CUDA-graph capture. |
| 108 | """ |
| 109 | global _fi_multi_ctas_kv_counter |
| 110 | if ( |
| 111 | _fi_multi_ctas_kv_counter is None |
| 112 | or _fi_multi_ctas_kv_counter.numel() < min_bytes |
| 113 | ): |
| 114 | _fi_multi_ctas_kv_counter = torch.zeros( |
| 115 | min_bytes, dtype=torch.uint8, device=device |
| 116 | ) |
| 117 | return _fi_multi_ctas_kv_counter |
| 118 | |
| 119 | |
| 120 | @dataclass |
| 121 | class FlashInferMLADecodeMetadata(MLACommonDecodeMetadata): |
| 122 | flattened_block_table: torch.Tensor | None = None |
| 123 | flattened_seq_lens: torch.Tensor | None = None |
| 124 | flattened_row_req: torch.Tensor | None = None |
| 125 | # Row count the flattened tensors were built for (0 = not built), not a query len. |
| 126 | query_len: int = 0 |
| 127 | query_start_loc: torch.Tensor | None = None |
| 128 | max_query_len: int = 1 |
| 129 | |
| 130 | |
| 131 | @dataclass |
| 132 | class FlashInferMLAMetadata(MLACommonMetadata[FlashInferMLADecodeMetadata]): |
| 133 | pass |
| 134 | |
| 135 | |
| 136 | class FlashInferMLAMetadataBuilder(MLACommonMetadataBuilder[FlashInferMLAMetadata]): |
| 137 | # Adaptive verification requires ALWAYS from every builder, matching upstream's |
| 138 | # DeepseekV4FlashMLAMetadataBuilder. The kernels tile ragged queries from the device |
| 139 | # query offsets (flashinfer #3238), so one k+1 graph replays any 1..k+1 mix. |
| 140 | _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS |
| 141 | query_len_support: ClassVar[QueryLenSupport] = QueryLenSupport.VARLEN |
| 142 | # Non-causal DSpark blocks are flattened to single-token rows in forward_mqa. |
| 143 | supports_non_causal_multi_token_decode: ClassVar[bool] = True |
| 144 | supports_non_causal_multi_token_dcp: ClassVar[bool] = True |
| 145 | |
| 146 | def __init__( |
| 147 | self, |
| 148 | kv_cache_spec: "AttentionSpec", |
| 149 | layer_names: list[str], |
| 150 | vllm_config: "VllmConfig", |
| 151 | device: torch.device, |
| 152 | ) -> None: |
| 153 | parallel_config = vllm_config.parallel_config |
| 154 | dcp_size = parallel_config.decode_context_parallel_size |
| 155 | interleave_size = parallel_config.cp_kv_cache_interleave_size |
| 156 | if dcp_size > 1 and interleave_size != 1: |
| 157 | raise ValueError( |
| 158 | "FlashInfer MLA native DCP requires " |
| 159 | "cp_kv_cache_interleave_size=1; got " |
| 160 | f"{interleave_size}." |
| 161 | ) |
| 162 | super().__init__( |
| 163 | kv_cache_spec, |
| 164 | layer_names, |
| 165 | vllm_config, |
| 166 | device, |
| 167 | FlashInferMLAMetadata, |
| 168 | supports_dcp_with_varlen=True, |
| 169 | ) |
| 170 | |
| 171 | def _build_decode( |
| 172 | self, |
| 173 | block_table_tensor: torch.Tensor, |
| 174 | seq_lens_device: torch.Tensor, |
| 175 | max_seq_len: int, |
| 176 | query_start_loc_cpu: torch.Tensor, |
| 177 | query_start_loc_device: torch.Tensor, |
| 178 | num_decode_tokens: int, |
| 179 | max_query_len: int, |
| 180 | dcp_tot_seq_lens_device: torch.Tensor | None, |
| 181 | ) -> FlashInferMLADecodeMetadata: |
| 182 | # Promised bound, not measured: a capture dummy puts one token/req, so |
| 183 | # measuring would bake an undersized graph. Ragged only when max_query_len>1. |
| 184 | return FlashInferMLADecodeMetadata( |
| 185 | block_table=block_table_tensor, |
| 186 | seq_lens=seq_lens_device, |
| 187 | dcp_tot_seq_lens=dcp_tot_seq_lens_device, |
| 188 | query_start_loc=query_start_loc_device, |
| 189 | max_query_len=max_query_len, |
| 190 | ) |
| 191 | |
| 192 | |
| 193 | class FlashInferMLABackend(MLACommonBackend): |
| 194 | supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16] |
| 195 | supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [ |
| 196 | "auto", |
| 197 | "float16", |
| 198 | "bfloat16", |
| 199 | "fp8", |
| 200 | "fp8_e4m3", |
| 201 | ] |
| 202 | |
| 203 | @staticmethod |
| 204 | def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]: |
| 205 | return [32, 64] |
| 206 | |
| 207 | @staticmethod |
| 208 | def get_name() -> str: |
| 209 | return "FLASHINFER_MLA" |
| 210 | |
| 211 | @staticmethod |
| 212 | def get_impl_cls() -> type["FlashInferMLAImpl"]: |
| 213 | return FlashInferMLAImpl |
| 214 | |
| 215 | @staticmethod |
| 216 | def get_builder_cls() -> type["FlashInferMLAMetadataBuilder"]: |
| 217 | return FlashInferMLAMetadataBuilder |
| 218 | |
| 219 | @classmethod |
| 220 | def supports_compute_capability(cls, capability: DeviceCapability) -> bool: |
| 221 | return capability.major == 10 |
| 222 | |
| 223 | @classmethod |
| 224 | def supports_non_causal(cls) -> bool: |
| 225 | return True |
| 226 | |
| 227 | @classmethod |
| 228 | def supports_combination( |
| 229 | cls, |
| 230 | head_size: int, |
| 231 | dtype: torch.dtype, |
| 232 | kv_cache_dtype: CacheDType | None, |
| 233 | block_size: int | None, |
| 234 | use_mla: bool, |
| 235 | has_sink: bool, |
| 236 | use_sparse: bool, |
| 237 | use_mm_prefix: bool, |
| 238 | device_capability: DeviceCapability, |
| 239 | ) -> str | None: |
| 240 | # FlashInfer MLA kernel requires qk_nope_head_dim in [64, 128, 192] |
| 241 | vllm_config = get_current_vllm_config() |
| 242 | if vllm_config.model_config is not None: |
| 243 | hf_text_config = vllm_config.model_config.hf_text_config |
| 244 | qk_nope_head_dim = getattr(hf_text_config, "qk_nope_head_dim", 1) |
| 245 | if qk_nope_head_dim not in [64, 128, 192]: |
| 246 | return ( |
| 247 | "FlashInfer MLA kernel requires qk_nope_head_dim " |
| 248 | f"in [64, 128, 192], but got {qk_nope_head_dim}" |
| 249 | ) |
| 250 | return None |
| 251 | |
| 252 | |
| 253 | class FlashInferMLAImpl(MLACommonImpl[FlashInferMLAMetadata]): |
| 254 | can_return_lse_for_decode: bool = True |
| 255 | supports_dcp: bool = True |
| 256 | # DCP is the only path that consumes LSE. It uses monolithic CuTeDSL, |
| 257 | # whose public LSE contract is natural-log. |
| 258 | lse_base_on_e: bool = True |
| 259 | |
| 260 | def __init__( |
| 261 | self, |
| 262 | num_heads: int, |
| 263 | head_size: int, |
| 264 | scale: float, |
| 265 | num_kv_heads: int, |
| 266 | alibi_slopes: list[float] | None, |
| 267 | sliding_window: int | None, |
| 268 | kv_cache_dtype: str, |
| 269 | logits_soft_cap: float | None, |
| 270 | attn_type: str, |
| 271 | kv_sharing_target_layer_name: str | None, |
| 272 | # MLA Specific Arguments |
| 273 | **mla_args, |
| 274 | ) -> None: |
| 275 | super().__init__( |
| 276 | num_heads, |
| 277 | head_size, |
| 278 | scale, |
| 279 | num_kv_heads, |
| 280 | alibi_slopes, |
| 281 | sliding_window, |
| 282 | kv_cache_dtype, |
| 283 | logits_soft_cap, |
| 284 | attn_type, |
| 285 | kv_sharing_target_layer_name, |
| 286 | **mla_args, |
| 287 | ) |
| 288 | |
| 289 | unsupported_features = [alibi_slopes, sliding_window, logits_soft_cap] |
| 290 | if any(unsupported_features): |
| 291 | raise NotImplementedError( |
| 292 | "FlashInferMLAImpl does not support one of the following: " |
| 293 | "alibi_slopes, sliding_window, logits_soft_cap" |
| 294 | ) |
| 295 | |
| 296 | if attn_type != AttentionType.DECODER: |
| 297 | raise NotImplementedError( |
| 298 | "Encoder self-attention and " |
| 299 | "encoder/decoder cross-attention " |
| 300 | "are not implemented for " |
| 301 | "FlashInferMLAImpl" |
| 302 | ) |
| 303 | |
| 304 | self.bmm1_scale: float | None = None |
| 305 | self.bmm2_scale: float | None = None |
| 306 | # Worst-case decode batch for the persistent trtllm-gen multi-CTA-KV |
| 307 | # counter buffer (see _get_multi_ctas_kv_counter_buffer). Captured here |
| 308 | # (config is in scope during construction) so the byte size can be |
| 309 | # resolved once on the first decode and never grows after CUDA-graph |
| 310 | # capture. The impl has no _vllm_config, hence get_current_vllm_config(). |
| 311 | _sched = get_current_vllm_config().scheduler_config |
| 312 | self._mla_counter_max_batch: int = ( |
| 313 | _sched.max_num_batched_tokens or _sched.max_num_seqs |
| 314 | ) |
| 315 | self._mla_counter_bytes: int | None = None |
| 316 | |
| 317 | def forward_mqa( |
| 318 | self, |
| 319 | q: torch.Tensor | tuple[torch.Tensor, torch.Tensor], |
| 320 | kv_c_and_k_pe_cache: torch.Tensor, |
| 321 | attn_metadata: FlashInferMLAMetadata, |
| 322 | layer: AttentionLayer, |
| 323 | ) -> tuple[torch.Tensor, torch.Tensor | None]: |
| 324 | assert kv_c_and_k_pe_cache.numel() > 0 |
| 325 | assert attn_metadata.decode is not None |
| 326 | |
| 327 | if isinstance(q, tuple): |
| 328 | q_nope, q_pe = q |
| 329 | q = torch.cat([q_nope, q_pe], dim=-1) |
| 330 | |
| 331 | block_table = attn_metadata.decode.block_table |
| 332 | seq_lens = attn_metadata.decode.seq_lens |
| 333 | |
| 334 | # Led by the promised bound (baked into the graph), not |
| 335 | # num_decode_tokens // num_decodes -- that average is a per-request length |
| 336 | # only for uniform batches. |
| 337 | multi_token_decode = ( |
| 338 | attn_metadata.decode.max_query_len > 1 |
| 339 | or attn_metadata.num_decode_tokens > attn_metadata.num_decodes |
| 340 | ) |
| 341 | cum_seq_lens_q: torch.Tensor | None = None |
| 342 | max_q_len: int | None = None |
| 343 | row_req: torch.Tensor | None = None |
| 344 | |
| 345 | if not attn_metadata.causal: |
| 346 | # FlashInfer decode has no causal flag. Flatten each non-causal |
| 347 | # query block into independent single-token rows. |
| 348 | q = q.unsqueeze(1) |
| 349 | if multi_token_decode: |
| 350 | block_table, seq_lens, row_req = self._flattened_decode_metadata( |
| 351 | attn_metadata, q.shape[0] |
| 352 | ) |
| 353 | elif attn_metadata.decode.max_query_len > 1 and self.dcp_world_size == 1: |
| 354 | # Causal spec decode: keep q compact and let the kernel tile each |
| 355 | # request's length (uniform 1+k and adaptive ragged). Mirrors vllm #52157. |
| 356 | # DCP keeps the uniform reshape below: it needs LSE, which the kernels |
| 357 | # do not return on the ragged path (flashinfer #3238). |
| 358 | cum_seq_lens_q = attn_metadata.decode.query_start_loc |
| 359 | max_q_len = attn_metadata.decode.max_query_len |
| 360 | # trtllm API requires extra dimension q_len_per_request for MTP |
| 361 | elif attn_metadata.num_decode_tokens % attn_metadata.num_decodes != 0: |
| 362 | logger.warning_once( |
| 363 | """FlashInferMLAImpl got a query of uneven length. |
| 364 | This usually indicates an issue in batch reordering |
| 365 | or incorrect setup in dummy_run.""" |
| 366 | ) |
| 367 | q = q.unsqueeze(1) |
| 368 | else: |
| 369 | q = q.view(attn_metadata.num_decodes, -1, q.shape[-2], q.shape[-1]) |
| 370 | |
| 371 | if self.bmm1_scale is None: |
| 372 | self.bmm1_scale = self.scale |
| 373 | if is_quantized_kv_cache(self.kv_cache_dtype): |
| 374 | self.bmm1_scale *= layer._q_scale_float * layer._k_scale_float |
| 375 | |
| 376 | if self.bmm2_scale is None: |
| 377 | self.bmm2_scale = 1.0 |
| 378 | if is_quantized_kv_cache(self.kv_cache_dtype): |
| 379 | self.bmm2_scale *= layer._k_scale_float |
| 380 | |
| 381 | return_lse = self.need_to_return_lse_for_decode |
| 382 | workspace_buffer = _get_workspace_buffer(return_lse) |
| 383 | # Parallel gathers can change the runtime Q heads from TP-local num_heads. |
| 384 | runtime_num_heads = q.shape[-2] |
| 385 | extra_kwargs: dict[str, Any] = {} |
| 386 | decode_backend: str | None |
| 387 | if self.dcp_world_size > 1: |
| 388 | causal_seqlens_kv_global = attn_metadata.decode.dcp_tot_seq_lens |
| 389 | assert causal_seqlens_kv_global is not None |
| 390 | if row_req is not None: |
| 391 | causal_seqlens_kv_global = causal_seqlens_kv_global[row_req] |
| 392 | extra_kwargs.update( |
| 393 | enable_dcp=True, |
| 394 | cp_world=self.dcp_world_size, |
| 395 | cp_rank=self.dcp_rank, |
| 396 | causal_seqlens_kv_global=causal_seqlens_kv_global, |
| 397 | ) |
| 398 | decode_backend = "cute-dsl" |
| 399 | else: |
| 400 | # trtllm-gen rejects MLA head counts it can't tile (e.g. 96); |
| 401 | # fall back to cute-dsl for those. |
| 402 | decode_backend = _select_mla_decode_backend(runtime_num_heads) |
| 403 | if cum_seq_lens_q is not None: |
| 404 | # Neither decode backend returns LSE on the ragged path |
| 405 | # (flashinfer #3238); DCP, the only LSE consumer, took the uniform |
| 406 | # branch above, so this only guards a future caller wiring the two. |
| 407 | assert not return_lse, ( |
| 408 | "FlashInferMLA ragged decode cannot return LSE; DCP and adaptive " |
| 409 | "variable-length decode are mutually exclusive." |
| 410 | ) |
| 411 | extra_kwargs["cum_seq_lens_q"] = cum_seq_lens_q |
| 412 | extra_kwargs["max_q_len"] = max_q_len |
| 413 | if decode_backend: |
| 414 | extra_kwargs["backend"] = decode_backend |
| 415 | elif kv_c_and_k_pe_cache.shape[-2] in (32, 64): |
| 416 | # The auto path can dispatch to trtllm-gen, whose multi-CTA-KV decode |
| 417 | # kernel self-resets its semaphore counter after each launch (so it |
| 418 | # only needs zeroing once). Pass a persistent counter buffer to skip |
| 419 | # the per-step re-allocate + re-zero the public entry point would |
| 420 | # otherwise do. Guarded to configs where a trtllm-gen runner is |
| 421 | # eligible (page/block size in {32, 64}); the arg is rejected when |
| 422 | # only a cute-dsl runner can run. |
| 423 | if self._mla_counter_bytes is None: |
| 424 | self._mla_counter_bytes = get_trtllm_gen_multi_ctas_kv_counter_bytes( |
| 425 | self._mla_counter_max_batch, |
| 426 | runtime_num_heads, |
| 427 | get_device_sm_count(q.device), |
| 428 | ) |
| 429 | extra_kwargs["multi_ctas_kv_counter_buffer"] = ( |
| 430 | _get_multi_ctas_kv_counter_buffer(self._mla_counter_bytes, q.device) |
| 431 | ) |
| 432 | kernel_out = trtllm_batch_decode_with_kv_cache_mla( |
| 433 | query=q, |
| 434 | kv_cache=kv_c_and_k_pe_cache.unsqueeze(1), |
| 435 | workspace_buffer=workspace_buffer, |
| 436 | qk_nope_head_dim=self.qk_nope_head_dim, |
| 437 | kv_lora_rank=self.kv_lora_rank, |
| 438 | qk_rope_head_dim=self.qk_rope_head_dim, |
| 439 | block_tables=block_table, |
| 440 | seq_lens=seq_lens, |
| 441 | max_seq_len=attn_metadata.max_seq_len, |
| 442 | bmm1_scale=self.bmm1_scale, |
| 443 | bmm2_scale=self.bmm2_scale, |
| 444 | return_lse=return_lse, |
| 445 | **extra_kwargs, |
| 446 | ) |
| 447 | if return_lse: |
| 448 | o, lse = kernel_out |
| 449 | lse = lse.view(-1, lse.shape[-1]) |
| 450 | else: |
| 451 | o, lse = kernel_out, None |
| 452 | |
| 453 | # Flatten the output for consistent shape |
| 454 | o = o.view(-1, o.shape[-2], o.shape[-1]) |
| 455 | |
| 456 | return o, lse |
| 457 | |
| 458 | def _flattened_decode_metadata( |
| 459 | self, |
| 460 | attn_metadata: FlashInferMLAMetadata, |
| 461 | num_rows: int, |
| 462 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| 463 | """Expand per-request decode tensors to one row per query token. |
| 464 | |
| 465 | Cached across the layers of a group, which all see the same batch. |
| 466 | """ |
| 467 | decode = attn_metadata.decode |
| 468 | assert decode is not None |
| 469 | if decode.query_len != num_rows: |
| 470 | cu = decode.query_start_loc |
| 471 | assert cu is not None |
| 472 | # searchsorted on the device offsets, not repeat_interleave(uniform_len): |
| 473 | # the latter only lines up for uniform batches and hands ragged rows |
| 474 | # another request's KV (silent garbage). |
| 475 | rows = torch.arange(num_rows, device=cu.device, dtype=cu.dtype) |
| 476 | row_req = torch.searchsorted(cu[1:], rows, right=True).clamp_( |
| 477 | max=decode.block_table.shape[0] - 1 |
| 478 | ) |
| 479 | decode.flattened_row_req = row_req |
| 480 | decode.flattened_block_table = decode.block_table[row_req] |
| 481 | decode.flattened_seq_lens = decode.seq_lens[row_req] |
| 482 | decode.query_len = num_rows |
| 483 | assert decode.flattened_block_table is not None |
| 484 | assert decode.flattened_seq_lens is not None |
| 485 | assert decode.flattened_row_req is not None |
| 486 | return ( |
| 487 | decode.flattened_block_table, |
| 488 | decode.flattened_seq_lens, |
| 489 | decode.flattened_row_req, |
| 490 | ) |