| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """DeepSeek V4.1 attention on FlashMLA's mega-attention kernel (SM100). |
| 4 | |
| 5 | One kernel does Q RoPE, sparse attention, the inverse RoPE of the output and |
| 6 | its FP8 cast, and writes straight into the buffer ``wo_a`` consumes. So the |
| 7 | layer declares ``accepts_unnormed_unroped_query`` -- the kernel RoPEs Q itself, |
| 8 | leaving the fused KV insert only the zero-pad to the kernel's head count -- and |
| 9 | its ``_alloc_attn_out`` / ``_o_proj`` pair speaks QuantizedActivation instead |
| 10 | of bf16, since the output needs no rotation or quantization here. |
| 11 | |
| 12 | ``wq_b`` rows and ``wo_a`` columns are permuted once at load so the surrounding |
| 13 | GEMMs speak the kernel's chunk-interleaved layouts directly. A step's prefill |
| 14 | and decode segments write disjoint token ranges of one output buffer pair, so |
| 15 | a single ``wo_a`` einsum covers the whole step. |
| 16 | """ |
| 17 | |
| 18 | from dataclasses import replace |
| 19 | from typing import TYPE_CHECKING, cast |
| 20 | |
| 21 | import torch |
| 22 | |
| 23 | from vllm.config import VllmConfig |
| 24 | from vllm.config.cache import CacheDType |
| 25 | from vllm.forward_context import get_forward_context |
| 26 | from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation |
| 27 | from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp8Dynamic |
| 28 | from vllm.models.deepseek_v41.common.ops import ( |
| 29 | combine_topk_swa_indices, |
| 30 | compute_global_topk_indices_and_lens, |
| 31 | dequantize_and_gather_k_cache, |
| 32 | ) |
| 33 | from vllm.models.deepseek_v41.common.ops.fused_layout import ( |
| 34 | WV_GROUP_SIZE, |
| 35 | permute_wo_a_, |
| 36 | permute_wq_b_, |
| 37 | ) |
| 38 | from vllm.models.deepseek_v41.nvidia.flashmla import DeepseekV4FlashMLAAttention |
| 39 | from vllm.models.deepseek_v41.sparse_mla import ( |
| 40 | DeepseekV4FlashMLAMetadata, |
| 41 | FlashMLAMegaAttnBackend, |
| 42 | ) |
| 43 | from vllm.platforms import current_platform |
| 44 | from vllm.utils.deep_gemm import fp8_einsum, get_tma_aligned_size |
| 45 | from vllm.utils.math_utils import round_up |
| 46 | from vllm.v1.attention.ops.flashmla import is_flashmla_sparse_supported |
| 47 | from vllm.v1.worker.workspace import current_workspace_manager |
| 48 | |
| 49 | if TYPE_CHECKING: |
| 50 | from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata |
| 51 | |
| 52 | _HEAD_DIM_V = 512 |
| 53 | _ROPE_DIM = 64 |
| 54 | # The kernel always emits per-32 scales, whatever num_per_channels says. |
| 55 | _QUANT_GROUP = 32 |
| 56 | # Fixed kernel conventions: GPT-J (non-neox) RoPE, TMA-aligned col-major |
| 57 | # scales, round-to-nearest, packed ue8m0 -- i.e. what DeepGEMM's fp8_einsum |
| 58 | # consumes. Q norm is off: the V4.1 layer norms qr before wq_b. |
| 59 | _ROPE_ARGS = (False, _ROPE_DIM) |
| 60 | _SF_ARGS = (_QUANT_GROUP, True, True, True) |
| 61 | |
| 62 | # Decode always goes through the mega kernel: there is no batch size at which |
| 63 | # falling back to the split-KV decode pays, so there is no minimum-token |
| 64 | # threshold and no fallback path to maintain. What decides whether this layer |
| 65 | # is worth selecting at all is TP, not the batch size. |
| 66 | # |
| 67 | # Measured on one GB300 with benchmarks/kernels/benchmark_dsv41_mega_attn.py, |
| 68 | # CUDA graphs on, counting every per-step op each path runs (64-head model, |
| 69 | # topk_swa=128, topk_extra=512, NVFP4 compressed cache), as mega / split-KV |
| 70 | # microseconds: |
| 71 | # |
| 72 | # s_q 1 32 128 256 512 |
| 73 | # TP1 (64 live heads) 29/29 31/31 31/40 53/72 96/140 |
| 74 | # TP2 (32 live) 29/29 31/30 35/35 61/63 110/119 |
| 75 | # TP4 (16 live) 29/29 31/29 34/34 59/58 109/109 |
| 76 | # TP8 (8 live) 29/29 31/29 33/33 59/57 107/104 |
| 77 | # |
| 78 | # The kernel absorbs the inverse RoPE and the FP8 cast, whose cost scales with |
| 79 | # *live* heads, while its own cost scales with *padded* heads -- always 64, |
| 80 | # since FlashMLA's FP8 decode takes only h_q in {64, 128}. So it wins at TP1 |
| 81 | # (up to 1.45x, and TP1 is also where the fused insert skips the Q pad |
| 82 | # entirely) and from TP2 down the two cancel to within a few percent, with a |
| 83 | # fixed ~2 us penalty at small s_q where the padded-head cost dominates. |
| 84 | # Measure before assuming a win at high TP -- and measure under CUDA graphs: |
| 85 | # in eager mode launch overhead swamps both paths and overstates the fused one |
| 86 | # several-fold. |
| 87 | |
| 88 | |
| 89 | def is_flashmla_mega_attn_supported() -> tuple[bool, str | None]: |
| 90 | """Whether FlashMLA's fused mega-attention kernels are usable here.""" |
| 91 | is_available, maybe_reason = is_flashmla_sparse_supported() |
| 92 | if not is_available: |
| 93 | return False, maybe_reason |
| 94 | if not current_platform.is_device_capability_family(100): |
| 95 | return False, "FlashMLA mega attention requires sm_10x GPUs." |
| 96 | if not hasattr(torch.ops._flashmla_C, "fused_norm_rope_attn_rope_cast_decode"): |
| 97 | return False, "vllm._flashmla_C was built without the mega attention ops." |
| 98 | if not hasattr(torch.ops._flashmla_C, "permute_q_b_proj"): |
| 99 | return False, "vllm._flashmla_C predates the mega attention weight permutes." |
| 100 | return True, None |
| 101 | |
| 102 | |
| 103 | def alloc_mega_attn_output( |
| 104 | num_tokens: int, |
| 105 | n_wv_group: int, |
| 106 | device: torch.device, |
| 107 | ) -> QuantizedActivation: |
| 108 | """Allocate the output buffer pair the mega-attention kernels write into. |
| 109 | |
| 110 | One pair per forward step: the prefill and decode segments fill disjoint |
| 111 | token ranges, so a single ``fp8_einsum`` over all ``num_tokens`` can |
| 112 | consume the result instead of one call per segment. |
| 113 | |
| 114 | The scale buffer is MN-major -- stride 1 along tokens, head-dim stride |
| 115 | ``ceil4(num_tokens)`` -- which is both what the kernel requires of a |
| 116 | caller-provided buffer and what DeepGEMM expects for this N. |
| 117 | """ |
| 118 | d = WV_GROUP_SIZE * _HEAD_DIM_V |
| 119 | data = torch.empty( |
| 120 | (num_tokens, n_wv_group, d), dtype=torch.float8_e4m3fn, device=device |
| 121 | ) |
| 122 | aligned = get_tma_aligned_size(num_tokens, torch.int32.itemsize) |
| 123 | scale = torch.empty( |
| 124 | (n_wv_group, d // (_QUANT_GROUP * 4), aligned), |
| 125 | dtype=torch.int32, |
| 126 | device=device, |
| 127 | ).permute(2, 0, 1)[:num_tokens] |
| 128 | return QuantizedActivation( |
| 129 | data=data, |
| 130 | scale=scale, |
| 131 | orig_dtype=torch.bfloat16, |
| 132 | orig_shape=data.shape, |
| 133 | quant_key=kMxfp8Dynamic, |
| 134 | ) |
| 135 | |
| 136 | |
| 137 | def _token_slice(out: QuantizedActivation, start: int, end: int) -> QuantizedActivation: |
| 138 | """The ``[start, end)`` token slice of a mega-attention output buffer.""" |
| 139 | data = out.data[start:end] |
| 140 | return replace(out, data=data, scale=out.scale[start:end], orig_shape=data.shape) |
| 141 | |
| 142 | |
| 143 | class DeepseekV4MegaAttnAttention(DeepseekV4FlashMLAAttention): |
| 144 | """FlashMLA mega-attention layer for DeepSeek V4.1 (SM100).""" |
| 145 | |
| 146 | backend_cls = FlashMLAMegaAttnBackend |
| 147 | |
| 148 | @classmethod |
| 149 | def is_available_for(cls, vllm_config: VllmConfig) -> bool: |
| 150 | """Whether this layer can serve the configured model on this device. |
| 151 | |
| 152 | The kernel is SM100-only and wants ``WV_GROUP_SIZE`` heads per ``wo_a`` |
| 153 | group, which stops holding once TP divides the head count far enough |
| 154 | (TP16 on a 64-head model leaves 4). ``__init__`` raises on both, so the |
| 155 | default-backend selector asks here first rather than turning an |
| 156 | unsupported topology into a startup crash. |
| 157 | """ |
| 158 | if not is_flashmla_mega_attn_supported()[0]: |
| 159 | return False |
| 160 | config = vllm_config.model_config.hf_config |
| 161 | config = getattr(config, "text_config", config) |
| 162 | n_heads = getattr(config, "num_attention_heads", 0) or 0 |
| 163 | n_groups = getattr(config, "o_groups", 0) or 0 |
| 164 | tp_size = vllm_config.parallel_config.tensor_parallel_size |
| 165 | if n_heads % tp_size or n_groups % tp_size: |
| 166 | return False |
| 167 | n_local_heads, n_local_groups = n_heads // tp_size, n_groups // tp_size |
| 168 | if not n_local_groups or n_local_heads % WV_GROUP_SIZE: |
| 169 | return False |
| 170 | return n_local_heads // n_local_groups == WV_GROUP_SIZE |
| 171 | |
| 172 | def __init__(self, vllm_config: VllmConfig, *args, **kwargs) -> None: |
| 173 | super().__init__(vllm_config, *args, **kwargs) |
| 174 | if self.n_local_heads % WV_GROUP_SIZE: |
| 175 | raise ValueError( |
| 176 | f"{self.prefix}: mega attention needs the local head count " |
| 177 | f"({self.n_local_heads}) to be a multiple of {WV_GROUP_SIZE}." |
| 178 | ) |
| 179 | if self.n_local_heads // self.n_local_groups != WV_GROUP_SIZE: |
| 180 | raise ValueError( |
| 181 | f"{self.prefix}: mega attention needs {WV_GROUP_SIZE} heads per " |
| 182 | "wo_a group." |
| 183 | ) |
| 184 | self.n_wv_group = self.padded_heads // WV_GROUP_SIZE |
| 185 | self._fused_layouts_ready = False |
| 186 | |
| 187 | # ---- interface contract ------------------------------------------------ |
| 188 | |
| 189 | @property |
| 190 | def accepts_unnormed_unroped_query(self) -> bool: |
| 191 | return True |
| 192 | |
| 193 | @property |
| 194 | def packed_kv_cache_dtype(self) -> CacheDType: |
| 195 | # This kernel is the only one that reads an NVFP4 compressed cache, so |
| 196 | # it is what an unspecific --kv-cache-dtype resolves to here. |
| 197 | return "nvfp4_ds_mla" |
| 198 | |
| 199 | def _alloc_attn_out( |
| 200 | self, num_tokens: int, hidden_states: torch.Tensor |
| 201 | ) -> QuantizedActivation: |
| 202 | return alloc_mega_attn_output(num_tokens, self.n_wv_group, hidden_states.device) |
| 203 | |
| 204 | def _o_proj( |
| 205 | self, attn_out: QuantizedActivation, positions: torch.Tensor |
| 206 | ) -> torch.Tensor: |
| 207 | """Grouped ``wo_a`` then ``wo_b`` over the kernel's quantized output. |
| 208 | |
| 209 | The inverse RoPE and the FP8 cast happened inside the attention |
| 210 | kernel, and ``wo_a`` is permuted for its output layout, so there is |
| 211 | nothing to rotate or quantize here and ``positions`` is unused. The |
| 212 | scale already is DeepGEMM's packed-ue8m0 MN-major layout, so the |
| 213 | einsum consumes the kernel's output with no repacking. |
| 214 | """ |
| 215 | groups = self.n_local_groups |
| 216 | z = torch.empty( |
| 217 | (attn_out.data.shape[0], groups, self.o_lora_rank), |
| 218 | dtype=torch.bfloat16, |
| 219 | device=attn_out.data.device, |
| 220 | ) |
| 221 | fp8_einsum( |
| 222 | "bhr,hdr->bhd", |
| 223 | (attn_out.data[:, :groups], attn_out.scale[:, :groups]), |
| 224 | (self.wo_a.weight, self.wo_a.weight_scale), |
| 225 | z, |
| 226 | recipe=self._einsum_recipe, |
| 227 | ) |
| 228 | return self._wo_b_proj(z.flatten(1)) |
| 229 | |
| 230 | # ---- weights ----------------------------------------------------------- |
| 231 | |
| 232 | def finalize_loaded_weights(self) -> None: |
| 233 | """Permute wq_b rows / wo_a columns into the kernel's layouts. |
| 234 | |
| 235 | Idempotent: a second post-load pass must not permute twice. |
| 236 | """ |
| 237 | if self._fused_layouts_ready: |
| 238 | return |
| 239 | permute_wq_b_( |
| 240 | self.wq_b.weight.data, self.wq_b.weight_scale.data, self.n_local_heads |
| 241 | ) |
| 242 | permute_wo_a_( |
| 243 | self.wo_a.weight.data, |
| 244 | self.wo_a.weight_scale.data, |
| 245 | self.n_local_heads // self.n_local_groups, |
| 246 | ) |
| 247 | self._fused_layouts_ready = True |
| 248 | |
| 249 | # ---- forward ----------------------------------------------------------- |
| 250 | |
| 251 | def forward_mqa( |
| 252 | self, |
| 253 | q: torch.Tensor, |
| 254 | kv: torch.Tensor, |
| 255 | positions: torch.Tensor, |
| 256 | output: QuantizedActivation, |
| 257 | ) -> None: |
| 258 | if not self._fused_layouts_ready: |
| 259 | raise RuntimeError( |
| 260 | f"{self.prefix}: wq_b / wo_a were never permuted for the mega " |
| 261 | "attention kernel; refusing to run with mismatched layouts." |
| 262 | ) |
| 263 | attn_metadata = get_forward_context().attn_metadata |
| 264 | if attn_metadata is None: |
| 265 | # Warmup dummy run: reserve the prefill workspace, produce zeros. |
| 266 | self._reserve_prefill_workspace(q) |
| 267 | output.data.zero_() |
| 268 | output.scale.zero_() |
| 269 | return |
| 270 | |
| 271 | assert isinstance(attn_metadata, dict) |
| 272 | # q was zero-padded to the kernel's head count by the fused KV insert. |
| 273 | assert q.shape[1] == self.padded_heads |
| 274 | flashmla_metadata = cast( |
| 275 | DeepseekV4FlashMLAMetadata | None, |
| 276 | attn_metadata.get(self.compressed_cache_prefix) |
| 277 | if self.compressed_cache_prefix is not None |
| 278 | else None, |
| 279 | ) |
| 280 | swa_metadata = cast( |
| 281 | "DeepseekSparseSWAMetadata", attn_metadata[self.swa_cache_layer.prefix] |
| 282 | ) |
| 283 | # The kernel takes int32 RoPE positions. |
| 284 | positions_int32 = positions.to(torch.int32) |
| 285 | num_decode_tokens = swa_metadata.num_decode_tokens |
| 286 | |
| 287 | if swa_metadata.num_prefills > 0: |
| 288 | self._forward_prefill_mega( |
| 289 | q[num_decode_tokens:], |
| 290 | positions_int32[num_decode_tokens:], |
| 291 | flashmla_metadata, |
| 292 | swa_metadata, |
| 293 | output, |
| 294 | num_decode_tokens, |
| 295 | ) |
| 296 | if swa_metadata.num_decodes > 0: |
| 297 | self._forward_decode_mega( |
| 298 | q[:num_decode_tokens], |
| 299 | positions_int32[:num_decode_tokens], |
| 300 | flashmla_metadata, |
| 301 | swa_metadata, |
| 302 | _token_slice(output, 0, num_decode_tokens), |
| 303 | ) |
| 304 | |
| 305 | def _reserve_prefill_workspace(self, q: torch.Tensor) -> None: |
| 306 | swa_only = self.compress_ratio == 0 |
| 307 | n = 0 if swa_only else -(-self.max_model_len // self.compress_ratio) |
| 308 | m = n + self.window_size + self.max_num_batched_tokens |
| 309 | if swa_only: |
| 310 | top_k = 0 |
| 311 | else: |
| 312 | assert self.topk_indices_buffer is not None |
| 313 | top_k = self.topk_indices_buffer.shape[-1] |
| 314 | combined_topk = round_up(top_k + self.window_size, 128) |
| 315 | current_workspace_manager().get_simultaneous( |
| 316 | ((self.PREFILL_CHUNK_SIZE, m, q.shape[-1]), torch.bfloat16), |
| 317 | ((self.max_num_batched_tokens, combined_topk), torch.int32), |
| 318 | ((self.max_num_batched_tokens,), torch.int32), |
| 319 | ) |
| 320 | |
| 321 | def _decode_compressed_kv_and_topk( |
| 322 | self, |
| 323 | flashmla_metadata: DeepseekV4FlashMLAMetadata | None, |
| 324 | swa_metadata: "DeepseekSparseSWAMetadata", |
| 325 | ) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]: |
| 326 | """The compressed cache and its ``[n, topk]`` slot indices / lengths.""" |
| 327 | if self.compress_ratio == 0: |
| 328 | return None, None, None |
| 329 | assert flashmla_metadata is not None |
| 330 | assert swa_metadata.is_valid_token is not None |
| 331 | assert self.topk_indices_buffer is not None |
| 332 | num_decode_tokens = swa_metadata.num_decode_tokens |
| 333 | indices, lens = compute_global_topk_indices_and_lens( |
| 334 | self.topk_indices_buffer[:num_decode_tokens], |
| 335 | swa_metadata.token_to_req_indices, |
| 336 | flashmla_metadata.block_table[: swa_metadata.num_decodes], |
| 337 | flashmla_metadata.block_size // self.compress_ratio, |
| 338 | swa_metadata.is_valid_token[:num_decode_tokens], |
| 339 | ) |
| 340 | return self._compressed_kv_cache().unsqueeze(-2), indices, lens |
| 341 | |
| 342 | def _forward_decode_mega( |
| 343 | self, |
| 344 | q: torch.Tensor, |
| 345 | positions_int32: torch.Tensor, |
| 346 | flashmla_metadata: DeepseekV4FlashMLAMetadata | None, |
| 347 | swa_metadata: "DeepseekSparseSWAMetadata", |
| 348 | out: QuantizedActivation, |
| 349 | ) -> None: |
| 350 | """Mega attention over the paged quantized caches, in place into ``out``. |
| 351 | |
| 352 | The caches are ``[num_blocks, page, 1, bytes]``, the record taken from |
| 353 | ``bytes`` (528 V4.1 fp8, 288 V4.1 NVFP4; NVFP4 only as the compressed |
| 354 | cache beside an fp8 SWA one), and the indices ``[s_q, topk]`` int32 |
| 355 | slot ids (``block * page + offset``, ``-1`` invalid). |
| 356 | """ |
| 357 | extra_cache, extra_idx, extra_len = self._decode_compressed_kv_and_topk( |
| 358 | flashmla_metadata, swa_metadata |
| 359 | ) |
| 360 | assert swa_metadata.decode_swa_indices is not None |
| 361 | torch.ops._flashmla_C.fused_norm_rope_attn_rope_cast_decode( |
| 362 | q, |
| 363 | self.swa_cache_layer.kv_cache.unsqueeze(-2), |
| 364 | swa_metadata.decode_swa_indices.view(q.shape[0], -1), |
| 365 | self.scale, |
| 366 | _HEAD_DIM_V, |
| 367 | self.attn_sink, |
| 368 | swa_metadata.decode_swa_lens, |
| 369 | extra_cache, |
| 370 | extra_idx, |
| 371 | extra_len, |
| 372 | False, # enable_q_norm |
| 373 | 0.0, # rms_norm_eps |
| 374 | positions_int32, |
| 375 | *_ROPE_ARGS, |
| 376 | self.rotary_emb.cos_sin_cache, |
| 377 | self.n_wv_group, |
| 378 | *_SF_ARGS, |
| 379 | out.data, |
| 380 | out.scale, |
| 381 | ) |
| 382 | |
| 383 | def _forward_prefill_mega( |
| 384 | self, |
| 385 | q: torch.Tensor, |
| 386 | positions_int32: torch.Tensor, |
| 387 | flashmla_metadata: DeepseekV4FlashMLAMetadata | None, |
| 388 | swa_metadata: "DeepseekSparseSWAMetadata", |
| 389 | out: QuantizedActivation, |
| 390 | token_base: int, |
| 391 | ) -> None: |
| 392 | swa_only = self.compress_ratio == 0 |
| 393 | num_decodes = swa_metadata.num_decodes |
| 394 | num_decode_tokens = swa_metadata.num_decode_tokens |
| 395 | seq_lens = swa_metadata.prefill_seq_lens |
| 396 | gather_lens = swa_metadata.prefill_gather_lens |
| 397 | query_start_loc_cpu = swa_metadata.query_start_loc_cpu |
| 398 | query_start_loc = swa_metadata.query_start_loc |
| 399 | assert seq_lens is not None and gather_lens is not None |
| 400 | assert query_start_loc_cpu is not None and query_start_loc is not None |
| 401 | prefill_token_base = query_start_loc_cpu[num_decodes] |
| 402 | assert self.topk_indices_buffer is not None |
| 403 | topk_indices = self.topk_indices_buffer[num_decode_tokens:][ |
| 404 | : swa_metadata.num_prefill_tokens |
| 405 | ] |
| 406 | top_k = 0 if swa_only else topk_indices.shape[-1] |
| 407 | chunk_plan = swa_metadata.get_prefill_chunk_plan( |
| 408 | compress_ratio=self.compress_ratio, |
| 409 | prefill_chunk_size=self.PREFILL_CHUNK_SIZE, |
| 410 | has_compressed=not swa_only, |
| 411 | ) |
| 412 | assert chunk_plan, "prefill chunk plan must be non-empty when num_prefills > 0" |
| 413 | workspace_manager = current_workspace_manager() |
| 414 | combined_topk = round_up(top_k + self.window_size, 128) |
| 415 | for chunk_start, chunk_end, chunk_n, chunk_m in chunk_plan: |
| 416 | chunk_size = chunk_end - chunk_start |
| 417 | kv_ws, idx_ws, lens_ws = workspace_manager.get_simultaneous( |
| 418 | ((chunk_size, chunk_m, q.shape[-1]), torch.bfloat16), |
| 419 | ((self.max_num_batched_tokens, combined_topk), torch.int32), |
| 420 | ((self.max_num_batched_tokens,), torch.int32), |
| 421 | ) |
| 422 | if not swa_only: |
| 423 | assert flashmla_metadata is not None |
| 424 | dequantize_and_gather_k_cache( |
| 425 | kv_ws[:chunk_size, :chunk_n], |
| 426 | self._compressed_kv_cache(), |
| 427 | seq_lens=seq_lens[chunk_start:chunk_end] // self.compress_ratio, |
| 428 | gather_lens=None, |
| 429 | block_table=flashmla_metadata.block_table[num_decodes:][ |
| 430 | chunk_start:chunk_end |
| 431 | ], |
| 432 | block_size=flashmla_metadata.block_size // self.compress_ratio, |
| 433 | offset=0, |
| 434 | ) |
| 435 | dequantize_and_gather_k_cache( |
| 436 | kv_ws[:chunk_size], |
| 437 | self.swa_cache_layer.kv_cache, |
| 438 | seq_lens=seq_lens[chunk_start:chunk_end], |
| 439 | gather_lens=gather_lens[chunk_start:chunk_end], |
| 440 | block_table=swa_metadata.block_table[num_decodes:][ |
| 441 | chunk_start:chunk_end |
| 442 | ], |
| 443 | block_size=swa_metadata.block_size, |
| 444 | offset=chunk_n, |
| 445 | ) |
| 446 | first, last = num_decodes + chunk_start, num_decodes + chunk_end |
| 447 | qs = int(query_start_loc_cpu[first] - prefill_token_base) |
| 448 | qe = int(query_start_loc_cpu[last] - prefill_token_base) |
| 449 | combined_indices, combined_lens = combine_topk_swa_indices( |
| 450 | topk_indices[qs:qe], |
| 451 | query_start_loc[first : last + 1], |
| 452 | seq_lens[chunk_start:chunk_end], |
| 453 | gather_lens[chunk_start:chunk_end], |
| 454 | self.window_size, |
| 455 | self.compress_ratio, |
| 456 | top_k, |
| 457 | chunk_m, |
| 458 | chunk_n, |
| 459 | out=(idx_ws[: qe - qs], lens_ws[: qe - qs]), |
| 460 | ) |
| 461 | chunk_out = _token_slice(out, token_base + qs, token_base + qe) |
| 462 | # Mega attention over the gathered non-paged bf16 KV (RoPE already |
| 463 | # applied), in place into this chunk's token range. `indices` |
| 464 | # entries outside [0, s_kv) skip. |
| 465 | torch.ops._flashmla_C.fused_norm_rope_attn_rope_cast_fwd( |
| 466 | q[qs:qe], |
| 467 | kv_ws.view(-1, 1, q.shape[-1]), |
| 468 | combined_indices.unsqueeze(1), |
| 469 | self.scale, |
| 470 | _HEAD_DIM_V, |
| 471 | self.attn_sink, |
| 472 | combined_lens, |
| 473 | False, # enable_q_norm |
| 474 | 0.0, # rms_norm_eps |
| 475 | positions_int32[qs:qe], |
| 476 | *_ROPE_ARGS, |
| 477 | self.rotary_emb.cos_sin_cache, |
| 478 | self.n_wv_group, |
| 479 | *_SF_ARGS, |
| 480 | chunk_out.data, |
| 481 | chunk_out.scale, |
| 482 | ) |