| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from dataclasses import dataclass |
| 4 | from typing import TYPE_CHECKING, ClassVar |
| 5 | |
| 6 | import numpy as np |
| 7 | import torch |
| 8 | import torch.distributed as dist |
| 9 | |
| 10 | from vllm import _custom_ops as ops |
| 11 | from vllm.config import VllmConfig, get_current_vllm_config |
| 12 | from vllm.config.cache import CacheDType |
| 13 | from vllm.distributed.parallel_state import get_dcp_group |
| 14 | from vllm.logger import init_logger |
| 15 | from vllm.model_executor.layers.attention.sparse_mla_attention import ( |
| 16 | SparseMLACommonImpl, |
| 17 | SparseMLACommonMetadata, |
| 18 | SparseMLACommonMetadataBuilder, |
| 19 | ) |
| 20 | from vllm.platforms import current_platform |
| 21 | from vllm.platforms.interface import DeviceCapability |
| 22 | from vllm.utils.math_utils import round_up |
| 23 | from vllm.utils.platform_utils import num_compute_units |
| 24 | from vllm.utils.torch_utils import ( |
| 25 | async_tensor_h2d, |
| 26 | current_stream, |
| 27 | is_quantized_kv_cache, |
| 28 | ) |
| 29 | from vllm.v1.attention.backend import ( |
| 30 | AttentionBackend, |
| 31 | AttentionCGSupport, |
| 32 | AttentionLayer, |
| 33 | CommonAttentionMetadata, |
| 34 | MultipleOf, |
| 35 | ) |
| 36 | from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup |
| 37 | from vllm.v1.attention.backends.mla.sparse_utils import ( |
| 38 | flat_kv_row_view, |
| 39 | request_row_bounds, |
| 40 | triton_convert_req_index_to_global_index, |
| 41 | triton_filter_and_convert_dcp_index, |
| 42 | ) |
| 43 | from vllm.v1.attention.backends.utils import ( |
| 44 | get_dcp_local_seq_lens, |
| 45 | reshape_attn_output_for_spec_decode, |
| 46 | reshape_query_for_spec_decode, |
| 47 | split_prefill_chunks, |
| 48 | ) |
| 49 | from vllm.v1.attention.ops.flashmla import ( |
| 50 | FlashMLASchedMeta, |
| 51 | flash_mla_sparse_fwd, |
| 52 | flash_mla_with_kvcache, |
| 53 | get_mla_metadata, |
| 54 | ) |
| 55 | from vllm.v1.kv_cache_interface import AttentionSpec |
| 56 | from vllm.v1.worker.workspace import current_workspace_manager |
| 57 | |
| 58 | if TYPE_CHECKING: |
| 59 | from vllm.model_executor.models.deepseek_v2 import Indexer |
| 60 | |
| 61 | logger = init_logger(__name__) |
| 62 | |
| 63 | # For FP8 sparse attention we have two implementations: |
| 64 | # 1. Mixed batch mode: use the FP8 decode kernel for both prefill and decode this is |
| 65 | # done by treating all tokens as single batch. |
| 66 | # 2. Separate prefill and decode mode: use the BF16 prefill kernel for prefill |
| 67 | # (upconverting the FP8 cache to BF16 then calling the prefill kernel) and using |
| 68 | # the FP8 decode kernel for decode. |
| 69 | # Currently we use #1 when the number of heads per rank is low (i.e. TP) since the BF16 |
| 70 | # prefill kernel requires padding the number of heads to 128 while the decode does not |
| 71 | # so when the per-rank head count is below MIN_HEADS_FOR_BF16_PREFILL we use the mixed |
| 72 | # batch mode (#1). |
| 73 | MIN_HEADS_FOR_BF16_PREFILL = 32 |
| 74 | |
| 75 | """ |
| 76 | NOTE: FlashMLA Sparse uses an fp8 cache with the following format |
| 77 | |
| 78 | For DeepSeek V3.2, in the "FP8 with scale" format, each token's KV cache is 656 |
| 79 | Bytes, structured as: |
| 80 | - **First 512 bytes:** The "quantized NoPE" part, containing 512 |
| 81 | `float8_e4m3` values. |
| 82 | - **Next 16 bytes:** Scale factors, containing 4 `float32` values. |
| 83 | The first `float32` is the scale for the first 128 `float8_e4m3` values, |
| 84 | the second for the next 128, and so on. |
| 85 | - **Last 128 bytes:** The "RoPE" part, containing 64 `bfloat16` values. This |
| 86 | part is not quantized for accuracy. |
| 87 | |
| 88 | For DeepSeek V4, in the "FP8 with scale" format, each token's KV cache is 584 |
| 89 | Bytes, structured as: |
| 90 | - **First 448 bytes:** The "quantized NoPE" part, containing 448 |
| 91 | `float8_e4m3` values. |
| 92 | - **Next 128 bytes:** The "RoPE" part, containing 64 `bfloat16` values. This |
| 93 | part is not quantized for accuracy. |
| 94 | - **Last 8 bytes:** Scale factors, containing 7 `ue8m0` values + 1B pad. |
| 95 | The first `ue8m0` is the scale for the first 64 `float8_e4m3` values, |
| 96 | the second for the next 64, and so on. |
| 97 | |
| 98 | For DeepSeek V4.1, each token's KV cache is 528 Bytes, structured as: |
| 99 | - **First 512 bytes:** all 512 dims as `float8_e4m3`. Unlike V4 the 64 RoPE |
| 100 | dims are quantized too, so there is no `bfloat16` part. |
| 101 | - **Last 16 bytes:** Scale factors, containing 16 `ue8m0` values, one per 32 |
| 102 | consecutive `float8_e4m3` values (i.e. MXFP8). |
| 103 | |
| 104 | The V4 and V4.1 records are not laid out token-by-token within a page: a page |
| 105 | holds all its data rows first and all its scale rows after, so a page is |
| 106 | `block_size * bytes_per_token` bytes rounded up to the decode kernel's TMA |
| 107 | stride (576 B for V4, 512 B for V4.1). |
| 108 | |
| 109 | In the "nvfp4_ds_mla" format (SM100 only, DeepSeek V3.2 geometry), each |
| 110 | token's KV cache is 352 Bytes, structured as: |
| 111 | - **First 256 bytes:** 512 `e2m1` NoPE values packed 2/byte (low nibble = |
| 112 | even element). |
| 113 | - **Next 64 bytes:** 64 `float8_e4m3` RoPE values. These carry no scale |
| 114 | factor: `e4m3`'s 4 exponent bits span the RoPE magnitude range unaided. |
| 115 | - **Last 32 bytes:** 32 `float8_e4m3` NoPE scale factors, one per 16 |
| 116 | elements, stored permuted (an 8x4 -> 4x8 transpose: the scale for element |
| 117 | block `s` lives at byte `8 * (s & 3) + (s >> 2)`) so that the 8 scales one |
| 118 | FlashMLA dequant thread needs are contiguous. See the layout comment in |
| 119 | `csrc/libtorch_stable/cache_kernels.cu`. |
| 120 | |
| 121 | """ |
| 122 | |
| 123 | # Quantized DS-MLA cache formats served by the FP8/NVFP4 sparse decode kernel |
| 124 | # path (as opposed to the plain bf16 cache). FlashMLA infers which of these the |
| 125 | # cache holds from its bytes-per-token, so nothing else needs to be passed down. |
| 126 | QUANTIZED_DS_MLA_CACHE_FORMATS: frozenset[str] = frozenset( |
| 127 | {"fp8_ds_mla", "nvfp4_ds_mla"} |
| 128 | ) |
| 129 | |
| 130 | |
| 131 | class FlashMLASparseBackend(AttentionBackend): |
| 132 | supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16] |
| 133 | supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [ |
| 134 | "auto", |
| 135 | "bfloat16", |
| 136 | "fp8_ds_mla", |
| 137 | "fp8", # alias for fp8_ds_mla |
| 138 | "nvfp4_ds_mla", # NVFP4 NoPE + FP8 RoPE (SM100 only) |
| 139 | ] |
| 140 | |
| 141 | @staticmethod |
| 142 | def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]: |
| 143 | return [64] |
| 144 | |
| 145 | @staticmethod |
| 146 | def get_name() -> str: |
| 147 | return "FLASHMLA_SPARSE" |
| 148 | |
| 149 | @staticmethod |
| 150 | def get_builder_cls() -> type["FlashMLASparseMetadataBuilder"]: |
| 151 | return FlashMLASparseMetadataBuilder |
| 152 | |
| 153 | @staticmethod |
| 154 | def get_impl_cls() -> type["FlashMLASparseImpl"]: |
| 155 | return FlashMLASparseImpl |
| 156 | |
| 157 | @classmethod |
| 158 | def get_supported_head_sizes(cls) -> list[int]: |
| 159 | # DeepSeek V3.2: 512 NoPE + 64 RoPE = 576; GLM5Next NoPE: 512. |
| 160 | return [576, 512] |
| 161 | |
| 162 | @classmethod |
| 163 | def is_mla(cls) -> bool: |
| 164 | return True |
| 165 | |
| 166 | @classmethod |
| 167 | def is_sparse(cls) -> bool: |
| 168 | return True |
| 169 | |
| 170 | @classmethod |
| 171 | def supports_compute_capability(cls, capability: DeviceCapability) -> bool: |
| 172 | return capability.major in [9, 10] |
| 173 | |
| 174 | @classmethod |
| 175 | def supports_combination( |
| 176 | cls, |
| 177 | head_size: int, |
| 178 | dtype: torch.dtype, |
| 179 | kv_cache_dtype: "CacheDType | None", |
| 180 | block_size: int | None, |
| 181 | use_mla: bool, |
| 182 | has_sink: bool, |
| 183 | use_sparse: bool, |
| 184 | use_mm_prefix: bool, |
| 185 | device_capability: DeviceCapability, |
| 186 | ) -> str | None: |
| 187 | if head_size == 512: |
| 188 | # GLM5Next NoPE (qk_rope_head_dim == 0, kv_lora_rank == 512) has |
| 189 | # head_size 512 and is served here by the direct 512-wide bf16 |
| 190 | # cache on SM90. Quantized DS-MLA caches need the zero-padded |
| 191 | # 576/656B envelope, which is wired up separately; plain fp8, |
| 192 | # SM100 bf16, and rope-carrying 512 models must fall through to |
| 193 | # FlashInfer/TRITON. |
| 194 | if ( |
| 195 | kv_cache_dtype in (None, "auto", "bfloat16", "float16") |
| 196 | and device_capability.major == 9 |
| 197 | ): |
| 198 | # Direct bf16 NoPE-512 is only correct for rope-free models. |
| 199 | # Precedent for reading hf_text_config in supports_combination: |
| 200 | # flashinfer_mla_sparse.py. |
| 201 | from vllm.config import get_current_vllm_config |
| 202 | |
| 203 | vllm_config = get_current_vllm_config() |
| 204 | if vllm_config.model_config is not None: |
| 205 | hf_text_config = vllm_config.model_config.hf_text_config |
| 206 | if getattr(hf_text_config, "qk_rope_head_dim", 64) != 0: |
| 207 | return ( |
| 208 | "FLASHMLA_SPARSE supports head_size 512 only for " |
| 209 | "rope-free (NoPE) models" |
| 210 | ) |
| 211 | else: |
| 212 | return ( |
| 213 | "FLASHMLA_SPARSE supports head_size 512 only with bf16 " |
| 214 | "kv-cache on SM90 (NoPE), got " |
| 215 | f"kv_cache_dtype={kv_cache_dtype}, " |
| 216 | f"capability={device_capability}" |
| 217 | ) |
| 218 | if kv_cache_dtype == "nvfp4_ds_mla" and device_capability.major != 10: |
| 219 | return ( |
| 220 | f"FLASHMLA_SPARSE only supports the {kv_cache_dtype} kv-cache " |
| 221 | "dtype on SM100 (Blackwell)" |
| 222 | ) |
| 223 | return None |
| 224 | |
| 225 | |
| 226 | def gathered_prefill_shards( |
| 227 | row_req_idx: np.ndarray, row_shard_rows: np.ndarray |
| 228 | ) -> tuple[np.ndarray, np.ndarray]: |
| 229 | """Under PCP+DCP, group this rank's prefill rows by request and take each |
| 230 | request's slice of this rank's KV shard. |
| 231 | |
| 232 | PCP gives a rank two adjacent chunk rows of a split prefill; they share |
| 233 | one context, so the workspace holds it once. ``row_shard_rows`` is |
| 234 | ``CommonAttentionMetadata.dcp_local_seq_lens_cpu_upper_bound``: the |
| 235 | largest DCP shard of the whole request on every row, identical on every |
| 236 | PCP rank, so the all-gather is shaped the same everywhere. |
| 237 | |
| 238 | Returns the request row bounds and, per request, its shard rows. |
| 239 | """ |
| 240 | row_bounds = request_row_bounds(row_req_idx) |
| 241 | shard_rows = row_shard_rows[row_bounds[:-1]].astype(np.int64) |
| 242 | assert np.all(shard_rows > 0), ( |
| 243 | f"PCP+DCP prefill got an empty context: {shard_rows.tolist()}" |
| 244 | ) |
| 245 | return row_bounds, shard_rows |
| 246 | |
| 247 | |
| 248 | @dataclass |
| 249 | class FlashMLASparseMetadata(SparseMLACommonMetadata): |
| 250 | @dataclass |
| 251 | class FP8KernelMetadata: |
| 252 | scheduler_metadata: FlashMLASchedMeta |
| 253 | dummy_block_table: torch.Tensor |
| 254 | cache_lens: torch.Tensor |
| 255 | |
| 256 | @dataclass |
| 257 | class FP8SeparatePrefillDecode: |
| 258 | @dataclass |
| 259 | class Decode: |
| 260 | seq_lens: torch.Tensor |
| 261 | kernel_metadata: "FlashMLASparseMetadata.FP8KernelMetadata" |
| 262 | decode_query_len: int # needed for reshape in spec decode |
| 263 | |
| 264 | @dataclass |
| 265 | class Prefill: |
| 266 | # Request ID for each token: -1 for decode tokens, request index |
| 267 | # (0, 1, 2, ...) for prefill tokens. |
| 268 | # Shape: [num_actual_tokens] |
| 269 | request_ids: torch.Tensor |
| 270 | |
| 271 | # Workspace start offsets for all prefill requests |
| 272 | # Shape: [num_prefill_reqs], adjusted in-place per chunk to be |
| 273 | # 0-indexed within each chunk. Used to map prefill tokens to workspace |
| 274 | # offsets in convert_logical_index_to_physical_index |
| 275 | workspace_starts: torch.Tensor |
| 276 | |
| 277 | @dataclass |
| 278 | class Chunk: |
| 279 | """Metadata for a chunk of prefill requests. |
| 280 | |
| 281 | Prefill requests may be chunked to fit within the fixed workspace size. |
| 282 | """ |
| 283 | |
| 284 | tokens_slice: slice |
| 285 | block_table: torch.Tensor |
| 286 | req_start_idx: int |
| 287 | workspace_starts: torch.Tensor |
| 288 | chunk_tot_seqlen: int |
| 289 | seq_lens: torch.Tensor | None = None |
| 290 | |
| 291 | chunks: list[Chunk] |
| 292 | |
| 293 | num_prefills: int = 0 |
| 294 | num_decodes: int = 0 |
| 295 | num_prefill_tokens: int = 0 |
| 296 | num_decode_tokens: int = 0 |
| 297 | |
| 298 | decode: Decode | None = None |
| 299 | prefill: Prefill | None = None |
| 300 | |
| 301 | fp8_extra_metadata: FP8SeparatePrefillDecode | FP8KernelMetadata | None = None |
| 302 | fp8_use_mixed_batch: bool = False |
| 303 | |
| 304 | |
| 305 | def get_prefill_workspace_size(max_model_len: int): |
| 306 | # NOTE(Lucas): 5 is a magic number for controlling the prefill buffer size. |
| 307 | # May be tuned later. |
| 308 | # Memory usage: 5 * max_model_len * 576 * 2 bytes |
| 309 | # Example: DeepSeek-V3.2 with max_model_len=163840 -> |
| 310 | # 5 * 163840 * 576 * 2 = ~900 MB |
| 311 | # This fits nicely below the typical MoE workspace size of >2GB so this is "free" |
| 312 | return max_model_len * 5 |
| 313 | |
| 314 | |
| 315 | class FlashMLASparseMetadataBuilder( |
| 316 | SparseMLACommonMetadataBuilder[FlashMLASparseMetadata] |
| 317 | ): |
| 318 | _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH |
| 319 | require_uniform_decodes: ClassVar[bool] = True |
| 320 | metadata_cls = FlashMLASparseMetadata |
| 321 | |
| 322 | def __init__( |
| 323 | self, |
| 324 | kv_cache_spec: AttentionSpec, |
| 325 | layer_names: list[str], |
| 326 | vllm_config: VllmConfig, |
| 327 | device: torch.device, |
| 328 | ) -> None: |
| 329 | super().__init__(kv_cache_spec, layer_names, vllm_config, device) |
| 330 | cache_config = vllm_config.cache_config |
| 331 | parallel_config = vllm_config.parallel_config |
| 332 | |
| 333 | num_q_heads = self.model_config.get_num_attention_heads(parallel_config) |
| 334 | if current_platform.is_device_capability_family(100): |
| 335 | threshold = {8: 128, 16: 128, 32: 128, 64: 256, 128: 1024}.get( |
| 336 | num_q_heads, 1024 |
| 337 | ) |
| 338 | else: |
| 339 | threshold = {16: 128, 32: 128, 64: 256, 128: 256}.get(num_q_heads, 256) |
| 340 | self.use_hisparse = vllm_config.attention_config.hisparse_config is not None |
| 341 | if self.use_hisparse: |
| 342 | threshold = 1 |
| 343 | # Varlen decodes are safe under DCP: causality comes from the |
| 344 | # indexer's top-k indices, not from the kernel metadata. |
| 345 | self._init_reorder_batch_threshold( |
| 346 | threshold, |
| 347 | supports_spec_as_decode=True, |
| 348 | supports_dcp_with_varlen=(parallel_config.cp_kv_cache_interleave_size == 1), |
| 349 | ) |
| 350 | |
| 351 | sm_count = num_compute_units(device.index) |
| 352 | |
| 353 | self.num_heads = self.model_config.get_num_attention_heads(parallel_config) |
| 354 | # FP8 decode kernel only supports h_q = 64 or 128, so we need to pad |
| 355 | self.fp8_decode_padded_heads = ( |
| 356 | FlashMLASparseImpl._compute_fp8_decode_padded_heads(self.num_heads) |
| 357 | ) |
| 358 | |
| 359 | self.use_fp8_kv_cache = ( |
| 360 | cache_config.cache_dtype in QUANTIZED_DS_MLA_CACHE_FORMATS |
| 361 | ) |
| 362 | max_num_seqs = vllm_config.scheduler_config.max_num_seqs |
| 363 | # Shape: [max_num_seqs], all elements = topk_tokens (constant for full-CG) |
| 364 | self.topk_tokens_tensor = torch.full( |
| 365 | (max_num_seqs,), self.topk_tokens, device=device, dtype=torch.int32 |
| 366 | ) |
| 367 | # Shape: [max_num_seqs], all elements = max_model_len |
| 368 | self.max_model_len_tensor = torch.full( |
| 369 | (max_num_seqs,), |
| 370 | self.model_config.max_model_len, |
| 371 | device=device, |
| 372 | dtype=torch.int32, |
| 373 | ) |
| 374 | # this is ignored by `flash_mla_with_kvcache` if indices not None |
| 375 | self.dummy_block_table = torch.empty( |
| 376 | (max_num_seqs, 1), dtype=torch.int32, device=self.device |
| 377 | ) |
| 378 | |
| 379 | # Equation taken from FlashMLA/csrc/api/sparse_decode.h |
| 380 | # For sparse FP8 decode, the formula depends on architecture: |
| 381 | # - SM90 (Hopper): num_sm_parts = num_sms / s_q / (h_q/64) |
| 382 | # - SM100 (Blackwell head64/head64x2): num_sm_parts = num_sms / s_q |
| 383 | # - SM100 (Blackwell head128): num_sm_parts = num_sms / s_q / 2 |
| 384 | # For max buffer size, use s_q = 1 (the case that produces largest output) |
| 385 | # Use padded head count since that's what will be passed to the kernel |
| 386 | h_q = self.fp8_decode_padded_heads |
| 387 | if current_platform.is_device_capability_family(100): |
| 388 | # SM100 head64 or head64x2 uses full SM count |
| 389 | max_num_sm_parts = sm_count |
| 390 | else: |
| 391 | # SM90 uses h_q/64 divisor |
| 392 | max_num_sm_parts = sm_count // max(1, h_q // 64) |
| 393 | self.tile_scheduler_metadata_buffer = torch.empty( |
| 394 | # TileSchedulerMetaDataSize = 8 |
| 395 | # see: FlashMLA/csrc/params.h |
| 396 | (max_num_sm_parts, 8), |
| 397 | dtype=torch.int32, |
| 398 | device=device, |
| 399 | ) |
| 400 | # Sized for per-request batching (num_decodes + 1) |
| 401 | self.num_splits_buffer = torch.empty( |
| 402 | (max_num_seqs + 1,), |
| 403 | dtype=torch.int32, |
| 404 | device=device, |
| 405 | ) |
| 406 | |
| 407 | # PCP+DCP attends prefill rows over the DCP-gathered KV, which only the |
| 408 | # separate prefill/decode path can do. |
| 409 | self.pcp_dcp_kv_gather = self.use_pcp and self.dcp_world_size > 1 |
| 410 | self.fp8_use_mixed_batch = ( |
| 411 | self.num_heads < MIN_HEADS_FOR_BF16_PREFILL |
| 412 | and not self.use_hisparse |
| 413 | and not self.pcp_dcp_kv_gather |
| 414 | ) |
| 415 | |
| 416 | if parallel_config.decode_context_parallel_size > 1: |
| 417 | if parallel_config.dcp_comm_backend != "ag_rs": |
| 418 | raise NotImplementedError( |
| 419 | "DCP for FlashMLA sparse is only validated with the " |
| 420 | "default 'ag_rs' DCP comm backend; got " |
| 421 | f"'{parallel_config.dcp_comm_backend}'" |
| 422 | ) |
| 423 | if self.pcp_dcp_kv_gather and cache_config.cache_dtype != "fp8_ds_mla": |
| 424 | raise NotImplementedError( |
| 425 | "PCP+DCP sparse prefill gathers the KV through the fp8_ds_mla " |
| 426 | f"upconvert; got a {cache_config.cache_dtype} cache" |
| 427 | ) |
| 428 | if not self.fp8_use_mixed_batch and not self.pcp_dcp_kv_gather: |
| 429 | raise NotImplementedError( |
| 430 | "DCP for FlashMLA sparse is only supported on the " |
| 431 | "mixed-batch fp8 path (num_heads < " |
| 432 | f"{MIN_HEADS_FOR_BF16_PREFILL}); the separate " |
| 433 | "prefill/decode path returns the LSE for decode tokens " |
| 434 | "only, while the DCP merge needs it for every token" |
| 435 | ) |
| 436 | # Head padding (and the tile-scheduler metadata sized from it) is |
| 437 | # computed from the local head count, but the kernel runs on the |
| 438 | # DCP-gathered heads. |
| 439 | if self.use_pcp: |
| 440 | gathered_num_heads = ( |
| 441 | self.num_heads * parallel_config.tensor_parallel_size |
| 442 | if self.dcp_world_size > self.pcp_world_size |
| 443 | else self.num_heads |
| 444 | ) |
| 445 | else: |
| 446 | gathered_num_heads = ( |
| 447 | self.num_heads * parallel_config.decode_context_parallel_size |
| 448 | ) |
| 449 | gathered_padded_heads = FlashMLASparseImpl._compute_fp8_decode_padded_heads( |
| 450 | gathered_num_heads |
| 451 | ) |
| 452 | if self.fp8_decode_padded_heads != gathered_padded_heads: |
| 453 | raise NotImplementedError( |
| 454 | "DCP for FlashMLA sparse requires the local and " |
| 455 | "DCP-gathered head counts to pad to the same fp8 decode " |
| 456 | f"kernel envelope; got {self.num_heads} local heads " |
| 457 | f"(pad to {self.fp8_decode_padded_heads}) vs " |
| 458 | f"{gathered_num_heads} gathered heads (pad to " |
| 459 | f"{gathered_padded_heads})" |
| 460 | ) |
| 461 | |
| 462 | def _build_fp8_mixed_decode_prefill( |
| 463 | self, |
| 464 | ) -> FlashMLASparseMetadata.FP8KernelMetadata: |
| 465 | """Build FP8 metadata treating MQA tokens as one batch. |
| 466 | |
| 467 | The scheduler initializes lazily from the runtime query shape, which may |
| 468 | be the full batch or only decodes when prefills use dense MHA. This avoids |
| 469 | the BF16 prefill kernel's head-padding overhead at high TP. |
| 470 | """ |
| 471 | scheduler_metadata, _ = get_mla_metadata() |
| 472 | return FlashMLASparseMetadata.FP8KernelMetadata( |
| 473 | scheduler_metadata=scheduler_metadata, |
| 474 | cache_lens=self.max_model_len_tensor[:1], |
| 475 | dummy_block_table=self.dummy_block_table[:1], |
| 476 | ) |
| 477 | |
| 478 | def _build_fp8_separate_prefill_decode( |
| 479 | self, |
| 480 | common_attn_metadata: CommonAttentionMetadata, |
| 481 | metadata: FlashMLASparseMetadata, |
| 482 | ) -> "FlashMLASparseMetadata.FP8SeparatePrefillDecode": |
| 483 | num_tokens = common_attn_metadata.num_actual_tokens |
| 484 | |
| 485 | (num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens) = ( |
| 486 | metadata.num_decodes, |
| 487 | metadata.num_prefills, |
| 488 | metadata.num_decode_tokens, |
| 489 | num_tokens - metadata.num_decode_tokens, |
| 490 | ) |
| 491 | |
| 492 | FP8Meta = FlashMLASparseMetadata.FP8SeparatePrefillDecode |
| 493 | # PCP decode sharding can leave a rank with only its collective-padding |
| 494 | # row when the global decode batch is smaller than the PCP world size. |
| 495 | # The row has no actual query tokens and must not be interpreted as a |
| 496 | # zero-length decode request by the sparse FP8 metadata builder. |
| 497 | if num_tokens == 0: |
| 498 | return FP8Meta() |
| 499 | |
| 500 | decode_query_len = 0 |
| 501 | active_num_decodes = num_decodes |
| 502 | if num_decodes > 0: |
| 503 | query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu |
| 504 | decode_query_len = (query_start_loc_cpu[1] - query_start_loc_cpu[0]).item() |
| 505 | assert decode_query_len > 0 |
| 506 | active_num_decodes = num_decode_tokens // decode_query_len |
| 507 | assert active_num_decodes * decode_query_len == num_decode_tokens |
| 508 | |
| 509 | fp8_metadata = FP8Meta( |
| 510 | num_decodes=active_num_decodes, |
| 511 | num_prefills=num_prefills, |
| 512 | num_decode_tokens=num_decode_tokens, |
| 513 | num_prefill_tokens=num_prefill_tokens, |
| 514 | ) |
| 515 | |
| 516 | # Extract prefill sequence lengths (context + query, not just query) |
| 517 | # Decode requests come first in the batch, prefill requests follow |
| 518 | prefill_request_id = None |
| 519 | prefill_workspace_starts = None |
| 520 | prefill_chunks = None |
| 521 | |
| 522 | # For pure decode batches, prefill_request_id will be None |
| 523 | # For mixed batches, it will have -1 for decode and request_id for prefill |
| 524 | if num_prefills > 0: |
| 525 | # Upper bound is exact for prefill rows (the `[num_decodes:]` |
| 526 | # slice below), so no D2H sync is needed. |
| 527 | seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound |
| 528 | assert seq_lens_cpu is not None |
| 529 | query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu |
| 530 | |
| 531 | prefill_seq_lens_cpu = seq_lens_cpu[num_decodes:] |
| 532 | # Chunk prefill requests to fit within workspace size |
| 533 | max_prefill_buffer_size = get_prefill_workspace_size( |
| 534 | self.vllm_config.model_config.max_model_len |
| 535 | ) |
| 536 | # Workspace rows per entry, and the prefill rows each entry covers. |
| 537 | workspace_rows = prefill_seq_lens_cpu |
| 538 | row_bounds = np.arange(num_prefills + 1) |
| 539 | prefill_block_table = common_attn_metadata.block_table_tensor[num_decodes:] |
| 540 | prefill_seq_lens = common_attn_metadata.seq_lens[num_decodes:] |
| 541 | if self.pcp_dcp_kv_gather: |
| 542 | # One entry per request, holding this rank's KV shard: the |
| 543 | # all-gather of every rank's shard is the whole context. |
| 544 | # A dummy batch bypasses the PCP manager: one row per request, |
| 545 | # so the row's own extent is the request's. |
| 546 | row_req_idx = common_attn_metadata.req_idx |
| 547 | if row_req_idx is None: |
| 548 | row_req_idx = np.arange(common_attn_metadata.num_reqs) |
| 549 | shard_rows_cpu = common_attn_metadata.dcp_local_seq_lens_cpu_upper_bound |
| 550 | if shard_rows_cpu is None: |
| 551 | shard_rows_cpu = get_dcp_local_seq_lens( |
| 552 | seq_lens_cpu, |
| 553 | self.dcp_world_size, |
| 554 | 0, |
| 555 | self.cp_kv_cache_interleave_size, |
| 556 | ) |
| 557 | row_bounds, rows_per_rank = gathered_prefill_shards( |
| 558 | row_req_idx[num_decodes:], shard_rows_cpu[num_decodes:].numpy() |
| 559 | ) |
| 560 | workspace_rows = torch.from_numpy(rows_per_rank.astype(np.int32)) |
| 561 | max_prefill_buffer_size //= self.dcp_world_size |
| 562 | entry_rows = async_tensor_h2d(row_bounds[:-1], device=self.device) |
| 563 | prefill_block_table = prefill_block_table.index_select(0, entry_rows) |
| 564 | prefill_seq_lens = prefill_seq_lens.index_select(0, entry_rows) |
| 565 | num_entries = len(workspace_rows) |
| 566 | |
| 567 | # Build prefill_request_id: -1 for decode, request index for |
| 568 | # prefill. This enables a single |
| 569 | # convert_logical_index_to_physical_index call for all tokens |
| 570 | prefill_request_id = torch.full( |
| 571 | (num_tokens,), -1, dtype=torch.int32, device=self.device |
| 572 | ) |
| 573 | # Map prefill tokens to their request IDs (0, 1, 2, ...) |
| 574 | for req_idx in range(num_entries): |
| 575 | # Get query token range for this prefill request |
| 576 | req_query_start = query_start_loc_cpu[ |
| 577 | num_decodes + int(row_bounds[req_idx]) |
| 578 | ] |
| 579 | req_query_end = query_start_loc_cpu[ |
| 580 | num_decodes + int(row_bounds[req_idx + 1]) |
| 581 | ] |
| 582 | prefill_request_id[req_query_start:req_query_end] = req_idx |
| 583 | |
| 584 | # will be adjusted by chunk loop |
| 585 | prefill_workspace_starts_cpu = torch.zeros( |
| 586 | num_entries, dtype=torch.int32, pin_memory=True |
| 587 | ) |
| 588 | prefill_workspace_starts_cpu[1:] = torch.cumsum(workspace_rows[:-1], dim=0) |
| 589 | # populated by non-blocking copy after prefill_workspace_starts_cpu is |
| 590 | # updated by each chunk |
| 591 | prefill_workspace_starts = torch.empty( |
| 592 | num_entries, dtype=torch.int32, device=self.device |
| 593 | ) |
| 594 | |
| 595 | chunk_bounds = split_prefill_chunks(workspace_rows, max_prefill_buffer_size) |
| 596 | |
| 597 | prefill_chunks = [] |
| 598 | for chunk_start, chunk_end in chunk_bounds: |
| 599 | # Adjust workspace_starts in-place per chunk to be |
| 600 | # 0-indexed within each chunk |
| 601 | # Example: seq_lens=[10,15,20,5], chunks=[[0,2],[2,4]] |
| 602 | # Initial: workspace_starts=[0,10,25,45] |
| 603 | # After: workspace_starts=[0,10,0,20] |
| 604 | # (chunk 0 starts at 0, chunk 1 starts at 0) |
| 605 | offset = prefill_workspace_starts_cpu[chunk_start].item() |
| 606 | prefill_workspace_starts_cpu[chunk_start:chunk_end] -= offset |
| 607 | |
| 608 | chunk_tot_seqlen = workspace_rows[chunk_start:chunk_end].sum() |
| 609 | token_start = query_start_loc_cpu[ |
| 610 | num_decodes + int(row_bounds[chunk_start]) |
| 611 | ].item() |
| 612 | token_end = query_start_loc_cpu[ |
| 613 | num_decodes + int(row_bounds[chunk_end]) |
| 614 | ].item() |
| 615 | tokens_slice = slice(token_start, token_end) |
| 616 | |
| 617 | # Create chunk view of gpu tensor |
| 618 | chunk_workspace_starts = prefill_workspace_starts[chunk_start:chunk_end] |
| 619 | chunk_block_table = prefill_block_table[chunk_start:chunk_end] |
| 620 | chunk_seq_lens = prefill_seq_lens[chunk_start:chunk_end] |
| 621 | |
| 622 | prefill_chunks.append( |
| 623 | FP8Meta.Prefill.Chunk( |
| 624 | tokens_slice=tokens_slice, |
| 625 | block_table=chunk_block_table, |
| 626 | req_start_idx=chunk_start, |
| 627 | workspace_starts=chunk_workspace_starts, |
| 628 | chunk_tot_seqlen=chunk_tot_seqlen, |
| 629 | seq_lens=chunk_seq_lens, |
| 630 | ) |
| 631 | ) |
| 632 | |
| 633 | prefill_workspace_starts.copy_( |
| 634 | prefill_workspace_starts_cpu, non_blocking=True |
| 635 | ) |
| 636 | |
| 637 | fp8_metadata.prefill = FP8Meta.Prefill( |
| 638 | request_ids=prefill_request_id, |
| 639 | workspace_starts=prefill_workspace_starts, |
| 640 | chunks=prefill_chunks, |
| 641 | ) |
| 642 | |
| 643 | if num_decodes > 0: |
| 644 | if self.pcp_dcp_kv_gather: |
| 645 | # Decode rows take the mixed-batch kernel call, which returns |
| 646 | # the LSE the DCP merge needs. |
| 647 | kernel_meta = self._build_fp8_mixed_decode_prefill() |
| 648 | else: |
| 649 | # Use padded head count since that's what the kernel will see |
| 650 | scheduler_metadata, _ = get_mla_metadata() |
| 651 | |
| 652 | kernel_meta = FlashMLASparseMetadata.FP8KernelMetadata( |
| 653 | scheduler_metadata=scheduler_metadata, |
| 654 | dummy_block_table=self.dummy_block_table[:active_num_decodes], |
| 655 | cache_lens=self.max_model_len_tensor[:active_num_decodes], |
| 656 | ) |
| 657 | fp8_metadata.decode = FP8Meta.Decode( |
| 658 | seq_lens=common_attn_metadata.seq_lens[:active_num_decodes], |
| 659 | kernel_metadata=kernel_meta, |
| 660 | decode_query_len=decode_query_len, |
| 661 | ) |
| 662 | |
| 663 | return fp8_metadata |
| 664 | |
| 665 | def build( |
| 666 | self, |
| 667 | common_prefix_len: int, |
| 668 | common_attn_metadata: CommonAttentionMetadata, |
| 669 | fast_build: bool = False, |
| 670 | ) -> FlashMLASparseMetadata: |
| 671 | metadata = super().build(common_prefix_len, common_attn_metadata, fast_build) |
| 672 | |
| 673 | metadata.fp8_use_mixed_batch = self.fp8_use_mixed_batch |
| 674 | if self.use_fp8_kv_cache: |
| 675 | if self.fp8_use_mixed_batch: |
| 676 | metadata.fp8_extra_metadata = self._build_fp8_mixed_decode_prefill() |
| 677 | else: |
| 678 | metadata.fp8_extra_metadata = self._build_fp8_separate_prefill_decode( |
| 679 | common_attn_metadata, metadata |
| 680 | ) |
| 681 | |
| 682 | return metadata |
| 683 | |
| 684 | |
| 685 | class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]): |
| 686 | can_return_lse_for_decode: bool = True |
| 687 | supports_dcp: bool = True |
| 688 | supports_mtp_with_cp_non_trivial_interleave_size: bool = True |
| 689 | |
| 690 | @staticmethod |
| 691 | def _compute_fp8_decode_padded_heads(num_heads: int) -> int: |
| 692 | # FP8 decode kernel only supports h_q = 64 or 128 |
| 693 | # Compute padded head count for decode |
| 694 | return 64 if num_heads <= 64 else 128 |
| 695 | |
| 696 | def __init__( |
| 697 | self, |
| 698 | num_heads: int, |
| 699 | head_size: int, |
| 700 | scale: float, |
| 701 | num_kv_heads: int, |
| 702 | alibi_slopes: list[float] | None, |
| 703 | sliding_window: int | None, |
| 704 | kv_cache_dtype: str, |
| 705 | logits_soft_cap: float | None, |
| 706 | attn_type: str, |
| 707 | kv_sharing_target_layer_name: str | None, |
| 708 | # MLA Specific Arguments |
| 709 | topk_indices_buffer: torch.Tensor | None = None, |
| 710 | indexer: "Indexer | None" = None, |
| 711 | **mla_args, |
| 712 | ) -> None: |
| 713 | super().__init__( |
| 714 | num_heads, |
| 715 | head_size, |
| 716 | scale, |
| 717 | num_kv_heads, |
| 718 | alibi_slopes, |
| 719 | sliding_window, |
| 720 | kv_cache_dtype, |
| 721 | logits_soft_cap, |
| 722 | attn_type, |
| 723 | kv_sharing_target_layer_name, |
| 724 | indexer=indexer, |
| 725 | topk_indices_buffer=topk_indices_buffer, |
| 726 | **mla_args, |
| 727 | ) |
| 728 | self.softmax_scale = scale |
| 729 | # Prefill BF16 kernel requires 64 on Hopper, 128 on Blackwell |
| 730 | self.prefill_padding = ( |
| 731 | 128 if current_platform.is_device_capability_family(100) else 64 |
| 732 | ) |
| 733 | self.fp8_decode_padded_heads = self._compute_fp8_decode_padded_heads(num_heads) |
| 734 | |
| 735 | vllm_config = get_current_vllm_config() |
| 736 | max_tokens = vllm_config.scheduler_config.max_num_batched_tokens |
| 737 | q_concat_heads = num_heads |
| 738 | if not is_quantized_kv_cache(kv_cache_dtype): |
| 739 | q_concat_heads = ( |
| 740 | (num_heads + self.prefill_padding - 1) |
| 741 | // self.prefill_padding |
| 742 | * self.prefill_padding |
| 743 | ) |
| 744 | q_concat_shape = (max_tokens, q_concat_heads, head_size) |
| 745 | if is_quantized_kv_cache(kv_cache_dtype): |
| 746 | assert kv_cache_dtype in QUANTIZED_DS_MLA_CACHE_FORMATS, ( |
| 747 | "FlashMLA Sparse Attention backend only supports the " |
| 748 | f"{sorted(QUANTIZED_DS_MLA_CACHE_FORMATS)} quantized kv-cache " |
| 749 | f"dtypes, got {kv_cache_dtype}" |
| 750 | ) |
| 751 | |
| 752 | if self.need_to_return_lse_for_decode and not is_quantized_kv_cache( |
| 753 | kv_cache_dtype |
| 754 | ): |
| 755 | raise NotImplementedError( |
| 756 | "DCP for FlashMLA sparse requires an fp8_ds_mla kv-cache; " |
| 757 | "the bf16 sparse path is not supported under DCP." |
| 758 | ) |
| 759 | |
| 760 | self.pcp_dcp_kv_gather = False |
| 761 | self.workspace_specs: list[tuple[tuple[int, ...], torch.dtype]] = [ |
| 762 | (q_concat_shape, torch.bfloat16) |
| 763 | ] |
| 764 | if kv_cache_dtype in QUANTIZED_DS_MLA_CACHE_FORMATS: |
| 765 | # Reserve workspace during initialization |
| 766 | assert vllm_config is not None and vllm_config.model_config is not None |
| 767 | prefill_workspace_size = get_prefill_workspace_size( |
| 768 | vllm_config.model_config.max_model_len |
| 769 | ) |
| 770 | parallel_config = vllm_config.parallel_config |
| 771 | self.pcp_dcp_kv_gather = ( |
| 772 | parallel_config.prefill_context_parallel_size > 1 |
| 773 | and parallel_config.decode_context_parallel_size > 1 |
| 774 | ) |
| 775 | shard_rows = prefill_workspace_size |
| 776 | if self.pcp_dcp_kv_gather: |
| 777 | # PCP+DCP upconverts this rank's KV shard, then all-gathers the |
| 778 | # shards into a workspace of the full prefill size. |
| 779 | shard_rows //= parallel_config.decode_context_parallel_size |
| 780 | self.prefill_workspace_shape = (shard_rows, head_size) |
| 781 | self.workspace_specs.append((self.prefill_workspace_shape, torch.bfloat16)) |
| 782 | if self.pcp_dcp_kv_gather: |
| 783 | self.workspace_specs.append( |
| 784 | ((prefill_workspace_size, head_size), torch.bfloat16) |
| 785 | ) |
| 786 | prefill_query_heads = num_heads |
| 787 | if self.pcp_dcp_kv_gather and self.dcp_world_size > self.pcp_world_size: |
| 788 | prefill_query_heads *= parallel_config.tensor_parallel_size |
| 789 | padded_prefill_query_heads = round_up( |
| 790 | prefill_query_heads, self.prefill_padding |
| 791 | ) |
| 792 | self.workspace_specs.extend( |
| 793 | (shape, torch.bfloat16) |
| 794 | for shape in ( |
| 795 | (max_tokens, padded_prefill_query_heads, head_size), |
| 796 | (max_tokens, padded_prefill_query_heads, self.kv_lora_rank), |
| 797 | (max_tokens, prefill_query_heads, self.kv_lora_rank), |
| 798 | ) |
| 799 | ) |
| 800 | # Reserve capacity without retaining views that prevent old storage |
| 801 | # from being released when another layer grows the shared workspace. |
| 802 | current_workspace_manager().get_simultaneous(*self.workspace_specs) |
| 803 | |
| 804 | def _forward_bf16_kv( |
| 805 | self, |
| 806 | q: torch.Tensor, |
| 807 | kv_c_and_k_pe_cache: torch.Tensor, |
| 808 | topk_indices: torch.Tensor, |
| 809 | attn_metadata: FlashMLASparseMetadata, |
| 810 | actual_num_heads: int, |
| 811 | ) -> tuple[torch.Tensor, torch.Tensor | None]: |
| 812 | index_group = self.index_group |
| 813 | if isinstance(index_group, HiSparseMLAIndexGroup): |
| 814 | cache = index_group.cache(self.index_group_index) |
| 815 | else: |
| 816 | cache = None |
| 817 | block_table = attn_metadata.block_table |
| 818 | # req_id_per_token covers the whole batch; slice it to the MQA tokens |
| 819 | # (q may exclude prefill tokens routed to dense MHA). |
| 820 | req_id_per_token = attn_metadata.req_id_per_token[: topk_indices.shape[0]] |
| 821 | decode_out: torch.Tensor | None = None |
| 822 | if cache is not None: |
| 823 | assert isinstance(index_group, HiSparseMLAIndexGroup) |
| 824 | num_decode_tokens = attn_metadata.num_decode_tokens |
| 825 | if num_decode_tokens > 0: |
| 826 | decode_topk, decode_lengths = ( |
| 827 | index_group.convert_logical_to_physical_topk( |
| 828 | self.index_group_index, |
| 829 | topk_indices[:num_decode_tokens], |
| 830 | attn_metadata, |
| 831 | block_stride_rows=None, |
| 832 | return_valid_counts=True, |
| 833 | ) |
| 834 | ) |
| 835 | decode_out, _ = self._bf16_flash_mla_kernel( |
| 836 | q[:num_decode_tokens], |
| 837 | index_group.physical_kv_cache(self.index_group_index), |
| 838 | decode_topk, |
| 839 | decode_lengths, |
| 840 | actual_num_heads, |
| 841 | ) |
| 842 | if num_decode_tokens == q.shape[0]: |
| 843 | return decode_out, None |
| 844 | q = q[num_decode_tokens:] |
| 845 | topk_indices = topk_indices[num_decode_tokens:] |
| 846 | kv_c_and_k_pe_cache, block_table, req_id_per_token = ( |
| 847 | index_group.stage_prefill_rows( |
| 848 | self.index_group_index, kv_c_and_k_pe_cache, attn_metadata |
| 849 | ) |
| 850 | ) |
| 851 | # Convert per-request indices to global slots (decode) or workspace offsets. |
| 852 | kv_rows, block_stride_rows = flat_kv_row_view( |
| 853 | kv_c_and_k_pe_cache, attn_metadata.block_size |
| 854 | ) |
| 855 | decode_only = ( |
| 856 | attn_metadata.num_decode_tokens |
| 857 | == attn_metadata.num_actual_tokens |
| 858 | == topk_indices.shape[0] |
| 859 | ) |
| 860 | uses_host_cache = isinstance(index_group, HiSparseMLAIndexGroup) |
| 861 | if not uses_host_cache and decode_only: |
| 862 | topk_indices, topk_length = self._convert_logical_to_physical_topk( |
| 863 | topk_indices, |
| 864 | attn_metadata, |
| 865 | block_stride_rows=block_stride_rows, |
| 866 | return_valid_counts=True, |
| 867 | ) |
| 868 | else: |
| 869 | topk_indices, topk_length = triton_convert_req_index_to_global_index( |
| 870 | req_id_per_token, |
| 871 | block_table, |
| 872 | topk_indices, |
| 873 | BLOCK_SIZE=attn_metadata.block_size, |
| 874 | BLOCK_STRIDE_ROWS=block_stride_rows, |
| 875 | NUM_TOPK_TOKENS=topk_indices.shape[1], |
| 876 | return_valid_counts=True, |
| 877 | ) |
| 878 | |
| 879 | attn_out, lse = self._bf16_flash_mla_kernel( |
| 880 | q, |
| 881 | kv_rows, |
| 882 | topk_indices, |
| 883 | topk_length, |
| 884 | actual_num_heads, |
| 885 | ) |
| 886 | if decode_out is None: |
| 887 | return attn_out, lse |
| 888 | return torch.cat([decode_out, attn_out], dim=0), None |
| 889 | |
| 890 | def _gather_prefill_chunk( |
| 891 | self, |
| 892 | chunk: "FlashMLASparseMetadata.FP8SeparatePrefillDecode.Prefill.Chunk", |
| 893 | shard: torch.Tensor, |
| 894 | topk_indices: torch.Tensor, |
| 895 | attn_metadata: FlashMLASparseMetadata, |
| 896 | prefill_meta: "FlashMLASparseMetadata.FP8SeparatePrefillDecode.Prefill", |
| 897 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| 898 | """All-gather this rank's upconverted KV shard so the chunk's rows attend |
| 899 | the whole context, and map their top-k onto the rank-major result.""" |
| 900 | shard_rows = int(chunk.chunk_tot_seqlen) |
| 901 | _, _, gathered_kv_workspace, *_ = current_workspace_manager().get_simultaneous( |
| 902 | *self.workspace_specs |
| 903 | ) |
| 904 | gathered_kv = gathered_kv_workspace[: self.dcp_world_size * shard_rows] |
| 905 | dist.all_gather_into_tensor( |
| 906 | gathered_kv, shard, group=get_dcp_group().device_group |
| 907 | ) |
| 908 | |
| 909 | tokens = chunk.tokens_slice |
| 910 | topk_indices, topk_length = triton_convert_req_index_to_global_index( |
| 911 | attn_metadata.req_id_per_token[tokens], |
| 912 | attn_metadata.block_table, |
| 913 | topk_indices, |
| 914 | BLOCK_SIZE=attn_metadata.block_size, |
| 915 | NUM_TOPK_TOKENS=topk_indices.shape[1], |
| 916 | HAS_PREFILL_WORKSPACE=True, |
| 917 | prefill_workspace_request_ids=prefill_meta.request_ids[tokens], |
| 918 | prefill_workspace_starts=prefill_meta.workspace_starts, |
| 919 | prefill_workspace_rank_stride=shard_rows, |
| 920 | dcp_size=self.dcp_world_size, |
| 921 | dcp_rank=self.dcp_rank, |
| 922 | cp_kv_cache_interleave_size=attn_metadata.cp_kv_cache_interleave_size, |
| 923 | return_valid_counts=True, |
| 924 | ) |
| 925 | return gathered_kv, topk_indices, topk_length |
| 926 | |
| 927 | def _forward_fp8_kv_separate_prefill_decode( |
| 928 | self, |
| 929 | q: torch.Tensor, |
| 930 | kv_c_and_k_pe_cache: torch.Tensor, |
| 931 | topk_indices: torch.Tensor, |
| 932 | attn_metadata: FlashMLASparseMetadata, |
| 933 | ) -> tuple[torch.Tensor, torch.Tensor | None]: |
| 934 | """The lse covers the decode rows and is only returned under PCP+DCP.""" |
| 935 | fp8_metadata = attn_metadata.fp8_extra_metadata |
| 936 | assert isinstance(fp8_metadata, FlashMLASparseMetadata.FP8SeparatePrefillDecode) |
| 937 | num_decodes = fp8_metadata.num_decodes |
| 938 | num_mqa_tokens = q.shape[0] |
| 939 | num_decode_tokens = fp8_metadata.num_decode_tokens |
| 940 | num_prefill_tokens = num_mqa_tokens - num_decode_tokens |
| 941 | assert num_prefill_tokens in (0, fp8_metadata.num_prefill_tokens), ( |
| 942 | "FP8 sparse MLA expects either the decode subset or the full batch" |
| 943 | ) |
| 944 | |
| 945 | decode_topk: torch.Tensor | None = None |
| 946 | index_group = self.index_group |
| 947 | uses_host_cache = isinstance(index_group, HiSparseMLAIndexGroup) |
| 948 | if uses_host_cache and num_decode_tokens > 0: |
| 949 | decode_topk = topk_indices[:num_decode_tokens] |
| 950 | |
| 951 | prefill_ready = None |
| 952 | if num_prefill_tokens > 0 and uses_host_cache: |
| 953 | assert fp8_metadata.prefill is not None |
| 954 | first_chunk = fp8_metadata.prefill.chunks[0] |
| 955 | assert isinstance(index_group, HiSparseMLAIndexGroup) |
| 956 | _, prefill_bf16_workspace, *_ = ( |
| 957 | current_workspace_manager().get_simultaneous(*self.workspace_specs) |
| 958 | ) |
| 959 | prefill_ready = index_group.gather_fp8_prefill( |
| 960 | self.index_group_index, |
| 961 | kv_c_and_k_pe_cache, |
| 962 | prefill_bf16_workspace[: first_chunk.chunk_tot_seqlen], |
| 963 | first_chunk.block_table, |
| 964 | first_chunk.workspace_starts, |
| 965 | len(first_chunk.block_table), |
| 966 | attn_metadata, |
| 967 | first_chunk.req_start_idx, |
| 968 | ) |
| 969 | |
| 970 | prefill_request_ids = None |
| 971 | prefill_workspace_starts = None |
| 972 | has_prefill_workspace = False |
| 973 | if num_prefill_tokens > 0: |
| 974 | assert fp8_metadata.prefill is not None |
| 975 | prefill_request_ids = fp8_metadata.prefill.request_ids |
| 976 | prefill_workspace_starts = fp8_metadata.prefill.workspace_starts |
| 977 | has_prefill_workspace = True |
| 978 | |
| 979 | # Convert per-request indices to global slots (decode) or workspace |
| 980 | # offsets (prefill). |
| 981 | # For FP8 cache: prefill uses workspace mapping (upconverted to BF16) |
| 982 | # For BF16 cache: always use global cache slots (no workspace) |
| 983 | # prefill_workspace_starts has been adjusted in-place per chunk so |
| 984 | # prefill indices automatically come out chunk-local |
| 985 | # Under PCP+DCP the decode rows are converted by the mixed-batch call |
| 986 | # and the prefill rows per chunk, against the gathered workspace. |
| 987 | lse: torch.Tensor | None = None |
| 988 | topk_length: torch.Tensor | None = None |
| 989 | if self.pcp_dcp_kv_gather: |
| 990 | pass |
| 991 | elif num_prefill_tokens == 0 and not uses_host_cache: |
| 992 | topk_indices, topk_length = self._convert_logical_to_physical_topk( |
| 993 | topk_indices, |
| 994 | attn_metadata, |
| 995 | block_stride_rows=None, |
| 996 | return_valid_counts=True, |
| 997 | ) |
| 998 | elif num_prefill_tokens > 0: |
| 999 | topk_indices, topk_length = triton_convert_req_index_to_global_index( |
| 1000 | attn_metadata.req_id_per_token[: topk_indices.shape[0]], |
| 1001 | attn_metadata.block_table, |
| 1002 | topk_indices, |
| 1003 | BLOCK_SIZE=attn_metadata.block_size, |
| 1004 | NUM_TOPK_TOKENS=topk_indices.shape[1], |
| 1005 | HAS_PREFILL_WORKSPACE=has_prefill_workspace, |
| 1006 | prefill_workspace_request_ids=prefill_request_ids, |
| 1007 | prefill_workspace_starts=prefill_workspace_starts, |
| 1008 | return_valid_counts=True, |
| 1009 | ) |
| 1010 | |
| 1011 | fp8_metadata = attn_metadata.fp8_extra_metadata |
| 1012 | assert isinstance(fp8_metadata, FlashMLASparseMetadata.FP8SeparatePrefillDecode) |
| 1013 | |
| 1014 | def _fp8_decode( |
| 1015 | q: torch.Tensor, |
| 1016 | topk_indices: torch.Tensor, |
| 1017 | ) -> torch.Tensor: |
| 1018 | assert fp8_metadata.decode is not None |
| 1019 | if self.pcp_dcp_kv_gather: |
| 1020 | # Attend this rank's KV shard; the LSE feeds the DCP merge. |
| 1021 | nonlocal lse |
| 1022 | out, lse = self._forward_fp8_kv_mixed_batch( |
| 1023 | q, |
| 1024 | kv_c_and_k_pe_cache, |
| 1025 | topk_indices, |
| 1026 | attn_metadata, |
| 1027 | kernel_metadata=fp8_metadata.decode.kernel_metadata, |
| 1028 | ) |
| 1029 | return out |
| 1030 | if uses_host_cache: |
| 1031 | return self._host_backed_fp8_decode( |
| 1032 | q, |
| 1033 | topk_indices, |
| 1034 | attn_metadata, |
| 1035 | fp8_metadata.decode.kernel_metadata, |
| 1036 | num_decodes, |
| 1037 | ) |
| 1038 | # Reshape q: (num_decode_tokens, num_heads, head_dim) |
| 1039 | # -> (num_decodes, seq_len, num_heads, head_dim) |
| 1040 | q = reshape_query_for_spec_decode(q, num_decodes) |
| 1041 | seq_len = q.shape[1] |
| 1042 | # Reshape topk_indices: (num_decode_tokens, topk) |
| 1043 | # -> (num_decodes, seq_len, topk) |
| 1044 | topk_indices = topk_indices.view(num_decodes, seq_len, -1) |
| 1045 | attn_out, _ = self._fp8_flash_mla_kernel( |
| 1046 | q=q, |
| 1047 | kv_c_and_k_pe_cache=kv_c_and_k_pe_cache, |
| 1048 | topk_indices=topk_indices, |
| 1049 | kernel_metadata=fp8_metadata.decode.kernel_metadata, |
| 1050 | ) |
| 1051 | # Reshape output: (num_decodes, seq_len, num_heads, head_dim_v) |
| 1052 | # -> (num_decode_tokens, num_heads, head_dim_v) |
| 1053 | return reshape_attn_output_for_spec_decode(attn_out) |
| 1054 | |
| 1055 | # Pure decode: direct call without allocation |
| 1056 | if num_decode_tokens > 0 and num_prefill_tokens == 0: |
| 1057 | assert fp8_metadata.decode is not None |
| 1058 | attn_out = _fp8_decode( |
| 1059 | q, decode_topk if decode_topk is not None else topk_indices |
| 1060 | ) |
| 1061 | else: |
| 1062 | _, prefill_bf16_workspace, *_, padded_out, compact_out = ( |
| 1063 | current_workspace_manager().get_simultaneous(*self.workspace_specs) |
| 1064 | ) |
| 1065 | attn_out = compact_out[:num_mqa_tokens] |
| 1066 | |
| 1067 | if num_decode_tokens > 0: |
| 1068 | attn_out[:num_decode_tokens] = _fp8_decode( |
| 1069 | q[:num_decode_tokens], |
| 1070 | decode_topk |
| 1071 | if decode_topk is not None |
| 1072 | else topk_indices[:num_decode_tokens], |
| 1073 | ) |
| 1074 | |
| 1075 | assert fp8_metadata.prefill is not None |
| 1076 | for chunk_index, chunk in enumerate(fp8_metadata.prefill.chunks): |
| 1077 | chunk_workspace = prefill_bf16_workspace[: chunk.chunk_tot_seqlen] |
| 1078 | if uses_host_cache and chunk_index > 0: |
| 1079 | assert isinstance(index_group, HiSparseMLAIndexGroup) |
| 1080 | prefill_ready = index_group.gather_fp8_prefill( |
| 1081 | self.index_group_index, |
| 1082 | kv_c_and_k_pe_cache, |
| 1083 | chunk_workspace, |
| 1084 | chunk.block_table, |
| 1085 | chunk.workspace_starts, |
| 1086 | len(chunk.block_table), |
| 1087 | attn_metadata, |
| 1088 | chunk.req_start_idx, |
| 1089 | ) |
| 1090 | if uses_host_cache: |
| 1091 | assert prefill_ready is not None |
| 1092 | current_stream().wait_event(prefill_ready) |
| 1093 | elif self.kv_cache_dtype == "fp8_ds_mla": |
| 1094 | ops.cp_gather_and_upconvert_fp8_kv_cache( |
| 1095 | kv_c_and_k_pe_cache, |
| 1096 | chunk_workspace, |
| 1097 | chunk.block_table, |
| 1098 | chunk.workspace_starts, |
| 1099 | len(chunk.block_table), |
| 1100 | ) |
| 1101 | else: |
| 1102 | ops.cp_gather_and_upconvert_nvfp4_kv_cache( |
| 1103 | kv_c_and_k_pe_cache.view(torch.uint8), |
| 1104 | chunk_workspace, |
| 1105 | chunk.block_table, |
| 1106 | chunk.workspace_starts, |
| 1107 | len(chunk.block_table), |
| 1108 | ) |
| 1109 | |
| 1110 | chunk_q = q[chunk.tokens_slice] |
| 1111 | if self.pcp_dcp_kv_gather: |
| 1112 | chunk_workspace, chunk_topk_indices_workspace, chunk_topk_length = ( |
| 1113 | self._gather_prefill_chunk( |
| 1114 | chunk, |
| 1115 | chunk_workspace, |
| 1116 | topk_indices[chunk.tokens_slice], |
| 1117 | attn_metadata, |
| 1118 | fp8_metadata.prefill, |
| 1119 | ) |
| 1120 | ) |
| 1121 | else: |
| 1122 | chunk_topk_indices_workspace = topk_indices[chunk.tokens_slice] |
| 1123 | assert topk_length is not None |
| 1124 | chunk_topk_length = topk_length[chunk.tokens_slice] |
| 1125 | |
| 1126 | chunk_out, chunk_lse = self._bf16_flash_mla_kernel( |
| 1127 | chunk_q, |
| 1128 | chunk_workspace, |
| 1129 | chunk_topk_indices_workspace, |
| 1130 | chunk_topk_length, |
| 1131 | out=padded_out[chunk.tokens_slice], |
| 1132 | ) |
| 1133 | attn_out[chunk.tokens_slice].copy_(chunk_out) |
| 1134 | del chunk_out, chunk_lse |
| 1135 | |
| 1136 | if self.pcp_dcp_kv_gather and lse is None: |
| 1137 | # No decode rows: the DCP merge still expects an LSE, of no rows. |
| 1138 | lse = q.new_empty((0, q.shape[1]), dtype=torch.float32) |
| 1139 | return attn_out, lse |
| 1140 | |
| 1141 | def _forward_fp8_kv_mixed_batch( |
| 1142 | self, |
| 1143 | q: torch.Tensor, |
| 1144 | kv_c_and_k_pe_cache: torch.Tensor, |
| 1145 | topk_indices: torch.Tensor, |
| 1146 | attn_metadata: FlashMLASparseMetadata, |
| 1147 | kernel_metadata: "FlashMLASparseMetadata.FP8KernelMetadata | None" = None, |
| 1148 | ) -> tuple[torch.Tensor, torch.Tensor | None]: |
| 1149 | """Mixed batch FP8 forward path that treats all tokens as one batch. |
| 1150 | |
| 1151 | This is equivalent to main branch's approach and avoids the BF16 |
| 1152 | prefill kernel which has head padding overhead when num_heads is small. |
| 1153 | Used when use_mixed_batch is True. |
| 1154 | |
| 1155 | The lse is only returned when DCP needs it, otherwise None. |
| 1156 | """ |
| 1157 | fp8_metadata = kernel_metadata or attn_metadata.fp8_extra_metadata |
| 1158 | assert isinstance(fp8_metadata, FlashMLASparseMetadata.FP8KernelMetadata) |
| 1159 | |
| 1160 | block_table = attn_metadata.block_table |
| 1161 | # req_id_per_token covers the whole batch; slice it to the MQA tokens |
| 1162 | # (q may exclude prefill tokens routed to dense MHA). |
| 1163 | req_id_per_token = attn_metadata.req_id_per_token[: topk_indices.shape[0]] |
| 1164 | if self.dcp_world_size > 1: |
| 1165 | # The indexer emits global token ids; keep this rank's shard and |
| 1166 | # convert to local slots. compact_valid_to_front=False keeps the |
| 1167 | # scattered -1s, which the fp8 kernel masks natively and the |
| 1168 | # empty-row neutralization below relies on. req_id is sliced to |
| 1169 | # topk_indices rows (the converter grids from req_id). |
| 1170 | topk_indices = triton_filter_and_convert_dcp_index( |
| 1171 | req_id_per_token, |
| 1172 | block_table, |
| 1173 | topk_indices, |
| 1174 | dcp_size=self.dcp_world_size, |
| 1175 | dcp_rank=self.dcp_rank, |
| 1176 | cp_kv_cache_interleave_size=attn_metadata.cp_kv_cache_interleave_size, |
| 1177 | BLOCK_SIZE=attn_metadata.block_size, |
| 1178 | NUM_TOPK_TOKENS=topk_indices.shape[1], |
| 1179 | compact_valid_to_front=False, |
| 1180 | ) |
| 1181 | else: |
| 1182 | # Convert per-request indices to global slots (decode) or workspace |
| 1183 | # offsets (prefill). |
| 1184 | decode_only = ( |
| 1185 | attn_metadata.num_decode_tokens |
| 1186 | == attn_metadata.num_actual_tokens |
| 1187 | == topk_indices.shape[0] |
| 1188 | ) |
| 1189 | if decode_only: |
| 1190 | topk_indices = self._convert_logical_to_physical_topk( |
| 1191 | topk_indices, |
| 1192 | attn_metadata, |
| 1193 | block_stride_rows=None, |
| 1194 | return_valid_counts=False, |
| 1195 | ) |
| 1196 | else: |
| 1197 | topk_indices = triton_convert_req_index_to_global_index( |
| 1198 | req_id_per_token, |
| 1199 | block_table, |
| 1200 | topk_indices, |
| 1201 | BLOCK_SIZE=attn_metadata.block_size, |
| 1202 | NUM_TOPK_TOKENS=topk_indices.shape[1], |
| 1203 | ) |
| 1204 | |
| 1205 | _attn_out, _lse = self._fp8_flash_mla_kernel( |
| 1206 | q=q.unsqueeze(0), # unsqueeze to add batch_dim: (T, H, D) -> (1, T, H, D) |
| 1207 | kv_c_and_k_pe_cache=kv_c_and_k_pe_cache, |
| 1208 | topk_indices=topk_indices.unsqueeze(0), # (T, topk) -> (1, T, topk) |
| 1209 | kernel_metadata=fp8_metadata, |
| 1210 | ) |
| 1211 | # Output is (1, T, H, D_v), squeeze back to (T, H, D_v) |
| 1212 | out = _attn_out.squeeze(0) |
| 1213 | |
| 1214 | if not self.need_to_return_lse_for_decode: |
| 1215 | return out, None |
| 1216 | |
| 1217 | # Kernel LSE is (1, H, T); the DCP merge consumes (T, H). |
| 1218 | lse = _lse.squeeze(0).transpose(0, 1) |
| 1219 | # Rows where this rank owns none of the selected tokens (all indices |
| 1220 | # -1) have undefined out/lse; (0, -inf) is the identity element of the |
| 1221 | # cross-rank LSE merge, so it drops this rank from those rows. |
| 1222 | empty_rows = (topk_indices == -1).all(dim=-1) |
| 1223 | out.masked_fill_(empty_rows.view(-1, 1, 1), 0.0) |
| 1224 | lse.masked_fill_(empty_rows.view(-1, 1), float("-inf")) |
| 1225 | # The head-padding slice above can leave `out` non-contiguous, and the |
| 1226 | # merge feeds it to reduce_scatter. |
| 1227 | return out.contiguous(), lse |
| 1228 | |
| 1229 | def _fp8_flash_mla_kernel( |
| 1230 | self, |
| 1231 | q: torch.Tensor, |
| 1232 | kv_c_and_k_pe_cache: torch.Tensor, |
| 1233 | topk_indices: torch.Tensor, |
| 1234 | kernel_metadata: FlashMLASparseMetadata.FP8KernelMetadata, |
| 1235 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 1236 | # q shape: (batch, seq_len, num_heads, head_dim) |
| 1237 | actual_num_heads = q.size(2) |
| 1238 | padded_num_heads = self.fp8_decode_padded_heads |
| 1239 | |
| 1240 | # Pad query if needed (kernel only supports h_q = 64 or 128) |
| 1241 | if actual_num_heads < padded_num_heads: |
| 1242 | logger.warning_once( |
| 1243 | f"Padding num_heads from {actual_num_heads} to " |
| 1244 | f"{padded_num_heads} for FP8 sparse decode kernel" |
| 1245 | ) |
| 1246 | q_padded = q.new_zeros((q.size(0), q.size(1), padded_num_heads, q.size(3))) |
| 1247 | q_padded[:, :, :actual_num_heads, :] = q |
| 1248 | q = q_padded |
| 1249 | |
| 1250 | out, lse = flash_mla_with_kvcache( |
| 1251 | q=q, |
| 1252 | k_cache=kv_c_and_k_pe_cache.view(torch.uint8).unsqueeze(-2), |
| 1253 | block_table=kernel_metadata.dummy_block_table, |
| 1254 | head_dim_v=512, |
| 1255 | cache_seqlens=kernel_metadata.cache_lens, |
| 1256 | tile_scheduler_metadata=kernel_metadata.scheduler_metadata, |
| 1257 | is_fp8_kvcache=True, |
| 1258 | indices=topk_indices, |
| 1259 | softmax_scale=self.softmax_scale, |
| 1260 | ) |
| 1261 | |
| 1262 | # Slice output and lse back to actual head count if we padded |
| 1263 | if actual_num_heads < padded_num_heads: |
| 1264 | out = out[:, :, :actual_num_heads, :] |
| 1265 | lse = lse[:, :actual_num_heads, :] |
| 1266 | |
| 1267 | return out, lse |
| 1268 | |
| 1269 | def _host_backed_fp8_decode( |
| 1270 | self, |
| 1271 | q: torch.Tensor, |
| 1272 | topk_indices: torch.Tensor, |
| 1273 | attn_metadata: FlashMLASparseMetadata, |
| 1274 | kernel_metadata: FlashMLASparseMetadata.FP8KernelMetadata, |
| 1275 | num_decodes: int, |
| 1276 | ) -> torch.Tensor: |
| 1277 | assert isinstance(self.index_group, HiSparseMLAIndexGroup) |
| 1278 | physical_topk = self.index_group.convert_logical_to_physical_topk( |
| 1279 | self.index_group_index, |
| 1280 | topk_indices, |
| 1281 | attn_metadata, |
| 1282 | block_stride_rows=None, |
| 1283 | return_valid_counts=False, |
| 1284 | ) |
| 1285 | assert isinstance(physical_topk, torch.Tensor) |
| 1286 | q = reshape_query_for_spec_decode(q, num_decodes) |
| 1287 | physical_topk = physical_topk.view(num_decodes, q.shape[1], -1) |
| 1288 | output, _ = self._fp8_flash_mla_kernel( |
| 1289 | q=q, |
| 1290 | kv_c_and_k_pe_cache=self.index_group.physical_kv_cache( |
| 1291 | self.index_group_index |
| 1292 | ), |
| 1293 | topk_indices=physical_topk, |
| 1294 | kernel_metadata=kernel_metadata, |
| 1295 | ) |
| 1296 | return reshape_attn_output_for_spec_decode(output) |
| 1297 | |
| 1298 | def _bf16_flash_mla_kernel( |
| 1299 | self, |
| 1300 | q: torch.Tensor, |
| 1301 | kv_c_and_k_pe_cache: torch.Tensor, |
| 1302 | topk_indices: torch.Tensor, |
| 1303 | topk_length: torch.Tensor | None = None, |
| 1304 | actual_num_heads: int | None = None, |
| 1305 | out: torch.Tensor | None = None, |
| 1306 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 1307 | num_tokens = q.shape[0] |
| 1308 | kv_c_and_k_pe_cache = kv_c_and_k_pe_cache.view( |
| 1309 | -1, 1, kv_c_and_k_pe_cache.shape[-1] |
| 1310 | ) |
| 1311 | |
| 1312 | # NOTE(Chen): kernel requires num_local_head to be a multiple of |
| 1313 | # 64 on hopper and 128 on blackwell. Pad from q's head count, not |
| 1314 | # self.num_heads: under DCP the heads are all-gathered before this. |
| 1315 | if actual_num_heads is None: |
| 1316 | actual_num_heads = q.shape[1] |
| 1317 | padded_num_heads = ( |
| 1318 | (actual_num_heads + self.prefill_padding - 1) |
| 1319 | // self.prefill_padding |
| 1320 | * self.prefill_padding |
| 1321 | ) |
| 1322 | if q.shape[1] < padded_num_heads: |
| 1323 | logger.warning_once( |
| 1324 | f"Padding num_heads from {actual_num_heads} to " |
| 1325 | f"{padded_num_heads} for BF16 sparse prefill kernel" |
| 1326 | ) |
| 1327 | if out is None: |
| 1328 | q_padded = q.new_empty((q.shape[0], padded_num_heads, q.shape[2])) |
| 1329 | else: |
| 1330 | *_, q_padded, _, _ = current_workspace_manager().get_simultaneous( |
| 1331 | *self.workspace_specs |
| 1332 | ) |
| 1333 | q_padded = q_padded[:num_tokens] |
| 1334 | q_padded[:, actual_num_heads:].zero_() |
| 1335 | q_padded[:, :actual_num_heads, :] = q |
| 1336 | q = q_padded |
| 1337 | |
| 1338 | topk_indices = topk_indices.view(num_tokens, 1, -1) |
| 1339 | output, _, lse = flash_mla_sparse_fwd( |
| 1340 | q, |
| 1341 | kv_c_and_k_pe_cache, |
| 1342 | topk_indices, |
| 1343 | self.softmax_scale, |
| 1344 | topk_length=topk_length, |
| 1345 | out=out, |
| 1346 | ) |
| 1347 | |
| 1348 | output = output[:, :actual_num_heads, :] |
| 1349 | lse = lse[:, :actual_num_heads] |
| 1350 | return output, lse |
| 1351 | |
| 1352 | def forward_mqa( |
| 1353 | self, |
| 1354 | q: torch.Tensor | tuple[torch.Tensor, torch.Tensor], |
| 1355 | kv_c_and_k_pe_cache: torch.Tensor, |
| 1356 | attn_metadata: FlashMLASparseMetadata, |
| 1357 | layer: AttentionLayer, |
| 1358 | ) -> tuple[torch.Tensor, torch.Tensor | None]: |
| 1359 | # NOTE(lucas): for the sparse FlashMLA kernels the kernels want to use |
| 1360 | # MQA 576/512 approach for both prefill and decode |
| 1361 | |
| 1362 | # Concatenate q if it's a tuple (ql_nope, q_pe) |
| 1363 | actual_num_heads = self.num_heads |
| 1364 | if isinstance(q, tuple): |
| 1365 | ql_nope, q_pe = q |
| 1366 | q_concat_buffer, *_ = current_workspace_manager().get_simultaneous( |
| 1367 | *self.workspace_specs |
| 1368 | ) |
| 1369 | q = q_concat_buffer[: ql_nope.shape[0]] |
| 1370 | if q_pe.size(-1) == 0: |
| 1371 | # NoPE (GLM5Next): concat_mla_q requires rope_dim == 64, |
| 1372 | # copy directly into the head-padded buffer instead. |
| 1373 | q[:, : ql_nope.shape[1]].copy_(ql_nope) |
| 1374 | else: |
| 1375 | ops.concat_mla_q(ql_nope, q_pe, q) |
| 1376 | else: |
| 1377 | actual_num_heads = q.shape[1] |
| 1378 | |
| 1379 | num_actual_toks = q.shape[0] |
| 1380 | |
| 1381 | # Get topk indices |
| 1382 | assert self.topk_indices_buffer is not None |
| 1383 | topk_indices = self.topk_indices_buffer[:num_actual_toks] |
| 1384 | |
| 1385 | use_fp8_cache = self.kv_cache_dtype in QUANTIZED_DS_MLA_CACHE_FORMATS |
| 1386 | |
| 1387 | lse: torch.Tensor | None = None |
| 1388 | |
| 1389 | if not use_fp8_cache: |
| 1390 | attn_out, bf16_lse = self._forward_bf16_kv( |
| 1391 | q, |
| 1392 | kv_c_and_k_pe_cache, |
| 1393 | topk_indices, |
| 1394 | attn_metadata, |
| 1395 | actual_num_heads, |
| 1396 | ) |
| 1397 | if self.need_to_return_lse_for_decode: |
| 1398 | lse = bf16_lse |
| 1399 | elif attn_metadata.fp8_use_mixed_batch: |
| 1400 | attn_out, lse = self._forward_fp8_kv_mixed_batch( |
| 1401 | q, kv_c_and_k_pe_cache, topk_indices, attn_metadata |
| 1402 | ) |
| 1403 | else: |
| 1404 | attn_out, lse = self._forward_fp8_kv_separate_prefill_decode( |
| 1405 | q, kv_c_and_k_pe_cache, topk_indices, attn_metadata |
| 1406 | ) |
| 1407 | |
| 1408 | return attn_out, lse |