A high-throughput and memory-efficient inference and serving engine for LLMs
[ROCm][DSV4][Perf] Fuse inverse-RoPE and cache bf16 wo_a in o-projection (#45103)
Signed-off-by: Fangzhou Ai <fangzhouai@gmail.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Fangzhou Ai committed
Jun 12, 2026 at 16:57 UTC
badddd254f744d26b6523b464c596f19015370f1
2 files changed
+341
-59
tests/kernels/attention/test_rocm_triton_attn_dsv4.py
+215
@@ -515,3 +515,218 @@ def test_sparse_attn_decode_split_k_kernel(
515
)
516
517
torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2)
518
+
519
+
520
+# ---------------------------------------------------------------------------
521
+# o-projection: fused inverse-RoPE + cached bf16 wo_a (rocm_inv_rope_einsum)
522
+# ---------------------------------------------------------------------------
523
+
524
+
525
+# Cache rows = max_position_embeddings * scaling_factor.
526
+_ROTARY_MAX_POS = 1024
527
+_ROTARY_SCALING_FACTOR = 4.0
528
+_ROTARY_CACHE_LEN = int(_ROTARY_MAX_POS * _ROTARY_SCALING_FACTOR)
529
+
530
+
531
+def _make_dsv4_rotary(device: torch.device):
532
+ """The official DSv4 rotary embedding, sized down for unit tests."""
533
+ from vllm.model_executor.layers.rotary_embedding.deepseek_scaling_rope import (
534
+ DeepseekV4ScalingRotaryEmbedding,
535
+ )
536
+
537
+ # The model loader constructs layers under a default-device context;
538
+ # mirror that so the fp32 cos_sin_cache lands on the GPU.
539
+ with torch.device(device):
540
+ rotary_emb = DeepseekV4ScalingRotaryEmbedding(
541
+ head_size=ROPE_HEAD_DIM,
542
+ rotary_dim=ROPE_HEAD_DIM,
543
+ max_position_embeddings=_ROTARY_MAX_POS,
544
+ base=10000,
545
+ is_neox_style=False,
546
+ scaling_factor=_ROTARY_SCALING_FACTOR,
547
+ dtype=torch.bfloat16,
548
+ mscale=1.0,
549
+ mscale_all_dim=1.0,
550
+ )
551
+ rotary_emb = rotary_emb.to(device)
552
+ assert rotary_emb.cos_sin_cache.shape == (_ROTARY_CACHE_LEN, ROPE_HEAD_DIM)
553
+ return rotary_emb
554
+
555
+
556
+def _inv_rope_via_rotary_native(
557
+ rotary_emb: torch.nn.Module,
558
+ o: torch.Tensor,
559
+ positions: torch.Tensor,
560
+) -> torch.Tensor:
561
+ """Reference: the official ``forward_native(inverse=True)`` path."""
562
+ expected, _ = rotary_emb.forward_native(positions, o.clone(), None, inverse=True)
563
+ return expected.to(torch.bfloat16)
564
+
565
+
566
+class _FakeWoA(torch.nn.Module):
567
+ """Stand-in for the wo_a linear layer holding the (optionally fp8) weight."""
568
+
569
+ def __init__(
570
+ self, weight: torch.Tensor, weight_scale_inv: torch.Tensor | None = None
571
+ ) -> None:
572
+ super().__init__()
573
+ self.weight = weight
574
+ if weight_scale_inv is not None:
575
+ self.weight_scale_inv = weight_scale_inv
576
+
577
+
578
+@pytest.mark.parametrize("num_tokens", [1, 7, 64])
579
+@pytest.mark.parametrize("num_heads", [1, 8])
580
+@pytest.mark.parametrize("pos_dtype", [torch.int32, torch.int64])
581
+@torch.inference_mode()
582
+def test_fused_inverse_rope_gptj_matches_rotary_native(
583
+ num_tokens: int, num_heads: int, pos_dtype: torch.dtype, default_vllm_config
584
+) -> None:
585
+ from vllm.v1.attention.ops.rocm_aiter_mla_sparse import _fused_inverse_rope_gptj
586
+
587
+ device = torch.device("cuda")
588
+ torch.manual_seed(0)
589
+ rotary_emb = _make_dsv4_rotary(device)
590
+ o = torch.randn(
591
+ num_tokens, num_heads, HEAD_DIM, dtype=torch.bfloat16, device=device
592
+ )
593
+ positions = torch.randint(
594
+ 0, _ROTARY_CACHE_LEN, (num_tokens,), dtype=pos_dtype, device=device
595
+ )
596
+
597
+ actual = _fused_inverse_rope_gptj(
598
+ o, positions, rotary_emb.cos_sin_cache, ROPE_HEAD_DIM
599
+ )
600
+ expected = _inv_rope_via_rotary_native(rotary_emb, o, positions)
601
+
602
+ assert actual.dtype == torch.bfloat16
603
+ assert actual.shape == o.shape
604
+ # NoPE lanes are a pure bf16 passthrough -> must be bit-exact.
605
+ assert torch.equal(actual[..., :NOPE_HEAD_DIM], expected[..., :NOPE_HEAD_DIM])
606
+ # RoPE lanes: tolerate at most ~1 bf16 ulp from fp32 fma ordering.
607
+ torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2)
608
+
609
+
610
+@torch.inference_mode()
611
+def test_fused_inverse_rope_gptj_empty(default_vllm_config) -> None:
612
+ from vllm.v1.attention.ops.rocm_aiter_mla_sparse import _fused_inverse_rope_gptj
613
+
614
+ device = torch.device("cuda")
615
+ rotary_emb = _make_dsv4_rotary(device)
616
+ o = torch.empty(0, 8, HEAD_DIM, dtype=torch.bfloat16, device=device)
617
+ positions = torch.empty(0, dtype=torch.int32, device=device)
618
+
619
+ out = _fused_inverse_rope_gptj(
620
+ o, positions, rotary_emb.cos_sin_cache, ROPE_HEAD_DIM
621
+ )
622
+ assert out.shape == (0, 8, HEAD_DIM)
623
+ assert out.dtype == torch.bfloat16
624
+
625
+
626
+@torch.inference_mode()
627
+def test_rocm_inv_rope_einsum_matches_rotary_native(default_vllm_config) -> None:
628
+ from vllm.v1.attention.ops.rocm_aiter_mla_sparse import rocm_inv_rope_einsum
629
+
630
+ device = torch.device("cuda")
631
+ torch.manual_seed(2)
632
+ num_tokens, num_heads = 5, 8
633
+ n_local_groups = num_heads
634
+ o_lora_rank = 16
635
+ hidden_dim = num_heads * HEAD_DIM // n_local_groups # 512
636
+
637
+ rotary_emb = _make_dsv4_rotary(device)
638
+ o = (
639
+ torch.randn(
640
+ num_tokens, num_heads, HEAD_DIM, dtype=torch.bfloat16, device=device
641
+ )
642
+ * 0.125
643
+ )
644
+ positions = torch.randint(
645
+ 0, _ROTARY_CACHE_LEN, (num_tokens,), dtype=torch.int32, device=device
646
+ )
647
+ weight = (
648
+ torch.randn(n_local_groups * o_lora_rank, hidden_dim, device=device) * 0.125
649
+ ).to(torch.bfloat16)
650
+ wo_a = _FakeWoA(weight)
651
+
652
+ actual = rocm_inv_rope_einsum(
653
+ rotary_emb, o, positions, ROPE_HEAD_DIM, n_local_groups, o_lora_rank, wo_a
654
+ )
655
+
656
+ o_ref = _inv_rope_via_rotary_native(rotary_emb, o, positions)
657
+ o_ref = o_ref.view(num_tokens, n_local_groups, -1)
658
+ wo_a_ref = weight.view(n_local_groups, o_lora_rank, hidden_dim).to(torch.bfloat16)
659
+ expected = torch.einsum("tgd,grd->tgr", o_ref, wo_a_ref)
660
+
661
+ assert actual.shape == (num_tokens, n_local_groups, o_lora_rank)
662
+ torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2)
663
+
664
+
665
+@torch.inference_mode()
666
+def test_get_cached_wo_a_bf16_plain_caches() -> None:
667
+ from vllm.v1.attention.ops.rocm_aiter_mla_sparse import _get_cached_wo_a_bf16
668
+
669
+ device = torch.device("cuda")
670
+ torch.manual_seed(4)
671
+ n_local_groups, o_lora_rank, hidden_dim = 2, 4, 8
672
+ weight = torch.randn(
673
+ n_local_groups * o_lora_rank, hidden_dim, dtype=torch.bfloat16, device=device
674
+ )
675
+ wo_a = _FakeWoA(weight)
676
+
677
+ out1 = _get_cached_wo_a_bf16(wo_a, n_local_groups, o_lora_rank, hidden_dim)
678
+ expected = weight.view(n_local_groups, o_lora_rank, hidden_dim).to(torch.bfloat16)
679
+ assert out1.shape == (n_local_groups, o_lora_rank, hidden_dim)
680
+ torch.testing.assert_close(out1, expected, atol=0, rtol=0)
681
+ assert hasattr(wo_a, "_dsv4_wo_a_bf16")
682
+
683
+ # Mutate the source weight: the cached tensor must be returned unchanged
684
+ # (proving the dequant is not recomputed per call).
685
+ wo_a.weight.zero_()
686
+ out2 = _get_cached_wo_a_bf16(wo_a, n_local_groups, o_lora_rank, hidden_dim)
687
+ assert out2 is out1
688
+ torch.testing.assert_close(out2, expected, atol=0, rtol=0)
689
+
690
+
691
+@torch.inference_mode()
692
+def test_get_cached_wo_a_bf16_fp8_blockscale_caches() -> None:
693
+ from vllm.v1.attention.ops.rocm_aiter_mla_sparse import _get_cached_wo_a_bf16
694
+
695
+ device = torch.device("cuda")
696
+ torch.manual_seed(5)
697
+ n_local_groups, o_lora_rank, hidden_dim = 2, 4, 8
698
+ row_block, col_block = 2, 2
699
+ row_blocks = o_lora_rank // row_block
700
+ col_blocks = hidden_dim // col_block
701
+
702
+ fp8_dtype = current_platform.fp8_dtype()
703
+ weight_f32 = (
704
+ torch.randn(
705
+ n_local_groups, o_lora_rank, hidden_dim, dtype=torch.float32, device=device
706
+ )
707
+ * 0.1
708
+ )
709
+ weight_fp8 = weight_f32.to(fp8_dtype)
710
+ scale = (
711
+ torch.rand(
712
+ n_local_groups, row_blocks, col_blocks, dtype=torch.float32, device=device
713
+ )
714
+ * 0.5
715
+ + 0.5
716
+ )
717
+ wo_a = _FakeWoA(
718
+ weight_fp8.reshape(n_local_groups * o_lora_rank, hidden_dim),
719
+ weight_scale_inv=scale.reshape(n_local_groups * row_blocks, col_blocks),
720
+ )
721
+
722
+ out = _get_cached_wo_a_bf16(wo_a, n_local_groups, o_lora_rank, hidden_dim)
723
+
724
+ scale_full = scale.repeat_interleave(row_block, dim=-2).repeat_interleave(
725
+ col_block, dim=-1
726
+ )
727
+ expected = (weight_fp8.to(torch.float32) * scale_full).to(torch.bfloat16)
728
+ assert out.shape == (n_local_groups, o_lora_rank, hidden_dim)
729
+ torch.testing.assert_close(out, expected, atol=0, rtol=0)
730
+
731
+ # Second call returns the same cached object.
732
+ assert _get_cached_wo_a_bf16(wo_a, n_local_groups, o_lora_rank, hidden_dim) is out
vllm/v1/attention/ops/rocm_aiter_mla_sparse.py
+126
-59
@@ -874,72 +874,113 @@ def _expand_2d_block_scales(
874
return scale
875
876
877
-def _apply_gptj_inv_rope_ref(
878
- x: torch.Tensor,
877
+@triton.jit
878
+def _inverse_rope_gptj_kernel(
879
+ o_ptr, # [T, H, D] input
880
+ out_ptr, # [T, H, D] bf16 output
881
+ pos_ptr, # [T] positions
882
+ cos_sin_ptr, # [P, rope_dim] fp32 (cos[:half] | sin[half:])
883
+ s_t,
884
+ s_h, # input row strides (last dim contiguous)
885
+ os_t,
886
+ os_h, # output row strides
887
+ cs_stride, # cos_sin_cache row stride
888
+ NOPE: tl.constexpr, # non-rope head dims (passed through)
889
+ HALF: tl.constexpr, # rope_dim // 2
890
+ BLOCK_NOPE: tl.constexpr,
891
+ BLOCK_HALF: tl.constexpr,
892
+):
893
+ """Fused inverse GPT-J RoPE on the trailing rope_dim of each (token, head).
894
+
895
+ Mirrors ``DeepseekV4ScalingRotaryEmbedding.forward_native(inverse=True)``
896
+ for the GPT-J (non-neox) layout, writing bf16 directly. Replaces the
897
+ clone + index_select + repeat_interleave + neg + stack + cat + cast chain
898
+ (~10 small kernels) with a single launch.
899
+ """
900
+ t = tl.program_id(0)
901
+ h = tl.program_id(1)
902
+ in_base = t * s_t + h * s_h
903
+ out_base = t * os_t + h * os_h
904
+
905
+ # NoPE lanes pass through unchanged (only cast to bf16).
906
+ n = tl.arange(0, BLOCK_NOPE)
907
+ nmask = n < NOPE
908
+ vals = tl.load(o_ptr + in_base + n, mask=nmask)
909
+ tl.store(out_ptr + out_base + n, vals.to(tl.bfloat16), mask=nmask)
910
+
911
+ # RoPE lanes: out_even = a*cos + b*sin, out_odd = b*cos - a*sin
912
+ # (a = even lane, b = odd lane; sin negated for the inverse rotation).
913
+ pos = tl.load(pos_ptr + t).to(tl.int64)
914
+ k = tl.arange(0, BLOCK_HALF)
915
+ kmask = k < HALF
916
+ a = tl.load(o_ptr + in_base + NOPE + 2 * k, mask=kmask).to(tl.float32)
917
+ b = tl.load(o_ptr + in_base + NOPE + 2 * k + 1, mask=kmask).to(tl.float32)
918
+ cos = tl.load(cos_sin_ptr + pos * cs_stride + k, mask=kmask)
919
+ sin = tl.load(cos_sin_ptr + pos * cs_stride + HALF + k, mask=kmask)
920
+ out_even = a * cos + b * sin
921
+ out_odd = b * cos - a * sin
922
+ tl.store(out_ptr + out_base + NOPE + 2 * k, out_even.to(tl.bfloat16), mask=kmask)
923
+ tl.store(out_ptr + out_base + NOPE + 2 * k + 1, out_odd.to(tl.bfloat16), mask=kmask)
924
+
925
+
926
+def _fused_inverse_rope_gptj(
927
+ o: torch.Tensor,
928
positions: torch.Tensor,
929
cos_sin_cache: torch.Tensor,
881
- rope_dim: int,
882
-) -> torch.Tensor:
883
- if rope_dim == 0 or x.numel() == 0:
884
- return x
885
- half_rot = rope_dim // 2
886
- nope_dim = x.shape[-1] - rope_dim
887
- dtype = x.dtype
888
- x = x.to(torch.float32)
889
- cache = cos_sin_cache.index_select(0, positions.to(torch.long))
890
- cos = cache[:, :half_rot].to(torch.float32)
891
- sin = cache[:, half_rot : 2 * half_rot].to(torch.float32)
892
- view_shape = (positions.shape[0],) + (1,) * (x.dim() - 2) + (half_rot,)
893
- cos = cos.view(view_shape)
894
- sin = sin.view(view_shape)
895
- rope = x[..., nope_dim:]
896
- y_even = rope[..., 0::2]
897
- y_odd = rope[..., 1::2]
898
- rope_out = torch.stack(
899
- (y_even * cos + y_odd * sin, y_odd * cos - y_even * sin),
900
- dim=-1,
901
- ).flatten(-2)
902
- x = x.clone()
903
- x[..., nope_dim:] = rope_out
904
- return x.to(dtype)
905
-
906
-
907
-def _apply_inv_rope_ref(
908
- rotary_emb: torch.nn.Module,
909
- x: torch.Tensor,
910
- positions: torch.Tensor,
911
- rope_dim: int,
930
+ rope_head_dim: int,
931
) -> torch.Tensor:
913
- if hasattr(rotary_emb, "forward_native"):
914
- try:
915
- query, _ = rotary_emb.forward_native(
916
- positions,
917
- x.clone(),
918
- None,
919
- inverse=True,
920
- )
921
- return query
922
- except TypeError:
923
- pass
924
- return _apply_gptj_inv_rope_ref(x, positions, rotary_emb.cos_sin_cache, rope_dim)
932
+ """bf16 inverse GPT-J RoPE via a single fused Triton kernel."""
933
+ assert o.dim() == 3 and o.stride(-1) == 1, (
934
+ "_fused_inverse_rope_gptj expects a [T, H, D] input with a contiguous last dim"
935
+ )
936
+ assert rope_head_dim > 0 and rope_head_dim % 2 == 0, (
937
+ f"_fused_inverse_rope_gptj expects an even rope_head_dim, got {rope_head_dim}"
938
+ )
939
+ assert cos_sin_cache.shape[-1] == rope_head_dim, (
940
+ "_fused_inverse_rope_gptj expects cos_sin_cache laid out as "
941
+ f"[P, {rope_head_dim}] = cos | sin, got {tuple(cos_sin_cache.shape)}"
942
+ )
943
+ num_tokens, num_heads, head_dim = o.shape
944
+ out = torch.empty(
945
+ (num_tokens, num_heads, head_dim), dtype=torch.bfloat16, device=o.device
946
+ )
947
+ if num_tokens == 0:
948
+ return out
949
+ _inverse_rope_gptj_kernel[(num_tokens, num_heads)](
950
+ o,
951
+ out,
952
+ positions,
953
+ cos_sin_cache,
954
+ o.stride(0),
955
+ o.stride(1),
956
+ out.stride(0),
957
+ out.stride(1),
958
+ cos_sin_cache.stride(0),
959
+ NOPE=head_dim - rope_head_dim,
960
+ HALF=rope_head_dim // 2,
961
+ BLOCK_NOPE=triton.next_power_of_2(head_dim - rope_head_dim),
962
+ BLOCK_HALF=triton.next_power_of_2(rope_head_dim // 2),
963
+ )
964
+ return out
965
966
927
-def rocm_inv_rope_einsum(
928
- rotary_emb: torch.nn.Module,
929
- o: torch.Tensor,
930
- positions: torch.Tensor,
931
- rope_head_dim: int,
967
+def _get_cached_wo_a_bf16(
968
+ wo_a: torch.nn.Module,
969
n_local_groups: int,
970
o_lora_rank: int,
934
- wo_a: torch.nn.Module,
971
+ hidden_dim: int,
972
) -> torch.Tensor:
936
- """Reference inverse-RoPE + WO_A einsum path used on ROCm."""
937
- o_ref = _apply_inv_rope_ref(rotary_emb, o, positions, rope_head_dim).to(
938
- torch.bfloat16
939
- )
940
- o_ref = o_ref.view(o.shape[0], n_local_groups, -1)
973
+ """Dequantize wo_a to bf16 once and cache it on the module.
974
942
- hidden_dim = o_ref.shape[-1]
975
+ wo_a weights are static, so the fp8 -> fp32 -> (* block scale) -> bf16
976
+ dequant only needs to run once. Recomputing it every decode step shows up
977
+ in the profile as the largest copy/mul kernels (``direct_copy float`` ~55us
978
+ and ``MulFunctor float`` ~31us per two layers). SGLang / ATOM keep wo_a in
979
+ bf16 and feed a plain bf16 GEMM; this mirrors that.
980
+ """
981
+ cached = getattr(wo_a, "_dsv4_wo_a_bf16", None)
982
+ if cached is not None:
983
+ return cached
984
if hasattr(wo_a, "weight_scale_inv"):
985
wo_a_weight = wo_a.weight.view(n_local_groups, o_lora_rank, hidden_dim).to(
986
torch.float32
@@ -951,11 +992,37 @@ def rocm_inv_rope_einsum(
992
o_lora_rank,
993
hidden_dim,
994
)
954
- wo_a_weight = (wo_a_weight * wo_a_scale).to(torch.bfloat16)
995
+ cached = (wo_a_weight * wo_a_scale).to(torch.bfloat16)
996
else:
956
- wo_a_weight = wo_a.weight.view(n_local_groups, o_lora_rank, hidden_dim).to(
997
+ cached = wo_a.weight.view(n_local_groups, o_lora_rank, hidden_dim).to(
998
torch.bfloat16
999
)
1000
+ wo_a._dsv4_wo_a_bf16 = cached
1001
+ return cached
1002
+
1003
+
1004
+def rocm_inv_rope_einsum(
1005
+ rotary_emb: torch.nn.Module,
1006
+ o: torch.Tensor,
1007
+ positions: torch.Tensor,
1008
+ rope_head_dim: int,
1009
+ n_local_groups: int,
1010
+ o_lora_rank: int,
1011
+ wo_a: torch.nn.Module,
1012
+) -> torch.Tensor:
1013
+ """Inverse-RoPE + WO_A bmm path used on ROCm.
1014
+
1015
+ Fuses the inverse GPT-J RoPE into one Triton kernel and caches the bf16
1016
+ wo_a weight so the per-step dequant disappears.
1017
+ """
1018
+ o_ref = _fused_inverse_rope_gptj(
1019
+ o, positions, rotary_emb.cos_sin_cache, rope_head_dim
1020
+ )
1021
+ o_ref = o_ref.view(o.shape[0], n_local_groups, -1)
1022
+
1023
+ wo_a_weight = _get_cached_wo_a_bf16(
1024
+ wo_a, n_local_groups, o_lora_rank, o_ref.shape[-1]
1025
+ )
1026
1027
return torch.einsum("tgd,grd->tgr", o_ref, wo_a_weight)
1028