vllm.models.minimax_m3.amd.indexer_aiter ¶
AITER (ROCm) indexer impl for MiniMax M3.
Scores index blocks and selects the top-k with AITER's fp8 MFMA kernels on both sides of the batch: pa_sparse_block_score_decode for the uniform-query-length decode rows and pa_sparse_block_score_prefill for the ragged prefill rows, then pa_sparse_block_topk for each. The two scoring passes share one tile body in AITER, so they agree block for block, and both encode the forced init/local blocks as the same sentinel scores the Triton indexer uses.
The top-k also emits the attend's page table. The winners are already in its workgroup's LDS, so resolving them through the block table there costs one wave and replaces the separate Triton pass in ops.sparse_pa; the rows it writes are one per (token, kv head) with the page ids folded head-minor, which is the layout pa_decode_gluon reads after it flattens the cache.
The score kernels are built on v_mfma_f32_16x16x32_fp8_fp8, so this impl requires an fp8 (e4m3) index cache and an fp8 index query -- the fused QK-norm/RoPE kernel emits both directly when the index cache is e4m3. See aiter_indexer_unsupported_reason for the full set of limits; select_aiter_indexer_impl_cls refuses to pick this impl unless they all hold, and the model falls back to the platform-neutral MiniMaxM3Indexer.
Classes:
-
MiniMaxM3AiterIndexer–MiniMaxM3Indexer's surface over the AITER impl. -
MiniMaxM3IndexerAiterBackend–Indexer side-cache backend selecting the AITER builder.
-
MiniMaxM3IndexerAiterImpl–AITER fp8 score + top-k for both prefill and decode.
-
MiniMaxM3IndexerAiterMetadata–Adds the per-row shape the ragged top-k needs.
-
MiniMaxM3IndexerAiterMetadataBuilder–The Triton indexer's metadata plus the prefill rows' causal shape.
Functions:
-
aiter_indexer_max_decode_query_len–Longest query a decode row can carry, which spec decode is what sets.
-
aiter_indexer_unsupported_reason–Return why this config cannot use the AITER indexer, or None if it can.
-
aiter_msa_kernels_unavailable_reason–Return why the AITER MSA score/top-k ops cannot be imported, or None.
-
score_block_width–Block-axis width the top-k requires of the score buffer.
-
select_aiter_indexer_impl_cls–The AITER indexer impl if this config can use it, else None.
MiniMaxM3AiterIndexer ¶
Bases: Module
MiniMaxM3Indexer's surface over the AITER impl.
The platform-neutral wrapper picks its impl through common's selector and forwards a Triton-only set of fused-table arguments, neither of which can reach this impl without editing common. This holds the same three members the attention layer uses -- impl, index_cache, num_index_heads -- and forwards the one argument AITER needs instead.
Source code in vllm/models/minimax_m3/amd/indexer_aiter.py
MiniMaxM3IndexerAiterBackend ¶
Bases: MiniMaxM3IndexerBackend
Indexer side-cache backend selecting the AITER builder.
Source code in vllm/models/minimax_m3/amd/indexer_aiter.py
MiniMaxM3IndexerAiterImpl ¶
Bases: MiniMaxM3IndexerImpl
AITER fp8 score + top-k for both prefill and decode.
Attributes:
-
pages_per_block(int) –Physical pages one selected block expands into for the attend.
Source code in vllm/models/minimax_m3/amd/indexer_aiter.py
361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499 500 501 502 503 504 505 506 507 508 509 510 511 512 513 514 515 516 517 518 519 520 521 522 523 524 525 526 527 528 529 530 531 532 533 534 535 536 537 538 539 540 541 542 543 544 545 546 547 548 549 550 551 552 553 554 555 556 557 558 559 560 | |
pages_per_block property ¶
Physical pages one selected block expands into for the attend.
_new_score(rows, max_seq_len) ¶
Score buffer for rows query rows.
Left uninitialized on purpose: the score pass writes every block up to the longest row it covers, and the top-k reads only the blocks its own row can see, so nothing downstream observes the padded tail. Filling it would cost a write over the whole [H, rows, width] extent, which at long context is the largest tensor in the indexer.
Source code in vllm/models/minimax_m3/amd/indexer_aiter.py
_table_rows(lo, hi) ¶
The page table and context rows covering score rows [lo, hi).
One table row per (token, kv head), head minor, which is the order pa_decode_gluon reads once it flattens the cache. Slices of the shared buffers rather than copies, so the attend sees the writes; a model that reserved no buffers gets throwaway ones, and the attend rebuilds the table itself in that case.
Source code in vllm/models/minimax_m3/amd/indexer_aiter.py
MiniMaxM3IndexerAiterMetadata dataclass ¶
Bases: MiniMaxM3IndexerMetadata
Adds the per-row shape the ragged top-k needs.
The uniform decode rows are recovered inside the kernel from seq_lens and the shared query length, which also clamps cudagraph padding rows to nothing. Prefill rows have no such shape, so it is materialized here once per forward and shared by every layer -- which is also where the emitted page table's tail block comes from, since a block count alone does not say how many tokens the last block holds.
Source code in vllm/models/minimax_m3/amd/indexer_aiter.py
MiniMaxM3IndexerAiterMetadataBuilder ¶
Bases: MiniMaxM3IndexerMetadataBuilder
The Triton indexer's metadata plus the prefill rows' causal shape.
Source code in vllm/models/minimax_m3/amd/indexer_aiter.py
229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 | |
aiter_indexer_max_decode_query_len(vllm_config) ¶
Longest query a decode row can carry, which spec decode is what sets.
Mirrors _init_reorder_batch_threshold(1, supports_spec_as_decode=True), since that is what the builder splits the batch on and therefore what the decode kernel will actually be handed.
Source code in vllm/models/minimax_m3/amd/indexer_aiter.py
aiter_indexer_unsupported_reason(*, topk_blocks, sparse_block_size, num_index_heads, index_head_dim, indexer_kv_dtype, max_model_len, max_decode_query_len=1, score_type='max') ¶
Return why this config cannot use the AITER indexer, or None if it can.
Checks platform (ROCm/gfx950), the AITER sparse PA attend, index-cache dtype, the compiled score/top-k contract, MFMA column limits, max context in blocks, and whether AITER exposes the MSA kernels. select_aiter_indexer_impl_cls logs the string and falls back when it is not None.
Source code in vllm/models/minimax_m3/amd/indexer_aiter.py
aiter_msa_kernels_unavailable_reason() cached ¶
Return why the AITER MSA score/top-k ops cannot be imported, or None.
They are a recent addition, so an AITER that predates them imports fine while these three names do not exist, and the failure would otherwise surface as an ImportError from the middle of a forward. compile_ops binds them lazily, so this costs the module import only -- the kernel build itself still happens on the first call.
Source code in vllm/models/minimax_m3/amd/indexer_aiter.py
score_block_width(max_seq_len, block_size) ¶
Block-axis width the top-k requires of the score buffer.
Lanes read whole wave-wide strips with no tail guard and each holds a power-of-two count of them, so the axis is padded past the block count.
Source code in vllm/models/minimax_m3/amd/indexer_aiter.py
select_aiter_indexer_impl_cls(*, topk_blocks, sparse_block_size, num_index_heads, index_head_dim, indexer_kv_dtype, score_type='max') ¶
The AITER indexer impl if this config can use it, else None.
None sends the caller to the platform-neutral MiniMaxM3Indexer, which on ROCm means the Triton indexer -- and a bf16-only one, so an fp8 index cache that lands here has nowhere to go.