@tduncan / Nvidia-Optimization-Stack-VLLM / commits / ca831d1c55

[ROCm][DSv4.1][Perf] Fuse the inverse RoPE into the sparse decode reduce (#57435)

Signed-off-by: Fangzhou Ai <fangzhou.ai@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Fangzhou Ai committed Sep 22, 2026 at 14:05 UTC ca831d1c55714aa10acd826389674ce8f6a4f872

3 files changed +184 -13

tests/kernels/attention/test_rocm_triton_attn_dsv4.py

+49

@@ -259,7 +259,10 @@ def _launch_sparse_decode_reduce(

259 part_l: torch.Tensor,

260 part_acc: torch.Tensor,

261 adaptive_splits: bool,

262 + positions: torch.Tensor | None = None,

263 + cos_sin_cache: torch.Tensor | None = None,

264 ) -> torch.Tensor:

265 + """Run the reduce kernel, with the inverse-RoPE epilogue if given positions."""

266 from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod

267

268 num_queries, num_splits, num_heads = part_m.shape

@@ -275,6 +278,8 @@ def _launch_sparse_decode_reduce(

278 part_acc,

279 attn_sink,

280 out,

281 + positions,

282 + cos_sin_cache,

283 out.stride(0),

284 out.stride(1),

285 part_m.stride(0),

@@ -282,6 +287,7 @@ def _launch_sparse_decode_reduce(

287 part_acc.stride(0),

288 part_acc.stride(1),

289 part_acc.stride(2),

290 + cos_sin_cache.stride(0) if cos_sin_cache is not None else 0,

291 num_heads,

292 HAS_ATTN_SINK=False,

293 ADAPTIVE_SPLITS=adaptive_splits,

@@ -289,6 +295,9 @@ def _launch_sparse_decode_reduce(

295 BLOCK_H=1,

296 NUM_SPLITS=num_splits,

297 SPLITS_PAD=1 << (num_splits - 1).bit_length(),

298 + FUSE_INV_ROPE=positions is not None,

299 + NOPE=NOPE_HEAD_DIM,

300 + HALF=ROPE_HEAD_DIM // 2,

301 num_warps=4,

302 )

303 return out

@@ -1008,6 +1017,46 @@ def test_sparse_attn_decode_gfx950_adaptive_reduce_ignores_stale_scratch() -> No

1017 assert torch.equal(actual, torch.full_like(actual, 2))

1018

1019

1020 +@requires_gfx950

1021 +@pytest.mark.parametrize("adaptive_splits", [False, True])

1022 +@torch.inference_mode()

1023 +def test_sparse_attn_decode_reduce_inverse_rope_epilogue(adaptive_splits: bool) -> None:

1024 + """FUSE_INV_ROPE must land where the standalone rotation would.

1025 +

1026 + The epilogue rotates a whole [H, nope + rope] row with one expression by

1027 + feeding cos=1/sin=0 to the NoPE lanes, which only holds if the pair split

1028 + lands on ``nope_head_dim // 2``; the NoPE lanes therefore have to come

1029 + back untouched, not merely close.

1030 + """

1031 + from vllm.v1.attention.ops.rocm_aiter_mla_sparse import _fused_inverse_rope_gptj

1032 +

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

1034 + torch.manual_seed(7)

1035 + num_queries, num_splits, num_heads, max_pos = 5, 4, 8, 64

1036 + shape = (num_queries, num_splits, num_heads)

1037 + part_m = torch.randn(shape, dtype=torch.float32, device=device)

1038 + part_l = torch.rand(shape, dtype=torch.float32, device=device) + 0.5

1039 + part_acc = torch.randn((*shape, HEAD_DIM), dtype=torch.float32, device=device)

1040 + positions = torch.randint(

1041 + 0, max_pos, (num_queries,), dtype=torch.int64, device=device

1042 + )

1043 + angle = torch.randn(max_pos, ROPE_HEAD_DIM // 2, device=device)

1044 + cos_sin_cache = torch.cat((angle.cos(), angle.sin()), dim=-1).contiguous()

1045 +

1046 + unfused = _launch_sparse_decode_reduce(part_m, part_l, part_acc, adaptive_splits)

1047 + expected = _fused_inverse_rope_gptj(

1048 + unfused, positions, cos_sin_cache, ROPE_HEAD_DIM

1049 + )

1050 + actual = _launch_sparse_decode_reduce(

1051 + part_m, part_l, part_acc, adaptive_splits, positions, cos_sin_cache

1052 + )

1053 +

1054 + torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2)

1055 + assert torch.equal(actual[..., :NOPE_HEAD_DIM], unfused[..., :NOPE_HEAD_DIM])

1056 + # An epilogue that quietly did nothing would satisfy everything above.

1057 + assert not torch.equal(actual[..., NOPE_HEAD_DIM:], unfused[..., NOPE_HEAD_DIM:])

1058 +

1059 +

1060 @requires_gfx950

1061 @pytest.mark.parametrize("extra_len", [0, 1, 31, 32, 33, 63, 64, 65])

1062 @torch.inference_mode()

vllm/models/deepseek_v41/amd/rocm.py

+22 -3

@@ -37,6 +37,7 @@ from vllm.v1.attention.backends.mla.sparse_swa import (

37 from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (

38 build_ragged_indices_from_dense,

39 rocm_inv_rope_einsum,

40 + rocm_inverse_rope_rows_,

41 rocm_sparse_attn_decode,

42 rocm_sparse_attn_prefill,

43 )

@@ -673,6 +674,7 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention):

674 self.n_local_groups,

675 self.o_lora_rank,

676 self.wo_a,

677 + inverse_rope=False,

678 )

679 zf = z.flatten(1)

680 if self._wo_b_scale is not None and zf.dim() == 2:

@@ -747,15 +749,28 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention):

749 attn_metadata=rocm_metadata,

750 swa_metadata=swa_metadata,

751 )

