| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Indexer metadata for the ROCm paged MXFP4 path. |
| 4 | |
| 5 | Extends the dense indexer metadata with what aiter's paged MXFP4 MQA-logits |
| 6 | kernel needs to read the cache in place: each prefill chunk as its requests, |
| 7 | packed through query_start_loc into one launch, and the decode rows as |
| 8 | next_n-row sequences on uniform steps. `RocmMxfp4IndexerMetadataBuilder` plans |
| 9 | this for any model's indexer, and a two-level indexer's candidate blocks once |
| 10 | a subclass sets their size. DeepSeek-V4.1's builder does, and with |
| 11 | ``AttentionConfig.indexer_sparse_logits`` adds the per-step state the |
| 12 | candidate consumers share. |
| 13 | """ |
| 14 | |
| 15 | import bisect |
| 16 | from dataclasses import dataclass, field, fields |
| 17 | |
| 18 | import torch |
| 19 | |
| 20 | from vllm.config import VllmConfig |
| 21 | from vllm.logger import init_logger |
| 22 | from vllm.v1.attention.backend import ( |
| 23 | AttentionCGSupport, |
| 24 | CommonAttentionMetadata, |
| 25 | MultipleOf, |
| 26 | ) |
| 27 | from vllm.v1.attention.backends.mla.indexer import ( |
| 28 | DeepseekV32IndexerMetadata, |
| 29 | DeepseekV32IndexerMetadataBuilder, |
| 30 | DeepseekV32IndexerPrefillChunkMetadata, |
| 31 | DeepseekV41IndexerBackend, |
| 32 | ) |
| 33 | from vllm.v1.attention.ops.rocm_paged_mxfp4_indexer import ( |
| 34 | build_rocm_mxfp4_decode_schedule, |
| 35 | rocm_mxfp4_consumer_rows, |
| 36 | rocm_mxfp4_decode_schedule_words, |
| 37 | rocm_paged_mxfp4_cache_layout, |
| 38 | ) |
| 39 | from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheSpec |
| 40 | |
| 41 | logger = init_logger(__name__) |
| 42 | |
| 43 | # The consumers walk the pool instead of the context once the context is this |
| 44 | # many pools long. Set from kernel timings, not end-to-end runs: the gather |
| 45 | # loads more per key than the dense walk, so it only wins on a large cut. |
| 46 | _GATHER_MIN_CUT_DECODE = 3.0 |
| 47 | _GATHER_MIN_CUT_PREFILL = 1.5 |
| 48 | |
| 49 | |
| 50 | @dataclass |
| 51 | class RocmMxfp4PrefillPlan: |
| 52 | """One prefill chunk as the requests the kernel launches on, all in one |
| 53 | launch: the kernel finds each row's request through query_start_loc, and |
| 54 | a request's rows share its KV tiles.""" |
| 55 | |
| 56 | requests: list[tuple[int, int, int]] |
| 57 | """(first row, end row, request) per request with rows in the chunk; rows |
| 58 | count from ``chunk.token_start``, requests index ``chunk.block_table``.""" |
| 59 | width: int |
| 60 | """Logits columns: an upper bound on the chunk's compressed contexts.""" |
| 61 | row_ends: torch.Tensor |
| 62 | """[rows] int32 exclusive compressed key bound of each row.""" |
| 63 | context_lens: torch.Tensor |
| 64 | """[chunk.num_reqs] int32 compressed context of each request.""" |
| 65 | query_start_loc: torch.Tensor |
| 66 | """[chunk.num_reqs + 1] int32 row offsets of the chunk's requests, local to |
| 67 | the chunk.""" |
| 68 | first_request: int |
| 69 | """The step's index of the chunk's first request.""" |
| 70 | block_ends: torch.Tensor | None = None |
| 71 | """[rows] int32 candidate blocks each row sees, for the source's pool.""" |
| 72 | use_gather: bool = False |
| 73 | |
| 74 | |
| 75 | @dataclass |
| 76 | class RocmMxfp4GatherLaunch: |
| 77 | """Query rows [token_start, token_end) of one request, which the candidate |
| 78 | consumers score against the pool in one launch.""" |
| 79 | |
| 80 | token_start: int |
| 81 | token_end: int |
| 82 | block_table: torch.Tensor |
| 83 | """[1, max_blocks] int32, the request's block table row.""" |
| 84 | context_len: torch.Tensor |
| 85 | """[1] int32 compressed context of the request.""" |
| 86 | row_ends: torch.Tensor |
| 87 | """[rows] int32 exclusive compressed key bound of each row.""" |
| 88 | pool: tuple[dict, torch.Tensor] | None = None |
| 89 | """The resolved candidate lists, built by the first consumer layer and |
| 90 | reused by the rest.""" |
| 91 | |
| 92 | |
| 93 | @dataclass |
| 94 | class RocmMxfp4NativeDecode: |
| 95 | """A decode step whose requests all have next_n query rows, as the |
| 96 | kernel's sequences: a workgroup then walks a KV tile once for up to |
| 97 | next_n rows instead of once per row.""" |
| 98 | |
| 99 | next_n: int |
| 100 | context_lens: torch.Tensor |
| 101 | """[requests] int32 compressed context of each request.""" |
| 102 | block_table: torch.Tensor |
| 103 | """[requests, max_blocks] int32, each request's first flattened row.""" |
| 104 | |
| 105 | |
| 106 | def native_decode( |
| 107 | row_lens: torch.Tensor, |
| 108 | row_block_table: torch.Tensor, |
| 109 | query_lens: list[int], |
| 110 | next_n: int, |
| 111 | context_lens: torch.Tensor, |
| 112 | ) -> RocmMxfp4NativeDecode | None: |
| 113 | """The step as next_n-row sequences, or None when a request has fewer rows. |
| 114 | |
| 115 | ``row_lens`` and ``row_block_table`` are the flattened rows' bounds and |
| 116 | block table, ``query_lens`` the decode requests' query lengths. Trailing |
| 117 | empty requests are cudagraph padding: they keep their next_n rows, with no |
| 118 | context. Only the step's shape decides, so a FULL graph captured on a |
| 119 | uniform batch replays the same launch on a padded one. |
| 120 | """ |
| 121 | num_reqs = len(query_lens) |
| 122 | num_full = next((i for i, n in enumerate(query_lens) if n != next_n), num_reqs) |
| 123 | if ( |
| 124 | next_n <= 1 |
| 125 | or row_lens.shape[0] != num_reqs * next_n |
| 126 | or any(query_lens[num_full:]) |
| 127 | ): |
| 128 | return None |
| 129 | return RocmMxfp4NativeDecode( |
| 130 | next_n, context_lens[:num_reqs], row_block_table[::next_n] |
| 131 | ) |
| 132 | |
| 133 | |
| 134 | @dataclass |
| 135 | class RocmMxfp4IndexerMetadata(DeepseekV32IndexerMetadata): |
| 136 | prefill_plans: list[RocmMxfp4PrefillPlan] = field(default_factory=list) |
| 137 | """Parallel to ``prefill.chunks``.""" |
| 138 | decode_row_lens: torch.Tensor | None = None |
| 139 | """[rows] int32 compressed context of each decode query row.""" |
| 140 | decode_block_ends: torch.Tensor | None = None |
| 141 | """[rows] int32 candidate blocks each decode row sees, for the source's |
| 142 | pool.""" |
| 143 | decode_native: RocmMxfp4NativeDecode | None = None |
| 144 | """The same rows as next_n-row sequences, for the dense launches.""" |
| 145 | decode_schedule: torch.Tensor | None = None |
| 146 | """The dense decode launches' work schedule, built once for the step.""" |
| 147 | decode_use_gather: bool = False |
| 148 | """Whether the candidate consumers gather for the decode rows.""" |
| 149 | decode_gather: tuple[dict, torch.Tensor] | None = None |
| 150 | """The pool the first consumer resolved for the decode rows.""" |
| 151 | gather_launches: list[RocmMxfp4GatherLaunch] = field(default_factory=list) |
| 152 | """The prefill rows of the chunks that gather, as the consumers launch |
| 153 | them.""" |
| 154 | |
| 155 | |
| 156 | def plan_prefill_chunks( |
| 157 | chunks: list[DeepseekV32IndexerPrefillChunkMetadata], |
| 158 | query_start_loc: list[int], |
| 159 | seq_lens_cpu: torch.Tensor, |
| 160 | context_lens: torch.Tensor, |
| 161 | compress_ratio: int, |
| 162 | candidate_block_size: int, |
| 163 | min_gather_width: float | None, |
| 164 | query_start_loc_device: torch.Tensor, |
| 165 | ) -> list[RocmMxfp4PrefillPlan]: |
| 166 | """Split each chunk into its requests. |
| 167 | |
| 168 | ``query_start_loc`` and ``seq_lens_cpu`` are the step's host copies (the |
| 169 | latter an upper bound), ``context_lens`` the device compressed lengths. |
| 170 | A chunk covers whole requests, or a query slice of one request; either |
| 171 | way it launches once, packed through the chunk's part of the device |
| 172 | ``query_start_loc``. |
| 173 | """ |
| 174 | block = candidate_block_size |
| 175 | plans = [] |
| 176 | for chunk in chunks: |
| 177 | t0, t1 = chunk.token_start, chunk.token_end |
| 178 | first = bisect.bisect_right(query_start_loc, t0) - 1 |
| 179 | requests = [] |
| 180 | for req in range(chunk.num_reqs): |
| 181 | lo = max(query_start_loc[first + req], t0) - t0 |
| 182 | hi = min(query_start_loc[first + req + 1], t1) - t0 |
| 183 | if hi > lo: |
| 184 | requests.append((lo, hi, req)) |
| 185 | reqs = slice(first, first + chunk.num_reqs) |
| 186 | packed = torch.clamp(query_start_loc_device[first : reqs.stop + 1], t0, t1) |
| 187 | packed -= t0 |
| 188 | width = int(seq_lens_cpu[reqs].max()) // compress_ratio |
| 189 | row_ends = chunk.cu_seqlen_ke - chunk.cu_seqlen_ks |
| 190 | plans.append( |
| 191 | RocmMxfp4PrefillPlan( |
| 192 | requests=requests, |
| 193 | width=width, |
| 194 | row_ends=row_ends, |
| 195 | context_lens=context_lens[reqs], |
| 196 | query_start_loc=packed, |
| 197 | first_request=first, |
| 198 | block_ends=(row_ends + block - 1) // block if block else None, |
| 199 | use_gather=min_gather_width is not None and width >= min_gather_width, |
| 200 | ) |
| 201 | ) |
| 202 | return plans |
| 203 | |
| 204 | |
| 205 | def split_prefill_chunks( |
| 206 | context_lens: torch.Tensor, |
| 207 | query_lens: torch.Tensor, |
| 208 | max_logits_bytes: int, |
| 209 | request_offset: int = 0, |
| 210 | ) -> list[tuple[slice, slice]]: |
| 211 | """(request slice, query slice) chunks, as the base chunker returns them. |
| 212 | |
| 213 | Nothing gathers K here, so the requests' contexts do not add up: a chunk's |
| 214 | logits are [rows, widest compressed context] fp32, at most |
| 215 | ``max_logits_bytes``. A request too wide to launch whole is cut on its |
| 216 | query rows. |
| 217 | """ |
| 218 | |
| 219 | def max_rows(width: int) -> int: |
| 220 | return max(1, max_logits_bytes // (4 * max(width, 1))) |
| 221 | |
| 222 | chunks: list[tuple[slice, slice]] = [] |
| 223 | end = 0 |
| 224 | while end < len(context_lens): |
| 225 | start, rows, width = end, 0, 0 |
| 226 | while end < len(context_lens): |
| 227 | q = int(query_lens[end]) |
| 228 | w = max(width, int(context_lens[end])) |
| 229 | if rows and rows + q > max_rows(w): |
| 230 | break |
| 231 | rows, width = rows + q, w |
| 232 | end += 1 |
| 233 | step = max_rows(width) |
| 234 | reqs = slice(start + request_offset, end + request_offset) |
| 235 | chunks.extend( |
| 236 | (reqs, slice(lo, min(lo + step, rows))) for lo in range(0, rows, step) |
| 237 | ) |
| 238 | return chunks |
| 239 | |
| 240 | |
| 241 | def plan_gather_launches( |
| 242 | chunks: list[DeepseekV32IndexerPrefillChunkMetadata], |
| 243 | plans: list[RocmMxfp4PrefillPlan], |
| 244 | max_rows: int, |
| 245 | ) -> list[RocmMxfp4GatherLaunch]: |
| 246 | """The candidate consumers' launches over the chunks that gather. |
| 247 | |
| 248 | A consumer's logits are [rows, pool] however long the context, so the |
| 249 | dense split, sized for the context, would launch it far more often than |
| 250 | its memory needs. Each request's rows go back together across the chunks |
| 251 | that cut them, then out again at most ``max_rows`` a launch. |
| 252 | """ |
| 253 | # Per request: (first token, [(plan, first row, end row)]), and the chunk |
| 254 | # and request index its block table and context come from. |
| 255 | runs: list[tuple[int, list, DeepseekV32IndexerPrefillChunkMetadata, int]] = [] |
| 256 | last: tuple[int, int] | None = None |
| 257 | for chunk, plan in zip(chunks, plans): |
| 258 | if not plan.use_gather: |
| 259 | last = None |
| 260 | continue |
| 261 | for lo, hi, req in plan.requests: |
| 262 | request, t_lo = plan.first_request + req, chunk.token_start + lo |
| 263 | if last != (request, t_lo): |
| 264 | runs.append((t_lo, [], chunk, req)) |
| 265 | runs[-1][1].append((plan, lo, hi)) |
| 266 | last = (request, chunk.token_start + hi) |
| 267 | launches = [] |
| 268 | for token_start, pieces, chunk, req in runs: |
| 269 | ends = [plan.row_ends[lo:hi] for plan, lo, hi in pieces] |
| 270 | row_ends = ends[0] if len(ends) == 1 else torch.cat(ends) |
| 271 | block_table = chunk.block_table[req : req + 1] |
| 272 | context_len = pieces[0][0].context_lens[req : req + 1] |
| 273 | for lo in range(0, row_ends.shape[0], max_rows): |
| 274 | hi = min(lo + max_rows, row_ends.shape[0]) |
| 275 | launches.append( |
| 276 | RocmMxfp4GatherLaunch( |
| 277 | token_start + lo, |
| 278 | token_start + hi, |
| 279 | block_table, |
| 280 | context_len, |
| 281 | row_ends[lo:hi], |
| 282 | ) |
| 283 | ) |
| 284 | return launches |
| 285 | |
| 286 | |
| 287 | class DeepseekV41RocmMxfp4IndexerBackend(DeepseekV41IndexerBackend): |
| 288 | @staticmethod |
| 289 | def get_name() -> str: |
| 290 | return "DEEPSEEK_V41_ROCM_MXFP4_INDEXER" |
| 291 | |
| 292 | @staticmethod |
| 293 | def get_supported_kernel_block_sizes( |
| 294 | kv_cache_spec: KVCacheSpec | None = None, |
| 295 | ) -> list[int | MultipleOf]: |
| 296 | # --block-size 128 preferred. |
| 297 | return [64, 128] |
| 298 | |
| 299 | @staticmethod |
| 300 | def get_builder_cls() -> type["DeepseekV41RocmMxfp4IndexerMetadataBuilder"]: |
| 301 | return DeepseekV41RocmMxfp4IndexerMetadataBuilder |
| 302 | |
| 303 | |
| 304 | class RocmMxfp4IndexerMetadataBuilder(DeepseekV32IndexerMetadataBuilder): |
| 305 | """Dense indexer metadata plus the in-place paged launches' layout. |
| 306 | |
| 307 | Requirements are checked here, at engine start. |
| 308 | """ |
| 309 | |
| 310 | @classmethod |
| 311 | def get_cudagraph_support( |
| 312 | cls, |
| 313 | vllm_config: VllmConfig, |
| 314 | kv_cache_spec: KVCacheSpec, |
| 315 | ) -> AttentionCGSupport: |
| 316 | # Decode rows are flattened into persistent buffers, and whether a |
| 317 | # step launches them as next_n-row sequences depends on its shape |
| 318 | # alone (see native_decode). |
| 319 | return AttentionCGSupport.ALWAYS |
| 320 | |
| 321 | def __init__( |
| 322 | self, |
| 323 | kv_cache_spec: AttentionSpec, |
| 324 | layer_names: list[str], |
| 325 | vllm_config: VllmConfig, |
| 326 | device: torch.device, |
| 327 | **kwargs, |
| 328 | ) -> None: |
| 329 | super().__init__(kv_cache_spec, layer_names, vllm_config, device, **kwargs) |
| 330 | name = "the ROCm MXFP4 indexer" |
| 331 | if self.dcp_world_size > 1 or self.use_pcp: |
| 332 | raise NotImplementedError( |
| 333 | f"{name} does not support decode/prefill context parallel." |
| 334 | ) |
| 335 | hf_config = vllm_config.model_config.hf_text_config |
| 336 | num_heads, head_dim = hf_config.index_n_heads, hf_config.index_head_dim |
| 337 | self.num_heads, self.head_dim = num_heads, head_dim |
| 338 | self.page_entries = kv_cache_spec.block_size // self.compress_ratio |
| 339 | self.cache_layout = rocm_paged_mxfp4_cache_layout( |
| 340 | num_heads, head_dim, self.page_entries |
| 341 | ) |
| 342 | # Ragged steps take a row per query token, as do the candidate |
| 343 | # consumers' gathers. Uniform steps also launch the dense layers on |
| 344 | # next_n-row sequences, see native_decode. |
| 345 | self.use_flattening = True |
| 346 | # A two-level (candidate block) indexer's builder sets these. |
| 347 | self.candidate_block_size = 0 |
| 348 | self.min_gather_width: float | None = None |
| 349 | |
| 350 | self.decode_block_ends_buffer = torch.zeros( |
| 351 | self.arange_buffer.shape[0], dtype=torch.int32, device=device |
| 352 | ) |
| 353 | self.context_lens_buffer = torch.zeros( |
| 354 | self.arange_buffer.shape[0], dtype=torch.int32, device=device |
| 355 | ) |
| 356 | spec_config = vllm_config.speculative_config |
| 357 | # Adaptive verification replays a FULL graph on drafts reallocated |
| 358 | # across requests, so a step's shape does not fix its query lengths. |
| 359 | adaptive = spec_config is not None and spec_config.enable_adaptive_verification |
| 360 | self.native_next_n = 1 if adaptive else self.num_speculative_tokens + 1 |
| 361 | # The indexer's logits width, which the decode schedule sizes slices by. |
| 362 | max_model_len = vllm_config.model_config.max_model_len |
| 363 | self.logits_width = max_model_len // self.compress_ratio |
| 364 | # Every dense layer of this group reads the same rows through the same |
| 365 | # page geometry, so the step builds their schedule once, not per layer. |
| 366 | self.decode_schedule_buffer = torch.empty( |
| 367 | rocm_mxfp4_decode_schedule_words( |
| 368 | num_heads, head_dim, self.page_entries, self.native_next_n |
| 369 | ), |
| 370 | dtype=torch.int32, |
| 371 | device=device, |
| 372 | ) |
| 373 | |
| 374 | def _split_indexer_prefill_chunks( # type: ignore[override] |
| 375 | self, |
| 376 | compressed_seq_lens_cpu: torch.Tensor, |
| 377 | prefill_query_lens_cpu: torch.Tensor, |
| 378 | workspace_size: int, |
| 379 | max_logits_bytes: int, |
| 380 | request_offset: int = 0, |
| 381 | ) -> list[tuple[slice, slice]]: |
| 382 | return split_prefill_chunks( |
| 383 | compressed_seq_lens_cpu, |
| 384 | prefill_query_lens_cpu, |
| 385 | max_logits_bytes, |
| 386 | request_offset, |
| 387 | ) |
| 388 | |
| 389 | def build( |
| 390 | self, |
| 391 | common_prefix_len: int, |
| 392 | common_attn_metadata: CommonAttentionMetadata, |
| 393 | fast_build: bool = False, |
| 394 | ) -> RocmMxfp4IndexerMetadata: |
| 395 | base = super().build(common_prefix_len, common_attn_metadata, fast_build) |
| 396 | metadata = RocmMxfp4IndexerMetadata( |
| 397 | **{f.name: getattr(base, f.name) for f in fields(base)} |
| 398 | ) |
| 399 | cm = common_attn_metadata |
| 400 | context_lens = self.context_lens_buffer[: cm.num_reqs] |
| 401 | torch.floor_divide( |
| 402 | cm.seq_lens[: cm.num_reqs], self.compress_ratio, out=context_lens |
| 403 | ) |
| 404 | if base.decode is not None: |
| 405 | lengths = base.decode.seq_lens.view(-1) |
| 406 | metadata.decode_row_lens = lengths |
| 407 | block = self.candidate_block_size |
| 408 | if block: |
| 409 | ends = self.decode_block_ends_buffer[: lengths.shape[0]] |
| 410 | torch.add(lengths, block - 1, out=ends) |
| 411 | ends.floor_divide_(block) |
| 412 | metadata.decode_block_ends = ends |
| 413 | query_start_loc = cm.query_start_loc_cpu[: base.num_decodes + 1] |
| 414 | metadata.decode_native = native_decode( |
| 415 | lengths, |
| 416 | base.decode.block_table, |
| 417 | torch.diff(query_start_loc).tolist(), |
| 418 | self.native_next_n, |
| 419 | context_lens, |
| 420 | ) |
| 421 | metadata.decode_schedule = build_rocm_mxfp4_decode_schedule( |
| 422 | lengths, |
| 423 | self.num_heads, |
| 424 | self.head_dim, |
| 425 | self.page_entries, |
| 426 | self.decode_schedule_buffer, |
| 427 | self.logits_width, |
| 428 | metadata.decode_native, |
| 429 | ) |
| 430 | if base.prefill is not None: |
| 431 | assert cm.seq_lens_cpu_upper_bound is not None |
| 432 | metadata.prefill_plans = plan_prefill_chunks( |
| 433 | base.prefill.chunks, |
| 434 | cm.query_start_loc_cpu.tolist(), |
| 435 | cm.seq_lens_cpu_upper_bound, |
| 436 | context_lens, |
| 437 | self.compress_ratio, |
| 438 | self.candidate_block_size, |
| 439 | self.min_gather_width, |
| 440 | cm.query_start_loc, |
| 441 | ) |
| 442 | return metadata |
| 443 | |
| 444 | |
| 445 | class DeepseekV41RocmMxfp4IndexerMetadataBuilder(RocmMxfp4IndexerMetadataBuilder): |
| 446 | """DeepSeek-V4.1's two-level indexer: sets the candidate block size the |
| 447 | base plans with and, with ``indexer_sparse_logits``, adds the consumers' |
| 448 | gathers.""" |
| 449 | |
| 450 | def __init__( |
| 451 | self, |
| 452 | kv_cache_spec: AttentionSpec, |
| 453 | layer_names: list[str], |
| 454 | vllm_config: VllmConfig, |
| 455 | device: torch.device, |
| 456 | **kwargs, |
| 457 | ) -> None: |
| 458 | super().__init__(kv_cache_spec, layer_names, vllm_config, device, **kwargs) |
| 459 | hf_config = vllm_config.model_config.hf_text_config |
| 460 | self.candidate_block_size = getattr(hf_config, "candidate_block_size", 0) |
| 461 | self.num_candidate_cols = 0 |
| 462 | if vllm_config.attention_config.indexer_sparse_logits: |
| 463 | prefix = "attention_config.indexer_sparse_logits" |
| 464 | block = self.candidate_block_size |
| 465 | topk_blocks = getattr(hf_config, "candidate_topk_blocks", 0) |
| 466 | layout = self.cache_layout |
| 467 | if block <= 0 or topk_blocks <= 0: |
| 468 | raise ValueError( |
| 469 | f"{prefix} requires a two-level (candidate block) indexer " |
| 470 | "model such as DeepSeek-V4.1." |
| 471 | ) |
| 472 | if topk_blocks & (topk_blocks - 1): |
| 473 | raise ValueError( |
| 474 | f"{prefix} on ROCm resolves the candidate pool with a sort " |
| 475 | f"that needs a power-of-two block count, not {topk_blocks}." |
| 476 | ) |
| 477 | if self.page_entries % block or block > layout.n_per_tile: |
| 478 | raise ValueError( |
| 479 | f"{prefix} on ROCm needs {block}-entry candidate blocks to " |
| 480 | f"tile a {self.page_entries}-entry indexer page and fit in " |
| 481 | f"one {layout.n_per_tile}-entry shuffle group." |
| 482 | ) |
| 483 | self.num_candidate_cols = topk_blocks * block |
| 484 | self.min_gather_width = self.num_candidate_cols * _GATHER_MIN_CUT_PREFILL |
| 485 | self.gather_rows = rocm_mxfp4_consumer_rows(self.num_candidate_cols) |
| 486 | logger.info_once( |
| 487 | "DeepSeek V4.1 indexer: candidate consumers score the %d " |
| 488 | "candidate positions with aiter's paged MXFP4 gather once the " |
| 489 | "context is %.1fx (decode) / %.1fx (prefill) as long.", |
| 490 | self.num_candidate_cols, |
| 491 | _GATHER_MIN_CUT_DECODE, |
| 492 | _GATHER_MIN_CUT_PREFILL, |
| 493 | ) |
| 494 | |
| 495 | def _gather_pays(self, context: int, cut: float) -> bool: |
| 496 | return self.num_candidate_cols > 0 and context >= cut * self.num_candidate_cols |
| 497 | |
| 498 | def build( |
| 499 | self, |
| 500 | common_prefix_len: int, |
| 501 | common_attn_metadata: CommonAttentionMetadata, |
| 502 | fast_build: bool = False, |
| 503 | ) -> RocmMxfp4IndexerMetadata: |
| 504 | metadata = super().build(common_prefix_len, common_attn_metadata, fast_build) |
| 505 | if metadata.decode is not None: |
| 506 | metadata.decode_use_gather = self._gather_pays( |
| 507 | metadata.max_seq_len // self.compress_ratio, _GATHER_MIN_CUT_DECODE |
| 508 | ) |
| 509 | if metadata.prefill is not None and self.num_candidate_cols: |
| 510 | metadata.gather_launches = plan_gather_launches( |
| 511 | metadata.prefill.chunks, metadata.prefill_plans, self.gather_rows |
| 512 | ) |
| 513 | return metadata |