@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / fcf5115c45
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

[ROCm][DSv4][Perf] Flash-decode split-K decode attention kernel (#44899)

Co-authored-by: vLLM Contributor <contributor@vllm.ai>

Fangzhou Ai committed Jun 11, 2026 at 23:17 UTC fcf5115c45b9acfe3a77052ddbb7dfb0f4d5ef18
2 files changed +675 -10
tests/kernels/attention/test_rocm_triton_attn_dsv4.py
+140
@@ -10,6 +10,25 @@ pytestmark = pytest.mark.skipif(
10 not current_platform.is_rocm(), reason="Only used by ROCm"
11 )
12
13 +
14 +def _on_gfx950() -> bool:
15 + if not current_platform.is_rocm():
16 + return False
17 + try:
18 + from vllm.platforms.rocm import _ON_GFX950
19 +
20 + return bool(_ON_GFX950)
21 + except Exception:
22 + return False
23 +
24 +
25 +# The flash-decode split-K decode path is only tuned for AMD gfx950; other
26 +# architectures take the fallback decode kernel, so its tests are skipped there.
27 +requires_gfx950 = pytest.mark.skipif(
28 + not _on_gfx950(),
29 + reason="split-K decode kernel is only tuned for AMD gfx950",
30 +)
31 +
32 NOPE_HEAD_DIM = 448
33 ROPE_HEAD_DIM = 64
34 HEAD_DIM = NOPE_HEAD_DIM + ROPE_HEAD_DIM
@@ -156,6 +175,20 @@ def _ref_sparse_decode_ragged(
175 return out.to(torch.bfloat16)
176
177
178 +def _ragged_from_rows(
179 + rows: list[list[int]], device: torch.device
180 +) -> tuple[torch.Tensor, torch.Tensor]:
181 + """Flatten per-query slot lists into ragged (indices, indptr) tensors."""
182 + flat = [slot for row in rows for slot in row]
183 + indptr = [0]
184 + for row in rows:
185 + indptr.append(indptr[-1] + len(row))
186 + return (
187 + torch.tensor(flat, dtype=torch.int32, device=device),
188 + torch.tensor(indptr, dtype=torch.int32, device=device),
189 + )
190 +
191 +
192 def _ref_combine_topk_swa_ragged(
193 device: torch.device,
194 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
@@ -375,3 +408,110 @@ def test_combine_topk_swa_indices_ragged() -> None:
408 )
409 torch.testing.assert_close(actual_indptr, expected_indptr)
410 torch.testing.assert_close(actual_lens, expected_lens)
411 +
412 +
413 +@requires_gfx950
414 +@torch.inference_mode()
415 +def test_decode_num_splits_heuristic(monkeypatch) -> None:
416 + """Split-count heuristic added with the flash-decode split-K decode path."""
417 + from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod
418 +
419 + # Pin the CU count so the heuristic is deterministic off-device.
420 + monkeypatch.setattr(mod, "_decode_cu_count", lambda: 256)
421 +
422 + # A batch that already fills the device should not be split.
423 + assert mod._decode_num_splits(256, 1, avg_main_len=128.0, avg_extra_len=0.0) == 1
424 + # A tiny batch on a large device should split to add parallelism.
425 + assert mod._decode_num_splits(2, 1, avg_main_len=256.0, avg_extra_len=0.0) > 1
426 +
427 + # The chosen count always stays within the searched [1, 16] range, and a
428 + # zero-length workload never splits (no work to parallelize).
429 + for num_queries in (1, 4, 24, 224, 1024):
430 + splits = mod._decode_num_splits(
431 + num_queries, 1, avg_main_len=512.0, avg_extra_len=128.0
432 + )
433 + assert 1 <= splits <= 16
434 + assert mod._decode_num_splits(2, 1, avg_main_len=0.0, avg_extra_len=0.0) >= 1
435 +
436 +
437 +@requires_gfx950
438 +@pytest.mark.parametrize("num_splits", [1, 2, 3, 4, 8])
439 +@pytest.mark.parametrize("with_extra", [True, False])
440 +@pytest.mark.parametrize("with_sink", [True, False])
441 +@torch.inference_mode()
442 +def test_sparse_attn_decode_split_k_kernel(
443 + monkeypatch, num_splits: int, with_extra: bool, with_sink: bool
444 +) -> None:
445 + """Flash-decode split-K decode path (partial + reduce kernels).
446 +
447 + This path is the gfx950 production path (``_ON_GFX950``), so the test only
448 + runs on gfx950. The split count is pinned so the partial/reduce kernels are
449 + exercised across split counts. ``num_splits=8`` drives splits past the
450 + shortest segment length, covering the empty-split edge case handled by the
451 + reduce kernel.
452 + """
453 + from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod
454 +
455 + device = torch.device("cuda")
456 + torch.manual_seed(7)
457 + block_size = 4
458 + num_heads = 3
459 +
460 + main_rows = [[0, 2, 4, 6, 1, 3, 7, 5], [4, 1, 6, 0, 2]]
461 + num_queries = len(main_rows)
462 + q = (
463 + torch.randn(
464 + num_queries, num_heads, HEAD_DIM, dtype=torch.bfloat16, device=device
465 + )
466 + * 0.125
467 + )
468 + main_kv = torch.randn(8, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125
469 + main_cache = _pack_fp8_ds_mla_cache(main_kv, block_size)
470 + main_indices, main_indptr = _ragged_from_rows(main_rows, device)
471 +
472 + extra_rows: list[list[int]] | None = None
473 + extra_cache: torch.Tensor | None = None
474 + extra_indices: torch.Tensor | None = None
475 + extra_indptr: torch.Tensor | None = None
476 + if with_extra:
477 + rows = [[1, 3, 0, 5, 2, 4], [3, 0, 6]]
478 + extra_kv = torch.randn(7, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125
479 + extra_rows = rows
480 + extra_cache = _pack_fp8_ds_mla_cache(extra_kv, block_size)
481 + extra_indices, extra_indptr = _ragged_from_rows(rows, device)
482 +
483 + attn_sink = (
484 + torch.tensor([-0.1, 0.0, 0.1], dtype=torch.float32, device=device)
485 + if with_sink
486 + else None
487 + )
488 + scale = HEAD_DIM**-0.5
489 +
490 + # Pin the split count so each parametrized value is exercised deterministically.
491 + monkeypatch.setattr(mod, "_decode_num_splits", lambda *args, **kwargs: num_splits)
492 +
493 + actual = mod._rocm_sparse_attn_decode_ragged_triton(
494 + q=q,
495 + main_cache=main_cache,
496 + main_indices=main_indices,
497 + main_indptr=main_indptr,
498 + scale=scale,
499 + attn_sink=attn_sink,
500 + nope_head_dim=NOPE_HEAD_DIM,
501 + rope_head_dim=ROPE_HEAD_DIM,
502 + extra_cache=extra_cache,
503 + extra_indices=extra_indices,
504 + extra_indptr=extra_indptr,
505 + )
506 + expected = _ref_sparse_decode_ragged(
507 + q=q,
508 + main_cache=main_cache,
509 + main_rows=main_rows,
510 + scale=scale,
511 + attn_sink=attn_sink,
512 + block_size=block_size,
513 + extra_cache=extra_cache,
514 + extra_rows=extra_rows,
515 + )
516 +
517 + torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2)
vllm/v1/attention/ops/rocm_aiter_mla_sparse.py
+535 -10
@@ -1406,6 +1406,348 @@ def _sparse_attn_decode_ragged_kernel(
1406 )
1407
1408
1409 +@triton.jit
1410 +def _sparse_attn_decode_partial_kernel(
1411 + q_ptr,
1412 + main_cache_ptr,
1413 + main_indices_ptr,
1414 + main_indptr_ptr,
1415 + extra_cache_ptr,
1416 + extra_indices_ptr,
1417 + extra_indptr_ptr,
1418 + part_m_ptr,
1419 + part_l_ptr,
1420 + part_acc_ptr,
1421 + q_stride0,
1422 + q_stride1,
1423 + main_cache_stride0,
1424 + extra_cache_stride0,
1425 + pm_stride0,
1426 + pm_stride_s,
1427 + pa_stride0,
1428 + pa_stride_s,
1429 + pa_stride_h,
1430 + main_num_rows,
1431 + extra_num_rows,
1432 + main_block_size,
1433 + extra_block_size,
1434 + scale,
1435 + num_heads,
1436 + HAS_EXTRA: tl.constexpr,
1437 + NOPE_DIM: tl.constexpr,
1438 + NOPE_BLOCK: tl.constexpr,
1439 + ROPE_DIM: tl.constexpr,
1440 + IS_FNUZ: tl.constexpr,
1441 + BLOCK_H: tl.constexpr,
1442 + BLOCK_K: tl.constexpr,
1443 + NUM_SPLITS: tl.constexpr,
1444 + NUM_STAGES: tl.constexpr,
1445 +):
1446 + query_idx = tl.program_id(0)
1447 + split_id = tl.program_id(1)
1448 + pid_h = tl.program_id(2)
1449 +
1450 + head_offsets = pid_h * BLOCK_H + tl.arange(0, BLOCK_H)
1451 + head_mask = head_offsets < num_heads
1452 + nope_offsets = tl.arange(0, NOPE_BLOCK)
1453 + nope_mask = nope_offsets < NOPE_DIM
1454 + rope_offsets = tl.arange(0, ROPE_DIM)
1455 +
1456 + q_row_ptr = q_ptr + query_idx * q_stride0 + head_offsets[:, None] * q_stride1
1457 + q_nope = tl.load(
1458 + q_row_ptr + nope_offsets[None, :],
1459 + mask=head_mask[:, None] & nope_mask[None, :],
1460 + other=0.0,
1461 + )
1462 + q_rope = tl.load(
1463 + q_row_ptr + NOPE_DIM + rope_offsets[None, :],
1464 + mask=head_mask[:, None],
1465 + other=0.0,
1466 + )
1467 +
1468 + neg_large = -3.4028234663852886e38
1469 + m_i = tl.full((BLOCK_H,), neg_large, dtype=tl.float32)
1470 + l_i = tl.zeros((BLOCK_H,), dtype=tl.float32)
1471 + acc_nope = tl.zeros((BLOCK_H, NOPE_BLOCK), dtype=tl.float32)
1472 + acc_rope = tl.zeros((BLOCK_H, ROPE_DIM), dtype=tl.float32)
1473 + k_offsets = tl.arange(0, BLOCK_K)
1474 +
1475 + zero_nope = tl.zeros((BLOCK_K, NOPE_BLOCK), dtype=tl.bfloat16)
1476 + zero_rope = tl.zeros((BLOCK_K, ROPE_DIM), dtype=tl.bfloat16)
1477 +
1478 + # Each split processes a contiguous slice of this query's main (SWA) and
1479 + # extra (topk) segments. Slices are handled independently so a block never
1480 + # straddles the main/extra boundary.
1481 + main_start = tl.load(main_indptr_ptr + query_idx)
1482 + main_end = tl.load(main_indptr_ptr + query_idx + 1)
1483 + main_len = main_end - main_start
1484 + main_chunk = (main_len + NUM_SPLITS - 1) // NUM_SPLITS
1485 + main_lo = split_id * main_chunk
1486 + main_hi = tl.minimum(main_lo + main_chunk, main_len)
1487 +
1488 + for k_start in tl.range(main_lo, main_hi, BLOCK_K, num_stages=NUM_STAGES):
1489 + k_pos = k_start + k_offsets
1490 + in_range = k_pos < main_hi
1491 + slot = tl.load(main_indices_ptr + main_start + k_pos, mask=in_range, other=-1)
1492 + valid = in_range & (slot >= 0) & (slot < main_num_rows)
1493 + safe_slot = tl.where(valid, slot, 0)
1494 +
1495 + block_idx = safe_slot // main_block_size
1496 + pos_in_block = safe_slot % main_block_size
1497 + cache_block_ptr = main_cache_ptr + block_idx.to(tl.int64) * main_cache_stride0
1498 + token_data_ptr = cache_block_ptr + pos_in_block * 576
1499 + token_scale_ptr = cache_block_ptr + main_block_size * 576 + pos_in_block * 8
1500 +
1501 + x_uint8 = tl.load(
1502 + token_data_ptr[:, None] + nope_offsets[None, :],
1503 + mask=valid[:, None] & nope_mask[None, :],
1504 + other=0,
1505 + )
1506 + if IS_FNUZ:
1507 + x_fp8 = x_uint8.to(tl.float8e4b15, bitcast=True)
1508 + else:
1509 + x_fp8 = x_uint8.to(tl.float8e4nv, bitcast=True)
1510 + encoded_scales = tl.load(
1511 + token_scale_ptr[:, None] + nope_offsets[None, :] // 64,
1512 + mask=valid[:, None] & nope_mask[None, :],
1513 + other=127,
1514 + )
1515 + scales = tl.exp2(encoded_scales.to(tl.float32) - 127.0)
1516 + k_nope = x_fp8.to(tl.bfloat16) * scales.to(tl.bfloat16)
1517 + k_nope = tl.where(valid[:, None] & nope_mask[None, :], k_nope, zero_nope)
1518 + k_nope = tl.where(k_nope == k_nope, k_nope, zero_nope)
1519 +
1520 + rope_ptr = (token_data_ptr + NOPE_DIM).to(tl.pointer_type(tl.bfloat16))
1521 + k_rope = tl.load(
1522 + rope_ptr[:, None] + rope_offsets[None, :],
1523 + mask=valid[:, None],
1524 + other=0.0,
1525 + )
1526 + k_rope = tl.where(valid[:, None], k_rope, zero_rope)
1527 + k_rope = tl.where(k_rope == k_rope, k_rope, zero_rope)
1528 +
1529 + scores = tl.dot(q_nope, tl.trans(k_nope)) + tl.dot(q_rope, tl.trans(k_rope))
1530 + scores *= scale
1531 + scores = tl.where(head_mask[:, None] & valid[None, :], scores, neg_large)
1532 +
1533 + m_block = tl.max(scores, axis=1)
1534 + m_new = tl.maximum(m_i, m_block)
1535 + alpha = tl.exp(m_i - m_new)
1536 + p = tl.exp(scores - m_new[:, None])
1537 + p = tl.where(head_mask[:, None] & valid[None, :], p, 0.0)
1538 + l_new = l_i * alpha + tl.sum(p, axis=1)
1539 +
1540 + acc_nope = acc_nope * alpha[:, None] + tl.dot(p.to(k_nope.dtype), k_nope)
1541 + acc_rope = acc_rope * alpha[:, None] + tl.dot(p.to(k_rope.dtype), k_rope)
1542 + m_i = m_new
1543 + l_i = l_new
1544 +
1545 + if HAS_EXTRA:
1546 + extra_start = tl.load(extra_indptr_ptr + query_idx)
1547 + extra_end = tl.load(extra_indptr_ptr + query_idx + 1)
1548 + extra_len = extra_end - extra_start
1549 + extra_chunk = (extra_len + NUM_SPLITS - 1) // NUM_SPLITS
1550 + extra_lo = split_id * extra_chunk
1551 + extra_hi = tl.minimum(extra_lo + extra_chunk, extra_len)
1552 +
1553 + for k_start in tl.range(extra_lo, extra_hi, BLOCK_K, num_stages=NUM_STAGES):
1554 + k_pos = k_start + k_offsets
1555 + in_range = k_pos < extra_hi
1556 + slot = tl.load(
1557 + extra_indices_ptr + extra_start + k_pos, mask=in_range, other=-1
1558 + )
1559 + valid = in_range & (slot >= 0) & (slot < extra_num_rows)
1560 + safe_slot = tl.where(valid, slot, 0)
1561 +
1562 + block_idx = safe_slot // extra_block_size
1563 + pos_in_block = safe_slot % extra_block_size
1564 + cache_block_ptr = (
1565 + extra_cache_ptr + block_idx.to(tl.int64) * extra_cache_stride0
1566 + )
1567 + token_data_ptr = cache_block_ptr + pos_in_block * 576
1568 + token_scale_ptr = (
1569 + cache_block_ptr + extra_block_size * 576 + pos_in_block * 8
1570 + )
1571 +
1572 + x_uint8 = tl.load(
1573 + token_data_ptr[:, None] + nope_offsets[None, :],
1574 + mask=valid[:, None] & nope_mask[None, :],
1575 + other=0,
1576 + )
1577 + if IS_FNUZ:
1578 + x_fp8 = x_uint8.to(tl.float8e4b15, bitcast=True)
1579 + else:
1580 + x_fp8 = x_uint8.to(tl.float8e4nv, bitcast=True)
1581 + encoded_scales = tl.load(
1582 + token_scale_ptr[:, None] + nope_offsets[None, :] // 64,
1583 + mask=valid[:, None] & nope_mask[None, :],
1584 + other=127,
1585 + )
1586 + scales = tl.exp2(encoded_scales.to(tl.float32) - 127.0)
1587 + k_nope = x_fp8.to(tl.bfloat16) * scales.to(tl.bfloat16)
1588 + k_nope = tl.where(valid[:, None] & nope_mask[None, :], k_nope, zero_nope)
1589 + k_nope = tl.where(k_nope == k_nope, k_nope, zero_nope)
1590 +
1591 + rope_ptr = (token_data_ptr + NOPE_DIM).to(tl.pointer_type(tl.bfloat16))
1592 + k_rope = tl.load(
1593 + rope_ptr[:, None] + rope_offsets[None, :],
1594 + mask=valid[:, None],
1595 + other=0.0,
1596 + )
1597 + k_rope = tl.where(valid[:, None], k_rope, zero_rope)
1598 + k_rope = tl.where(k_rope == k_rope, k_rope, zero_rope)
1599 +
1600 + scores = tl.dot(q_nope, tl.trans(k_nope)) + tl.dot(
1601 + q_rope,
1602 + tl.trans(k_rope),
1603 + )
1604 + scores *= scale
1605 + scores = tl.where(head_mask[:, None] & valid[None, :], scores, neg_large)
1606 +
1607 + m_block = tl.max(scores, axis=1)
1608 + m_new = tl.maximum(m_i, m_block)
1609 + alpha = tl.exp(m_i - m_new)
1610 + p = tl.exp(scores - m_new[:, None])
1611 + p = tl.where(head_mask[:, None] & valid[None, :], p, 0.0)
1612 + l_new = l_i * alpha + tl.sum(p, axis=1)
1613 +
1614 + acc_nope = acc_nope * alpha[:, None] + tl.dot(p.to(k_nope.dtype), k_nope)
1615 + acc_rope = acc_rope * alpha[:, None] + tl.dot(p.to(k_rope.dtype), k_rope)
1616 + m_i = m_new
1617 + l_i = l_new
1618 +
1619 + # Store raw (un-normalized) partial state for this split. Softmax sink and
1620 + # final normalization happen in the reduce kernel.
1621 + pm_base = query_idx * pm_stride0 + split_id * pm_stride_s + head_offsets
1622 + tl.store(part_m_ptr + pm_base, m_i, mask=head_mask)
1623 + tl.store(part_l_ptr + pm_base, l_i, mask=head_mask)
1624 + acc_base = (
1625 + part_acc_ptr
1626 + + query_idx * pa_stride0
1627 + + split_id * pa_stride_s
1628 + + head_offsets[:, None] * pa_stride_h
1629 + )
1630 + tl.store(
1631 + acc_base + nope_offsets[None, :],
1632 + acc_nope,
1633 + mask=head_mask[:, None] & nope_mask[None, :],
1634 + )
1635 + tl.store(
1636 + acc_base + NOPE_DIM + rope_offsets[None, :],
1637 + acc_rope,
1638 + mask=head_mask[:, None],
1639 + )
1640 +
1641 +
1642 +@triton.jit
1643 +def _sparse_attn_decode_reduce_kernel(
1644 + part_m_ptr,
1645 + part_l_ptr,
1646 + part_acc_ptr,
1647 + attn_sink_ptr,
1648 + out_ptr,
1649 + out_stride0,
1650 + out_stride1,
1651 + pm_stride0,
1652 + pm_stride_s,
1653 + pa_stride0,
1654 + pa_stride_s,
1655 + pa_stride_h,
1656 + num_heads,
1657 + HAS_ATTN_SINK: tl.constexpr,
1658 + COMB_DIM: tl.constexpr,
1659 + BLOCK_H: tl.constexpr,
1660 + NUM_SPLITS: tl.constexpr,
1661 + SPLITS_PAD: tl.constexpr,
1662 +):
1663 + query_idx = tl.program_id(0)
1664 + pid_h = tl.program_id(1)
1665 +
1666 + head_offsets = pid_h * BLOCK_H + tl.arange(0, BLOCK_H)
1667 + head_mask = head_offsets < num_heads
1668 + comb_offsets = tl.arange(0, COMB_DIM)
1669 + # SPLITS_PAD is NUM_SPLITS rounded up to a power of two so the parallel
1670 + # split-axis load is a legal arange for any split count; padding lanes are
1671 + # masked off.
1672 + split_offsets = tl.arange(0, SPLITS_PAD)
1673 + split_mask = split_offsets < NUM_SPLITS
1674 +
1675 + neg_large = -3.4028234663852886e38
1676 +
1677 + # Phase 1: load every split's running max/sum at once and reduce the max
1678 + # in parallel (tl.max over the split axis) instead of walking the splits
1679 + # serially. This breaks the long online-softmax dependency chain that made
1680 + # the reduce latency-bound.
1681 + load_mask = split_mask[:, None] & head_mask[None, :]
1682 + pm_split = (
1683 + part_m_ptr
1684 + + query_idx * pm_stride0
1685 + + split_offsets[:, None] * pm_stride_s
1686 + + head_offsets[None, :]
1687 + )
1688 + m_all = tl.load(pm_split, mask=load_mask, other=neg_large) # [S, H]
1689 + l_all = tl.load(
1690 + part_l_ptr
1691 + + query_idx * pm_stride0
1692 + + split_offsets[:, None] * pm_stride_s
1693 + + head_offsets[None, :],
1694 + mask=load_mask,
1695 + other=0.0,
1696 + )
1697 +
1698 + m_comb = tl.max(m_all, axis=0) # [H]
1699 + if HAS_ATTN_SINK:
1700 + sink = tl.load(
1701 + attn_sink_ptr + head_offsets, mask=head_mask, other=neg_large
1702 + ).to(tl.float32)
1703 + m_final = tl.maximum(m_comb, sink)
1704 + else:
1705 + m_final = m_comb
1706 +
1707 + w_all = tl.exp(m_all - m_final[None, :]) # [S, H]
1708 + w_all = tl.where(load_mask, w_all, 0.0)
1709 + l_final = tl.sum(w_all * l_all, axis=0) # [H]
1710 + if HAS_ATTN_SINK:
1711 + l_final = l_final + tl.exp(sink - m_final)
1712 + denom = tl.maximum(l_final, 1.0e-30)
1713 +
1714 + # Phase 2: weighted sum of the per-split accumulators. The combine weight
1715 + # for each split only depends on the (already known) global max, so the
1716 + # acc loads carry no cross-split dependency and the compiler can pipeline
1717 + # them; only the cheap FMA into `acc` is loop-carried.
1718 + acc = tl.zeros((BLOCK_H, COMB_DIM), dtype=tl.float32)
1719 + for s in tl.static_range(NUM_SPLITS):
1720 + m_s = tl.load(
1721 + part_m_ptr + query_idx * pm_stride0 + s * pm_stride_s + head_offsets,
1722 + mask=head_mask,
1723 + other=neg_large,
1724 + )
1725 + w_s = tl.exp(m_s - m_final)
1726 + acc_base = (
1727 + part_acc_ptr
1728 + + query_idx * pa_stride0
1729 + + s * pa_stride_s
1730 + + head_offsets[:, None] * pa_stride_h
1731 + )
1732 + acc_s = tl.load(
1733 + acc_base + comb_offsets[None, :],
1734 + mask=head_mask[:, None],
1735 + other=0.0,
1736 + )
1737 + acc += w_s[:, None] * acc_s
1738 +
1739 + out = tl.where(l_final[:, None] > 0.0, acc / denom[:, None], 0.0)
1740 +
1741 + out_row_ptr = (
1742 + out_ptr + query_idx * out_stride0 + head_offsets[:, None] * out_stride1
1743 + )
1744 + tl.store(
1745 + out_row_ptr + comb_offsets[None, :],
1746 + out,
1747 + mask=head_mask[:, None],
1748 + )
1749 +
1750 +
1751 def _rocm_sparse_attn_prefill_ragged_triton(
1752 q: torch.Tensor,
1753 kv: torch.Tensor,
@@ -1502,6 +1844,101 @@ def _rocm_sparse_attn_prefill_triton(
1844 )
1845
1846
1847 +@functools.lru_cache
1848 +def _decode_cu_count() -> int:
1849 + try:
1850 + return torch.cuda.get_device_properties(0).multi_processor_count
1851 + except Exception:
1852 + return 256 # For gfx950 arch, gated behind a fallback path for other archs.
1853 +
1854 +
1855 +def _decode_partial_iters(
1856 + avg_main_len: float, avg_extra_len: float, splits: int, block_k: int
1857 +) -> int:
1858 + """BLOCK_K iterations one partial workgroup walks for ``splits`` splits.
1859 +
1860 + Each split processes ``ceil(seg_len / splits)`` tokens of a segment, walked
1861 + ``BLOCK_K`` at a time, and the main/extra segments are handled separately.
1862 + """
1863 + main_iters = (
1864 + math.ceil(math.ceil(avg_main_len / splits) / block_k) if avg_main_len > 0 else 0
1865 + )
1866 + extra_iters = (
1867 + math.ceil(math.ceil(avg_extra_len / splits) / block_k)
1868 + if avg_extra_len > 0
1869 + else 0
1870 + )
1871 + return main_iters + extra_iters
1872 +
1873 +
1874 +def _decode_num_splits(
1875 + num_queries: int,
1876 + heads_blocks: int,
1877 + avg_main_len: float = 0.0,
1878 + avg_extra_len: float = 0.0,
1879 + block_k: int = 32,
1880 +) -> int:
1881 + """Pick a flash-decode split count to keep the GPU busy across batch sizes.
1882 +
1883 + Decode launches only ``num_queries * heads_blocks`` workgroups otherwise,
1884 + which severely under-fills the device for the low-concurrency regime that
1885 + dominates latency. Splitting the KV sequence adds parallelism.
1886 +
1887 + We model the relative partial-kernel latency for a given split count ``s``
1888 + as ``waves * (1/s + mu)`` where ``waves = ceil(base * s / CU)`` and ``mu``
1889 + is a small per-wave overhead penalty:
1890 +
1891 + - ``waves / s`` captures the partial compute: each wave walks roughly
1892 + ``total_tokens / s`` tokens and there are ``waves`` of them, so dividing
1893 + by ``s`` makes more splits cheaper *until* they spill into extra waves.
1894 + - ``mu * waves`` charges per-wave launch/tail overhead so we do not
1895 + over-split into many mostly-idle waves (e.g. batch 224 on 256 CUs is
1896 + best left at 1 split rather than 8 splits across 7 waves).
1897 +
1898 + The minimiser naturally prefers split counts that pack the device into full
1899 + waves (``base * s`` near a multiple of ``CU``) and falls back to 1 split
1900 + once the batch already fills the device. Ties favour the smaller split
1901 + count (less reduce work).
1902 +
1903 + Finally we "snap down" the chosen split count to the smallest value that
1904 + yields the same wave count *and* the same per-workgroup BLOCK_K iteration
1905 + count. Because latency tracks iteration count (not raw token count), extra
1906 + splits that do not lower the iteration count add only reduce/HBM overhead
1907 + for no parallelism gain (e.g. batch 24: s8 and s10 both walk 4 extra iters
1908 + in one wave, so s8 is strictly better). Snapping needs the average segment
1909 + lengths, which the caller derives sync-free from the ragged index sizes.
1910 + """
1911 + base = max(1, num_queries * heads_blocks)
1912 + # Target ~1 workgroup per CU: enough to fill the device while keeping the
1913 + # reduce cost (which grows with split count) small. Tuned on gfx950.
1914 + cu = max(1, _decode_cu_count())
1915 + # Per-wave overhead penalty: higher values discourage split counts that
1916 + # spill into extra GPU waves. Tuned on gfx950.
1917 + mu = 0.04
1918 + best_splits = 1
1919 + best_cost = None
1920 + # Search up to 16 splits; beyond that the reduce/HBM overhead dominates.
1921 + for splits in range(1, 17):
1922 + waves = (base * splits + cu - 1) // cu
1923 + cost = waves * (1.0 / splits + mu)
1924 + if best_cost is None or cost < best_cost - 1e-9:
1925 + best_splits = splits
1926 + best_cost = cost
1927 +
1928 + if best_splits > 1 and (avg_main_len > 0 or avg_extra_len > 0):
1929 + target_waves = (base * best_splits + cu - 1) // cu
1930 + target_iters = _decode_partial_iters(
1931 + avg_main_len, avg_extra_len, best_splits, block_k
1932 + )
1933 + for splits in range(1, best_splits):
1934 + waves = (base * splits + cu - 1) // cu
1935 + iters = _decode_partial_iters(avg_main_len, avg_extra_len, splits, block_k)
1936 + if waves == target_waves and iters == target_iters:
1937 + best_splits = splits
1938 + break
1939 + return best_splits
1940 +
1941 +
1942 def _rocm_sparse_attn_decode_ragged_triton(
1943 q: torch.Tensor,
1944 main_cache: torch.Tensor,
@@ -1575,9 +2012,70 @@ def _rocm_sparse_attn_decode_ragged_triton(
2012 extra_indptr = torch.zeros(num_queries + 1, device=q.device, dtype=torch.int32)
2013
2014 block_h = 16
1578 - block_k = 16 if head_dim >= 256 else 32
2015 out = torch.empty_like(q, dtype=torch.bfloat16)
1580 - _sparse_attn_decode_ragged_kernel[(num_queries, triton.cdiv(num_heads, block_h))](
2016 + heads_blocks = triton.cdiv(num_heads, block_h)
2017 + nope_block = triton.next_power_of_2(nope_head_dim)
2018 + comb_dim = nope_head_dim + rope_head_dim
2019 + is_fnuz = current_platform.is_fp8_fnuz()
2020 +
2021 + if not _ON_GFX950: # Fallback path for un-tuned architectures.
2022 + block_k = 16 if head_dim >= 256 else 32
2023 + _sparse_attn_decode_ragged_kernel[(num_queries, heads_blocks)](
2024 + q,
2025 + main_cache,
2026 + main_indices,
2027 + main_indptr,
2028 + extra_cache,
2029 + extra_indices,
2030 + extra_indptr,
2031 + attn_sink,
2032 + out,
2033 + q.stride(0),
2034 + q.stride(1),
2035 + out.stride(0),
2036 + out.stride(1),
2037 + main_cache.stride(0),
2038 + extra_cache.stride(0),
2039 + main_cache.shape[0] * main_cache.shape[1],
2040 + extra_cache.shape[0] * extra_cache.shape[1],
2041 + main_cache.shape[1],
2042 + extra_cache.shape[1],
2043 + scale,
2044 + num_heads,
2045 + HAS_ATTN_SINK=has_attn_sink,
2046 + HAS_EXTRA=has_extra,
2047 + NOPE_DIM=nope_head_dim,
2048 + NOPE_BLOCK=nope_block,
2049 + ROPE_DIM=rope_head_dim,
2050 + IS_FNUZ=is_fnuz,
2051 + BLOCK_H=block_h,
2052 + BLOCK_K=block_k,
2053 + num_warps=8,
2054 + )
2055 + return out
2056 +
2057 + block_k = 32 # KV tokens walked per split-K iteration. Tuned on gfx950.
2058 + # Average per-query segment lengths, read sync-free from the ragged index
2059 + # sizes, let the split heuristic avoid over-splitting
2060 + # main_indices/extra_indices are flat [nnz] int32.
2061 + inv_q = 1.0 / max(1, num_queries)
2062 + avg_main_len = main_indices.numel() * inv_q
2063 + avg_extra_len = (extra_indices.numel() * inv_q) if has_extra else 0.0
2064 + num_splits = _decode_num_splits(
2065 + num_queries, heads_blocks, avg_main_len, avg_extra_len, block_k
2066 + )
2067 +
2068 + part_m = torch.empty(
2069 + (num_queries, num_splits, num_heads), dtype=torch.float32, device=q.device
2070 + )
2071 + part_l = torch.empty_like(part_m)
2072 + part_acc = torch.empty(
2073 + (num_queries, num_splits, num_heads, comb_dim),
2074 + dtype=torch.float32,
2075 + device=q.device,
2076 + )
2077 +
2078 + _sparse_attn_decode_partial_kernel[(num_queries, num_splits, heads_blocks)](
2079 q,
2080 main_cache,
2081 main_indices,
@@ -1585,29 +2083,56 @@ def _rocm_sparse_attn_decode_ragged_triton(
2083 extra_cache,
2084 extra_indices,
2085 extra_indptr,
1588 - attn_sink,
1589 - out,
2086 + part_m,
2087 + part_l,
2088 + part_acc,
2089 q.stride(0),
2090 q.stride(1),
1592 - out.stride(0),
1593 - out.stride(1),
2091 main_cache.stride(0),
2092 extra_cache.stride(0),
2093 + part_m.stride(0),
2094 + part_m.stride(1),
2095 + part_acc.stride(0),
2096 + part_acc.stride(1),
2097 + part_acc.stride(2),
2098 main_cache.shape[0] * main_cache.shape[1],
2099 extra_cache.shape[0] * extra_cache.shape[1],
2100 main_cache.shape[1],
2101 extra_cache.shape[1],
2102 scale,
2103 num_heads,
1602 - HAS_ATTN_SINK=has_attn_sink,
2104 HAS_EXTRA=has_extra,
2105 NOPE_DIM=nope_head_dim,
1605 - NOPE_BLOCK=triton.next_power_of_2(nope_head_dim),
2106 + NOPE_BLOCK=nope_block,
2107 ROPE_DIM=rope_head_dim,
1607 - IS_FNUZ=current_platform.is_fp8_fnuz(),
2108 + IS_FNUZ=is_fnuz,
2109 BLOCK_H=block_h,
2110 BLOCK_K=block_k,
1610 - num_warps=8,
2111 + NUM_SPLITS=num_splits,
2112 + NUM_STAGES=1,
2113 + num_warps=4,
2114 + )
2115 +
2116 + _sparse_attn_decode_reduce_kernel[(num_queries, heads_blocks)](
2117 + part_m,
2118 + part_l,
2119 + part_acc,
2120 + attn_sink,
2121 + out,
2122 + out.stride(0),
2123 + out.stride(1),
2124 + part_m.stride(0),
2125 + part_m.stride(1),
2126 + part_acc.stride(0),
2127 + part_acc.stride(1),
2128 + part_acc.stride(2),
2129 + num_heads,
2130 + HAS_ATTN_SINK=has_attn_sink,
2131 + COMB_DIM=comb_dim,
2132 + BLOCK_H=block_h,
2133 + NUM_SPLITS=num_splits,
2134 + SPLITS_PAD=triton.next_power_of_2(num_splits),
2135 + num_warps=4,
2136 )
2137 return out
2138