752 + rotated = 0

753 if num_decodes > 0:

751 - self._forward_decode(

754 + rotated = self._forward_decode(

755 q=q[:num_decode_tokens],

756 + positions=positions[:num_decode_tokens],

757 kv_cache=self_kv_cache,

758 swa_metadata=swa_metadata,

759 attn_metadata=rocm_metadata,

760 swa_only=swa_only,

761 output=output[:num_decode_tokens],

762 )

763 + # Only the decode reduce rotates its own rows, and only the leading

764 + # `rotated` of them; prefill rows and any decode path that did not

765 + # fuse still owe the standalone pass. Settle that here rather than in

766 + # _o_proj: the split is batch-dependent and _o_proj runs compiled,

767 + # where such a value freezes at its trace-time value.

768 + rocm_inverse_rope_rows_(

769 + output[rotated:, : self.n_local_heads, :],

770 + positions[rotated:],

771 + self.rotary_emb.cos_sin_cache,

772 + self.rope_head_dim,

773 + )

774

775 def _decode_topk_ragged(

776 self,

@@ -796,12 +811,14 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention):

811 def _forward_decode(

812 self,

813 q: torch.Tensor,

814 + positions: torch.Tensor,

815 kv_cache: torch.Tensor | None,

816 swa_metadata: DeepseekV4ROCMAiterSparseSWAMetadata,

817 attn_metadata: DeepseekV4FlashMLAMetadata | None,

818 swa_only: bool,

819 output: torch.Tensor,

804 - ) -> None:

820 + ) -> int:

821 + """Returns how many leading rows the decode epilogue inverse-RoPE'd."""

822 num_decodes = swa_metadata.num_decodes

823 num_decode_tokens = swa_metadata.num_decode_tokens

824

@@ -820,7 +837,7 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention):

837 num_decode_tokens=num_decode_tokens,

838 )

839

823 - rocm_sparse_attn_decode(

840 + return rocm_sparse_attn_decode(

841 q=q,

842 kv_cache=kv_cache,

843 swa_k_cache=self.swa_cache_layer.kv_cache,

@@ -839,6 +856,8 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention):

856 nope_head_dim=self.nope_head_dim,

857 rope_head_dim=self.rope_head_dim,

858 output=output,

859 + inv_rope_positions=positions,

860 + inv_rope_cos_sin_cache=self.rotary_emb.cos_sin_cache,

861 extra_cache_nan_free=_trust_dsv4_extra_cache_nan_free(

862 self.kv_cache_dtype,

863 self._has_kv_transfer,

vllm/v1/attention/ops/rocm_aiter_mla_sparse.py

+113 -10

@@ -1572,8 +1572,13 @@ def _fused_inverse_rope_gptj(

1572 positions: torch.Tensor,

1573 cos_sin_cache: torch.Tensor,

1574 rope_head_dim: int,

1575 + out: torch.Tensor | None = None,

1576 ) -> torch.Tensor:

1576 - """bf16 inverse GPT-J RoPE via a single fused Triton kernel."""

1577 + """bf16 inverse GPT-J RoPE via a single fused Triton kernel.

1578 +

1579 + ``out`` may alias ``o``: the rotation is a per-row bijection whose kernel

1580 + reads both lanes of a pair before storing either.

1581 + """

1582 assert o.dim() == 3 and o.stride(-1) == 1, (

1583 "_fused_inverse_rope_gptj expects a [T, H, D] input with a contiguous last dim"

1584 )

@@ -1585,9 +1590,14 @@ def _fused_inverse_rope_gptj(

1590 f"[P, {rope_head_dim}] = cos | sin, got {tuple(cos_sin_cache.shape)}"

1591 )

1592 num_tokens, num_heads, head_dim = o.shape

1588 - out = torch.empty(

1589 - (num_tokens, num_heads, head_dim), dtype=torch.bfloat16, device=o.device

1590 - )

1593 + if out is None:

1594 + out = torch.empty(

1595 + (num_tokens, num_heads, head_dim), dtype=torch.bfloat16, device=o.device

1596 + )

1597 + else:

1598 + assert out.dtype == torch.bfloat16, (

1599 + f"inverse RoPE writes bf16, got an output buffer of {out.dtype}"

1600 + )

1601 if num_tokens == 0:

1602 return out

1603 _inverse_rope_gptj_kernel[(num_tokens, num_heads)](

@@ -1608,6 +1618,25 @@ def _fused_inverse_rope_gptj(

1618 return out

1619

1620

1621 +def rocm_inverse_rope_rows_(

1622 + o: torch.Tensor,

1623 + positions: torch.Tensor,

1624 + cos_sin_cache: torch.Tensor,

1625 + rope_head_dim: int,

1626 +) -> None:

1627 + """Inverse-RoPE attention output rows in place.

1628 +

1629 + For rows no attention kernel rotated in its epilogue. Call it from the

1630 + eager attention segment: which rows still owe a rotation depends on the

1631 + prefill/decode split, and the o_proj that used to do this runs inside the

1632 + compiled region, where a batch-dependent Python value would be frozen at

1633 + trace time.

1634 + """

1635 + if o.shape[0] == 0:

1636 + return

1637 + _fused_inverse_rope_gptj(o, positions, cos_sin_cache, rope_head_dim, out=o)

1638 +

1639 +

1640 def _get_cached_wo_a_bf16(

1641 wo_a: torch.nn.Module,

1642 n_local_groups: int,

@@ -1664,16 +1693,27 @@ def rocm_inv_rope_einsum(

1693 n_local_groups: int,

1694 o_lora_rank: int,

1695 wo_a: torch.nn.Module,

1696 + inverse_rope: bool = True,

1697 ) -> torch.Tensor:

1698 """Inverse-RoPE + WO_A bmm path used on ROCm.

1699

1700 Fuses the inverse GPT-J RoPE into one Triton kernel and caches the bf16

1671 - wo_a weight so the per-step dequant disappears.

1701 + wo_a weight so the per-step dequant disappears. Callers whose attention

1702 + already rotated every row pass ``inverse_rope=False``; that is a property

1703 + of the attention backend, not of the batch, so it stays constant across

1704 + steps and is safe to read from compiled code.

1705 """

1673 - o_ref = _fused_inverse_rope_gptj(

1674 - o, positions, rotary_emb.cos_sin_cache, rope_head_dim

1675 - )

1676 - o_ref = o_ref.view(o.shape[0], n_local_groups, -1)

1706 + if inverse_rope:

1707 + o_ref = _fused_inverse_rope_gptj(

1708 + o, positions, rotary_emb.cos_sin_cache, rope_head_dim

1709 + )

1710 + else:

1711 + assert o.dtype == torch.bfloat16, (

1712 + "a pre-rotated attention output feeds the wo_a bmm directly, so it "

1713 + f"must already be bf16, got {o.dtype}"

1714 + )

1715 + o_ref = o

1716 + o_ref = o_ref.reshape(o.shape[0], n_local_groups, -1)

1717

1718 wo_a_weight = _get_cached_wo_a_bf16(

1719 wo_a, n_local_groups, o_lora_rank, o_ref.shape[-1]

@@ -2889,6 +2929,8 @@ def _sparse_attn_decode_reduce_kernel(

2929 part_acc_ptr,

2930 attn_sink_ptr,

2931 out_ptr,

2932 + pos_ptr,

2933 + cos_sin_ptr,

2934 out_stride0,

2935 out_stride1,

2936 pm_stride0,

@@ -2896,6 +2938,7 @@ def _sparse_attn_decode_reduce_kernel(

2938 pa_stride0,

2939 pa_stride_s,

2940 pa_stride_h,

2941 + cs_stride,

2942 num_heads,

2943 HAS_ATTN_SINK: tl.constexpr,

2944 ADAPTIVE_SPLITS: tl.constexpr,

@@ -2903,6 +2946,9 @@ def _sparse_attn_decode_reduce_kernel(

2946 BLOCK_H: tl.constexpr,

2947 NUM_SPLITS: tl.constexpr,

2948 SPLITS_PAD: tl.constexpr,

2949 + FUSE_INV_ROPE: tl.constexpr,

2950 + NOPE: tl.constexpr,

2951 + HALF: tl.constexpr,

2952 ):

2953 query_idx = tl.program_id(0)

2954 pid_h = tl.program_id(1)

@@ -2992,6 +3038,26 @@ def _sparse_attn_decode_reduce_kernel(

3038

3039 out = tl.where(l_final[:, None] > 0.0, acc / denom[:, None], 0.0)

3040

3041 + if FUSE_INV_ROPE:

3042 + # Inverse GPT-J RoPE on the trailing rope lanes, straight out of the

3043 + # combine registers: out_even = a*cos + b*sin, out_odd = b*cos - a*sin.

3044 + # NoPE lanes take cos=1/sin=0 so one expression covers the whole row

3045 + # and the o_proj rotation pass disappears for these tokens.

3046 + pos = tl.load(pos_ptr + query_idx).to(tl.int64)

3047 + pair_idx = tl.arange(0, COMB_DIM // 2) - (NOPE // 2)

3048 + is_rope = pair_idx >= 0

3049 + k = tl.where(is_rope, pair_idx, 0)

3050 + cos = tl.where(is_rope, tl.load(cos_sin_ptr + pos * cs_stride + k), 1.0)

3051 + sin = tl.where(is_rope, tl.load(cos_sin_ptr + pos * cs_stride + HALF + k), 0.0)

3052 + even, odd = tl.split(tl.reshape(out, (BLOCK_H, COMB_DIM // 2, 2)))

3053 + out = tl.reshape(

3054 + tl.join(

3055 + even * cos[None, :] + odd * sin[None, :],

3056 + odd * cos[None, :] - even * sin[None, :],

3057 + ),

3058 + (BLOCK_H, COMB_DIM),

3059 + )

3060 +

3061 out_row_ptr = (

3062 out_ptr + query_idx * out_stride0 + head_offsets[:, None] * out_stride1

3063 )

@@ -3312,6 +3378,8 @@ def _rocm_sparse_attn_decode_ragged_triton(

3378 out: torch.Tensor | None = None,

3379 extra_cache_nan_free: bool = False,

3380 adaptive_splits: bool = False,

3381 + inv_rope_positions: torch.Tensor | None = None,

3382 + inv_rope_cos_sin_cache: torch.Tensor | None = None,

3383 ) -> torch.Tensor:

3384 assert q.ndim == 3, f"expected q=[b,h,d], got {q.shape}"

3385 assert main_cache.ndim == 3, (

@@ -3554,12 +3622,27 @@ def _rocm_sparse_attn_decode_ragged_triton(

3622 num_warps=4,

3623 )

3624

3625 + fuse_inv_rope = inv_rope_positions is not None

3626 + if fuse_inv_rope:

3627 + assert inv_rope_cos_sin_cache is not None

3628 + assert inv_rope_cos_sin_cache.shape[-1] == rope_head_dim, (

3629 + "fused inverse RoPE expects cos_sin_cache laid out as "

3630 + f"[P, {rope_head_dim}] = cos | sin, got "

3631 + f"{tuple(inv_rope_cos_sin_cache.shape)}"

3632 + )

3633 + assert nope_head_dim % 2 == 0 and rope_head_dim % 2 == 0, (

3634 + "fused inverse RoPE pairs adjacent lanes, so both head dims must "

3635 + f"be even, got nope={nope_head_dim} rope={rope_head_dim}"

3636 + )

3637 +

3638 _sparse_attn_decode_reduce_kernel[(num_queries, num_heads)](

3639 part_m,

3640 part_l,

3641 part_acc,

3642 attn_sink,

3643 out,

3644 + inv_rope_positions,

3645 + inv_rope_cos_sin_cache,

3646 out.stride(0),

3647 out.stride(1),

3648 part_m.stride(0),

@@ -3567,6 +3650,7 @@ def _rocm_sparse_attn_decode_ragged_triton(

3650 part_acc.stride(0),

3651 part_acc.stride(1),

3652 part_acc.stride(2),

3653 + inv_rope_cos_sin_cache.stride(0) if inv_rope_cos_sin_cache is not None else 0,

3654 num_heads,

3655 HAS_ATTN_SINK=has_attn_sink,

3656 ADAPTIVE_SPLITS=adaptive_splits,

@@ -3574,6 +3658,9 @@ def _rocm_sparse_attn_decode_ragged_triton(

3658 BLOCK_H=1,

3659 NUM_SPLITS=num_splits,

3660 SPLITS_PAD=triton.next_power_of_2(num_splits),

3661 + FUSE_INV_ROPE=fuse_inv_rope,

3662 + NOPE=nope_head_dim,

3663 + HALF=rope_head_dim // 2,

3664 num_warps=4,

3665 )

3666 return out

@@ -3598,6 +3685,8 @@ def _rocm_sparse_attn_decode_triton(

3685 out: torch.Tensor | None = None,

3686 extra_cache_nan_free: bool = False,

3687 adaptive_splits: bool = False,

3688 + inv_rope_positions: torch.Tensor | None = None,

3689 + inv_rope_cos_sin_cache: torch.Tensor | None = None,

3690 ) -> torch.Tensor:

3691 if main_ragged_indices is None or main_ragged_indptr is None:

3692 main_ragged_indices, main_ragged_indptr = build_ragged_indices_from_dense(

@@ -3636,6 +3725,8 @@ def _rocm_sparse_attn_decode_triton(

3725 out=out,

3726 extra_cache_nan_free=extra_cache_nan_free,

3727 adaptive_splits=adaptive_splits,

3728 + inv_rope_positions=inv_rope_positions,

3729 + inv_rope_cos_sin_cache=inv_rope_cos_sin_cache,

3730 )

3731

3732

@@ -3737,7 +3828,16 @@ def rocm_sparse_attn_decode(

3828 output: torch.Tensor,

3829 extra_cache_nan_free: bool = False,

3830 adaptive_splits: bool = False,

3740 -) -> None:

3831 + inv_rope_positions: torch.Tensor | None = None,

3832 + inv_rope_cos_sin_cache: torch.Tensor | None = None,

3833 +) -> int:

3834 + """Run sparse MLA decode into ``output``.

3835 +

3836 + Passing ``inv_rope_positions`` folds the inverse RoPE into the reduce

3837 + epilogue. Returns how many leading rows of ``output`` came back rotated,

3838 + so a caller mixing in a decode path that does not fuse still knows what it

3839 + owes the standalone pass. Read it from the eager attention segment only.

3840 + """

3841 assert swa_k_cache.dtype == torch.uint8, (

3842 "ROCm Triton sparse decode expects uint8 fp8_ds_mla SWA cache, "

3843 f"got {swa_k_cache.dtype}"

@@ -3786,6 +3886,9 @@ def rocm_sparse_attn_decode(

3886 out=direct_out,

3887 extra_cache_nan_free=extra_cache_nan_free,

3888 adaptive_splits=adaptive_splits,

3889 + inv_rope_positions=inv_rope_positions,

3890 + inv_rope_cos_sin_cache=inv_rope_cos_sin_cache,

3891 )

3892 if direct_out is None:

3893 output.copy_(attn_out.to(output.dtype))

3894 + return output.shape[0] if inv_rope_positions is not None else 0