| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Sparse attention indexer on aiter's paged MXFP4 MQA-logits kernel. |
| 4 | |
| 5 | gfx950 only. The kernel reads the preshuffled paged indexer K cache in place, |
| 6 | so no layer gathers K into a contiguous buffer, prefill included. Any model |
| 7 | whose indexer supplies the inputs `rocm_mxfp4_sparse_attn_indexer` documents |
| 8 | can run the dense path. DeepSeek-V4.1's two-level indexer adds to it: the |
| 9 | candidate source takes its block maxima from the same walk that writes its |
| 10 | logits, and the candidate consumers can walk the candidate pool instead of |
| 11 | the whole context: the pool is resolved once per step and shared by all of |
| 12 | them. |
| 13 | """ |
| 14 | |
| 15 | import functools |
| 16 | from collections.abc import Callable |
| 17 | from dataclasses import dataclass |
| 18 | from typing import TYPE_CHECKING, NamedTuple |
| 19 | |
| 20 | import torch |
| 21 | |
| 22 | import vllm.envs as envs |
| 23 | from vllm.compilation.breakable_cudagraph import eager_break_during_capture |
| 24 | from vllm.config import CUDAGraphMode |
| 25 | from vllm.forward_context import get_forward_context |
| 26 | from vllm.model_executor.kernels.attention.dsa.candidate_blocks import ( |
| 27 | apply_candidate_mask, |
| 28 | ) |
| 29 | from vllm.triton_utils import tl, triton |
| 30 | from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( |
| 31 | _apply_candidate_mask_strided, |
| 32 | _max_decode_logits_rows, |
| 33 | ) |
| 34 | from vllm.v1.worker.workspace import current_workspace_manager |
| 35 | |
| 36 | if TYPE_CHECKING: |
| 37 | from vllm.v1.attention.backends.mla.indexer import ( |
| 38 | DeepseekV32IndexerPrefillChunkMetadata, |
| 39 | ) |
| 40 | from vllm.v1.attention.backends.mla.rocm_paged_mxfp4_indexer import ( |
| 41 | RocmMxfp4GatherLaunch, |
| 42 | RocmMxfp4IndexerMetadata, |
| 43 | RocmMxfp4NativeDecode, |
| 44 | RocmMxfp4PrefillPlan, |
| 45 | ) |
| 46 | |
| 47 | MXFP4_BLOCK_SIZE = 32 |
| 48 | |
| 49 | |
| 50 | @functools.cache |
| 51 | def _aiter(): |
| 52 | """The aiter module with the paged MXFP4 MQA-logits launcher, its schedule |
| 53 | and cache_format.""" |
| 54 | from aiter.ops.triton.attention import pa_mqa_logits_mxfp4 |
| 55 | |
| 56 | return pa_mqa_logits_mxfp4 |
| 57 | |
| 58 | |
| 59 | @functools.cache |
| 60 | def _aiter_topk() -> Callable[..., None]: |
| 61 | """The aiter per-row top-k.""" |
| 62 | from aiter.ops.topk import top_k_per_row_decode |
| 63 | |
| 64 | return top_k_per_row_decode |
| 65 | |
| 66 | |
| 67 | class RocmPagedMxfp4CacheLayout(NamedTuple): |
| 68 | """The byte order of an indexer K page, as the kernel reads it, and the |
| 69 | shuffle pattern aiter's K cache op writes it with. |
| 70 | |
| 71 | A page is cut into runs of ``n_per_tile`` tokens. Inside a run, each token's |
| 72 | packed values are split into ``d_per_tile``-byte chunks along the head |
| 73 | dimension, and the run stores chunk 0 of every token, then chunk 1, and so |
| 74 | on: ``[chunk, token, byte]``. Its e8m0 scales are stored as |
| 75 | ``[scale % scale_lanes, token, scale // scale_lanes]``. |
| 76 | """ |
| 77 | |
| 78 | n_per_tile: int |
| 79 | d_per_tile: int |
| 80 | scale_lanes: int |
| 81 | |
| 82 | |
| 83 | # The scale order the kernel reads: aiter's mode 1, a lane's scales innermost. |
| 84 | # cache_format does not report the order or the lane split yet. |
| 85 | _SCALE_ORDER = 1 |
| 86 | _WAVE_SIZE = 64 |
| 87 | |
| 88 | |
| 89 | @functools.cache |
| 90 | def rocm_paged_mxfp4_cache_layout( |
| 91 | num_heads: int, head_dim: int, page_entries: int |
| 92 | ) -> RocmPagedMxfp4CacheLayout: |
| 93 | """The K page layout from aiter's ``cache_format``, the only place vLLM |
| 94 | reads it, so a change there cannot reach the writer unchecked. |
| 95 | |
| 96 | Raises: |
| 97 | ValueError: aiter rejects the page size, or describes an order this |
| 98 | layout cannot express. |
| 99 | |
| 100 | """ |
| 101 | fmt = _aiter().cache_format(num_heads, head_dim, page_entries) |
| 102 | try: |
| 103 | n_per_tile = int(fmt["n_per_tile"]) |
| 104 | d_per_tile = int(fmt["d_per_tile"]) |
| 105 | scale_order = fmt.get("scale_mode", _SCALE_ORDER) |
| 106 | scale_lanes = int(fmt.get("scale_lanes", _WAVE_SIZE // n_per_tile)) |
| 107 | except (KeyError, TypeError, ValueError, ZeroDivisionError) as e: |
| 108 | raise ValueError(f"aiter's cache_format() returned {fmt!r}") from e |
| 109 | if ( |
| 110 | scale_order != _SCALE_ORDER |
| 111 | or min(n_per_tile, d_per_tile, scale_lanes) <= 0 |
| 112 | or page_entries % n_per_tile |
| 113 | or (head_dim // 2) % d_per_tile |
| 114 | or (head_dim // MXFP4_BLOCK_SIZE) % scale_lanes |
| 115 | ): |
| 116 | raise ValueError( |
| 117 | f"aiter's cache_format() returned {fmt!r} for {num_heads} heads of " |
| 118 | f"{head_dim} in {page_entries}-entry pages, a K page order " |
| 119 | "RocmPagedMxfp4CacheLayout cannot express" |
| 120 | ) |
| 121 | return RocmPagedMxfp4CacheLayout(n_per_tile, d_per_tile, scale_lanes) |
| 122 | |
| 123 | |
| 124 | def rocm_mxfp4_decode_schedule_words( |
| 125 | num_heads: int, head_dim: int, page_entries: int, next_n: int = 1 |
| 126 | ) -> int: |
| 127 | """int32 words of the largest schedule a decode step can build, flattened or |
| 128 | not. The slice cap can take it past target_wgs, up to SCHED_SLOT_CAP.""" |
| 129 | words = 4 * _aiter().SCHED_SLOT_CAP |
| 130 | for rows in {1, next_n}: |
| 131 | config = _aiter().select_config(num_heads, head_dim, rows, page_entries) |
| 132 | words = max(words, 4 * config["target_wgs"]) |
| 133 | return words |
| 134 | |
| 135 | |
| 136 | def build_rocm_mxfp4_decode_schedule( |
| 137 | row_lens: torch.Tensor, |
| 138 | num_heads: int, |
| 139 | head_dim: int, |
| 140 | page_entries: int, |
| 141 | out: torch.Tensor, |
| 142 | logits_width: int, |
| 143 | native: "RocmMxfp4NativeDecode | None" = None, |
| 144 | ) -> torch.Tensor | None: |
| 145 | """Work descriptors that even out a decode step, or None where the static |
| 146 | grid already fills the machine. Depends only on the rows' lengths, the |
| 147 | cache geometry and the logits width, so one serves every layer of a group. |
| 148 | The width sizes the slices, capped the way the static grid caps them.""" |
| 149 | if native is None: |
| 150 | return _aiter().build_schedule( |
| 151 | row_lens, |
| 152 | 1, |
| 153 | num_heads, |
| 154 | head_dim, |
| 155 | page_entries, |
| 156 | out=out, |
| 157 | max_model_len=logits_width, |
| 158 | ) |
| 159 | return _aiter().build_schedule( |
| 160 | native.context_lens, |
| 161 | native.next_n, |
| 162 | num_heads, |
| 163 | head_dim, |
| 164 | page_entries, |
| 165 | out=out, |
| 166 | row_ends=row_lens, |
| 167 | max_model_len=logits_width, |
| 168 | ) |
| 169 | |
| 170 | |
| 171 | def _kv_view(kv_cache: torch.Tensor, head_dim: int) -> torch.Tensor: |
| 172 | """The indexer cache as [pages, entries, 1, bytes]. The page stride is the |
| 173 | block-major pool's, not the page's own size.""" |
| 174 | num_pages, entries, width = kv_cache.shape |
| 175 | assert kv_cache.dtype == torch.uint8 |
| 176 | assert width == head_dim // 2 + head_dim // MXFP4_BLOCK_SIZE, width |
| 177 | return torch.as_strided( |
| 178 | kv_cache, (num_pages, entries, 1, width), (kv_cache.stride(0), width, width, 1) |
| 179 | ) |
| 180 | |
| 181 | |
| 182 | def rocm_mxfp4_consumer_rows(num_candidate_cols: int) -> int: |
| 183 | """Query rows per candidate-consumer launch. Its logits are [rows, pool] |
| 184 | fp32 whatever the context, so the logits budget alone sizes it.""" |
| 185 | budget = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024 |
| 186 | return max(1, budget // (4 * num_candidate_cols)) |
| 187 | |
| 188 | |
| 189 | def reserve_rocm_mxfp4_indexer_workspace( |
| 190 | hidden_states: torch.Tensor, |
| 191 | logits_width: int, |
| 192 | candidate_block_size: int = 0, |
| 193 | gather_block_size: int = 0, |
| 194 | num_candidate_cols: int = 0, |
| 195 | ) -> None: |
| 196 | """Profiling run: claim the decode logits workspace and the peak prefill |
| 197 | logits, block scores included when the layer writes them, candidate lists |
| 198 | when it gathers.""" |
| 199 | rows = _max_decode_logits_rows(hidden_states.shape[0]) |
| 200 | specs = [((rows, logits_width), torch.float32)] |
| 201 | budget = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024 |
| 202 | if candidate_block_size: |
| 203 | nblocks = triton.cdiv(logits_width, candidate_block_size) |
| 204 | specs.append(((rows, nblocks), torch.float32)) |
| 205 | budget += budget // candidate_block_size |
| 206 | if gather_block_size: |
| 207 | # aiter allocates the candidate lists: each gathered row, 8 B per |
| 208 | # candidate block (int32 slot and position). The first consumer builds |
| 209 | # them and they live until the step ends, across the later layers, so |
| 210 | # one reservation is held for the rest of the forward. |
| 211 | held = get_forward_context().additional_kwargs |
| 212 | if "rocm_mxfp4_candidate_lists" not in held: |
| 213 | held["rocm_mxfp4_candidate_lists"] = torch.empty( |
| 214 | hidden_states.shape[0] * (num_candidate_cols // gather_block_size) * 8, |
| 215 | dtype=torch.uint8, |
| 216 | device=hidden_states.device, |
| 217 | ) |
| 218 | current_workspace_manager().get_simultaneous(*specs) |
| 219 | torch.empty(budget, dtype=torch.uint8, device=hidden_states.device) |
| 220 | |
| 221 | |
| 222 | def _topk( |
| 223 | logits: torch.Tensor, lengths: torch.Tensor, out: torch.Tensor, k: int |
| 224 | ) -> None: |
| 225 | """Top-k over each row's [0, lengths[row]); -1 pads rows shorter than k. |
| 226 | |
| 227 | aiter's kernel is faster on the candidate top-k and from 128 rows on; below |
| 228 | that, a wide decode step is faster on vLLM's. Shape alone decides, so a FULL |
| 229 | graph replays the same choice at every context length. |
| 230 | """ |
| 231 | rows, width = logits.shape |
| 232 | if rows == 0: |
| 233 | return |
| 234 | if width >= k and (k >= 2048 or rows >= 128): |
| 235 | _aiter_topk()( |
| 236 | logits, |
| 237 | 1, |
| 238 | lengths.reshape(-1), |
| 239 | out, |
| 240 | rows, |
| 241 | logits.stride(0), |
| 242 | logits.stride(1), |
| 243 | k=k, |
| 244 | ) |
| 245 | else: |
| 246 | torch.ops._C.top_k_per_row_decode( |
| 247 | logits, 1, lengths, out, rows, logits.stride(0), logits.stride(1), k |
| 248 | ) |
| 249 | |
| 250 | |
| 251 | @triton.jit |
| 252 | def _remap_compact_topk_kernel( |
| 253 | idx_ptr, |
| 254 | idx_stride, |
| 255 | pos_ptr, |
| 256 | pos_stride, |
| 257 | K: tl.constexpr, |
| 258 | BLOCK: tl.constexpr, |
| 259 | PADDED_K: tl.constexpr, |
| 260 | ): |
| 261 | row = tl.program_id(0).to(tl.int64) |
| 262 | cols = tl.arange(0, PADDED_K) |
| 263 | live = cols < K |
| 264 | slot = tl.load(idx_ptr + row * idx_stride + cols, mask=live, other=-1) |
| 265 | hit = slot >= 0 |
| 266 | slot = tl.where(hit, slot, 0) |
| 267 | start = tl.load( |
| 268 | pos_ptr + row * pos_stride + slot // BLOCK, mask=live & hit, other=0 |
| 269 | ) |
| 270 | tl.store( |
| 271 | idx_ptr + row * idx_stride + cols, |
| 272 | tl.where(hit, start + slot % BLOCK, -1).to(tl.int32), |
| 273 | mask=live, |
| 274 | ) |
| 275 | |
| 276 | |
| 277 | def _remap_compact_topk( |
| 278 | indices: torch.Tensor, positions: torch.Tensor, block: int |
| 279 | ) -> None: |
| 280 | """Candidate slots to context positions, in place: slot j sits in pool |
| 281 | block j // block, which starts at positions[j // block].""" |
| 282 | rows, k = indices.shape |
| 283 | if rows == 0: |
| 284 | return |
| 285 | _remap_compact_topk_kernel[(rows,)]( |
| 286 | indices, |
| 287 | indices.stride(0), |
| 288 | positions, |
| 289 | positions.stride(0), |
| 290 | K=k, |
| 291 | BLOCK=block, |
| 292 | PADDED_K=triton.next_power_of_2(k), |
| 293 | num_warps=4, |
| 294 | ) |
| 295 | |
| 296 | |
| 297 | @dataclass |
| 298 | class _Layer: |
| 299 | """One indexer layer's inputs and outputs for this step.""" |
| 300 | |
| 301 | metadata: "RocmMxfp4IndexerMetadata" |
| 302 | kv: torch.Tensor |
| 303 | q: torch.Tensor |
| 304 | q_scale: torch.Tensor |
| 305 | weights: torch.Tensor |
| 306 | topk_buffer: torch.Tensor |
| 307 | topk_tokens: int |
| 308 | compress_ratio: int |
| 309 | candidates: torch.Tensor | None |
| 310 | block: int |
| 311 | full_graph: bool |
| 312 | |
| 313 | @property |
| 314 | def num_heads(self) -> int: |
| 315 | return self.q.shape[1] |
| 316 | |
| 317 | @property |
| 318 | def head_dim(self) -> int: |
| 319 | return self.q.shape[2] * 2 |
| 320 | |
| 321 | def rows(self, lo: int, hi: int, seqs: int = 1) -> tuple[torch.Tensor, ...]: |
| 322 | """Q, its scales and the weights of token rows [lo, hi), as ``seqs`` |
| 323 | sequences of next_n rows each.""" |
| 324 | return ( |
| 325 | self.q[lo:hi].view(seqs, -1, *self.q.shape[1:]), |
| 326 | self.q_scale[lo:hi].view(seqs, -1, *self.q_scale.shape[1:]), |
| 327 | self.weights[lo:hi], |
| 328 | ) |
| 329 | |
| 330 | def decode_rows(self, n: int) -> tuple[torch.Tensor, ...]: |
| 331 | """The first n token rows, one sequence each.""" |
| 332 | return ( |
| 333 | self.q[:n].unsqueeze(1), |
| 334 | self.q_scale[:n].unsqueeze(1), |
| 335 | self.weights[:n], |
| 336 | ) |
| 337 | |
| 338 | |
| 339 | def _layer( |
| 340 | k_cache_prefix: str, |
| 341 | kv_cache: torch.Tensor, |
| 342 | q_values: torch.Tensor, |
| 343 | q_scale: torch.Tensor, |
| 344 | weights: torch.Tensor, |
| 345 | topk_indices_buffer: torch.Tensor, |
| 346 | topk_tokens: int, |
| 347 | head_dim: int, |
| 348 | compress_ratio: int, |
| 349 | candidate_blocks: torch.Tensor | None, |
| 350 | candidate_block_size: int, |
| 351 | ) -> _Layer: |
| 352 | from vllm.v1.attention.backends.mla.rocm_paged_mxfp4_indexer import ( |
| 353 | RocmMxfp4IndexerMetadata, |
| 354 | ) |
| 355 | |
| 356 | forward_context = get_forward_context() |
| 357 | attn_metadata = forward_context.attn_metadata |
| 358 | assert isinstance(attn_metadata, dict) |
| 359 | metadata = attn_metadata[k_cache_prefix] |
| 360 | assert isinstance(metadata, RocmMxfp4IndexerMetadata), ( |
| 361 | "the ROCm MXFP4 indexer needs RocmMxfp4IndexerMetadataBuilder metadata" |
| 362 | ) |
| 363 | kv = _kv_view(kv_cache, head_dim) |
| 364 | num_heads = q_values.shape[1] |
| 365 | return _Layer( |
| 366 | metadata=metadata, |
| 367 | kv=kv, |
| 368 | q=q_values, |
| 369 | # The fused Q kernel returns one int32 per head; the kernel reads the |
| 370 | # four ue8m0 bytes behind it. |
| 371 | q_scale=q_scale.view(torch.uint8).view(q_scale.shape[0], num_heads, -1), |
| 372 | weights=weights, |
| 373 | topk_buffer=topk_indices_buffer, |
| 374 | topk_tokens=topk_tokens, |
| 375 | compress_ratio=compress_ratio, |
| 376 | candidates=candidate_blocks, |
| 377 | block=candidate_block_size, |
| 378 | full_graph=forward_context.cudagraph_runtime_mode == CUDAGraphMode.FULL, |
| 379 | ) |
| 380 | |
| 381 | |
| 382 | def _block_scores(layer: _Layer, scores: torch.Tensor | None) -> dict: |
| 383 | if scores is None: |
| 384 | return {} |
| 385 | return dict( |
| 386 | calc_block_scores=True, block_scores=scores, candidate_block_size=layer.block |
| 387 | ) |
| 388 | |
| 389 | |
| 390 | def _dense_prefill( |
| 391 | layer: _Layer, |
| 392 | chunk: "DeepseekV32IndexerPrefillChunkMetadata", |
| 393 | plan: "RocmMxfp4PrefillPlan", |
| 394 | candidate_write: bool, |
| 395 | ) -> None: |
| 396 | pa = _aiter() |
| 397 | t0, t1 = chunk.token_start, chunk.token_end |
| 398 | logits = layer.q.new_empty((t1 - t0, plan.width), dtype=torch.float32) |
| 399 | scores = None |
| 400 | if candidate_write: |
| 401 | nblocks = triton.cdiv(plan.width, layer.block) |
| 402 | scores = layer.q.new_empty((t1 - t0, nblocks), dtype=torch.float32) |
| 403 | # One launch for the chunk: the kernel finds each row's request from |
| 404 | # query_start_loc on its own grid. |
| 405 | pa.paged_mxfp4_mqa_logits( |
| 406 | layer.q[t0:t1], |
| 407 | layer.q_scale[t0:t1], |
| 408 | layer.kv, |
| 409 | layer.weights[t0:t1], |
| 410 | plan.context_lens, |
| 411 | chunk.block_table, |
| 412 | plan.width, |
| 413 | out_logits=logits, |
| 414 | clean_logits=False, |
| 415 | row_ends=plan.row_ends, |
| 416 | query_start_loc=plan.query_start_loc, |
| 417 | **_block_scores(layer, scores), |
| 418 | ) |
| 419 | candidates = None if layer.candidates is None else layer.candidates[t0:t1] |
| 420 | if scores is not None: |
| 421 | assert candidates is not None and plan.block_ends is not None |
| 422 | _topk(scores, plan.block_ends, candidates, candidates.shape[1]) |
| 423 | elif candidates is not None: |
| 424 | apply_candidate_mask(logits, None, plan.row_ends, candidates, layer.block) |
| 425 | _topk( |
| 426 | logits, |
| 427 | plan.row_ends, |
| 428 | layer.topk_buffer[t0:t1, : layer.topk_tokens], |
| 429 | layer.topk_tokens, |
| 430 | ) |
| 431 | |
| 432 | |
| 433 | def _dense_decode(layer: _Layer, logits_width: int, candidate_write: bool) -> None: |
| 434 | metadata = layer.metadata |
| 435 | lengths = metadata.decode_row_lens |
| 436 | assert metadata.decode is not None and lengths is not None |
| 437 | rows = lengths.shape[0] |
| 438 | specs = [((rows, logits_width), torch.float32)] |
| 439 | if candidate_write: |
| 440 | specs.append(((rows, triton.cdiv(logits_width, layer.block)), torch.float32)) |
| 441 | logits, *scores = current_workspace_manager().get_simultaneous(*specs) |
| 442 | native = metadata.decode_native |
| 443 | if native is not None: |
| 444 | # A request's next_n rows go in as one sequence, so a workgroup walks |
| 445 | # each KV tile once for all of them. |
| 446 | q, q_scale, weights = layer.rows(0, rows, native.context_lens.shape[0]) |
| 447 | context_lens, block_table, row_ends = ( |
| 448 | native.context_lens, |
| 449 | native.block_table, |
| 450 | lengths, |
| 451 | ) |
| 452 | else: |
| 453 | q, q_scale, weights = layer.decode_rows(rows) |
| 454 | context_lens, block_table, row_ends = lengths, metadata.decode.block_table, None |
| 455 | assert block_table.shape[0] == rows, "one block-table row per query row" |
| 456 | _aiter().paged_mxfp4_mqa_logits( |
| 457 | q, |
| 458 | q_scale, |
| 459 | layer.kv, |
| 460 | weights, |
| 461 | context_lens, |
| 462 | block_table, |
| 463 | logits_width, |
| 464 | out_logits=logits, |
| 465 | clean_logits=False, |
| 466 | row_ends=row_ends, |
| 467 | schedule=metadata.decode_schedule, |
| 468 | **_block_scores(layer, scores[0] if scores else None), |
| 469 | ) |
| 470 | candidates = None if layer.candidates is None else layer.candidates[:rows] |
| 471 | if scores: |
| 472 | assert candidates is not None and metadata.decode_block_ends is not None |
| 473 | _topk(scores[0], metadata.decode_block_ends, candidates, candidates.shape[1]) |
| 474 | elif candidates is not None: |
| 475 | _apply_candidate_mask_strided(logits, None, lengths, candidates, layer.block) |
| 476 | _topk( |
| 477 | logits, |
| 478 | lengths, |
| 479 | layer.topk_buffer[:rows, : layer.topk_tokens], |
| 480 | layer.topk_tokens, |
| 481 | ) |
| 482 | |
| 483 | |
| 484 | def _gather_prefill( |
| 485 | layer: _Layer, launch: "RocmMxfp4GatherLaunch", num_cols: int |
| 486 | ) -> None: |
| 487 | pa = _aiter() |
| 488 | assert layer.candidates is not None |
| 489 | t0, t1 = launch.token_start, launch.token_end |
| 490 | if launch.pool is None: |
| 491 | launch.pool = pa.build_candidate_gather( |
| 492 | layer.candidates[t0:t1], |
| 493 | launch.row_ends, |
| 494 | launch.block_table.expand(t1 - t0, -1), |
| 495 | layer.kv, |
| 496 | layer.num_heads, |
| 497 | layer.head_dim, |
| 498 | layer.block, |
| 499 | ) |
| 500 | gather, slot_ends = launch.pool |
| 501 | compact = layer.q.new_empty((t1 - t0, num_cols), dtype=torch.float32) |
| 502 | q, q_scale, weights = layer.rows(t0, t1) |
| 503 | pa.paged_mxfp4_mqa_logits( |
| 504 | q, |
| 505 | q_scale, |
| 506 | layer.kv, |
| 507 | weights, |
| 508 | launch.context_len, |
| 509 | launch.block_table, |
| 510 | num_cols, |
| 511 | out_logits=compact, |
| 512 | clean_logits=False, |
| 513 | row_ends=slot_ends, |
| 514 | use_gather=True, |
| 515 | candidates=gather, |
| 516 | ) |
| 517 | out = layer.topk_buffer[t0:t1, : layer.topk_tokens] |
| 518 | _topk(compact, slot_ends, out, layer.topk_tokens) |
| 519 | _remap_compact_topk(out, gather["positions"], layer.block) |
| 520 | |
| 521 | |
| 522 | def _gather_decode(layer: _Layer, num_cols: int) -> None: |
| 523 | pa = _aiter() |
| 524 | assert layer.candidates is not None |
| 525 | metadata = layer.metadata |
| 526 | lengths = metadata.decode_row_lens |
| 527 | assert metadata.decode is not None and lengths is not None |
| 528 | block_table = metadata.decode.block_table |
| 529 | rows = lengths.shape[0] |
| 530 | if metadata.decode_gather is None: |
| 531 | metadata.decode_gather = pa.build_candidate_gather( |
| 532 | layer.candidates[:rows], |
| 533 | lengths, |
| 534 | block_table, |
| 535 | layer.kv, |
| 536 | layer.num_heads, |
| 537 | layer.head_dim, |
| 538 | layer.block, |
| 539 | ) |
| 540 | gather, slot_ends = metadata.decode_gather |
| 541 | (compact,) = current_workspace_manager().get_simultaneous( |
| 542 | ((rows, num_cols), torch.float32) |
| 543 | ) |
| 544 | q, q_scale, weights = layer.decode_rows(rows) |
| 545 | pa.paged_mxfp4_mqa_logits( |
| 546 | q, |
| 547 | q_scale, |
| 548 | layer.kv, |
| 549 | weights, |
| 550 | lengths, |
| 551 | block_table, |
| 552 | num_cols, |
| 553 | out_logits=compact, |
| 554 | clean_logits=False, |
| 555 | row_ends=slot_ends, |
| 556 | use_gather=True, |
| 557 | candidates=gather, |
| 558 | ) |
| 559 | out = layer.topk_buffer[:rows, : layer.topk_tokens] |
| 560 | _topk(compact, slot_ends, out, layer.topk_tokens) |
| 561 | _remap_compact_topk(out, gather["positions"], layer.block) |
| 562 | |
| 563 | |
| 564 | @eager_break_during_capture |
| 565 | def rocm_mxfp4_sparse_attn_indexer( |
| 566 | hidden_states: torch.Tensor, |
| 567 | k_cache_prefix: str, |
| 568 | kv_cache: torch.Tensor, |
| 569 | q_values: torch.Tensor, |
| 570 | q_scale: torch.Tensor, |
| 571 | weights: torch.Tensor, |
| 572 | topk_tokens: int, |
| 573 | head_dim: int, |
| 574 | max_model_len: int, |
| 575 | topk_indices_buffer: torch.Tensor, |
| 576 | compress_ratio: int = 1, |
| 577 | candidate_blocks: torch.Tensor | None = None, |
| 578 | candidate_block_size: int = 0, |
| 579 | candidate_write: bool = False, |
| 580 | ) -> torch.Tensor: |
| 581 | """Dense indexer: every layer that scores the whole context, the |
| 582 | candidate source included. With ``candidate_blocks`` and not |
| 583 | ``candidate_write`` the scores are masked to the pool first, as the |
| 584 | shared path does. |
| 585 | |
| 586 | The model supplies ``q_values`` [T, H, D // 2] packed e2m1, ``q_scale`` |
| 587 | holding each head's D // 32 ue8m0 bytes (one int32 per head at D = 128), |
| 588 | fp32 ``weights`` with the softmax and head scales folded in, and |
| 589 | ``kv_cache`` pages in `rocm_paged_mxfp4_cache_layout` order for H heads. |
| 590 | """ |
| 591 | if not isinstance(get_forward_context().attn_metadata, dict): |
| 592 | reserve_rocm_mxfp4_indexer_workspace( |
| 593 | hidden_states, |
| 594 | max_model_len, |
| 595 | candidate_block_size if candidate_write else 0, |
| 596 | ) |
| 597 | return topk_indices_buffer |
| 598 | layer = _layer( |
| 599 | k_cache_prefix, |
| 600 | kv_cache, |
| 601 | q_values, |
| 602 | q_scale, |
| 603 | weights, |
| 604 | topk_indices_buffer, |
| 605 | topk_tokens, |
| 606 | head_dim, |
| 607 | compress_ratio, |
| 608 | candidate_blocks, |
| 609 | candidate_block_size, |
| 610 | ) |
| 611 | metadata = layer.metadata |
| 612 | topk_indices_buffer[: hidden_states.shape[0]] = -1 |
| 613 | if metadata.prefill is not None: |
| 614 | for chunk, plan in zip(metadata.prefill.chunks, metadata.prefill_plans): |
| 615 | _dense_prefill(layer, chunk, plan, candidate_write) |
| 616 | if metadata.decode is not None: |
| 617 | _dense_decode(layer, max_model_len, candidate_write) |
| 618 | return topk_indices_buffer |
| 619 | |
| 620 | |
| 621 | @eager_break_during_capture |
| 622 | def rocm_mxfp4_sparse_mqa_indexer( |
| 623 | hidden_states: torch.Tensor, |
| 624 | k_cache_prefix: str, |
| 625 | kv_cache: torch.Tensor, |
| 626 | q_values: torch.Tensor, |
| 627 | q_scale: torch.Tensor, |
| 628 | weights: torch.Tensor, |
| 629 | topk_tokens: int, |
| 630 | head_dim: int, |
| 631 | max_model_len: int, |
| 632 | topk_indices_buffer: torch.Tensor, |
| 633 | compress_ratio: int, |
| 634 | candidate_blocks: torch.Tensor, |
| 635 | candidate_block_size: int, |
| 636 | num_candidate_cols: int, |
| 637 | ) -> torch.Tensor: |
| 638 | """Candidate consumer: score only the source's pool where the builder's |
| 639 | length gate says it pays, else the dense walk masked to the pool.""" |
| 640 | if not isinstance(get_forward_context().attn_metadata, dict): |
| 641 | reserve_rocm_mxfp4_indexer_workspace( |
| 642 | hidden_states, |
| 643 | max(max_model_len, num_candidate_cols), |
| 644 | gather_block_size=candidate_block_size, |
| 645 | num_candidate_cols=num_candidate_cols, |
| 646 | ) |
| 647 | return topk_indices_buffer |
| 648 | layer = _layer( |
| 649 | k_cache_prefix, |
| 650 | kv_cache, |
| 651 | q_values, |
| 652 | q_scale, |
| 653 | weights, |
| 654 | topk_indices_buffer, |
| 655 | topk_tokens, |
| 656 | head_dim, |
| 657 | compress_ratio, |
| 658 | candidate_blocks, |
| 659 | candidate_block_size, |
| 660 | ) |
| 661 | metadata = layer.metadata |
| 662 | topk_indices_buffer[: hidden_states.shape[0]] = -1 |
| 663 | if metadata.prefill is not None: |
| 664 | for chunk, plan in zip(metadata.prefill.chunks, metadata.prefill_plans): |
| 665 | if not plan.use_gather: |
| 666 | _dense_prefill(layer, chunk, plan, candidate_write=False) |
| 667 | for launch in metadata.gather_launches: |
| 668 | _gather_prefill(layer, launch, num_candidate_cols) |
| 669 | if metadata.decode is not None: |
| 670 | # A FULL graph cannot follow the per-step gate; the gather is the side |
| 671 | # that stays flat in the context length. |
| 672 | if metadata.decode_use_gather or layer.full_graph: |
| 673 | _gather_decode(layer, num_candidate_cols) |
| 674 | else: |
| 675 | _dense_decode(layer, max_model_len, candidate_write=False) |
| 676 | return topk_indices_buffer |