@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 71fc70d3ae

[ROCm][DSv4] Fix sparse-indexer logits collapse on gfx950/gfx942 (#50455)

Signed-off-by: Frida Andersson <fanderss@amd.com> Signed-off-by: frida-andersson <fanderss@amd.com>

frida-andersson committed Sep 18, 2026 at 18:57 UTC 71fc70d3ae1df53300a23ec69f6d97b7207a109f

2 files changed +373 -5

tests/kernels/attention/test_rocm_triton_attn_dsv4.py

+186

@@ -1430,3 +1430,189 @@ def test_get_cached_wo_a_bf16_fp8_blockscale_caches() -> None:

1430

1431 # Second call returns the same cached object.

1432 assert _get_cached_wo_a_bf16(wo_a, n_local_groups, o_lora_rank, hidden_dim) is out

1433 +

1434 +

1435 +@requires_split_decode_arch

1436 +@torch.inference_mode()

1437 +def test_fp8_paged_mqa_logits_triton_matches_torch_ref() -> None:

1438 + """Block-flat Triton decode logits vs the eager torch reference.

1439 +

1440 + Compare only valid key positions: the kernel skips the -inf pad that

1441 + the torch ref writes past seq_len.

1442 + """

1443 + from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod

1444 + from vllm.v1.worker.workspace import init_workspace_manager

1445 +

1446 + device = torch.device("cuda")

1447 + init_workspace_manager(device)

1448 + torch.manual_seed(0)

1449 +

1450 + fp8_dtype = current_platform.fp8_dtype()

1451 + batch_size = 2

1452 + next_n = 1

1453 + num_heads = 4

1454 + head_size = 128

1455 + block_size = 64

1456 + max_model_len = 256

1457 + seq_lens = [180, 64]

1458 + num_pages = [(seq_len + block_size - 1) // block_size for seq_len in seq_lens]

1459 + num_blocks = sum(num_pages)

1460 +

1461 + q = torch.randn(

1462 + batch_size, next_n, num_heads, head_size, device=device, dtype=torch.bfloat16

1463 + ).to(fp8_dtype)

1464 + weights = torch.rand(

1465 + batch_size * next_n, num_heads, device=device, dtype=torch.float32

1466 + )

1467 + context_lens = torch.tensor(seq_lens, device=device, dtype=torch.int32)

1468 + block_tables = torch.full(

1469 + (batch_size, max(num_pages)), 0, device=device, dtype=torch.int32

1470 + )

1471 + page = 0

1472 + for i, n_pages in enumerate(num_pages):

1473 + block_tables[i, :n_pages] = torch.arange(

1474 + page, page + n_pages, device=device, dtype=torch.int32

1475 + )

1476 + page += n_pages

1477 +

1478 + kv_cache = torch.empty(

1479 + num_blocks, block_size, 1, head_size + 4, device=device, dtype=torch.uint8

1480 + )

1481 + values = torch.randn(

1482 + num_blocks, block_size, head_size, device=device, dtype=torch.bfloat16

1483 + ).to(fp8_dtype)

1484 + scales = torch.rand(

1485 + num_blocks, block_size, device=device, dtype=torch.float32

1486 + ).clamp(min=1e-3)

1487 + kv_flat = kv_cache.view(num_blocks, -1)

1488 + scale_off = block_size * head_size

1489 + kv_flat[:, :scale_off] = values.reshape(num_blocks, -1).view(torch.uint8)

1490 + kv_flat[:, scale_off:] = scales.contiguous().view(torch.uint8).view(num_blocks, -1)

1491 +

1492 + got = mod.rocm_fp8_paged_mqa_logits_triton(

1493 + q, kv_cache, weights, context_lens, block_tables, max_model_len

1494 + )

1495 + ref = mod.fp8_paged_mqa_logits_torch(

1496 + q, kv_cache, weights, context_lens, block_tables, max_model_len

1497 + )

1498 + for i, seq_len in enumerate(seq_lens):

1499 + torch.testing.assert_close(

1500 + got[i, :seq_len], ref[i, :seq_len], atol=2e-3, rtol=2e-3

1501 + )

1502 +

1503 +

1504 +def test_indexer_k_is_c4a_block_flat_gate() -> None:

1505 + from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (

1506 + _indexer_k_is_c4a_block_flat,

1507 + )

1508 +

1509 + assert not _indexer_k_is_c4a_block_flat(1)

1510 + assert not _indexer_k_is_c4a_block_flat(2)

1511 + assert _indexer_k_is_c4a_block_flat(4)

1512 + assert not _indexer_k_is_c4a_block_flat(128)

1513 +

1514 +

1515 +@requires_split_decode_arch

1516 +@torch.inference_mode()

1517 +def test_paged_mqa_logits_gate_v41_shuffle_vs_c4a_flat() -> None:

1518 + """Ratio 1/2 stay on AITER; ratio 4 takes Triton."""

1519 + from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (

1520 + fp8_paged_mqa_logits_torch,

1521 + indexer_k_quant_and_cache_triton,

1522 + rocm_fp8_paged_mqa_logits,

1523 + rocm_fp8_paged_mqa_logits_triton,

1524 + )

1525 + from vllm.v1.worker.workspace import init_workspace_manager

1526 +

1527 + device = torch.device("cuda")

1528 + init_workspace_manager(device)

1529 + torch.manual_seed(1)

1530 +

1531 + fp8_dtype = current_platform.fp8_dtype()

1532 + batch_size = 1

1533 + next_n = 1

1534 + num_heads = 32

1535 + head_size = 128

1536 + block_size = 64

1537 + seq_len = 180

1538 + max_model_len = 256

1539 + num_pages = (seq_len + block_size - 1) // block_size

1540 + num_tokens = num_pages * block_size

1541 +

1542 + q = torch.randn(

1543 + batch_size, next_n, num_heads, head_size, device=device, dtype=torch.bfloat16

1544 + ).to(fp8_dtype)

1545 + weights = torch.rand(

1546 + batch_size * next_n, num_heads, device=device, dtype=torch.float32

1547 + )

1548 + context_lens = torch.tensor([seq_len], device=device, dtype=torch.int32)

1549 + block_tables = torch.arange(num_pages, device=device, dtype=torch.int32).view(1, -1)

1550 + dummy_sched = torch.zeros(8, 2, device=device, dtype=torch.int32)

1551 +

1552 + k_bf16 = torch.randn(num_tokens, head_size, device=device, dtype=torch.bfloat16)

1553 + slot = torch.arange(num_tokens, device=device, dtype=torch.int32)

1554 + cache_shuffle = torch.zeros(

1555 + num_pages, block_size, head_size + 4, device=device, dtype=torch.uint8

1556 + )

1557 + indexer_k_quant_and_cache_triton(k_bf16, cache_shuffle, slot, 128, "ue8m0")

1558 + cache_shuffle = cache_shuffle.unsqueeze(2)

1559 +

1560 + aiter = rocm_fp8_paged_mqa_logits(

1561 + q,

1562 + cache_shuffle,

1563 + weights,

1564 + context_lens,

1565 + block_tables,

1566 + dummy_sched,

1567 + max_model_len,

1568 + compress_ratio=1,

1569 + )

1570 + for ratio in (1, 2):

1571 + got = rocm_fp8_paged_mqa_logits(

1572 + q,

1573 + cache_shuffle,

1574 + weights,

1575 + context_lens,

1576 + block_tables,

1577 + dummy_sched,

1578 + max_model_len,

1579 + compress_ratio=ratio,

1580 + )

1581 + torch.testing.assert_close(got[:, :seq_len], aiter[:, :seq_len], atol=0, rtol=0)

1582 +

1583 + cache_flat = torch.empty(

1584 + num_pages, block_size, 1, head_size + 4, device=device, dtype=torch.uint8

1585 + )

1586 + values = torch.randn(

1587 + num_pages, block_size, head_size, device=device, dtype=torch.bfloat16

1588 + ).to(fp8_dtype)

1589 + scales = torch.rand(

1590 + num_pages, block_size, device=device, dtype=torch.float32

1591 + ).clamp(min=1e-3)

1592 + kv_flat = cache_flat.view(num_pages, -1)

1593 + scale_off = block_size * head_size

1594 + kv_flat[:, :scale_off] = values.reshape(num_pages, -1).view(torch.uint8)

1595 + kv_flat[:, scale_off:] = scales.contiguous().view(torch.uint8).view(num_pages, -1)

1596 +

1597 + gated_c4a = rocm_fp8_paged_mqa_logits(

1598 + q,

1599 + cache_flat,

1600 + weights,

1601 + context_lens,

1602 + block_tables,

1603 + dummy_sched,

1604 + max_model_len,

1605 + compress_ratio=4,

1606 + )

1607 + triton = rocm_fp8_paged_mqa_logits_triton(

1608 + q, cache_flat, weights, context_lens, block_tables, max_model_len

1609 + )

1610 + ref = fp8_paged_mqa_logits_torch(

1611 + q, cache_flat, weights, context_lens, block_tables, max_model_len

1612 + )

1613 + torch.testing.assert_close(

1614 + gated_c4a[:, :seq_len], triton[:, :seq_len], atol=0, rtol=0

1615 + )

1616 + torch.testing.assert_close(

1617 + gated_c4a[:, :seq_len], ref[:, :seq_len], atol=2e-3, rtol=2e-3

1618 + )

vllm/v1/attention/ops/rocm_aiter_mla_sparse.py

+187 -5

@@ -66,6 +66,11 @@ _GFX950_DSV4_NATIVE_MAX_COLUMNS = 1024 * 1024

66 _GFX950_AITER_SPARSE_PREFILL_OPUS_MIN_QUERIES = 1024

69 +def _indexer_k_is_c4a_block_flat(compress_ratio: int) -> bool:

70 + """V4.0 C4A is block-flat (NORMAL). Ratio 1 and 2 are 16×16 SHUFFLE."""

71 + return compress_ratio == 4

72 +

73 +

74 def _get_aiter_top_k_kernel(

75 *,

76 is_prefill: bool,

@@ -458,6 +463,7 @@ def cp_gather_indexer_k_quant_cache_triton(

463 token_to_seq: torch.Tensor,

464 block_tile_size: int = 16,

465 head_tile_size: int = 16,

466 + cache_layout: str | None = None,

467 ):

468 num_tokens = k_fp8.size(0)

469 block_size = k_cache.size(1)

@@ -470,7 +476,10 @@ def cp_gather_indexer_k_quant_cache_triton(

476 k_cache_scale = k_cache[:, block_size * head_dim :].view(torch.float32)

477 grid = (num_tokens,)

478 k_fp8_scale = k_fp8_scale.view(torch.float32)

473 - layout = "NORMAL" if block_size == 1 else "SHUFFLE"

479 + if cache_layout is None:

480 + layout = "NORMAL" if block_size == 1 else "SHUFFLE"

481 + else:

482 + layout = cache_layout

483 kernel_args = (

484 k_cache_value,

485 k_cache_scale,

@@ -607,6 +616,158 @@ def fp8_paged_mqa_logits_torch(

616 return logits

617

618

619 +@triton.jit

620 +def _fp8_paged_mqa_logits_decode_kernel(

621 + q_ptr, # fp8 [B, NEXT_N, H, D]

622 + kv_val_ptr, # fp8, block-flat: [num_blocks, block_size*(D+4)]

623 + kv_scale_ptr, # fp32, block-flat: [num_blocks, block_size*(D+4)//4]

624 + weights_ptr, # fp32 [B*NEXT_N, H]

625 + ctx_lens_ptr, # int32 [B*NEXT_N] if CTX_PER_ROW else [B]

626 + block_tables_ptr, # int32 [B, max_blocks]

627 + logits_ptr, # fp32 [B*NEXT_N, max_model_len]

628 + stride_q_b,

629 + stride_q_n,

630 + stride_q_h,

631 + stride_w_row,

632 + stride_kvblk_fp8,

633 + stride_kvblk_f32,

634 + scale_region_off, # block_size*D // 4 (fp32 offset of scale region within a block)

635 + stride_bt_b,

636 + stride_logits_row,

637 + max_blocks, # block_tables width; guards the block-table gather

638 + max_model_len,

639 + NUM_HEADS: tl.constexpr,

640 + HEAD_SIZE: tl.constexpr,

641 + BLOCK_SIZE: tl.constexpr, # paged-cache page size

642 + BLOCK_KV: tl.constexpr, # positions per tile; must divide BLOCK_SIZE

643 + N_SPLITS: tl.constexpr, # KV-tile parallelism factor (grid dim 1)

644 + NEXT_N: tl.constexpr, # query positions per batch (1 = decode; >1 = MTP verify)

645 + CTX_PER_ROW: tl.constexpr, # ctx_lens is already per-(b, n)

646 +):

647 + # scale_region_off = block_size*D//4 needs D % 4 == 0 (D=128).

648 + tl.static_assert(HEAD_SIZE % 4 == 0)

649 + # Grid (B*NEXT_N, N_SPLITS): disjoint KV tiles per program, no host sync.

650 + row = tl.program_id(0)

651 + split = tl.program_id(1)

652 + b = row // NEXT_N

653 + n = row % NEXT_N

654 + if CTX_PER_ROW:

655 + seq_len = tl.load(ctx_lens_ptr + row)

656 + else:

657 + seq_len = tl.load(ctx_lens_ptr + b) - NEXT_N + n + 1

658 + seq_len = tl.minimum(tl.maximum(seq_len, 0), max_model_len)

659 +

660 + h = tl.arange(0, NUM_HEADS)

661 + d = tl.arange(0, HEAD_SIZE)

662 + # keep q/kv fp8 -> fp8 MFMA (f32 accum), not the slow f32 path

663 + q = tl.load(

664 + q_ptr + b * stride_q_b + n * stride_q_n + h[:, None] * stride_q_h + d[None, :]

665 + )

666 + w = tl.load(weights_ptr + row * stride_w_row + h).to(tl.float32) # [H]

667 +

668 + kv_col = tl.arange(0, BLOCK_KV)

669 + for kv_start in tl.range(split * BLOCK_KV, seq_len, N_SPLITS * BLOCK_KV):

670 + pos = kv_start + kv_col

671 + # Tiles are page-aligned (BLOCK_KV | BLOCK_SIZE), so a tile is one page.

672 + logical_blk = kv_start // BLOCK_SIZE

673 + blk_ok = logical_blk < max_blocks # guard against a wild page -> fault

674 + mask_pos = (pos < seq_len) & blk_ok

675 + page = tl.load(

676 + block_tables_ptr + b * stride_bt_b + logical_blk, mask=blk_ok, other=0

677 + ).to(tl.int64)

678 + pos_in_blk = pos - logical_blk * BLOCK_SIZE # [BLOCK_KV] == pos % BLOCK_SIZE

679 +

680 + # block-flat layout: values region, then scales region.

681 + val_off = page * stride_kvblk_fp8 + pos_in_blk[:, None] * HEAD_SIZE + d[None, :]

682 + kv = tl.load(

683 + kv_val_ptr + val_off, mask=mask_pos[:, None], other=0.0

684 + ) # [BLOCK_KV, D] fp8

685 + sc = tl.load(

686 + kv_scale_ptr + page * stride_kvblk_f32 + scale_region_off + pos_in_blk,

687 + mask=mask_pos,

688 + other=0.0,

689 + ) # [BLOCK_KV]

690 +

691 + # fp8 upcasts exactly; bit-identical on gfx942, ~1e-5 rel on gfx950

692 + # (accum order), top-k unchanged.

693 + s = tl.dot(kv, tl.trans(q)) # [BLOCK_KV, H]

694 + s = tl.maximum(s, 0.0)

695 + s = s * w[None, :]

696 + s = tl.sum(s, axis=1) # [BLOCK_KV]

697 + s = s * sc

698 + tl.store(logits_ptr + row * stride_logits_row + pos, s, mask=mask_pos)

699 +

700 +

701 +def rocm_fp8_paged_mqa_logits_triton(

702 + q_fp8: torch.Tensor,

703 + kv_cache_fp8: torch.Tensor,

704 + weights: torch.Tensor,

705 + context_lens: torch.Tensor,

706 + block_tables: torch.Tensor,

707 + max_model_len: int,

708 +) -> torch.Tensor:

709 + """Triton paged MQA-logits for decode and MTP; matches the torch ref but

710 + has no host sync, so it is safe to capture under a full CUDA graph."""

711 + batch_size, next_n, num_heads, head_size = q_fp8.shape

712 + block_size = kv_cache_fp8.shape[1]

713 + BLOCK_KV = 64

714 + assert block_size % BLOCK_KV == 0

715 +

716 + fp8_dtype = current_platform.fp8_dtype()

717 + num_blocks = kv_cache_fp8.shape[0]

718 + kv_flat = kv_cache_fp8.reshape(

719 + num_blocks, -1

720 + ) # uint8 [num_blocks, block_size*(D+4)]

721 + kv_val = kv_flat.view(fp8_dtype) # [num_blocks, block_size*(D+4)] fp8

722 + kv_scale = kv_flat.view(torch.float32) # [num_blocks, block_size*(D+4)//4] fp32

723 +

724 + cl = context_lens.reshape(-1)

725 + ctx_per_row = not (next_n > 1 and cl.numel() == batch_size)

726 +

727 + max_blocks = block_tables.shape[1]

728 + (out_logits,) = current_workspace_manager().get_simultaneous(

729 + ((batch_size * next_n, max_model_len), torch.float32),

730 + )

731 +

732 + # Memory-bound over the KV range: split each row's keys across programs so

733 + # few-row / long-context launches still fill the GPU. Cap splits at the

734 + # device CU count (304 on gfx942, 256 on gfx950) rather than a gfx950-sized

735 + # constant. All terms are static at launch, so the grid stays CUDA-graph-safe.

736 + rows = batch_size * next_n

737 + tiles_cap = (max_model_len + BLOCK_KV - 1) // BLOCK_KV

738 + N_SPLITS = max(1, min(max(1, _decode_cu_count()), tiles_cap, 1024 // rows))

739 + _fp8_paged_mqa_logits_decode_kernel[(rows, N_SPLITS)](

740 + q_fp8,

741 + kv_val,

742 + kv_scale,

743 + weights,

744 + cl,

745 + block_tables,

746 + out_logits,

747 + q_fp8.stride(0),

748 + q_fp8.stride(1),

749 + q_fp8.stride(2),

750 + weights.stride(0),

751 + kv_val.stride(0),

752 + kv_scale.stride(0),

753 + (block_size * head_size) // 4,

754 + block_tables.stride(0),

755 + out_logits.stride(0),

756 + max_blocks,

757 + max_model_len,

758 + NUM_HEADS=num_heads,

759 + HEAD_SIZE=head_size,

760 + BLOCK_SIZE=block_size,

761 + BLOCK_KV=BLOCK_KV,

762 + N_SPLITS=N_SPLITS,

763 + NEXT_N=next_n,

764 + CTX_PER_ROW=ctx_per_row,

765 + num_warps=4,

766 + num_stages=2,

767 + )

768 + return out_logits

769 +

770 +

771 @functools.lru_cache

772 def paged_mqa_logits_module():

773 paged_mqa_logits_module_path = None

@@ -632,6 +793,8 @@ def rocm_fp8_paged_mqa_logits(

793 block_tables: torch.Tensor,

794 schedule_metadata: torch.Tensor,

795 max_model_len: int,

796 + *,

797 + compress_ratio: int = 1,

798 ) -> torch.Tensor:

799 """Compute FP8 MQA logits using paged KV-cache.

800

@@ -639,8 +802,7 @@ def rocm_fp8_paged_mqa_logits(

802 q_fp8: Query tensor of shape [B, next_n, H, D]. Casted to

803 `torch.float8_e4m3fn` by caller.

804 kv_cache_fp8: Paged KV-cache in packed FP8+scale layout with shape

642 - [num_blocks, block_size, 1, D+4], dtype `torch.uint8`. The last

643 - 4 bytes per (block,pos) store the `float` dequant scale.

805 + [num_blocks, block_size, 1, D+4], dtype `torch.uint8`.

806 weights: Tensor of shape [B * next_n, H], dtype `torch.float32`.

807 context_lens: Tensor of shape [B], dtype int32; effective context length

808 for each batch element.

@@ -649,6 +811,7 @@ def rocm_fp8_paged_mqa_logits(

811 schedule_metadata: Returned by `get_paged_mqa_logits_metadata`;

812 used to distribute work across SMs.

813 max_model_len: Maximum sequence length used to size the logits output.

814 + compress_ratio: C4A (4) takes block-flat Triton; 1 and 2 stay on AITER.

815

816 Returns:

817 Logits tensor of shape [B * next_n, max_model_len], dtype

@@ -657,11 +820,26 @@ def rocm_fp8_paged_mqa_logits(

820 """

821 from vllm._aiter_ops import rocm_aiter_ops

822

660 - aiter_paged_mqa_logits_module = None

661 - # if rocm_aiter_ops.is_enabled():

823 batch_size, next_n = q_fp8.shape[:2]

824 block_size = kv_cache_fp8.shape[1]

825

826 + # C4A only: Flash/DSv3.2 also skip insert but still write SHUFFLE.

827 + if (

828 + (_ON_GFX950 or _ON_GFX942)

829 + and _indexer_k_is_c4a_block_flat(compress_ratio)

830 + and block_size > 1

831 + ):

832 + if block_size % 64 == 0:

833 + return rocm_fp8_paged_mqa_logits_triton(

834 + q_fp8, kv_cache_fp8, weights, context_lens, block_tables, max_model_len

835 + )

836 + # Non 64-aligned page size (not used in prod): eager torch ref.

837 + return fp8_paged_mqa_logits_torch(

838 + q_fp8, kv_cache_fp8, weights, context_lens, block_tables, max_model_len

839 + )

840 +

841 + aiter_paged_mqa_logits_module = None

842 +

843 if rocm_aiter_ops.is_enabled() or rocm_aiter_ops.is_rdna_aiter_enabled():

844 aiter_paged_mqa_logits_module = paged_mqa_logits_module()

845

@@ -1128,6 +1306,9 @@ def rocm_aiter_sparse_attn_indexer(

1306 chunk.block_table,

1307 chunk.cu_seq_lens,

1308 token_to_seq=chunk.token_to_seq,

1309 + cache_layout=(

1310 + "NORMAL" if _indexer_k_is_c4a_block_flat(compress_ratio) else None

1311 + ),

1312 )

1313 logits = rocm_fp8_mqa_logits(

1314 q_fp8[chunk.token_start : chunk.token_end],

@@ -1225,6 +1406,7 @@ def rocm_aiter_sparse_attn_indexer(

1406 decode_metadata.block_table,

1407 decode_metadata.schedule_metadata,

1408 max_model_len=max_model_len,

1409 + compress_ratio=compress_ratio,

1410 )

1411

1412 if candidate_blocks is not None: